#include "../include/rcore.h"
#include "../core/allocations.h"
#include "../core/contingency.tables.h"
#include "../core/data.table.h"
#include "../core/sets.h"
#include "../math/sparse.amat.h"
#include "../minimal/common.h"

static int deff_node_root(int *n, int llx) {

int i = 0, nonzero = 0;

  for (i = 0; i < llx; i++)
    nonzero += (n[i] > 0);

  return nonzero;

}/*DEFF_NODE_ROOT*/

static int deff_node(int **n, int *nj, int llx, int lly) {

int i = 0, j = 0, nonzero = 0;

  for (i = 0; i < llx; i++)
    for (j = 0; j < lly; j++)
      nonzero += (n[i][j] > 0);
  for (j = 0; j < lly; j++)
    nonzero -= (nj[j] > 0);

  return nonzero;

}/*DEFF_NODE*/

static double entropy1d(int *n, int llx, int num) {

int i = 0;
double res = 0, alpha = 1 / (double)llx;

  for (i = 0; i < llx; i++)
      res += ((double)(n[i] + 1) / (double)(num + llx) - alpha) *
               log((double)(n[i] + 1) / (double)(num + llx));

  return res;

}/*ENTROPY1D*/

static double entropy2d(int **n, int *nj, int llx, int lly, int num) {

int i = 0, j = 0;
double res = 0, alpha = 1 / (double)(llx * lly);

  for (i = 0; i < llx; i++)
    for (j = 0; j < lly; j++)
      res += ((double)(n[i][j] + 1) / ((double)(num + llx * lly)) - alpha) *
               log((double)(n[i][j] + 1) / (double)(nj[j] + llx));

  return res;

}/*ENTROPY2D*/

/* Harald Steck optimal imaginary sample size estimator. */
SEXP alpha_star(SEXP x, SEXP data, SEXP debug) {

int i = 0, j = 0, nnodes = length(data), nobs = length(VECTOR_ELT(data, 0));
bool debugging = isTRUE(debug);
long double num = 0, den = 0, ratio = 0;
counts1d marginal = { 0 };
counts2d joint = { 0 };
tabular dt = tabular_from_SEXP(data, 0, 0, R_NilValue);
sparse_amat sam = sparse_amat_from_SEXP(x, TRUE);

  /* iterate over the nodes in the graph. */
  for (i = 0; i < nnodes; i++) {

    const char *label = sam.labels[i];
    int node = dt.map[tabular_column(dt, label)];
    int npar = sam.rowptr[i + 1] - sam.rowptr[i];

    if (debugging)
      Rprintf("* processing node %s.\n", label);

    if (npar == 0) {

      marginal = new_1d_table(dt.nlvl[node]);
      fill_1d_table(dt.dcol[node], &marginal, nobs);

      /* add the effective degrees of freedom to the numerator. */
      num += deff_node_root(marginal.n, marginal.llx);
      /* add the entropy delta to the denominator. */
      den += entropy1d(marginal.n, marginal.llx, marginal.nobs);

      Free1DTAB(marginal);

    }/*THEN*/
    else {

      /* gather the parent columns (and their levels) from the tabular data and
       * generate the (1-based, all-levels) parents' configurations. */
      int **par_cols = Calloc1D(npar, sizeof(int *));
      int *par_lvls = Calloc1D(npar, sizeof(int));
      int *config = Calloc1D(nobs, sizeof(int)), nconfig = 0;

      for (j = 0; j < npar; j++) {

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

      }/*FOR*/
      c_fast_config(par_cols, nobs, npar, par_lvls, config, &nconfig, 1);

      joint = new_2d_table(dt.nlvl[node], nconfig, TRUE);
      fill_2d_table(dt.dcol[node], config, &joint, nobs);

      /* add the effective degrees of freedom to the numerator. */
      num += deff_node(joint.n, joint.nj, joint.llx, joint.lly);
      /* add the entropy delta to the denominator. */
      den += entropy2d(joint.n, joint.nj, joint.llx, joint.lly, joint.nobs);

      Free2DTAB(joint);
      Free1D(par_cols);
      Free1D(par_lvls);
      Free1D(config);

    }/*ELSE*/

    if (debugging)
      Rprintf("  > numerator is now %lf, denominator is now %lf.\n",
        (double)num, (double)den);

  }/*FOR*/

  FreeSparseAMAT(sam);
  FreeTAB(dt);

  /* check that the proposed alpha is equal or greater than 1. */
  ratio = num / den;

  if (ratio < 1)
    warning("the estimated imaginary sample size is very small (%lf).",
      (double)ratio);

  return ScalarReal((double)ratio);

}/*ALPHA_STAR*/
