// Copyright (C) 2019 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)

/** \file Pybind11StOptReg.cpp
 * \brief Map regression  classes to python
 * \author Xavier Warin, Thomas Ouillon
 */

#include <Eigen/Dense>
#include <iostream>
#include <memory>
#include <vector>
#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <pybind11/stl_bind.h>
#include <pybind11/stl.h>
#include "StOpt/core/grids/SpaceGrid.h"
#include "StOpt/core/grids/FullGrid.h"
#include "StOpt/core/utils/Polynomials1D.h"
#include "StOpt/regression/LocalLinearRegression.h"
#include "StOpt/regression/LocalConstRegression.h"
#include "StOpt/regression/LocalSameSizeLinearRegression.h"
#include "StOpt/regression/LocalSameSizeConstRegression.h"
#include "StOpt/regression/LocalKMeansRegression.h"
#include "StOpt/regression/LocalGridKernelRegression.h"
#include "StOpt/regression/LocalConstDiscrLastDimRegression.h"
#include "StOpt/regression/LocalLinearDiscrLastDimRegression.h"
#include "StOpt/regression/LaplacianConstKernelRegression.h"
#include "StOpt/regression/LaplacianLinearKernelRegression.h"
#include "StOpt/regression/LaplacianGridKernelRegression.h"
#include "StOpt/regression/SparseRegression.h"
#include "StOpt/regression/GlobalRegression.h"
#include "StOpt/regression/ContinuationValue.h"
#include "StOpt/regression/ContinuationCuts.h"
#include "StOpt/regression/GridAndRegressedValueGeners.h"


namespace py = pybind11;

// Wrapper for regression methods
//******************************

// wrapper for BaseRegression
class PyBaseRegression : public StOpt::BaseRegression
{
public :

    using StOpt::BaseRegression::BaseRegression;

    void updateSimulations(const bool &p_bZeroDate, const Eigen::ArrayXXd &p_particles)   override
    {
        PYBIND11_OVERLOAD_PURE(void, StOpt::BaseRegression, updateSimulations, p_bZeroDate, p_particles);
    }

    Eigen::ArrayXd getCoordBasisFunction(const Eigen::ArrayXd &p_fToRegress) const  override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::ArrayXd, StOpt::BaseRegression, getCoordBasisFunction, p_fToRegress) ;
    }

    Eigen::ArrayXXd getCoordBasisFunctionMultiple(const Eigen::ArrayXXd &p_fToRegress) const  override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::ArrayXXd, StOpt::BaseRegression, getCoordBasisFunctionMultiple, p_fToRegress);
    }

    double reconstructionASim(const int &p_isim, const Eigen::ArrayXd &p_basisCoefficients) const  override
    {
        PYBIND11_OVERLOAD_PURE(double,  StOpt::BaseRegression, reconstructionASim,  p_isim, p_basisCoefficients);
    }

    Eigen::ArrayXd getAllSimulations(const Eigen::ArrayXd &p_fToRegress) const  override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::ArrayXd, StOpt::BaseRegression, getAllSimulations, p_fToRegress);
    }

    Eigen::ArrayXXd getAllSimulationsMultiple(const Eigen::ArrayXXd &p_fToRegress)  const  override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::ArrayXXd,  StOpt::BaseRegression, getAllSimulationsMultiple, p_fToRegress);
    }

    Eigen::ArrayXd reconstruction(const Eigen::ArrayXd &p_basisCoefficients) const  override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::ArrayXd, StOpt::BaseRegression, reconstruction,  p_basisCoefficients);
    }

    Eigen::ArrayXXd reconstructionMultiple(const Eigen::ArrayXXd &p_basisCoefficients) const  override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::ArrayXXd, StOpt::BaseRegression, reconstructionMultiple, p_basisCoefficients);
    }

    double getValue(const Eigen::ArrayXd &p_coordinates, const Eigen::ArrayXd &p_coordBasisFunction) const override
    {
        PYBIND11_OVERLOAD_PURE(double,   StOpt::BaseRegression,  getValue, p_coordinates, p_coordBasisFunction);
    }


    double getAValue(const Eigen::ArrayXd &p_coordinates, const Eigen::ArrayXd &p_ptOfStock,
                     const std::vector< std::shared_ptr<StOpt::InterpolatorSpectral> > &m_interpFuncBasis) const override
    {
        PYBIND11_OVERLOAD_PURE(double, StOpt::BaseRegression, getAValue, p_coordinates, p_ptOfStock, m_interpFuncBasis);
    }

    Eigen::ArrayXXd getParticles() const override
    {
        PYBIND11_OVERLOAD(Eigen::ArrayXXd, StOpt::BaseRegression, getParticles,);
    }

    int getNumberOfFunction() const override
    {
        PYBIND11_OVERLOAD_PURE(int, StOpt::BaseRegression, getNumberOfFunction,);
    }

    std::shared_ptr<StOpt::BaseRegression> clone() const override
    {
        PYBIND11_OVERLOAD_PURE(std::shared_ptr<StOpt::BaseRegression>, StOpt::BaseRegression, clone,);
    }

};


/// \brief Encapsulation for regression module
PYBIND11_MODULE(StOptReg, m)
{
    // Map regressors
    //****************
    py::class_<StOpt::BaseRegression, std::shared_ptr<StOpt::BaseRegression>, PyBaseRegression>(m, "BaseRegression")
    .def(py::init<>())
    .def("updateSimulations", &StOpt::BaseRegression::updateSimulations)
    .def("getCoordBasisFunction", &StOpt::BaseRegression::getCoordBasisFunction)
    .def("getCoordBasisFunctionMultiple", &StOpt::BaseRegression::getCoordBasisFunctionMultiple)
    .def("getAllSimulations", &StOpt::BaseRegression::getAllSimulations)
    .def("getAllSimulationsMultiple", &StOpt::BaseRegression::getAllSimulationsMultiple)
    .def("reconstruction", &StOpt::BaseRegression::reconstruction)
    .def("getValue", &StOpt::BaseRegression::getValue)
    .def("getValues", &StOpt::BaseRegression::getValues)
    .def("getAValue", &StOpt::BaseRegression::getAValue)
    .def("getNbSimul", &StOpt::BaseRegression::getNbSimul)
    .def("getParticles", &StOpt::BaseRegression::getParticles)
    .def("getDimension", &StOpt::BaseRegression::getDimension)
    .def("getNumberOfFunction", &StOpt::BaseRegression::getNumberOfFunction)
    .def("reconstructionMultiple", &StOpt::BaseRegression::reconstructionMultiple)
    .def("getBRotationAndRescale", &StOpt::BaseRegression::getBRotationAndRescale)
    .def("getMeanX", &StOpt::BaseRegression::getMeanX)
    .def("getEtypX", &StOpt::BaseRegression::getEtypX)
    .def("getSvdMatrix", &StOpt::BaseRegression::getSvdMatrix)
    .def("getSing", &StOpt::BaseRegression::getSing)
    ;


    py::class_<StOpt::LocalLinearRegression, std::shared_ptr< StOpt::LocalLinearRegression>, StOpt::BaseRegression> (m, "LocalLinearRegression")
    .def(py::init< const Eigen::ArrayXi &>())
    .def(py::init< const Eigen::ArrayXi &, bool>())
    .def(py::init< const bool &, const Eigen::ArrayXXd &, const Eigen::ArrayXi &> ())
    .def(py::init< const bool &, const Eigen::ArrayXXd &, const Eigen::ArrayXi &, bool> ())
    .def("getMesh1D",  [ ](StOpt::LocalLinearRegression & p_x)
    {
        auto ret = p_x.getMesh1DCopy();
        std::vector< Eigen::ArrayXd > result;
        for (auto &array : ret)
            result.push_back(*array);
        return result;
    })
    .def("getNbMesh", &StOpt::LocalRegression::getNbMeshCopy)
    .def("getSimToCell", &StOpt::LocalRegression::getSimToCellCopy)
    .def("getAllSimulationsConvex", &StOpt::LocalLinearRegression::getAllSimulationsConvex)
    ;

    py::class_<StOpt::LocalConstRegression, std::shared_ptr< StOpt::LocalConstRegression>, StOpt::BaseRegression> (m, "LocalConstRegression")
    .def(py::init< const Eigen::ArrayXi &>())
    .def(py::init< const Eigen::ArrayXi &, bool>())
    .def(py::init< const bool &, const Eigen::ArrayXXd &, const Eigen::ArrayXi &> ())
    .def(py::init< const bool &, const Eigen::ArrayXXd &, const Eigen::ArrayXi &, bool> ())
    .def("getMesh1D",  [ ](StOpt::LocalConstRegression & p_x)
    {
        auto ret = p_x.getMesh1DCopy();
        std::vector< Eigen::ArrayXd > result;
        for (auto &array : ret)
            result.push_back(*array);
        return result;
    })
    .def("getNbMesh", &StOpt::LocalRegression::getNbMeshCopy)
    .def("getSimToCell", &StOpt::LocalRegression::getSimToCellCopy)
    ;

    py::class_<StOpt::LocalLinearDiscrLastDimRegression, std::shared_ptr< StOpt::LocalLinearDiscrLastDimRegression>, StOpt::BaseRegression> (m, "LocalLinearDiscrLastDimRegression")
    .def(py::init< const Eigen::ArrayXi &>())
    .def(py::init< const Eigen::ArrayXi &, bool>())
    .def(py::init< const bool &, const Eigen::ArrayXXd &, const Eigen::ArrayXi &> ())
    .def(py::init< const bool &, const Eigen::ArrayXXd &, const Eigen::ArrayXi &, bool> ())
    .def("getMesh1D",  [ ](StOpt::LocalLinearDiscrLastDimRegression & p_x)
    {
        auto ret = p_x.getMesh1DCopy();
        std::vector< Eigen::ArrayXd > result;
        for (auto &array : ret)
            result.push_back(*array);
        return result;
    })
    .def("getNbMesh", &StOpt::LocalRegression::getNbMeshCopy)
    .def("getSimToCell", &StOpt::LocalRegression::getSimToCellCopy)
    ;

    py::class_<StOpt::LocalConstDiscrLastDimRegression, std::shared_ptr< StOpt::LocalConstDiscrLastDimRegression>, StOpt::BaseRegression> (m, "LocalConstDiscrLastDimRegression")
    .def(py::init< const Eigen::ArrayXi &>())
    .def(py::init< const Eigen::ArrayXi &, bool>())
    .def(py::init< const bool &, const Eigen::ArrayXXd &, const Eigen::ArrayXi &> ())
    .def(py::init< const bool &, const Eigen::ArrayXXd &, const Eigen::ArrayXi &, bool> ())
    .def("getMesh1D",  [ ](StOpt::LocalConstDiscrLastDimRegression & p_x)
    {
        auto ret = p_x.getMesh1DCopy();
        std::vector< Eigen::ArrayXd > result;
        for (auto &array : ret)
            result.push_back(*array);
        return result;
    })
    .def("getNbMesh", &StOpt::LocalRegression::getNbMeshCopy)
    .def("getSimToCell", &StOpt::LocalRegression::getSimToCellCopy)
    ;

    py::class_<StOpt::LocalKMeansRegression, std::shared_ptr< StOpt::LocalKMeansRegression>, StOpt::BaseRegression> (m, "LocalKMeansRegression")
    .def(py::init< const Eigen::ArrayXi &>())
    .def(py::init< const Eigen::ArrayXi &, bool>())
    .def(py::init< const bool &, const Eigen::ArrayXXd &, const Eigen::ArrayXi &> ())
    .def(py::init< const bool &, const Eigen::ArrayXXd &, const Eigen::ArrayXi &, bool> ())
    ;

    py::class_<StOpt::LocalSameSizeLinearRegression, std::shared_ptr< StOpt::LocalSameSizeLinearRegression>, StOpt::BaseRegression> (m, "LocalSameSizeLinearRegression")
    .def(py::init<const Eigen::ArrayXd &, const Eigen::ArrayXd &, const  Eigen::ArrayXi & >())
    .def(py::init<const bool &,   const Eigen::ArrayXXd &,   const Eigen::ArrayXd &,  const Eigen::ArrayXd &,  const  Eigen::ArrayXi & > ())
    ;

    py::class_<StOpt::LocalSameSizeConstRegression, std::shared_ptr< StOpt::LocalSameSizeConstRegression>, StOpt::BaseRegression> (m, "LocalSameSizeConstRegression")
    .def(py::init<const Eigen::ArrayXd &, const Eigen::ArrayXd &, const  Eigen::ArrayXi & >())
    .def(py::init<const bool &,   const Eigen::ArrayXXd &,   const Eigen::ArrayXd &,  const Eigen::ArrayXd &,  const  Eigen::ArrayXi & > ())
    ;

    py::class_<StOpt::SparseRegression, std::shared_ptr<StOpt::SparseRegression>, StOpt::BaseRegression > (m, "SparseRegression")
    .def(py::init<  const int &, const Eigen::ArrayXd &, const int &>())
    .def(py::init< const bool &, const Eigen::ArrayXXd &,  const int &, const Eigen::ArrayXd &, const int &> ())
    ;

    py::class_<StOpt::GlobalRegression<StOpt::Hermite>, std::shared_ptr<StOpt::GlobalRegression<StOpt::Hermite> >, StOpt::BaseRegression > (m, "GlobalHermiteRegression")
    .def(py::init< const int &, const int & >())
    .def(py::init< const bool &, Eigen::ArrayXXd &,  const int & > ())
    ;
    py::class_<StOpt::GlobalRegression<StOpt::Canonical>, std::shared_ptr<StOpt::GlobalRegression<StOpt::Canonical> >, StOpt::BaseRegression > (m, "GlobalCanonicalRegression")
    .def(py::init< const int &, const int & >())
    .def(py::init< const bool &, const  Eigen::ArrayXXd &,  const int & > ())
    ;
    py::class_<StOpt::GlobalRegression<StOpt::Tchebychev>, std::shared_ptr<StOpt::GlobalRegression<StOpt::Tchebychev> >, StOpt::BaseRegression > (m, "GlobalTchebychevRegression")
    .def(py::init< const int &, const int & >())
    .def(py::init< const int &, const int &, bool >())
    .def(py::init< const bool &, const Eigen::ArrayXXd &,  const int & > ())
    .def(py::init< const bool &, const Eigen::ArrayXXd &,  const int &, bool > ())
    ;

    py::class_<StOpt::LocalGridKernelRegression, std::shared_ptr<StOpt::LocalGridKernelRegression>, StOpt::BaseRegression > (m, "LocalGridKernelRegression")
    .def(py::init< const bool &, const Eigen::ArrayXXd &, const double &, const double &, const bool &> ())
    .def(py::init<  const double &, const double &, const bool &> ())
    ;
    py::class_<StOpt::LaplacianConstKernelRegression, std::shared_ptr<StOpt::LaplacianConstKernelRegression>, StOpt::BaseRegression > (m, "LaplacianConstKernelRegression")
    .def(py::init< const bool &, const  Eigen::ArrayXXd &, const Eigen::ArrayXd &> ())
    .def(py::init< const Eigen::ArrayXd &> ())
    .def(py::init<  const bool &, const Eigen::ArrayXXd &> ())
    ;
    py::class_<StOpt::LaplacianLinearKernelRegression, std::shared_ptr<StOpt::LaplacianLinearKernelRegression>, StOpt::BaseRegression > (m, "LaplacianLinearKernelRegression")
    .def(py::init< const bool &, const  Eigen::ArrayXXd &, const Eigen::ArrayXd &> ())
    .def(py::init< const Eigen::ArrayXd &> ())
    .def(py::init<  const bool &, const  Eigen::ArrayXXd &> ())
    ;
    py::class_<StOpt::LaplacianGridKernelRegression, std::shared_ptr<StOpt::LaplacianGridKernelRegression>, StOpt::BaseRegression > (m, "LaplacianGridKernelRegression")
    .def(py::init< const bool &, const Eigen::ArrayXXd &,  const Eigen::ArrayXd &,  const double &, const bool &> ())
    .def(py::init<const Eigen::ArrayXd &,  const double &,  const bool &  > ())
    ;

    // map contination values
    py::class_<StOpt::ContinuationValue, std::shared_ptr<StOpt::ContinuationValue> >(m, "ContinuationValue")
    .def(py::init <  const std::shared_ptr< StOpt::SpaceGrid >  &,  const std::shared_ptr< StOpt::BaseRegression> &,    const Eigen::ArrayXXd & > ())
    .def(py::init<   const std::shared_ptr< StOpt::BaseRegression> &>())
    .def("loadForSimulation", &StOpt::ContinuationValue::loadForSimulation)
    .def("loadCondExpForSimulation", &StOpt::ContinuationValue::loadCondExpForSimulation)
    .def("getAllSimulations", (Eigen::ArrayXd(StOpt::ContinuationValue::*)(const Eigen::ArrayXd &) const) &StOpt::ContinuationValue::getAllSimulations)
    .def("getAllSimulations", (Eigen::ArrayXd(StOpt::ContinuationValue::*)(const StOpt::Interpolator &) const) &StOpt::ContinuationValue::getAllSimulations)
    .def("getASimulation", (double (StOpt::ContinuationValue::*)(const int &, const StOpt::Interpolator &) const)  &StOpt::ContinuationValue::getASimulation, "From interpolator")
    .def("getValue", &StOpt::ContinuationValue::getValue)
    .def("getValues", &StOpt::ContinuationValue::getValuesCopy)
    .def("getGrid", &StOpt::ContinuationValue::getGrid)
    .def("getCondExp", &StOpt::ContinuationValue::getCondExp)
    .def("getNbSimul", &StOpt::ContinuationValue::getNbSimul)
    .def("getParticles", &StOpt::ContinuationValue::getParticles)
    ;

    // map grid and regressed values
    py::class_<StOpt::GridAndRegressedValue, std::shared_ptr<StOpt::GridAndRegressedValue> >(m, "GridAndRegressedValue")
    .def(py::init <  const std::shared_ptr< StOpt::SpaceGrid >  &,  const std::shared_ptr< StOpt::BaseRegression> &,    const Eigen::ArrayXXd & > ())
    .def("getValue", &StOpt::GridAndRegressedValue::getValue)
    .def("getGrid", &StOpt::GridAndRegressedValue::getGrid)
    .def("getRegressor", &StOpt::GridAndRegressedValue::getRegressor)
    ;


    // map continuation cut values
    py::class_<StOpt::ContinuationCuts, std::shared_ptr<StOpt::ContinuationCuts> >(m, "ContinuationCut")
    .def(py::init <  const std::shared_ptr< StOpt::SpaceGrid >  &,  const std::shared_ptr< StOpt::BaseRegression> &,  const Eigen::ArrayXXd & > ())
    .def(py::init<   const std::shared_ptr< StOpt::BaseRegression> &>())
    .def("loadForSimulation", [ ](StOpt::ContinuationCuts &  p_x, const  std::shared_ptr< StOpt::SpaceGrid > &p_grid,
                                  const std::shared_ptr< StOpt::BaseRegression >   &p_condExp, const py::list & ns)
    {
        Eigen::Array< Eigen::ArrayXXd, Eigen::Dynamic, 1  >  values(ns.size());
        for (size_t i = 0; i < ns.size(); ++i)
        {
            values(i) = ns[i].cast<Eigen::ArrayXXd>();
        }
        return p_x.loadForSimulation(p_grid, p_condExp, values);
        ;
    }
        )
    .def("getCutsAllSimulations", &StOpt::ContinuationCuts::getCutsAllSimulations)
    .def("getCutsASim", &StOpt::ContinuationCuts::getCutsASim)
    .def("getGrid", &StOpt::ContinuationCuts::getGrid)
    .def("getCondExp", &StOpt::ContinuationCuts::getCondExp)
    .def("getNbSimul", &StOpt::ContinuationCuts::getNbSimul)
    .def("getParticles", &StOpt::ContinuationCuts::getParticles)
    ;
}
