#include "../../include/rcore.h"
#include "../../core/allocations.h"
#include "../../core/data.table.h"
#include "../../fitted/fitted.h"
#include "../../include/globals.h"
#include "../../minimal/common.h"
#include "../loglikelihood/loglikelihood.h"

/* multi-purpose log-likelihood function for data and a fitted network. */
SEXP loglikelihood_function(SEXP fitted, SEXP data, SEXP by_sample,
    SEXP keep_nodes, SEXP propagate_missing, SEXP debug) {

int ndata = length(VECTOR_ELT(data, 0)), nparams = 0;
double *loglik = NULL;
fitted_bn bn = fitted_network_from_SEXP(fitted);
bool by = isTRUE(by_sample), propagate = isTRUE(propagate_missing);
bool debugging = isTRUE(debug);
SEXP keep, loglikelihood, nodes_in_fitted;

  /* allocate the return value: a vector with length equal to the sample size if
   * we are returning the log-likelihood of each observation, or a vector of
   * length 1 if we are returning the log-likelihood of the whole data set. */
  if (by) {

    PROTECT(loglikelihood = allocVector(REALSXP, ndata));
    loglik = REAL(loglikelihood);
    memset(loglik, '\0', sizeof(double) * ndata);

  }/*THEN*/
  else {

    PROTECT(loglikelihood = ScalarReal(0));
    loglik = Calloc1D(ndata, sizeof(double));

  }/*ELSE*/

  /* find out which nodes to compute the log-likelihood for, zero-indexed; this
   * is passed to tabular_from_SEXP() as the set of columns to keep, while the
   * completeness flags are read from the data's own metadata. */
  PROTECT(nodes_in_fitted = getAttrib(fitted, R_NamesSymbol));
  PROTECT(keep = match(keep_nodes, nodes_in_fitted, 0));

  if (bn.type == DNET || bn.type == ONET || bn.type == DONET) {

    if (debugging)
      Rprintf("> computing the log-likelihood of a discrete network.\n");

    tabular dt = tabular_from_SEXP(data, 0, 0, keep);

    if (by) {

      bysample_discrete_loglikelihood(bn, dt, loglik, debugging);

    }/*THEN*/
    else {

      NUM(loglikelihood) =
        data_discrete_loglikelihood(bn, dt, propagate, debugging);

    }/*ELSE*/

    FreeTAB(dt);

  }/*THEN*/
  else if (bn.type == GNET) {

    if (debugging)
      Rprintf("> computing the log-likelihood of a Gaussian network.\n");

    tabular dt = tabular_from_SEXP(data, 0, 0, keep);

    if (by) {

      bysample_gaussian_loglikelihood(bn, dt, loglik, FALSE, debugging);

    }/*THEN*/
    else {

      NUM(loglikelihood) =
        data_gaussian_loglikelihood(bn, dt, loglik, propagate, debugging);

    }/*ELSE*/

    FreeTAB(dt);

  }/*THEN*/
  else if (bn.type == CGNET) {

    if (debugging)
      Rprintf("> computing the log-likelihood of a conditional Gaussian network.\n");

    tabular dt = tabular_from_SEXP(data, 0, 0, keep);

    if (by) {

      bysample_clgaussian_loglikelihood(bn, dt, loglik, FALSE, debugging);

    }/*THEN*/
    else {

      NUM(loglikelihood) =
        data_clgaussian_loglikelihood(bn, dt, loglik, propagate, debugging);

    }/*ELSE*/

    FreeTAB(dt);

  }/*THEN*/
  else if (bn.type == ZINET) {

    if (debugging)
      Rprintf("> computing the log-likelihood of a zero-inflated network.\n");

    tabular dt = tabular_from_SEXP(data, 0, 0, keep);

    if (by) {

      bysample_zeroinflated_loglikelihood(bn, dt, loglik, FALSE, debugging);

    }/*THEN*/
    else {

      NUM(loglikelihood) =
        data_zeroinflated_loglikelihood(bn, dt, loglik, propagate, debugging);

    }/*ELSE*/

    FreeTAB(dt);

  }/*THEN*/
  else {

    error("unknown network type, unable to compute the log-likelihood.");

  }/*ELSE*/

  /* compute the number of parameters. */
  for (int i = 0; i < bn.nnodes; i++)
    nparams += nparams_fitted_node(bn.ldists[i], bn.node_types[i]);

  /* set the class and the expected attributes. */
  setAttrib(loglikelihood, R_ClassSymbol, mkString("logLik"));
  setAttrib(loglikelihood, BN_NobsSymbol, ScalarReal(ndata));
  setAttrib(loglikelihood, BN_DfSymbol, ScalarReal(nparams));

  if (!by)
    Free1D(loglik);
  FreeFittedBN(bn);
  UNPROTECT(3);

  return loglikelihood;

}/*LOGLIKELIHOOD_FUNCTION*/

