#include <R.h>
#include <Rinternals.h>
#include <Rmath.h>
#include <R_ext/Rdynload.h>
#include <math.h>

static double log1pexp_stable(double x) {
  if (x > 35.0) return x;
  if (x < -35.0) return exp(x);
  return log1p(exp(x));
}

static double logsumexp_vec(const double *x, int n) {
  double m = x[0];
  for (int i = 1; i < n; i++) {
    if (x[i] > m) m = x[i];
  }
  double s = 0.0;
  for (int i = 0; i < n; i++) {
    s += exp(x[i] - m);
  }
  return m + log(s);
}

static void estep_core(
    const double *W, const int *T, const int *G, int n,
    const double *par, int free_delta0,
    int grid_size, double grid_sd,
    double *post_mean, double *post_second, double *total_ll,
    double *xgrid_mat, double *wpost_mat) {

  double *log_kernel = (double *) R_alloc(grid_size, sizeof(double));
  double *x_grid = (double *) R_alloc(grid_size, sizeof(double));
  double *w_post = (double *) R_alloc(grid_size, sizeof(double));

  double mu0 = par[0];
  double tau = par[1];
  double sigmaX2 = par[2];
  double alpha0 = par[3];
  double alphaT = par[4];
  double alphaX = par[5];
  double delta0 = free_delta0 ? par[6] : 0.0;
  double delta1 = free_delta0 ? par[7] : par[6];
  double sigmaW0_2 = free_delta0 ? par[8] : par[7];
  double sigmaW1_2 = free_delta0 ? par[9] : par[8];

  *total_ll = 0.0;

  for (int i = 0; i < n; i++) {
    double Ti = (double) T[i];
    int Gi = G[i];
    double Wi = W[i];

    double mu_prior = mu0 + tau * Ti;
    double delta_g = (Gi == 0) ? delta0 : delta1;
    double sig_wg2 = (Gi == 0) ? sigmaW0_2 : sigmaW1_2;

    double prec = 1.0 / sigmaX2 + 1.0 / sig_wg2;
    double v0 = 1.0 / prec;
    double m0 = v0 * (mu_prior / sigmaX2 + (Wi - delta_g) / sig_wg2);
    double s0 = sqrt(v0);

    double xmin = m0 - grid_sd * s0;
    double xmax = m0 + grid_sd * s0;
    double dx = (xmax - xmin) / (grid_size - 1);

    for (int k = 0; k < grid_size; k++) {
      double x = xmin + dx * k;
      x_grid[k] = x;
      double eta = alpha0 + alphaT * Ti + alphaX * (x - mu0);
      double log_prior = dnorm4(x, mu_prior, sqrt(sigmaX2), 1);
      double log_meas = dnorm4(Wi, x + delta_g, sqrt(sig_wg2), 1);
      double log_resp = ((double) Gi) * eta - log1pexp_stable(eta);
      log_kernel[k] = log_prior + log_meas + log_resp;
    }

    double log_norm = logsumexp_vec(log_kernel, grid_size);
    double ex = 0.0;
    double ex2 = 0.0;
    for (int k = 0; k < grid_size; k++) {
      w_post[k] = exp(log_kernel[k] - log_norm);
      ex += w_post[k] * x_grid[k];
      ex2 += w_post[k] * x_grid[k] * x_grid[k];
      if (xgrid_mat != NULL) xgrid_mat[i + n * k] = x_grid[k];
      if (wpost_mat != NULL) wpost_mat[i + n * k] = w_post[k];
    }
    post_mean[i] = ex;
    post_second[i] = ex2;
    *total_ll += logsumexp_vec(log_kernel, grid_size) + log(dx);
  }
}

SEXP c_e_step_loglik_free(SEXP W_, SEXP T_, SEXP G_, SEXP par_, SEXP grid_size_, SEXP grid_sd_) {
  int n = LENGTH(W_);
  const double *W = REAL(W_);
  const int *T = INTEGER(T_);
  const int *G = INTEGER(G_);
  const double *par = REAL(par_);
  int grid_size = INTEGER(grid_size_)[0];
  double grid_sd = REAL(grid_sd_)[0];

  SEXP out = PROTECT(allocVector(VECSXP, 5));
  SEXP post_mean_ = PROTECT(allocVector(REALSXP, n));
  SEXP post_second_ = PROTECT(allocVector(REALSXP, n));
  SEXP loglik_ = PROTECT(allocVector(REALSXP, 1));
  SEXP xgrid_ = PROTECT(allocMatrix(REALSXP, n, grid_size));
  SEXP wpost_ = PROTECT(allocMatrix(REALSXP, n, grid_size));

  double total_ll = 0.0;
  estep_core(W, T, G, n, par, 1, grid_size, grid_sd,
             REAL(post_mean_), REAL(post_second_), &total_ll,
             REAL(xgrid_), REAL(wpost_));
  REAL(loglik_)[0] = total_ll;

  SET_VECTOR_ELT(out, 0, post_mean_);
  SET_VECTOR_ELT(out, 1, post_second_);
  SET_VECTOR_ELT(out, 2, loglik_);
  SET_VECTOR_ELT(out, 3, xgrid_);
  SET_VECTOR_ELT(out, 4, wpost_);
  SEXP names = PROTECT(allocVector(STRSXP, 5));
  SET_STRING_ELT(names, 0, mkChar("post_mean"));
  SET_STRING_ELT(names, 1, mkChar("post_second"));
  SET_STRING_ELT(names, 2, mkChar("loglik"));
  SET_STRING_ELT(names, 3, mkChar("xgrid"));
  SET_STRING_ELT(names, 4, mkChar("wpost"));
  setAttrib(out, R_NamesSymbol, names);

  UNPROTECT(7);
  return out;
}

SEXP c_e_step_loglik_fixed(SEXP W_, SEXP T_, SEXP G_, SEXP par_, SEXP grid_size_, SEXP grid_sd_) {
  int n = LENGTH(W_);
  const double *W = REAL(W_);
  const int *T = INTEGER(T_);
  const int *G = INTEGER(G_);
  const double *par = REAL(par_);
  int grid_size = INTEGER(grid_size_)[0];
  double grid_sd = REAL(grid_sd_)[0];

  SEXP out = PROTECT(allocVector(VECSXP, 5));
  SEXP post_mean_ = PROTECT(allocVector(REALSXP, n));
  SEXP post_second_ = PROTECT(allocVector(REALSXP, n));
  SEXP loglik_ = PROTECT(allocVector(REALSXP, 1));
  SEXP xgrid_ = PROTECT(allocMatrix(REALSXP, n, grid_size));
  SEXP wpost_ = PROTECT(allocMatrix(REALSXP, n, grid_size));

  double total_ll = 0.0;
  estep_core(W, T, G, n, par, 0, grid_size, grid_sd,
             REAL(post_mean_), REAL(post_second_), &total_ll,
             REAL(xgrid_), REAL(wpost_));
  REAL(loglik_)[0] = total_ll;

  SET_VECTOR_ELT(out, 0, post_mean_);
  SET_VECTOR_ELT(out, 1, post_second_);
  SET_VECTOR_ELT(out, 2, loglik_);
  SET_VECTOR_ELT(out, 3, xgrid_);
  SET_VECTOR_ELT(out, 4, wpost_);
  SEXP names = PROTECT(allocVector(STRSXP, 5));
  SET_STRING_ELT(names, 0, mkChar("post_mean"));
  SET_STRING_ELT(names, 1, mkChar("post_second"));
  SET_STRING_ELT(names, 2, mkChar("loglik"));
  SET_STRING_ELT(names, 3, mkChar("xgrid"));
  SET_STRING_ELT(names, 4, mkChar("wpost"));
  setAttrib(out, R_NamesSymbol, names);

  UNPROTECT(7);
  return out;
}

SEXP c_loglik_free(SEXP W_, SEXP T_, SEXP G_, SEXP par_, SEXP grid_size_, SEXP grid_sd_) {
  int n = LENGTH(W_);
  const double *W = REAL(W_);
  const int *T = INTEGER(T_);
  const int *G = INTEGER(G_);
  const double *par = REAL(par_);
  int grid_size = INTEGER(grid_size_)[0];
  double grid_sd = REAL(grid_sd_)[0];

  SEXP out = PROTECT(allocVector(REALSXP, 1));
  double *post_mean = (double *) R_alloc(n, sizeof(double));
  double *post_second = (double *) R_alloc(n, sizeof(double));
  double total_ll = 0.0;
  estep_core(W, T, G, n, par, 1, grid_size, grid_sd, post_mean, post_second, &total_ll, NULL, NULL);
  REAL(out)[0] = total_ll;
  UNPROTECT(1);
  return out;
}

SEXP c_loglik_fixed(SEXP W_, SEXP T_, SEXP G_, SEXP par_, SEXP grid_size_, SEXP grid_sd_) {
  int n = LENGTH(W_);
  const double *W = REAL(W_);
  const int *T = INTEGER(T_);
  const int *G = INTEGER(G_);
  const double *par = REAL(par_);
  int grid_size = INTEGER(grid_size_)[0];
  double grid_sd = REAL(grid_sd_)[0];

  SEXP out = PROTECT(allocVector(REALSXP, 1));
  double *post_mean = (double *) R_alloc(n, sizeof(double));
  double *post_second = (double *) R_alloc(n, sizeof(double));
  double total_ll = 0.0;
  estep_core(W, T, G, n, par, 0, grid_size, grid_sd, post_mean, post_second, &total_ll, NULL, NULL);
  REAL(out)[0] = total_ll;
  UNPROTECT(1);
  return out;
}

SEXP c_q_alpha(SEXP xgrid_, SEXP wpost_, SEXP T_, SEXP G_, SEXP alpha_, SEXP mu0_) {
  int n = LENGTH(T_);
  const int *T = INTEGER(T_);
  const int *G = INTEGER(G_);
  const double *xgrid = REAL(xgrid_);
  const double *wpost = REAL(wpost_);
  const double *alpha = REAL(alpha_);
  SEXP dim = getAttrib(xgrid_, R_DimSymbol);
  int grid_size = INTEGER(dim)[1];

  double alpha0 = alpha[0];
  double alphaT = alpha[1];
  double alphaX = alpha[2];
  double mu0 = REAL(mu0_)[0];
  SEXP out = PROTECT(allocVector(REALSXP, 1));
  double total = 0.0;
  for (int i = 0; i < n; i++) {
    double Ti = (double) T[i];
    int Gi = G[i];
    for (int k = 0; k < grid_size; k++) {
      double x = xgrid[i + n * k];
      double w = wpost[i + n * k];
      double eta = alpha0 + alphaT * Ti + alphaX * (x - mu0);
      total += w * (((double) Gi) * eta - log1pexp_stable(eta));
    }
  }
  REAL(out)[0] = total;
  UNPROTECT(1);
  return out;
}

static const R_CallMethodDef CallEntries[] = {
  {"c_e_step_loglik_free",  (DL_FUNC) &c_e_step_loglik_free,  6},
  {"c_e_step_loglik_fixed", (DL_FUNC) &c_e_step_loglik_fixed, 6},
  {"c_loglik_free",         (DL_FUNC) &c_loglik_free,         6},
  {"c_loglik_fixed",        (DL_FUNC) &c_loglik_fixed,        6},
  {"c_q_alpha",             (DL_FUNC) &c_q_alpha,             6},
  {NULL, NULL, 0}
};

void R_init_prome(DllInfo *dll) {
  R_registerRoutines(dll, NULL, CallEntries, NULL, NULL);
  R_useDynamicSymbols(dll, FALSE);
}
