#define PY_SSIZE_T_CLEAN
#include <stdexcept>
#include "PyWrap.h"
#include <string>
#include <map>
#include <iostream>
#include <sstream>
#include <boost/python/numpy.hpp>

namespace pywrap {
    namespace np = boost::python::numpy;
    PyWrap* PyWrap::wrapper = nullptr;
    std::mutex PyWrap::mutex;
    CPyInstance* PyWrap::pyInstance = nullptr;
    auto moduleClassMap = std::map<std::pair<std::string, std::string>, std::tuple<PyObject*, PyObject*, PyObject*>>();

    PyWrap* PyWrap::GetInstance()
    {
        std::lock_guard<std::mutex> lock(mutex);
        if (wrapper == nullptr) {
            wrapper = new PyWrap();
            pyInstance = new CPyInstance();
        }
        return wrapper;
    }
    void PyWrap::RemoveInstance()
    {
        if (wrapper != nullptr) {
            if (pyInstance != nullptr) {
                delete pyInstance;
            }
            pyInstance = nullptr;
            if (wrapper != nullptr) {
                delete wrapper;
            }
            wrapper = nullptr;
        }
    }
    void PyWrap::importClass(const std::string& moduleName, const std::string& className)
    {
        auto result = moduleClassMap.find({ moduleName, className });
        if (result != moduleClassMap.end()) {
            return;
        }
        CPyObject module = PyImport_ImportModule(moduleName.c_str());
        if (PyErr_Occurred()) {
            errorAbort("Couldn't import module " + moduleName);
        }
        CPyObject classObject = PyObject_GetAttrString(module, className.c_str());
        if (PyErr_Occurred()) {
            errorAbort("Couldn't find class " + className);
        }
        CPyObject instance = PyObject_CallObject(classObject, NULL);
        if (PyErr_Occurred()) {
            errorAbort("Couldn't create instance of class " + className);
        }
        std::lock_guard<std::mutex> lock(mutex);
        module.AddRef();
        classObject.AddRef();
        instance.AddRef();
        moduleClassMap.insert({ { moduleName, className }, { module.getObject(), classObject.getObject(), instance.getObject() } });
    }
    void PyWrap::clean(const std::string& moduleName, const std::string& className)
    {
        std::lock_guard<std::mutex> lock(mutex);
        auto result = moduleClassMap.find({ moduleName, className });
        if (result == moduleClassMap.end()) {
            return;
        }
        Py_DECREF(std::get<0>(result->second));
        Py_DECREF(std::get<1>(result->second));
        Py_DECREF(std::get<2>(result->second));
        moduleClassMap.erase(result);
        if (PyErr_Occurred()) {
            PyErr_Print();
            errorAbort("Error cleaning module " + moduleName + " and class " + className);
        }
        if (moduleClassMap.empty()) {
            RemoveInstance();
        }
    }
    void PyWrap::errorAbort(const std::string& message)
    {
        std::cerr << message << std::endl;
        PyErr_Print();
        RemoveInstance();
        exit(1);
    }
    PyObject* PyWrap::getClass(const std::string& moduleName, const std::string& className)
    {
        auto item = moduleClassMap.find({ moduleName, className });
        if (item == moduleClassMap.end()) {
            errorAbort("Module " + moduleName + " and class " + className + " not found");
        }
        return std::get<2>(item->second);
    }
    std::string PyWrap::callMethodString(const std::string& moduleName, const std::string& className, const std::string& method)
    {
        PyObject* instance = getClass(moduleName, className);
        PyObject* result;
        try {
            if (!(result = PyObject_CallMethod(instance, method.c_str(), NULL)))
                errorAbort("Couldn't call method " + method);
        }
        catch (const std::exception& e) {
            errorAbort(e.what());
        }
        std::string value = PyUnicode_AsUTF8(result);
        Py_XDECREF(result);
        return value;
    }
    std::string PyWrap::version(const std::string& moduleName, const std::string& className)
    {
        return callMethodString(moduleName, className, "version");
    }
    void PyWrap::setHyperparameters(const std::string& moduleName, const std::string& className, const json& hyperparameters)
    {
        // Set hyperparameters as attributes of the class
        std::cout << "Building dictionary of arguments" << std::endl;
        PyObject* pValue;
        PyObject* instance = getClass(moduleName, className);
        for (const auto& [key, value] : hyperparameters.items()) {
            std::stringstream oss;
            oss << value.type_name();
            if (oss.str() == "string") {
                pValue = Py_BuildValue("s", value.get<std::string>().c_str());
            } else {
                if (value.is_number_integer()) {
                    pValue = Py_BuildValue("i", value.get<int>());
                } else {
                    pValue = Py_BuildValue("f", value.get<double>());
                }
            }
            int res = PyObject_SetAttrString(instance, key.c_str(), pValue);
            if (res == -1 && PyErr_Occurred()) {
                cleanDictionary(args);
                errorAbort("Couldn't set attribute " + key + "=" + value.dump());
            }
            Py_XDECREF(pValue);
        }
    }
    void PyWrap::fit(const std::string& moduleName, const std::string& className, CPyObject& X, CPyObject& y)
    {
        PyObject* instance = getClass(moduleName, className);
        CPyObject result;
        std::string method = "fit";
        try {
            if (!(result = PyObject_CallMethodObjArgs(instance, PyUnicode_FromString(method.c_str()), X.getObject(), y.getObject(), NULL)))
                errorAbort("Couldn't call method fit");
        }
        catch (const std::exception& e) {
            errorAbort(e.what());
        }
        // Py_XDECREF(result);
    }

    PyObject* PyWrap::predict(const std::string& moduleName, const std::string& className, CPyObject& X)
    {
        PyObject* instance = getClass(moduleName, className);
        PyObject* result;
        std::string method = "predict";
        try {
            if (!(result = PyObject_CallMethodObjArgs(instance, PyUnicode_FromString(method.c_str()), X.getObject(), NULL)))
                errorAbort("Couldn't call method predict");
        }
        catch (const std::exception& e) {
            errorAbort(e.what());
        }
        Py_INCREF(result);
        return result; // Caller must free this object
    }
    double PyWrap::score(const std::string& moduleName, const std::string& className, CPyObject& X, CPyObject& y)
    {
        PyObject* instance = getClass(moduleName, className);
        CPyObject result;
        std::string method = "score";
        try {
            if (!(result = PyObject_CallMethodObjArgs(instance, PyUnicode_FromString(method.c_str()), X.getObject(), y.getObject(), NULL)))
                errorAbort("Couldn't call method score");
        }
        catch (const std::exception& e) {
            errorAbort(e.what());
        }
        double resultValue = PyFloat_AsDouble(result);
        // Py_XDECREF(result);
        return resultValue;
    }
}