#include "../include/rcore.h"
#include "../core/allocations.h"
#include "../core/contingency.tables.h"
#include "../core/data.table.h"
#include "../core/sets.h"
#include "../minimal/common.h"
#include "../minimal/data.frame.h"
#include "scores.h"

double pdnode(int *xx, int llx, int nobs, int *new_xx, int new_nobs,
    double *nparams) {

double res = 0;
counts1d train = { 0 }, test = { 0 };

  /* initialize the contingency tables for train and test data. */
  train = new_1d_table(llx);
  test = new_1d_table(train.llx);
  fill_1d_table(xx, &train, nobs);
  fill_1d_table(new_xx, &test, new_nobs);

  /* compute the entropy from the joint and marginal frequencies. */
  for (int i = 0; i < train.llx; i++)
    if (train.n[i] != 0)
      res += (double)test.n[i] * log((double)train.n[i] / test.nobs);

  /* we may want to store the number of parameters. */
  if (nparams)
    *nparams = train.llx - 1;

  Free1DTAB(train);
  Free1DTAB(test);

  return res;

}/*PDNODE*/

double cpdnode(int *xx, int llx, int *yy, int lly, int nobs, int *new_xx,
    int *new_yy, int new_nobs, double *nparams) {

double res = 0;
counts2d train = { 0 }, test = { 0 };

  /* initialize the contingency tables for train and test data. */
  train = new_2d_table(llx, lly, TRUE);
  test = new_2d_table(train.llx, train.lly, FALSE);
  fill_2d_table(xx, yy, &train, nobs);
  fill_2d_table(new_xx, new_yy, &test, new_nobs);

  /* compute the conditional entropy from the joint and marginal frequencies. */
  for (int i = 0; i < train.llx; i++)
    for (int j = 0; j < train.lly; j++)
      if (train.n[i][j] != 0)
        res += (double)test.n[i][j] * log((double)train.n[i][j] / train.nj[j]);

  /* we may want to store the number of parameters. */
  if (nparams)
    *nparams = (train.llx - 1) * train.lly;

  Free2DTAB(train);
  Free2DTAB(test);

  return res;

}/*CPDNODE*/

double predictive_loglik_dnode(SEXP target, sparse_amat sam, tabular dt,
    tabular newdt, double *nparams, bool debugging) {

double loglik = 0;
char *t = (char *)CHAR(STRING_ELT(target, 0));
int i = 0, npar = 0, node = 0, new_node = 0, nconfig = 0, it = 0;
int **par_cols = NULL, **new_par_cols = NULL, *par_lvls = NULL;
int *config = NULL, *new_config = NULL;

  /* 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 training and test data. */
  node = dt.map[tabular_column(dt, t)];
  new_node = newdt.map[tabular_column(newdt, t)];

  if (npar == 0) {

    loglik = pdnode(dt.dcol[node], dt.nlvl[node], dt.m.nobs, newdt.dcol[new_node],
               newdt.m.nobs, nparams);

  }/*THEN*/
  else {

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

      const char *pname = sam.labels[sam.colidx[sam.rowptr[it] + i]];
      int pc = dt.map[tabular_column(dt, pname)];

      par_cols[i] = dt.dcol[pc];
      par_lvls[i] = dt.nlvl[pc];
      new_par_cols[i] = newdt.dcol[newdt.map[tabular_column(newdt, pname)]];

    }/*FOR*/

    /* generate the (1-based, all-levels) configurations of the parents. */
    config = Calloc1D(dt.m.nobs, sizeof(int));
    new_config = Calloc1D(newdt.m.nobs, sizeof(int));
    c_fast_config(par_cols, dt.m.nobs, npar, par_lvls, config, &nconfig, 1);
    c_fast_config(new_par_cols, newdt.m.nobs, npar, par_lvls, new_config, NULL, 1);

    /* compute the log-likelihood. */
    loglik = cpdnode(dt.dcol[node], dt.nlvl[node], config, nconfig, dt.m.nobs,
               newdt.dcol[new_node], new_config, newdt.m.nobs, nparams);

    Free1D(par_cols);
    Free1D(new_par_cols);
    Free1D(par_lvls);
    Free1D(config);
    Free1D(new_config);

  }/*ELSE*/

  if (debugging)
    Rprintf("  > log-likelihood is %lf.\n", loglik);

  return loglik;

}/*PREDICTIVE_LOGLIK_DNODE*/
