#include "../include/rcore.h"
#include "../core/allocations.h"
#include "../core/sets.h"
#include "../include/globals.h"
#include "../math/linear.algebra.h"
#include "../minimal/common.h"
#include "../minimal/data.frame.h"
#include "scores.h"

static double c_pcgnode(double *xx, double *xx2, double **gp, double **gp2,
    int ngp, int nobs, int nobs2, int *config, int *config2, int nconfig) {

int i = 0, j = 0;
double res = 0, fitted = 0, *beta = NULL, *sd = NULL;

  /* allocate the fitted values and the standard error. */
  sd = Calloc1D((!config) ? 1 : nconfig, sizeof(double));

  /* if the regression is conditional on config, iterate over its values,
   * otherwise fit using the whole sample. */
  if (!config) {

    beta = Calloc1D(ngp + 1, sizeof(double));

    c_ols(gp, xx, nobs, ngp, NULL, NULL, beta, sd, NULL, FALSE);

    /* compute the log-likelihood (singular models haze zero density). */
    if (*sd < MACHINE_TOL)
      res = R_NegInf;
    else
      for (j = 0; j < nobs2; j++) {

        for (i = 1, fitted = beta[0]; i < ngp + 1; i++)
          fitted += beta[j] * gp2[i - 1][j];

        res += dnorm(xx2[j], fitted, *sd, TRUE);

      }/*FOR*/

  }/*THEN*/
  else {

    beta = Calloc1D((ngp + 1) * nconfig, sizeof(double));

    c_cls(gp, xx, config, nobs, ngp, nconfig, NULL, NULL, beta, sd, NULL, FALSE);

    /* if the model is not singular compute the log-likelihood. */
    for (j = 0; j < nobs2; j++) {

      if (sd[config2[j] - 1] < MACHINE_TOL) {

        res = R_NegInf;
        goto end;

      }/*THEN*/

      for (i = 1, fitted = beta[0 + (config2[j] - 1) * (ngp + 1)]; i < ngp + 1; i++)
         fitted += beta[i + (config2[j] - 1) * (ngp + 1)] * gp2[i - 1][j];

      res += dnorm(xx2[j], fitted, sd[config2[j] - 1], TRUE);

    }/*FOR*/

  }/*ELSE*/

end:

  Free1D(beta);
  Free1D(sd);

  return res;

}/*C_PCGNODE*/

static double pcgnode(double *xx, int nobs, double *xx2, int nobs2, int **dp,
    int **dp2, int *nlevels, int ndp, double **gp, double **gp2, int ngp,
    double *nparams) {

int nconfig = 0;
int *config = NULL, *config2 = NULL;
double res = 0;

  /* build the configurations of the discrete parents for both data sets. */
  config = Calloc1D(nobs, sizeof(int));
  config2 = Calloc1D(nobs2, sizeof(int));
  c_fast_config(dp, nobs, ndp, nlevels, config, &nconfig, 1);
  c_fast_config(dp2, nobs2, ndp, nlevels, config2, NULL, 1);

  res = c_pcgnode(xx, xx2, gp, gp2, ngp, nobs, nobs2, config, config2, nconfig);

  /* we may want to store the number of parameters. */
  if (nparams)
    *nparams = nconfig * (ngp + 1);

  Free1D(config);
  Free1D(config2);

  return res;

}/*PCGNODE*/

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

double loglik = 0, **gp = NULL, **gp2 = NULL;
char *t = (char *)CHAR(STRING_ELT(target, 0));
int i = 0, npar = 0, node_col = 0, node = 0, new_node = 0, it = 0;
int ndp = 0, ngp = 0, nconfig = 0;
int **dp = NULL, **dp2 = NULL, *nlevels = NULL, *config = NULL, *config2 = NULL;
bool node_discrete = FALSE;

  /* 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 column in the training and test data. */
  node_col = tabular_column(dt, t);
  node = dt.map[node_col];
  new_node = newdt.map[tabular_column(newdt, t)];
  node_discrete = dt.m.flag[node_col].discrete;

  if (npar == 0) {

    /* no parents, reuse the marginal likelihoods. */
    if (node_discrete)
      loglik = pdnode(dt.dcol[node], dt.nlvl[node], dt.m.nobs,
                 newdt.dcol[new_node], newdt.m.nobs, nparams);
    else
      loglik = pgnode(dt.ccol[node], dt.m.nobs, newdt.ccol[new_node],
                 newdt.m.nobs, nparams);

  }/*THEN*/
  else {

    /* gather the discrete and continuous parents' columns separately, from
     * both the training and the test data. */
    dp = Calloc1D(npar, sizeof(int *));
    dp2 = Calloc1D(npar, sizeof(int *));
    nlevels = Calloc1D(npar, sizeof(int));
    gp = Calloc1D(npar, sizeof(double *));
    gp2 = Calloc1D(npar, sizeof(double *));

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

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

      if (dt.m.flag[pc].discrete) {

        dp[ndp] = dt.dcol[slot];
        dp2[ndp] = newdt.dcol[new_slot];
        nlevels[ndp++] = dt.nlvl[slot];

      }/*THEN*/
      else {

        gp[ngp] = dt.ccol[slot];
        gp2[ngp++] = newdt.ccol[new_slot];

      }/*ELSE*/

    }/*FOR*/

    if (node_discrete) {

      /* discrete nodes are not allowed to have continuous parents; then reuse
       * the conditional discrete likelihood. */
      if (ngp != 0) {

        loglik = R_NegInf;

      }/*THEN*/
      else {

        config = Calloc1D(dt.m.nobs, sizeof(int));
        config2 = Calloc1D(newdt.m.nobs, sizeof(int));
        c_fast_config(dp, dt.m.nobs, ndp, nlevels, config, &nconfig, 1);
        c_fast_config(dp2, newdt.m.nobs, ndp, nlevels, config2, NULL, 1);
        loglik = cpdnode(dt.dcol[node], dt.nlvl[node], config, nconfig,
                   dt.m.nobs, newdt.dcol[new_node], config2, newdt.m.nobs,
                   nparams);
        Free1D(config);
        Free1D(config2);

      }/*ELSE*/

    }/*THEN*/
    else {

      if (ndp == 0)
        loglik = cpgnode(dt.ccol[node], dt.m.nobs, newdt.ccol[new_node],
                   newdt.m.nobs, gp, gp2, ngp, nparams);
      else
        loglik = pcgnode(dt.ccol[node], dt.m.nobs, newdt.ccol[new_node],
                   newdt.m.nobs, dp, dp2, nlevels, ndp, gp, gp2, ngp, nparams);

    }/*ELSE*/

    Free1D(dp);
    Free1D(dp2);
    Free1D(nlevels);
    Free1D(gp);
    Free1D(gp2);

  }/*ELSE*/

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

  return loglik;

}/*PREDICTIVE_LOGLIK_CGNODE*/
