// Copyright (C) 2019 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#ifdef USE_MPI
#include <boost/mpi.hpp>
#endif
#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <pybind11/stl_bind.h>
#include <pybind11/stl.h>
#include "StOpt/core/utils/version.h"
#include "StOpt/sddp/SimulatorSDDPBase.h"
#include "StOpt/sddp/OptimizerSDDPBase.h"
#include "StOpt/sddp/LocalLinearRegressionForSDDP.h"
#include "StOpt/sddp/LocalLinearRegressionForSDDPGeners.h"
#include "StOpt/sddp/SDDPFinalCut.h"
#include "StOpt/sddp/backwardForwardSDDP.h"
#include "StOpt/sddp/backwardSDDP.h"
#include "StOpt/sddp/forwardSDDP.h"
#include "StOpt/sddp/SDDPLocalCut.h"
#include "StOpt/python/BinaryFileArchiveStOpt.h"

namespace py = pybind11;


/** \file Pybind11SOptSDDP.cpp
 * \brief Mapping SDDP optimizer and simulator  base classes to python
 * \author Xavier Warin
 */

/// \brief backwardForward  wrapper
/// \return backward, forward valorization ,  number of iterations used in SDDP, and acciracy achieved
py::dict pyBackwardForwardSDDP(std::shared_ptr<StOpt::OptimizerSDDPBase> p_optimizer,
                               const int   &p_nbSimulCheckForSimu,
                               const Eigen::ArrayXd &p_initialState,
                               const StOpt::SDDPFinalCut &p_finalCut,
                               const Eigen::ArrayXd &p_dates,
                               const Eigen::ArrayXi &p_meshForReg,
                               const std::string &p_nameRegressor,
                               const std::string &p_nameCut,
                               const std::string &p_nameVisitedStates,
                               int p_iter,
                               double p_accuracy,
                               const int &p_nStepConv)
{
#ifdef USE_MPI
    boost::mpi::communicator world;
#endif
    std::ostringstream  stringStream ;
    std::pair<double, double> retLoc = StOpt::backwardForwardSDDP<StOpt::LocalLinearRegressionForSDDP>(p_optimizer, p_nbSimulCheckForSimu, p_initialState,
                                       p_finalCut, p_dates, p_meshForReg, p_nameRegressor, p_nameCut,
                                       p_nameVisitedStates, p_iter, p_accuracy,
                                       p_nStepConv,
                                       stringStream
#ifdef USE_MPI
                                       , world
#endif
                                                                                                      );
    py::dict dicReturn;
    dicReturn["BackWardEstimation"] = retLoc.first;
    dicReturn["ForwardEStimation"] = retLoc.second;
    dicReturn["NumberOfSDDPIterations"] = p_iter;
    dicReturn["AccuracyReached"] = p_accuracy;
    dicReturn["message"] = stringStream.str();
    return dicReturn;

}

double pyBackwardSDDP(std::shared_ptr<StOpt::OptimizerSDDPBase> p_optimizer,
                      std::shared_ptr<StOpt::SimulatorSDDPBase> p_simulator,
                      const Eigen::ArrayXd   &p_dates,
                      const Eigen::ArrayXd &p_initialState,
                      const StOpt::SDDPFinalCut &p_finalCut,
                      const std::shared_ptr<BinaryFileArchiveStOpt> &p_archiveRegresssor,
                      const std::string &p_nameVisitedStates,
                      const std::shared_ptr<BinaryFileArchiveStOpt> &p_archiveCut)
{
#ifdef USE_MPI
    boost::mpi::communicator world;
#endif
    return StOpt::backwardSDDP<StOpt::LocalLinearRegressionForSDDP>(p_optimizer, p_simulator, p_dates,
            p_initialState, p_finalCut, p_archiveRegresssor, p_nameVisitedStates, p_archiveCut
#ifdef USE_MPI
            , world
#endif
                                                                   );
}

double	pyForwardSDDP(std::shared_ptr<StOpt::OptimizerSDDPBase>    p_optimizer,
                      std::shared_ptr<StOpt::SimulatorSDDPBase>  p_simulator,
                      const Eigen::ArrayXd &p_dates,
                      const Eigen::ArrayXd &p_initialState,
                      const StOpt::SDDPFinalCut &p_finalCut,
                      const bool   &p_bIncreaseCut,
                      const std::shared_ptr<BinaryFileArchiveStOpt > &p_archiveRegresssor,
                      const std::shared_ptr<BinaryFileArchiveStOpt > &p_archiveCutToRead,
                      const std::string &p_nameVisitedStates)
{
#ifdef USE_MPI
    boost::mpi::communicator world;
#endif
    return StOpt::forwardSDDP<StOpt::LocalLinearRegressionForSDDP>(p_optimizer, p_simulator, p_dates, p_initialState,
            p_finalCut, p_bIncreaseCut, p_archiveRegresssor, p_archiveCutToRead, p_nameVisitedStates
#ifdef USE_MPI
            , world
#endif
                                                                  );
}

// Wrapper for SDDP simulator
//***************************
class PySimulatorSDDPBase : StOpt::SimulatorSDDPBase
{
public:

    using StOpt::SimulatorSDDPBase::SimulatorSDDPBase;

    int getNbSimul() const override
    {
        PYBIND11_OVERLOAD_PURE(int, StOpt::SimulatorSDDPBase, getNbSimul);
    }
    int getNbSample() const override
    {
        PYBIND11_OVERLOAD_PURE(int, StOpt::SimulatorSDDPBase, getNbSample);
    }
    void updateDateIndex(const int &p_idate) override
    {
        PYBIND11_OVERLOAD_PURE(void, StOpt::SimulatorSDDPBase, updateDateIndex, p_idate) ;
    }

    Eigen::VectorXd getOneParticle(const int &p_isim) const override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::VectorXd, StOpt::SimulatorSDDPBase,  getOneParticle, p_isim);
    }
    Eigen::MatrixXd getParticles() const override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::MatrixXd, StOpt::SimulatorSDDPBase, getParticles);
    }
    void resetTime() override
    {
        PYBIND11_OVERLOAD_PURE(void, StOpt::SimulatorSDDPBase, resetTime) ;
    }

    void updateSimulationNumberAndResetTime(const int &p_nbSimul) override
    {
        PYBIND11_OVERLOAD_PURE(void, StOpt::SimulatorSDDPBase, updateSimulationNumberAndResetTime, p_nbSimul);
    }
};


/// Wrap SDDP optimizer
//*********************
class  PyOptimizerSDDPBase : StOpt::OptimizerSDDPBase
{
public :

    using StOpt::OptimizerSDDPBase::OptimizerSDDPBase;

    Eigen::ArrayXd oneStepBackward(const StOpt::SDDPCutOptBase &, const std::tuple< std::shared_ptr<Eigen::ArrayXd>, int, int > &, const Eigen::ArrayXd &, const int &) const override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::ArrayXd, StOpt::OptimizerSDDPBase, oneStepBackward);
    }

    double oneStepForward(const Eigen::ArrayXd &, Eigen::ArrayXd &,  Eigen::ArrayXd &, const StOpt::SDDPCutOptBase &, const int &) const override
    {
        PYBIND11_OVERLOAD_PURE(double, StOpt::OptimizerSDDPBase, oneStepForward);
    }
    void updateDates(const double &p_date, const double &p_dateNext) override
    {
        PYBIND11_OVERLOAD_PURE(void, StOpt::OptimizerSDDPBase, updateDates, p_date, p_dateNext);
    }
    Eigen::ArrayXd oneAdmissibleState(const double   &p_date)override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::ArrayXd, StOpt::OptimizerSDDPBase, oneAdmissibleState, p_date);
    }
    int getStateSize() const override
    {
        PYBIND11_OVERLOAD_PURE(int, StOpt::OptimizerSDDPBase, getStateSize);
    }
    std::shared_ptr< StOpt::SimulatorSDDPBase > getSimulatorBackward() const override
    {
        using shared_ptrSimTemp = std::shared_ptr< StOpt::SimulatorSDDPBase >;
        PYBIND11_OVERLOAD_PURE(shared_ptrSimTemp, StOpt::OptimizerSDDPBase, getSimulatorBackward);
    }
    std::shared_ptr< StOpt::SimulatorSDDPBase > getSimulatorForward() const override
    {
        using shared_ptrSimTemp = std::shared_ptr< StOpt::SimulatorSDDPBase >;
        PYBIND11_OVERLOAD_PURE(shared_ptrSimTemp, StOpt::OptimizerSDDPBase, getSimulatorForward);
    }
};



/// \brief Encapsulation for regression module
PYBIND11_MODULE(StOptSDDP, m)
{

    // version
    m.def("getVersion", StOpt::getStOptVersion);

    /// map sddp base simulator
    py::class_<StOpt::SimulatorSDDPBase, std::shared_ptr<StOpt::SimulatorSDDPBase>, PySimulatorSDDPBase>(m, "SimulatorSDDPBase")
    .def("getNbSimul", &StOpt::SimulatorSDDPBase::getNbSimul)
    .def("getNbSample", &StOpt::SimulatorSDDPBase::getNbSample)
    .def("updateDateIndex", &StOpt::SimulatorSDDPBase::updateDateIndex)
    .def("getOneParticle", &StOpt::SimulatorSDDPBase::getOneParticle)
    .def("getParticles", &StOpt::SimulatorSDDPBase::getParticles)
    .def("resetTime", &StOpt::SimulatorSDDPBase::resetTime)
    .def("updateSimulationNumberAndResetTime", &StOpt::SimulatorSDDPBase::updateSimulationNumberAndResetTime)
    ;

    // map sddp base optimizer
    py::class_<StOpt::OptimizerSDDPBase, std::shared_ptr<StOpt::OptimizerSDDPBase>, PyOptimizerSDDPBase>(m, "OptimizerSDDPBase")
    .def("updateDates", &StOpt::OptimizerSDDPBase::updateDates)
    .def("getStateSize", &StOpt::OptimizerSDDPBase::getStateSize)
    .def("oneAdmissibleState", &StOpt::OptimizerSDDPBase::oneAdmissibleState)
    ;

    // map SDDP final cut constructor
    py::class_<StOpt::SDDPFinalCut, std::shared_ptr<StOpt::SDDPFinalCut>  >(m, "SDDPFinalCut")
    .def(py::init< >())
    .def(py::init< const Eigen::ArrayXXd &>())
    ;

    // map LocalLinearRegressionForSDDP
    py::class_<StOpt::LocalLinearRegressionForSDDP, std::shared_ptr<StOpt::LocalLinearRegressionForSDDP> >(m, "LocalLinearRegressionForSDDP")
    .def(py::init<const bool &,  const Eigen::ArrayXXd &,  const Eigen::ArrayXi &>())
    .def("getNbMeshTotal", &StOpt::LocalLinearRegressionForSDDP::getNbMeshTotal)
    ;

    /// map SDDPVisitedStates
    py::class_<StOpt::SDDPVisitedStates, std::shared_ptr<StOpt::SDDPVisitedStates> >(m, "SDDPVisitedStates")
    .def(py::init<>())
    .def(py::init<const int & >())
    .def("addVisitedState", [  ](StOpt::SDDPVisitedStates  & p_x, const Eigen::Ref<Eigen::ArrayXd> &p_state, const Eigen::Ref<Eigen::ArrayXd> &p_particle, const StOpt::LocalLinearRegressionForSDDP   & p_regressor)
    {
        std::shared_ptr< Eigen::ArrayXd > state = std::make_shared<Eigen::ArrayXd>(p_state);
        p_x.addVisitedState(state, p_particle, p_regressor);
        ;
    }
        )
    .def("addVisitedStateForAll", [  ](StOpt::SDDPVisitedStates  & p_x, const Eigen::Ref<Eigen::ArrayXd> &p_state,  const StOpt::LocalLinearRegressionForSDDP   & p_regressor)
    {
        std::shared_ptr< Eigen::ArrayXd > state = std::make_shared<Eigen::ArrayXd>(p_state);
        p_x.addVisitedStateForAll(state, p_regressor);
        ;
    }
        )
    ;

    // SDDP functions
    m.def("backwardForwardSDDP", pyBackwardForwardSDDP);
    m.def("backwardSDDP", pyBackwardSDDP);
    m.def("forwardSDDP", pyForwardSDDP);

}
