#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <iostream>

#include <gdichotomy/utils.hpp>
#include <gdichotomy/basesolver.hpp>
#include <gdichotomy/general_dichotomy.hpp>
#include <gdichotomy/gdexception.hpp>

/*!
 * \brief wrapper to the general dichotomy abstract solver class
 */
class PySolver: public BaseSolver {

    public:

        PySolver(size_t nobj, PyObject *py_solve_problem, PyObject *py_is_new) : _nobj(nobj), _py_solve_problem(py_solve_problem) ,_py_is_new(py_is_new) {
            
        }

        virtual bool solve(dichotomy::WeightVec& weight, std::set<size_t> adj_solutions, dichotomy::CostVec& solution, dichotomy::COST& obj_value) {

            // create a python list of the weights
            PyObject *pyweights = PyList_New(_nobj);
            for(Py_ssize_t obj_ind = 0; obj_ind < _nobj; obj_ind ++) {
                PyList_SetItem(pyweights, obj_ind, PyFloat_FromDouble(weight[obj_ind]));
            }

            // create a python list of the solution adjacency
            PyObject *pyadj = PyList_New(adj_solutions.size());
            Py_ssize_t adj_ind = 0;
            for(auto& elt: adj_solutions) {
                PyList_SetItem(pyadj, adj_ind, PyLong_FromLong(static_cast<long>(elt)));
                adj_ind ++;
            }

            // passing the current weights to the pysolver as a python list
            PyObject *arglist = Py_BuildValue("(OO)", pyweights, pyadj);

            // call the python solve method
            PyObject *result = PyObject_CallObject(_py_solve_problem, arglist);

            // dereference the arguments
            Py_DECREF(arglist);

            if(result == NULL) { // python exception occurred
                throw Dichotomy::GDException("Error when calling solve(): solution of the weighted problem failed");

            }

            // extraction of the returned tuple
            if(!PyTuple_Check(result) || PyTuple_Size(result) != 2) {
                Py_DECREF(result);
                throw Dichotomy::GDException("Error when calling solve(): return value is invalid");
            }

            PyObject *output_list = PyTuple_GetItem(result, 0);
            
            if(output_list == NULL || !PyList_Check(output_list) || PyList_Size(output_list) != _nobj) {
                std::cout << "output list size: " << PyList_Size(output_list) << std::endl;
                throw Dichotomy::GDException("Error when calling solve(): extraction of the objective cost vector (0) failed");
            }

            PyObject *output_value = PyTuple_GetItem(result, 1);
            if(output_value == NULL || !PyFloat_Check(output_value)) {
                throw Dichotomy::GDException("Error when calling solve(): extraction the weighted objective value (1) failed");
            }

            obj_value = PyFloat_AsDouble(output_value);

            // extract the resulting cost vector
            solution.resize(_nobj);
            for(Py_ssize_t obj_ind = 0; obj_ind < _nobj; obj_ind ++) {
                solution[obj_ind] = PyFloat_AsDouble(PyList_GetItem(output_list, obj_ind));
            }

            // dereference the result
            Py_DECREF(result);

            return true;

        }

        virtual void isNew(bool is_new) {
            
            // call py_is_new

            // passing the current weights to the pysolver as a python list
            PyObject *arglist = Py_BuildValue("(I)", static_cast<unsigned int>(is_new));

            // call the python solve method
            PyObject *result = PyObject_CallObject(_py_is_new, arglist);

            // dereference the arguments
            Py_DECREF(arglist);

            if(result == NULL) {
                throw Dichotomy::GDException();
            } else {
                // dereference the result
                Py_DECREF(result);
            }

        }

        /*!
         * \brief get the number of objectives of the problem
         */
        inline std::size_t getNObj() {
            return _nobj;
        }

    private:

        std::size_t _nobj;

        // user callback functions
        PyObject *_py_solve_problem = NULL;
        PyObject *_py_is_new = NULL;
 
};

extern "C" {

/*!
 * \brief call the general dichotomy method
 */
static PyObject* run_algorithm(PyObject* self, PyObject* args) {
    
    std::cout << "running the general dichotomy scheme" << std::endl;

    PyObject *result = NULL;

    // to arguments are expected : number of objectives and the function that solved the problem
    PyObject *solving_func = NULL;
    PyObject *py_is_new = NULL;

    // parameters
    unsigned int verbose;
    unsigned int nobj;
    PyObject *directions = NULL;
    unsigned int backend;
    double wnorm;
    unsigned int max_iter;
    unsigned int max_depth;
    double cost_scale;
    unsigned int export_weights;
    unsigned int export_weight_graph;
    unsigned int search_scheme;
    PyObject *init_weight = NULL; 
    
    if (PyArg_ParseTuple(args, "IOdIIdOIIIIIOO:run_algorithm", &backend, &directions, &wnorm, &max_iter, &max_depth, &cost_scale, &init_weight, &search_scheme, &export_weights, &export_weight_graph, &verbose, &nobj, &solving_func, &py_is_new)) {

        if (!PyCallable_Check(solving_func) || !PyCallable_Check(py_is_new)) {
            PyErr_SetString(PyExc_TypeError, "Error when executing the GeneralDichotomy algorithm: callback functions parameters must be callable");
            return NULL;
        }

        Py_XINCREF(solving_func);
        Py_XINCREF(py_is_new);

        Py_XINCREF(directions);
        Py_XINCREF(init_weight);

        GeneralDichotomy::Parameters params;
        params.verbose = verbose;
        params.wnorm = wnorm;
        params.max_iter = static_cast<size_t>(max_iter);
        params.max_depth = static_cast<size_t>(max_depth);
        params.cost_scale = static_cast<dichotomy::COST>(cost_scale);
        params.export_weight_graph = export_weight_graph;
        if(!export_weights && params.export_weight_graph) {
            params.export_weight_graph = false;
            if(params.verbose >= 1) {
                std::cerr << "Warning: export_weights parameters must be set to true to export the weight graph (export_weight_graph=True)" << std::endl;
            }
        }
        if(backend == 1) {
            params.backend = GeneralDichotomy::Backend::CDD; 
        } else if(backend == 2) {
            params.backend = GeneralDichotomy::Backend::PPL; 
        } else {
            PyErr_SetString(PyExc_TypeError, "Error, unknown polytope enumeration backend");
            return NULL;
        }
        if(search_scheme == 0) {
            params.search_scheme = GeneralDichotomy::DIVERSIFICATION;
        } else { // search_scheme == 1
            params.search_scheme = GeneralDichotomy::INTENSIFICATION;
        }

        // parameter directions
        if(!PyTuple_Check(directions)) {
            PyErr_SetString(PyExc_TypeError, "Error, param directions is not a tuple");
            return NULL;
        }
        params.directions.resize(static_cast<size_t>(PyTuple_Size(directions)));
        if(!params.directions.empty()) {
            for(size_t obj_ind = 0; obj_ind < nobj; obj_ind ++) {
                long value = PyLong_AsLong(PyTuple_GetItem(directions, static_cast<Py_ssize_t>(obj_ind)));
                if(value == 0) {
                    params.directions[obj_ind] = dichotomy::MIN;
                } else {
                    params.directions[obj_ind] = dichotomy::MAX;
                }
            }
        }
        
        // parameter init_weight
        if(!PyList_Check(init_weight)) {
            PyErr_SetString(PyExc_TypeError, "Error, param init_weight is not a list");
            return NULL;
        }
        if(PyList_Size(init_weight) > 0) {
            params.initial_weight.resize(static_cast<size_t>(PyList_Size(init_weight)));
            for(size_t w_ind = 0; w_ind < params.initial_weight.size(); w_ind ++) {
                double value = PyFloat_AsDouble(PyList_GetItem(init_weight, static_cast<Py_ssize_t>(w_ind)));
                params.initial_weight[w_ind] = value;
            }
        }
        
        if(params.verbose > 0) {
            Utils::printParameters(params, std::cout);
            std::cout << "\t-export_weights: " << export_weights << std::endl;
            std::cout << std::endl;
        }

        // create the pysolver (takes the solving func as parameter)
        PySolver py_solver(nobj, solving_func, py_is_new);

        // create the general dichotomy solver
        GeneralDichotomy gdsolver(py_solver);

        try {

            // solve the problem
            dichotomy::Result dichotomy_result = gdsolver.run(params);

            Py_ssize_t result_size = 1;
            if(export_weights) {
                result_size ++;
            } if(params.export_weight_graph) {
                result_size ++;
            }
            result = PyTuple_New(result_size);

            // export the solutions
            PyObject* sol_list = PyList_New(static_cast<Py_ssize_t>(dichotomy_result.solutions.size())); 
            for(size_t sol_ind = 0; sol_ind < dichotomy_result.solutions.size(); sol_ind ++) {
                // objective cost vector
                PyObject *current_sol = PyList_New(static_cast<Py_ssize_t>(nobj));
                for(size_t obj_ind = 0; obj_ind < nobj; obj_ind ++) {
                    PyList_SetItem(current_sol, static_cast<Py_ssize_t>(obj_ind), PyFloat_FromDouble(dichotomy_result.solutions[sol_ind][obj_ind]));
                } 
                PyList_SetItem(sol_list, static_cast<Py_ssize_t>(sol_ind), current_sol);
            }
            PyTuple_SetItem(result, 0, sol_list);

            // export the weights
            if(export_weights) {
                PyObject* weight_list = PyList_New(static_cast<Py_ssize_t>(dichotomy_result.weights.size()));
                for(size_t weight_ind = 0; weight_ind < dichotomy_result.weights.size(); weight_ind ++) {
                    PyObject* weight_tuple = PyTuple_New(static_cast<Py_ssize_t>(3));
                    // weight vector
                    PyObject* weight_vector = PyList_New(static_cast<Py_ssize_t>(nobj));
                    for(size_t obj_ind = 0; obj_ind < nobj; obj_ind ++) {
                        PyList_SetItem(weight_vector, static_cast<Py_ssize_t>(obj_ind), PyFloat_FromDouble(static_cast<double>(dichotomy_result.weights[weight_ind].weight[obj_ind])));
                    }
                    PyTuple_SetItem(weight_tuple, 0, weight_vector);
                    PyTuple_SetItem(weight_tuple, 1, PyFloat_FromDouble(static_cast<double>(dichotomy_result.weights[weight_ind].z)));
                    // adjacency list
                    PyObject* adj_list = PyList_New(static_cast<Py_ssize_t>(dichotomy_result.weights[weight_ind].bnd.size()+dichotomy_result.weights[weight_ind].adj.size()));
                    Py_ssize_t adj_ind = 0;
                    for(auto elt: dichotomy_result.weights[weight_ind].bnd) {
                        PyList_SetItem(adj_list, adj_ind, PyLong_FromLong(static_cast<long>(elt)));
                        adj_ind ++;
                    }
                    for(auto elt: dichotomy_result.weights[weight_ind].adj) {
                        PyList_SetItem(adj_list, adj_ind, PyLong_FromLong(static_cast<long>(elt)));
                        adj_ind ++;
                    }
                    PyTuple_SetItem(weight_tuple, 2, adj_list);
                    PyList_SetItem(weight_list, static_cast<Py_ssize_t>(weight_ind), weight_tuple);
                }
                PyTuple_SetItem(result, 1, weight_list);
            }

            // export the weight graph
            if(params.export_weight_graph) {
                PyObject* weight_graph = PyList_New(static_cast<Py_ssize_t>(dichotomy_result.weights.size()));
                for(size_t weight_ind = 0; weight_ind < dichotomy_result.weightGraph.size(); weight_ind ++) {
                    PyObject* neighbours = PyList_New(static_cast<Py_ssize_t>(dichotomy_result.weightGraph[weight_ind].size()));
                    Py_ssize_t neighbour_ind = 0;
                    for(auto neighbour: dichotomy_result.weightGraph[weight_ind]) {
                        PyList_SetItem(neighbours, neighbour_ind, PyLong_FromLong(static_cast<long>(neighbour)));
                        neighbour_ind ++;
                    }
                    PyList_SetItem(weight_graph, static_cast<Py_ssize_t>(weight_ind), neighbours);
                }
                PyTuple_SetItem(result, 2, weight_graph);
            }

        } catch(Dichotomy::GDException& exception) {
            PyErr_PrintEx(0);
            if(exception.msg().length() > 0) {
                PyErr_SetString(PyExc_Exception, exception.msg().c_str());
            } else {
                PyErr_SetString(PyExc_Exception, "An exception occurred during the general dichotomy search!");
            }
            Py_XDECREF(directions);
            Py_XDECREF(init_weight);

            Py_XDECREF(solving_func);
            Py_XDECREF(py_is_new);

            return NULL;
        }

        Py_XDECREF(directions);
        Py_XDECREF(init_weight);

        Py_XDECREF(solving_func);
        Py_XDECREF(py_is_new);

    } else {
        throw Dichotomy::GDException();
    }

    // return the result
    return result;
}




static PyMethodDef pygdichotomy_methods[] = {
    {"run_algorithm", run_algorithm, METH_VARARGS, "Run the general dichotomy scheme."},
    {NULL, NULL, 0, NULL}
};



static struct PyModuleDef libpygdichotomy = {
    PyModuleDef_HEAD_INIT,
    "libpygdichotomy",
    NULL,
    -1,
    pygdichotomy_methods
};



PyMODINIT_FUNC PyInit_libpygdichotomy(void) {
    return PyModule_Create(&libpygdichotomy);
}

}