#include "../include/rcore.h"
#include "../core/allocations.h"
#include "../core/sets.h"
#include "../core/sort.h"
#include "../minimal/common.h"
#include "../minimal/data.frame.h"
#include "scores.h"

/* relabel raw configuration codes to a dense 1-based range, mirroring the
 * sparse factor produced by int2fac(); return the number of distinct codes.
 * assumes complete data, which the Dirichlet scores require. */
static int densify_config(int *config, int nobs) {

int i = 0, k = 0;
int *sorted = NULL;

  if (nobs == 0)
    return 0;

  /* collect the distinct configuration codes in ascending order. */
  sorted = Calloc1D(nobs, sizeof(int));
  memcpy(sorted, config, nobs * sizeof(int));
  i_sort(sorted, NULL, nobs);
  for (i = 1; i < nobs; i++)
    if (sorted[i] != sorted[k])
      sorted[++k] = sorted[i];
  k++;

  /* remap every configuration to its 1-based rank by binary search. */
  for (i = 0; i < nobs; i++) {

    int lo = 0, hi = k - 1, mid = 0;

    while (lo <= hi) {

      mid = (lo + hi) / 2;
      if (sorted[mid] < config[i])
        lo = mid + 1;
      else if (sorted[mid] > config[i])
        hi = mid - 1;
      else
        break;

    }/*WHILE*/

    config[i] = mid + 1;

  }/*FOR*/

  Free1D(sorted);

  return k;

}/*DENSIFY_CONFIG*/

/* posterior Dirichlet probability (covers BD and K2 scores). */
double dpost(int *xx, int llx, int num, SEXP iss, int per_cell, SEXP exp) {

int i = 0, k = 0, *n = NULL;
double imaginary = 0, alpha = 0, res = 0;

  /* the correct vaules for the hyperparameters alpha are documented in
   * "Learning Bayesian Networks: The Combination of Knowledge and Statistical
   * Data" by Heckerman, Geiger & Chickering (1995), page 17. */

  if (per_cell) {

    /* this is for K2 and BDJ, which does not define an imaginary sample size;
     * all hyperparameters are set to 1 or 1/2 in the prior distribution. */
    imaginary = llx;
    alpha = NUM(iss);

  }/*THEN*/
  else {

    /* this is for the BDe and BDs scores. */
    imaginary = NUM(iss);
    alpha = imaginary / llx;

  }/*ELSE*/

  /* initialize the contingency table. */
  n = Calloc1D(llx, sizeof(int));

  /* compute the frequency table of x, disregarding experimental data. */
  if (exp == R_NilValue) {

    for (i = 0; i < num; i++)
      n[xx[i] - 1]++;

  }/*THEN*/
  else {

    int nexp = length(exp), *e = INTEGER(exp);

    for (i = 0, k = 0; i < num; i++) {
      if ((k >= nexp) || (i != e[k] - 1))
        n[xx[i] - 1]++;
      else
        k++;

    }/*FOR*/

    /* adjust the sample size to match the number of observational data. */
    num -= nexp;

  }/*ELSE*/

  /* compute the posterior probability. */
  for (i = 0; i < llx; i++)
    res += lgammafn(n[i] + alpha) - lgammafn(alpha);
  res += lgammafn(imaginary) - lgammafn(imaginary + num);

  Free1D(n);

  return res;

}/*DPOST*/

/* conditional posterior Dirichlet probability (covers BD and K2 scores). */
double cdpost(int *xx, int llx, int *yy, int lly, int num, SEXP iss,
    int per_cell, SEXP exp) {

int i = 0, j = 0, k = 0, p = llx * lly;
int **n = NULL, *nj = NULL;
double imaginary = 0, alpha = 0, res = 0;

  if (per_cell) {

    /* this is for K2 and BDJ, which does not define an imaginary sample size;
     * all hyperparameters are set to 1 or 1/2 in the prior distribution. */
    alpha = NUM(iss);
    imaginary = alpha * p;

  }/*THEN*/
  else {

    /* this is for the BDe and BDs scores. */
    imaginary = NUM(iss);
    alpha = imaginary / p;

  }/*ELSE*/

  /* initialize the contingency table. */
  n = (int **) Calloc2D(llx, lly, sizeof(int));
  nj = Calloc1D(lly, sizeof(int));

  /* compute the joint frequency of x and y. */
  if (exp == R_NilValue) {

    for (i = 0; i < num; i++) {

      n[xx[i] - 1][yy[i] - 1]++;
      nj[yy[i] - 1]++;

    }/*FOR*/

  }/*THEN*/
  else {

    int nexp = length(exp), *e = INTEGER(exp);

    for (i = 0, k = 0; i < num; i++) {

      if ((k >= nexp) || (i != e[k] - 1)) {

        n[xx[i] - 1][yy[i] - 1]++;
        nj[yy[i] - 1]++;

      }/*THEN*/
      else {

        k++;

      }/*ELSE*/

    }/*FOR*/

    /* adjust the sample size to match the number of observational data. */
    num -= nexp;

  }/*ELSE*/

  /* compute the conditional posterior probability. */
  for (i = 0; i < llx; i++)
    for (j = 0; j < lly; j++)
      res += lgammafn(n[i][j] + alpha) - lgammafn(alpha);
  for (j = 0; j < lly; j++)
    res += lgammafn(imaginary / lly) - lgammafn(nj[j] + imaginary / lly);

  Free1D(nj);
  Free2D(n, llx);

  return res;

}/*CDPOST*/

/* Dirichlet posterior probabilities (covers BD and K2 scores). */
double dirichlet_node(SEXP target, sparse_amat sam, tabular dt, SEXP iss,
    int per_cell, SEXP prior, SEXP beta, SEXP experimental, int sparse,
    bool debugging) {

char *t = (char *)CHAR(STRING_ELT(target, 0));
double prob = 0, prior_prob = 0;
int i = 0, node = 0, npar = 0, nconfig = 0, it = 0;
int **par_cols = NULL, *par_lvls = NULL, *config = NULL;
SEXP exp_data;

  /* get the parents of the node from the sparse adjacency matrix. */
  it = sparse_amat_node(sam, t);
  npar = sam.rowptr[it + 1] - sam.rowptr[it];

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

  /* the experimental data are not part of the tabular structure. */
  PROTECT(exp_data = c_dataframe_column(experimental, target, TRUE, FALSE));
  /* compute the prior probability component for the node. */
  prior_prob = graph_prior_prob(prior, target, beta, sam, debugging);

  if (npar == 0) {

    prob = dpost(dt.dcol[node], dt.nlvl[node], dt.m.nobs, iss, per_cell,
             exp_data);

  }/*THEN*/
  else {

    /* gather the parent columns and their numbers of levels. */
    par_cols = Calloc1D(npar, sizeof(int *));
    par_lvls = Calloc1D(npar, sizeof(int));
    for (i = 0; i < npar; i++) {

      int pc = dt.map[tabular_column(dt, sam.labels[sam.colidx[sam.rowptr[it] + i]])];
      par_cols[i] = dt.dcol[pc];
      par_lvls[i] = dt.nlvl[pc];

    }/*FOR*/

    /* generate the configurations of the parents: dense over the observed
     * values for the sparse score, over the full grid otherwise. */
    config = Calloc1D(dt.m.nobs, sizeof(int));
    if (sparse) {

      c_fast_config(par_cols, dt.m.nobs, npar, par_lvls, config, &nconfig, 0);
      nconfig = densify_config(config, dt.m.nobs);

    }/*THEN*/
    else {

      c_fast_config(par_cols, dt.m.nobs, npar, par_lvls, config, &nconfig, 1);

    }/*ELSE*/

    /* compute the marginal likelihood. */
    prob = cdpost(dt.dcol[node], dt.nlvl[node], config, nconfig, dt.m.nobs,
             iss, per_cell, exp_data);

    Free1D(par_cols);
    Free1D(par_lvls);
    Free1D(config);

  }/*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(1);

  return prob;

}/*DIRICHLET_NODE*/

