#include "../include/rcore.h"
#include "../include/globals.h"
#include "../fitted/fitted.h"
#include "common.h"
#include "data.frame.h"
#include "strings.h"

/* read the per-node completeness flags stored in a data frame's metadata. */
SEXP dataframe_complete_nodes(SEXP data) {

  return getListElement(getAttrib(data, BN_MetaDataSymbol), "complete.nodes");

}/*DATAFRAME_COMPLETE_NODES*/

/* read the per-node latent (fully-unobserved) flags from a data frame's
 * metadata. */
SEXP dataframe_latent_nodes(SEXP data) {

  return getListElement(getAttrib(data, BN_MetaDataSymbol), "latent.nodes");

}/*DATAFRAME_LATENT_NODES*/

/* in-place conversion of a list into a data frame. */
SEXP minimal_data_frame(SEXP obj, int nrows) {

int *row = NULL;
SEXP rownames;

  /* generate and set the row names. */
  if (nrows > 0) {

    PROTECT(rownames = allocVector(INTSXP, 2));
    row = INTEGER(rownames);

    row[0] = NA_INTEGER;
    row[1] = -nrows;

  }/*THEN*/
  else {

    PROTECT(rownames = allocVector(INTSXP, 0));

  }/*ELSE*/

  setAttrib(obj, R_RowNamesSymbol, rownames);

  /* set the class name. */
  setAttrib(obj, R_ClassSymbol, mkString("data.frame"));

  UNPROTECT(1);

  return obj;

}/*MINIMAL_DATA_FRAME*/

/* efficient column extraction from data frames. */
SEXP c_dataframe_column(SEXP dataframe, SEXP name, bool drop, bool keep_names) {

SEXP try, result, rownames, colnames = getAttrib(dataframe, R_NamesSymbol);
int *idx = NULL, *row = NULL, nrows = 0;
int nnames = length(name), name_type = TYPEOF(name);

  if (dataframe == R_NilValue)
    return R_NilValue;

  switch(name_type) {

    case STRSXP:

      /* column names passed as strings; match the corresponding indexes. */
      PROTECT(try = match(colnames, name, 0));
      idx = INTEGER(try);
      break;

    case REALSXP:

      /* these are almost good enough, coerce them to integers. */
      PROTECT(try = coerceVector(name, INTSXP));
      idx = INTEGER(try);
      break;

    case INTSXP:

      /* these are already indexes, nothing to do. */
      idx = INTEGER(name);
      break;

    default:

      error("this SEXP type is not handled in c_dataframe_column().");

  }/*SWITCH*/

  if ((nnames > 1) || !drop) {

    PROTECT(result = allocVector(VECSXP, nnames));

    for (int i = 0; i < nnames; i++)
      SET_VECTOR_ELT(result, i, VECTOR_ELT(dataframe, idx[i] - 1));

    /* set the row names. */
    if (nnames >= 1)
      nrows = length(VECTOR_ELT(dataframe, 0));
    else {

      SEXP dim = getAttrib(dataframe, R_RowNamesSymbol);
      if (isInteger(dim) && length(dim) == 2 && INTEGER(dim)[0] == NA_INTEGER)
        nrows = abs(INTEGER(dim)[1]);
      else
        nrows = length(dim);

    }/*ELSE*/

    if (nrows > 0) {

      PROTECT(rownames = allocVector(INTSXP, 2));
      row = INTEGER(rownames);

      row[0] = NA_INTEGER;
      row[1] = -nrows;

    }/*THEN*/
    else {

      PROTECT(rownames = allocVector(INTSXP, 0));

    }/*ELSE*/

    setAttrib(result, R_RowNamesSymbol, rownames);

    /* set the class name. */
    setAttrib(result, R_ClassSymbol, mkString("data.frame"));

    if (keep_names)
      setAttrib(result, R_NamesSymbol, name);

    /* carry a subsetted copy of the completeness/latent metadata so the
     * per-column flags survive extraction and can be read straight back by
     * tabular_from_SEXP(), rather than being restored by name afterwards. */
    SEXP metadata = getAttrib(dataframe, BN_MetaDataSymbol);

    if (metadata != R_NilValue) {

      SEXP complete = getListElement(metadata, "complete.nodes");
      SEXP latent = getListElement(metadata, "latent.nodes");
      SEXP submeta, subc, subl;

      PROTECT(subc = allocVector(LGLSXP, nnames));
      PROTECT(subl = allocVector(LGLSXP, nnames));
      for (int i = 0; i < nnames; i++) {

        LOGICAL(subc)[i] =
          (complete == R_NilValue) ? TRUE : LOGICAL(complete)[idx[i] - 1];
        LOGICAL(subl)[i] =
          (latent == R_NilValue) ? FALSE : LOGICAL(latent)[idx[i] - 1];

      }/*FOR*/

      PROTECT(submeta = allocVector(VECSXP, 3));
      setAttrib(submeta, R_NamesSymbol,
        mkStringVec(3, "type", "complete.nodes", "latent.nodes"));
      SET_VECTOR_ELT(submeta, 0, getListElement(metadata, "type"));
      SET_VECTOR_ELT(submeta, 1, subc);
      SET_VECTOR_ELT(submeta, 2, subl);
      setAttrib(result, BN_MetaDataSymbol, submeta);

      UNPROTECT(3);

    }/*THEN*/

    UNPROTECT(2);

  }/*THEN*/
  else {

    if (*idx != 0)
      result = VECTOR_ELT(dataframe, *idx - 1);
    else
      result = R_NilValue;

  }/*ELSE*/

  if (name_type != INTSXP)
    UNPROTECT(1);

  return result;

}/*C_DATAFRAME_COLUMN*/

/* extract the union of up to three sets of named columns (any of which may be
 * R_NilValue) as one data frame, preserving the given order and the names. Used
 * to pull the variables under test and their conditioning set together, so they
 * occupy the leading columns without reserved slots or manual insertion. */
SEXP c_dataframe_column_set(SEXP dataframe, SEXP a, SEXP b, SEXP c) {

int na = (a == R_NilValue) ? 0 : length(a);
int nb = (b == R_NilValue) ? 0 : length(b);
int nc = (c == R_NilValue) ? 0 : length(c);
int k = 0;
SEXP names, result;

  PROTECT(names = allocVector(STRSXP, na + nb + nc));
  for (int j = 0; j < na; j++)
    SET_STRING_ELT(names, k++, STRING_ELT(a, j));
  for (int j = 0; j < nb; j++)
    SET_STRING_ELT(names, k++, STRING_ELT(b, j));
  for (int j = 0; j < nc; j++)
    SET_STRING_ELT(names, k++, STRING_ELT(c, j));

  result = c_dataframe_column(dataframe, names, FALSE, TRUE);

  UNPROTECT(1);

  return result;

}/*C_DATAFRAME_COLUMN_SET*/

/* create a data column from a standalone bn.fit.*node object. */
SEXP node2df(SEXP target, int n) {

fitted_node_e node_type = fitted_node_to_enum(target);
SEXP result, res_levels;

  switch(node_type) {

    case GNODE:
    case CGNODE:
    case ZIHPNODE:
    case ZINBNODE:

      return allocVector(REALSXP, n);

    case DNODE:
    case ONODE:

      PROTECT(result = allocVector(INTSXP, n));
      memset(INTEGER(result), '\0', n * sizeof(int));
      if (node_type == ONODE)
        setAttrib(result, R_ClassSymbol, mkStringVec(2, "ordered", "factor"));
      else if (node_type == DNODE)
        setAttrib(result, R_ClassSymbol, mkString("factor"));
      res_levels = getAttrib(getListElement(target, "prob"), R_DimNamesSymbol);
      setAttrib(result, R_LevelsSymbol, VECTOR_ELT(res_levels, 0));

      UNPROTECT(1);

      return result;

    default:
      error("unknown node type (class: %s).",
         CHAR(STRING_ELT(getAttrib(target, R_ClassSymbol), 0)));

  }/*SWITCH*/

}/*NODE2DF*/

/* create a data column from a node in a fitted network. */
SEXP fitnode2df(SEXP fitted, SEXP node, int n) {

SEXP target = getListElement(fitted, (char *)CHAR(node));

  return node2df(target, n);

}/*FITNODE2DF*/

/* create a data frame from a fitted network. */
SEXP fit2df(SEXP fitted, int n) {

int i = 0, nnodes = length(fitted);
SEXP df, nodes;

  /* allocate the data frame. */
  PROTECT(nodes = getAttrib(fitted, R_NamesSymbol));
  PROTECT(df = allocVector(VECSXP, nnodes));
  /* fill the columns. */
  for (i = 0; i < nnodes; i++)
    SET_VECTOR_ELT(df, i, fitnode2df(fitted, STRING_ELT(nodes, i), n));
  /* set the column names. */
  setAttrib(df, R_NamesSymbol, nodes);
  /* add the labels to the return value. */
  minimal_data_frame(df, n);

  UNPROTECT(2);

  return df;

}/*FIT2DF*/

