#include "../include/rcore.h"
#include "../core/allocations.h"
#include "../core/covariance.matrix.h"
#include "../core/data.table.h"
#include "../core/moments.h"
#include "../math/linear.algebra.h"
#include "../minimal/common.h"
#include "../minimal/data.frame.h"
#include "scores.h"

double wpost(double *xx, int data_ncols, int nobs, double alpha_mu, double *nu,
    double alpha_w) {

double r = 0, t = 0, xm = 0, xsse = 0;
long double logprob = 0;

  /* first term. */
  logprob = 0.5 * (log(alpha_mu) - log(nobs + alpha_mu));

  /* Gamma_l ratio in the second term. */
  logprob += lgammafn(0.5 * (nobs + alpha_w - data_ncols + 1)) -
             lgammafn(0.5 * (alpha_w - data_ncols+ 1));

  /* leftover from the second term. */
  logprob -= 0.5 * nobs * log(M_PI);

  /* third term, numerator. */
  t = alpha_mu * (alpha_w - data_ncols - 1) / (alpha_mu + 1);
  logprob += 0.5 * (alpha_w - data_ncols + 1) * log(t);

  /* third term, denominator. */
  xm = c_mean(xx, nobs);
  xsse = c_sse(xx, xm, nobs);

  r = t + xsse + (nobs * alpha_mu) / (nobs + alpha_mu) * (xm - *nu) * (xm - *nu);

  logprob -= 0.5 * (nobs + alpha_w - data_ncols + 1) * log(r);

  return (double)logprob;

}/*WPOST*/

double cwpost(double *xx, double **parents, int p, int data_ncols, int nobs,
    double alpha_mu, double *nu, double alpha_w) {

int i = 0, j = 0;
double t = 0;
long double logprob = 0;
tabular dtx = { 0 };
covariance R = { 0 }, Rtilde = { 0 };

  /* first term. */
  logprob = 0.5 * (log(alpha_mu) - log(nobs + alpha_mu));

  /* Gamma_l ratio in the second term. */
  logprob += lgammafn(0.5 * (nobs + alpha_w - data_ncols + p + 1)) -
             lgammafn(0.5 * (alpha_w - data_ncols + p + 1));

  /* leftover from the second term. */
  logprob -= 0.5 * nobs * log(M_PI);

  /* third term, ratio of the determinants of the prior T matrices. */
  t = alpha_mu * (alpha_w - data_ncols - 1) / (alpha_mu + 1);

  logprob += 0.5 * (alpha_w - data_ncols + p + 1) * (p + 1) * log(t) -
             0.5 * (alpha_w - data_ncols + p) * (p) * log(t);

  /* third term, ratio of the determinants of the posterior R matrices; the
   * node fills the leading continuous slot, the parents the following ones. */
  dtx = empty_tabular(nobs, 0, p + 1);
  dtx.ccol[0] = xx;
  for (i = 0; i < p; i++)
    dtx.ccol[i + 1] = parents[i];
  tabular_cache_means(&dtx, 0);
  R = new_covariance(dtx.m.ncols, FALSE);
  Rtilde = new_covariance(dtx.m.ncols - 1, FALSE);

  /* compute the rescaled covariance matrix. */
  c_covmat(dtx.ccol, dtx.mean, dtx.m.nobs, dtx.m.ncols, R, 0);
  for (i = 0; i < R.dim * R.dim; i++)
    R.mat[i] *= nobs - 1;

  /* add the prior matrix, diagonal with elements t. */
  for (i = 0; i < R.dim; i++)
    R.mat[CMC(i, i, R.dim)] += t;

  /* add the outer product of the difference in observed and prior means. */
  for (i = 0; i < R.dim; i++)
    for (j = 0; j < R.dim; j++) {

      R.mat[CMC(i, j, R.dim)] += (nobs * alpha_mu) / (nobs + alpha_mu) *
                                  (dtx.mean[i] - nu[i]) * (dtx.mean[j] - nu[j]);

  }/*FOR*/

  /* subset the rescaled covariance matrix to include only the parents. */
  covariance_drop_variable(&R, &Rtilde, 0);

  logprob += 0.5 * (nobs + alpha_w - data_ncols + p) *
               c_logdet(Rtilde.mat, Rtilde.dim);

  logprob -= 0.5 * (nobs + alpha_w - data_ncols + p + 1) *
               c_logdet(R.mat, R.dim);

  FreeTAB(dtx);
  FreeCOV(R);
  FreeCOV(Rtilde);

  return (double)logprob;

}/*CWPOST*/

double wishart_node(SEXP target, sparse_amat sam, tabular dt, SEXP iss_mu,
    SEXP nu, SEXP iss_w, SEXP prior, SEXP beta, bool debugging) {

int i = 0, node = 0, npar = 0, it = 0, ncols = dt.m.ncols, nobs = dt.m.nobs;
char *t = (char *)CHAR(STRING_ELT(target, 0));
double prob = 0, prior_prob = 0, **par_cols = NULL;
SEXP parents, nu_subset;

  /* get the parents of the node from the sparse adjacency matrix, as labels so
   * they can index the prior mean vector by name. */
  it = sparse_amat_node(sam, t);
  npar = sam.rowptr[it + 1] - sam.rowptr[it];
  PROTECT(parents = allocVector(STRSXP, npar));
  for (i = 0; i < npar; i++)
    SET_STRING_ELT(parents, i, mkChar(sam.labels[sam.colidx[sam.rowptr[it] + i]]));

  /* locate the node's (continuous) column in the tabular data. */
  node = dt.map[tabular_column(dt, t)];

  /* compute the prior probability component for the node. */
  prior_prob = graph_prior_prob(prior, target, beta, sam, debugging);
  /* extract the elements of the prior mean vector. */
  PROTECT(nu_subset = subset_by_name(nu, 2, target, parents));

  if (npar == 0) {

    prob = wpost(dt.ccol[node], ncols, nobs, NUM(iss_mu), REAL(nu_subset),
             NUM(iss_w));

  }/*THEN*/
  else {

    /* gather the parents' columns. */
    par_cols = Calloc1D(npar, sizeof(double *));
    for (i = 0; i < npar; i++)
      par_cols[i] = dt.ccol[dt.map[tabular_column(dt, CHAR(STRING_ELT(parents, i)))]];

    /* compute the marginal likelihood. */
    prob = cwpost(dt.ccol[node], par_cols, npar, ncols, nobs, NUM(iss_mu),
             REAL(nu_subset), NUM(iss_w));

    Free1D(par_cols);

  }/*ELSE*/

  if (debugging) {

    Rprintf("  > (log)prior probability is %lf.\n", prior_prob);
    Rprintf("  > (log)posterior density is %lf.\n", prob);

  }/*THEN*/

  /* add the (log)prior to the marginal (log)likelihood to get the (log)posterior. */
  prob += prior_prob;

  UNPROTECT(2);

  return prob;

}/*WISHART_NODE*/
