#include <Python.h>
#include <iostream>
#include <Eigen/Dense>

using namespace Eigen;

typedef struct
{
 PyObject_HEAD
 MatrixXd *matrix=nullptr;
} PyMatrixObject;

static PyObject *
PyMatrix_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
 PyMatrixObject *self;
 self = (PyMatrixObject *)type->tp_alloc(type, 0);

 char *kwlist[] = {"width", "height", NULL};
 int width = 0;
 int height = 0;

 if (!PyArg_ParseTupleAndKeywords(args, kwds, "ii", kwlist,
 &width, &height))
 {
 Py_DECREF(self);
 return NULL;
 }
 if (width <= 0 or height <= 0)
 {
 PyErr_SetString(PyExc_ValueError, "The height and width must be greater than 0.");
 return NULL;
 }

 self->matrix = new MatrixXd(width, height);
 return (PyObject *)self;
}

static void
*PyMatrix_dealloc(PyObject *obj)
{
 delete ((PyMatrixObject *)obj)->matrix; 
 Py_TYPE(obj)->tp_free(obj);
}


inline MatrixXd *ParseMatrix(PyObject *obj){
 return ((PyMatrixObject *)obj)->matrix;
}

inline PyObject *ReturnMatrix(MatrixXd *m, PyTypeObject *type){
 PyMatrixObject *obj = PyObject_NEW(PyMatrixObject, type);
 obj->matrix = m;
 return (PyObject *)obj;
}

static PyObject *
PyMatrix_add(PyObject *a, PyObject *b)
{
 MatrixXd *matrix_a = ParseMatrix(a);
 MatrixXd *matrix_b = ParseMatrix(b);

 if (matrix_a->cols() != matrix_b->cols() or matrix_a->rows() != matrix_b->rows()){
 PyErr_SetString(PyExc_ValueError, "The input matrix must be the same shape.");
 return NULL;
 }

 MatrixXd *matrix_c = new MatrixXd(matrix_a->cols(), matrix_b->rows());
 *matrix_c = *matrix_a + *matrix_b;

 return ReturnMatrix(matrix_c, a->ob_type);
}

static PyObject *
PyMatrix_minus(PyObject *a, PyObject *b)
{
 MatrixXd *matrix_a = ParseMatrix(a);
 MatrixXd *matrix_b = ParseMatrix(b);

 if (matrix_a->cols() != matrix_b->cols() or matrix_a->rows() != matrix_b->rows()){
 PyErr_SetString(PyExc_ValueError, "The input matrix must be the same shape.");
 return NULL;
 }

 MatrixXd *matrix_c = new MatrixXd(matrix_a->cols(), matrix_b->rows());
 *matrix_c = *matrix_a + *matrix_b;
 return ReturnMatrix(matrix_c, a->ob_type);
}

static PyObject *
PyMatrix_multiply(PyObject *a, PyObject *b)
{
 MatrixXd *matrix_a = ParseMatrix(a);
 MatrixXd *matrix_b = ParseMatrix(b);

 if (matrix_a->cols() != matrix_b->rows()){
 PyErr_SetString(PyExc_ValueError, "The colonm rank of matrix A must be the same as the row rank of matrix B.");
 return NULL;
 }
 MatrixXd *matrix_c = new MatrixXd(matrix_a->rows(), matrix_b->cols());
 *matrix_c = (*matrix_a) * (*matrix_b);
 return ReturnMatrix(matrix_c, a->ob_type);
}

static PyObject *PyMatrix_str(PyObject *a)
{
 MatrixXd *matrix = ParseMatrix(a);
 std::stringstream ss;
 ss << *matrix;
 return Py_BuildValue("s", ss.str().c_str());
}

static PyNumberMethods numberMethods = {
 PyMatrix_add, //nb_add
 PyMatrix_minus, //nb_subtract;
 PyMatrix_multiply, //nb_multiply
 nullptr, //nb_remainder;
 nullptr, //nb_divmod;
 nullptr, // nb_power;
 nullptr, // nb_negative;
 nullptr, // nb_positive;
 nullptr, // nb_absolute;
 nullptr, // nb_bool;
 nullptr, // nb_invert;
 nullptr, // nb_lshift;
 nullptr, // nb_rshift;
 nullptr, // nb_and;
 nullptr, // nb_xor;
 nullptr, // nb_or;
 nullptr, // nb_int;
 nullptr, // nb_reserved;
 nullptr, // nb_float;

 nullptr, // nb_inplace_add;
 nullptr, // nb_inplace_subtract;
 nullptr, // nb_inplace_multiply;
 nullptr, // nb_inplace_remainder;
 nullptr, // nb_inplace_power;
 nullptr, // nb_inplace_lshift;
 nullptr, // nb_inplace_rshift;
 nullptr, // nb_inplace_and;
 nullptr, // nb_inplace_xor;
 nullptr, // nb_inplace_or;

 nullptr, // nb_floor_divide;
 nullptr, // nb_true_divide;
 nullptr, // nb_inplace_floor_divide;
 nullptr, // nb_inplace_true_divide;

 nullptr, // nb_index;

 nullptr, //nb_matrix_multiply;
 nullptr //nb_inplace_matrix_multiply;

};

PyObject *PyMatrix_data(PyObject *self, void *closure)
{

 PyMatrixObject *obj = (PyMatrixObject *)self;
 Py_ssize_t width = obj->matrix->cols();
 Py_ssize_t height = obj->matrix->rows();

 PyObject *list = PyList_New(height);
 for (int i = 0; i < height; i++)
 {
 PyObject *internal = PyList_New(width);

 for (int j = 0; j < width; j++)
 {
 PyObject *value = PyFloat_FromDouble((*obj->matrix)(i, j));
 PyList_SetItem(internal, j, value);
 }

 PyList_SetItem(list, i, internal);
 }
 return list;
}

PyObject *PyMatrix_rows(PyObject *self, void *closure)
{
 PyMatrixObject *obj = (PyMatrixObject *)self;
 return Py_BuildValue("i", obj->matrix->rows());
}

PyObject *PyMatrix_cols(PyObject *self, void *closure)
{
 PyMatrixObject *obj = (PyMatrixObject *)self;
 return Py_BuildValue("i", obj->matrix->cols());
}

static PyGetSetDef MatrixGetSet[] = {
 {"data", (getter)PyMatrix_data, nullptr, nullptr},
 {"row", (getter)PyMatrix_rows, nullptr, nullptr},
 {"colunm", (getter)PyMatrix_cols, nullptr, nullptr},
 {nullptr}};

PyObject *PyMatrix_tolist(PyObject *self, PyObject *args)
{
 return PyMatrix_data(self, nullptr);
}

static PyMethodDef MatrixMethods[] = {
 {"to_list", (PyCFunction)PyMatrix_tolist, METH_VARARGS, "Return the matrix data to a list object."},
 {nullptr}};

static PyTypeObject MatrixType = {
 PyVarObject_HEAD_INIT(nullptr, 0) "matrix.Matrix", /* tp_name */
 sizeof(PyMatrixObject), /* tp_basicsize */
 0, /* tp_itemsize */
 (destructor)PyMatrix_dealloc, /* tp_dealloc */
 nullptr, /* tp_print */
 nullptr, /* tp_getattr */
 nullptr, /* tp_setattr */
 nullptr, /* tp_reserved */
 nullptr, /* tp_repr */
 &numberMethods, /* tp_as_number */
 nullptr, /* tp_as_sequence */
 nullptr, /* tp_as_mapping */
 nullptr, /* tp_hash */
 nullptr, /* tp_call */
 PyMatrix_str, /* tp_str */
 nullptr, /* tp_getattro */
 nullptr, /* tp_setattro */
 nullptr, /* tp_as_buffer */
 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
 "Coustom matrix class.", /* tp_doc */
 nullptr, /* tp_traverse */
 nullptr, /* tp_clear */
 nullptr, /* tp_richcompare */
 0, /* tp_weaklistoffset */
 nullptr, /* tp_iter */
 nullptr, /* tp_iternext */
 MatrixMethods, /* tp_methods */
 nullptr, /* tp_members */
 MatrixGetSet, /* tp_getset */
 nullptr, /* tp_base */
 nullptr, /* tp_dict */
 nullptr, /* tp_descr_get */
 nullptr, /* tp_descr_set */
 0, /* tp_dictoffset */
 nullptr, /* tp_init */
 nullptr, /* tp_alloc */
 PyMatrix_new /* tp_new */
};

static PyObject *PyMatrix_ones(PyObject *self, PyObject *args, PyObject *kwargs)
{
 PyMatrixObject *m = (PyMatrixObject *)PyMatrix_new(&MatrixType, args, kwargs);
 m->matrix->setOnes();
 return (PyObject *)m;
}

static PyObject *PyMatrix_zeros(PyObject *self, PyObject *args, PyObject *kwargs)
{
 PyMatrixObject *m = (PyMatrixObject *)PyMatrix_new(&MatrixType, args, kwargs);
 m->matrix->setZero();
 return (PyObject *)m;
}

static PyObject *PyMatrix_random(PyObject *self, PyObject *args, PyObject *kwargs)
{
 PyMatrixObject *m = (PyMatrixObject *)PyMatrix_new(&MatrixType, args, kwargs);
 m->matrix->setRandom();
 return (PyObject *)m;
}

static PyObject *PyMatrix_matrix(PyObject *self, PyObject *args)
{
 PyObject *data = nullptr;
 if (!PyArg_ParseTuple(args, "O", &data))
 {
 PyErr_SetString(PyExc_ValueError, "Please pass a 2 dimensions list object. 1");
 return nullptr;
 }
 if (!PyList_Check(data))
 {
 PyErr_SetString(PyExc_ValueError, "Please pass a 2 dimensions list object. 2");
 return nullptr;
 }
 int height = PyList_GET_SIZE(data);
 if (height <= 0)
 {
 PyErr_SetString(PyExc_ValueError, "Please pass a 2 dimensions list object. 2");
 return nullptr;
 }
 PyObject *list = PyList_GET_ITEM(data, 0);
 if (!PyList_Check(list))
 {
 PyErr_SetString(PyExc_ValueError, "Please pass a 2 dimensions list object. 3");
 return nullptr;
 }
 int width = PyList_GET_SIZE(list);
 MatrixXd *p_mat = new MatrixXd(width, height);
 for (int i = 0; i < height; i++)
 {
 PyObject *list = PyList_GET_ITEM(data, i);
 if (!PyList_Check(list))
 {
 PyErr_SetString(PyExc_ValueError, "Please pass a 2 dimensions list object. 3");
 return nullptr;
 }
 int tmp = PyList_GET_SIZE(list);
 if (width != tmp)
 {
 PyErr_SetString(PyExc_ValueError, "Please pass a 2 dimensions list object. Each elements of it must be the same length.");
 return nullptr;
 }
 width = tmp;

 for (int j = 0; j < width; j++)
 {
 PyObject *num = PyList_GET_ITEM(list, j);
 if (!PyFloat_Check(num))
 {
 PyErr_SetString(PyExc_ValueError, "Every elements of the matrix must float.");
 return nullptr;
 }
 (*p_mat)(i, j) = ((PyFloatObject *)num)->ob_fval;
 }
 }

 return ReturnMatrix(p_mat, &MatrixType);
}

static PyMethodDef matrixMethods[] = {
 {"ones", (PyCFunction)PyMatrix_ones, METH_VARARGS | METH_KEYWORDS, "Return a new matrix with initial values one."},
 {"zeros", (PyCFunction)PyMatrix_zeros, METH_VARARGS | METH_KEYWORDS, "Return a new matrix with initial values zero."},
 {"random", (PyCFunction)PyMatrix_random, METH_VARARGS | METH_KEYWORDS, "Return a new matrix with random values"},
 {"matrix", (PyCFunction)PyMatrix_matrix, METH_VARARGS, "Return a new matrix with given values"},
 {nullptr}};

static struct PyModuleDef module = {
 PyModuleDef_HEAD_INIT,
 "matrix",
 "Python interface for Matrix calculation",
 -1,
 matrixMethods};

PyObject *initModule(void)
{
 PyObject *m;
 if (PyType_Ready(&MatrixType) < 0)
 return NULL;

 m = PyModule_Create(&module);
 if (m == NULL)
 return NULL;

 Py_INCREF(&MatrixType);
 if (PyModule_AddObject(m, "Matrix", (PyObject *)&MatrixType) < 0)
 {
 Py_DECREF(&MatrixType);
 Py_DECREF(m);
 return NULL;
 }

 return m;
}