#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION

#include <Python.h>
#include <numpy/arrayobject.h>
#include "tree_shap.h"
#include <iostream>

typedef double tfloat;

static PyObject *_cext_tree_shap(PyObject *self, PyObject *args);
static PyObject *_cext_compute_expectations(PyObject *self, PyObject *args);

static PyMethodDef module_methods[] = {
  {"tree_shap", _cext_tree_shap, METH_VARARGS, "C implementation of Tree SHAP."},
  {"compute_expectations", _cext_compute_expectations, METH_VARARGS, "Compute expectations of internal nodes."},
  {NULL, NULL, 0, NULL}
};

PyMODINIT_FUNC PyInit__cext(void)
{
  PyObject *module;
  static struct PyModuleDef moduledef = {
    PyModuleDef_HEAD_INIT,
    "_cext",
    "This module provides an interface for a fast Tree SHAP implementation.",
    -1,
    module_methods,
    NULL,
    NULL,
    NULL,
    NULL
  };
  module = PyModule_Create(&moduledef);
  if (!module) return NULL;

  /* Load `numpy` functionality. */
  import_array();

  return module;
}

static PyObject *_cext_compute_expectations(PyObject *self, PyObject *args)
{
  PyObject *children_left_obj;
  PyObject *children_right_obj;
  PyObject *node_sample_weight_obj;
  PyObject *values_obj;

  /* Parse the input tuple */
  if (!PyArg_ParseTuple(
    args, "OOOO", &children_left_obj, &children_right_obj, &node_sample_weight_obj, &values_obj
  )) return NULL;

  /* Interpret the input objects as numpy arrays. */
  PyArrayObject *children_left_array = (PyArrayObject*)PyArray_FROM_OTF(children_left_obj, NPY_INT, NPY_ARRAY_IN_ARRAY);
  PyArrayObject *children_right_array = (PyArrayObject*)PyArray_FROM_OTF(children_right_obj, NPY_INT, NPY_ARRAY_IN_ARRAY);
  PyArrayObject *node_sample_weight_array = (PyArrayObject*)PyArray_FROM_OTF(node_sample_weight_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
  PyArrayObject *values_array = (PyArrayObject*)PyArray_FROM_OTF(values_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);

  /* If that didn't work, throw an exception. */
  if (children_left_array == NULL || children_right_array == NULL ||
      values_array == NULL || node_sample_weight_array == NULL) {
    Py_XDECREF(children_left_array);
    Py_XDECREF(children_right_array);
    Py_XDECREF(values_array);
    Py_XDECREF(node_sample_weight_array);
    return NULL;
  }

  /* Get pointers to the data as C-types. */
  int *children_left = (int*)PyArray_DATA(children_left_array);
  int *children_right = (int*)PyArray_DATA(children_right_array);
  tfloat *values = (tfloat*)PyArray_DATA(values_array);
  tfloat *node_sample_weight = (tfloat*)PyArray_DATA(node_sample_weight_array);

  const int max_depth = compute_expectations(children_left, children_right, node_sample_weight, values, 0, 0);

  // clean up the created python objects
  Py_XDECREF(children_left_array);
  Py_XDECREF(children_right_array);
  Py_XDECREF(values_array);
  Py_XDECREF(node_sample_weight_array);

  PyObject *ret = Py_BuildValue("i", max_depth);
  return ret;
}

static PyObject *_cext_tree_shap(PyObject *self, PyObject *args)
{
  int max_depth;
  PyObject *children_left_obj;
  PyObject *children_right_obj;
  PyObject *children_default_obj;
  PyObject *features_obj;
  PyObject *thresholds_obj;
  PyObject *values_obj;
  PyObject *node_sample_weight_obj;
  PyObject *x_obj;
  PyObject *x_missing_obj;
  PyObject *out_contribs_obj;
  int condition;
  int condition_feature;

  /* Parse the input tuple */
  if (!PyArg_ParseTuple(
    args, "iOOOOOOOOOOii", &max_depth, &children_left_obj, &children_right_obj, &children_default_obj,
    &features_obj, &thresholds_obj, &values_obj, &node_sample_weight_obj, &x_obj,
    &x_missing_obj, &out_contribs_obj, &condition, &condition_feature
  )) return NULL;

  /* Interpret the input objects as numpy arrays. */
  PyArrayObject *children_left_array = (PyArrayObject*)PyArray_FROM_OTF(children_left_obj, NPY_INT, NPY_ARRAY_IN_ARRAY);
  PyArrayObject *children_right_array = (PyArrayObject*)PyArray_FROM_OTF(children_right_obj, NPY_INT, NPY_ARRAY_IN_ARRAY);
  PyArrayObject *children_default_array = (PyArrayObject*)PyArray_FROM_OTF(children_default_obj, NPY_INT, NPY_ARRAY_IN_ARRAY);
  PyArrayObject *features_array = (PyArrayObject*)PyArray_FROM_OTF(features_obj, NPY_INT, NPY_ARRAY_IN_ARRAY);
  PyArrayObject *thresholds_array = (PyArrayObject*)PyArray_FROM_OTF(thresholds_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
  PyArrayObject *values_array = (PyArrayObject*)PyArray_FROM_OTF(values_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
  PyArrayObject *node_sample_weight_array = (PyArrayObject*)PyArray_FROM_OTF(node_sample_weight_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
  PyArrayObject *x_array = (PyArrayObject*)PyArray_FROM_OTF(x_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
  PyArrayObject *x_missing_array = (PyArrayObject*)PyArray_FROM_OTF(x_missing_obj, NPY_BOOL, NPY_ARRAY_IN_ARRAY);
  PyArrayObject *out_contribs_array = (PyArrayObject*)PyArray_FROM_OTF(out_contribs_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);

  /* If that didn't work, throw an exception. */
  if (children_left_array == NULL || children_right_array == NULL ||
      children_default_array == NULL || features_array == NULL || thresholds_array == NULL ||
      values_array == NULL || node_sample_weight_array == NULL || x_array == NULL ||
      x_missing_array == NULL || out_contribs_array == NULL) {
    Py_XDECREF(children_left_array);
    Py_XDECREF(children_right_array);
    Py_XDECREF(children_default_array);
    Py_XDECREF(features_array);
    Py_XDECREF(thresholds_array);
    Py_XDECREF(values_array);
    Py_XDECREF(node_sample_weight_array);
    Py_XDECREF(x_array);
    Py_XDECREF(x_missing_array);
    Py_XDECREF(out_contribs_array);
    return NULL;
  }

  // number of features
  const unsigned M = PyArray_DIM(x_array, 0);

  // Get pointers to the data as C-types
  int *children_left = (int*)PyArray_DATA(children_left_array);
  int *children_right = (int*)PyArray_DATA(children_right_array);
  int *children_default = (int*)PyArray_DATA(children_default_array);
  int *features = (int*)PyArray_DATA(features_array);
  tfloat *thresholds = (tfloat*)PyArray_DATA(thresholds_array);
  tfloat *values = (tfloat*)PyArray_DATA(values_array);
  tfloat *node_sample_weight = (tfloat*)PyArray_DATA(node_sample_weight_array);
  tfloat *x = (tfloat*)PyArray_DATA(x_array);
  bool *x_missing = (bool*)PyArray_DATA(x_missing_array);
  tfloat *out_contribs = (tfloat*)PyArray_DATA(out_contribs_array);

  //const int max_depth = compute_expectations(children_left, children_right, node_sample_weight, values, 0, 0);

  tree_shap(
    M, max_depth, children_left, children_right, children_default, features,
    thresholds, values, node_sample_weight, x, x_missing, out_contribs,
    condition, condition_feature
  );

  // clean up the created python objects
  Py_XDECREF(children_left_array);
  Py_XDECREF(children_right_array);
  Py_XDECREF(children_default_array);
  Py_XDECREF(features_array);
  Py_XDECREF(thresholds_array);
  Py_XDECREF(values_array);
  Py_XDECREF(node_sample_weight_array);
  Py_XDECREF(x_array);
  Py_XDECREF(x_missing_array);
  Py_XDECREF(out_contribs_array);

                                      // unsigned M
                                      // unsigned max_depth
  // double m, b;
  // PyObject *x_obj, *y_obj, *yerr_obj;
  //
  // /* Parse the input tuple */
  // if (!PyArg_ParseTuple(args, "ddOOO", &m, &b, &x_obj, &y_obj,
  //                                     &yerr_obj))
  //     return NULL;
  //
  // /* Interpret the input objects as numpy arrays. */
  // PyObject *x_array = PyArray_FROM_OTF(x_obj, NPY_DOUBLE, NPY_IN_ARRAY);
  // PyObject *y_array = PyArray_FROM_OTF(y_obj, NPY_DOUBLE, NPY_IN_ARRAY);
  // PyObject *yerr_array = PyArray_FROM_OTF(yerr_obj, NPY_DOUBLE,
  //                                         NPY_IN_ARRAY);
  //
  // /* If that didn't work, throw an exception. */
  // if (x_array == NULL || y_array == NULL || yerr_array == NULL) {
  //     Py_XDECREF(x_array);
  //     Py_XDECREF(y_array);
  //     Py_XDECREF(yerr_array);
  //     return NULL;
  // }
  //
  // /* How many data points are there? */
  // int N = (int)PyArray_DIM(x_array, 0);
  //
  // /* Get pointers to the data as C-types. */
  // double *x    = (double*)PyArray_DATA(x_array);
  // double *y    = (double*)PyArray_DATA(y_array);
  // double *yerr = (double*)PyArray_DATA(yerr_array);
  //
  // /* Call the external C function to compute the chi-squared. */
  // double value = chi2(m, b, x, y, yerr, N);
  //
  // /* Clean up. */
  // Py_DECREF(x_array);
  // Py_DECREF(y_array);
  // Py_DECREF(yerr_array);
  //
  // if (value < 0.0) {
  //     PyErr_SetString(PyExc_RuntimeError,
  //                 "Chi-squared returned an impossible value.");
  //     return NULL;
  // }

  // const int max_depth = compute_expectations(children_left, children_right, node_sample_weight, values, 0, 0);
  //
  // calculate_contributions(const unsigned M, const unsigned max_depth, const int *children_left,
  //                                     const int *children_right,
  //                                     const int *children_default, const int *features,
  //                                     const tfloat *thresholds, const tfloat *values,
  //                                     const tfloat *node_sample_weight,
  //                                     const tfloat *x, const bool *x_missing, unsigned root_id,
  //                                     tfloat *out_contribs, int condition,
  //                                     unsigned condition_feature) const {
    // find the expected value of the tree's predictions
    //const int max_depth = compute_expectations(children_left, children_right, node_sample_weight, values, i, depth=0);


  /* Build the output tuple */
  PyObject *ret = Py_BuildValue("d", (double)values[0]);
  return ret;
}
