// Copyright (C) 2016 Fime
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#ifndef USE_MPI
#define BOOST_TEST_MODULE testGasStorageSDDP
#endif
#define BOOST_TEST_DYN_LINK
#ifdef USE_MPI
#include <boost/mpi.hpp>
#endif
#include <functional>
#include <array>
#include <memory>
#define _USE_MATH_DEFINES
#include <math.h>
#include <boost/test/unit_test.hpp>
#include <Eigen/Dense>
#include "StOpt/core/grids/OneDimRegularSpaceGrid.h"
#include "StOpt/core/grids/OneDimData.h"
#include "StOpt/core/grids/RegularSpaceGridGeners.h"
#include "StOpt/sddp/LocalLinearRegressionForSDDPGeners.h"
#include "StOpt/sddp/LocalConstRegressionForSDDPGeners.h"
#include "StOpt/sddp/SDDPFinalCut.h"
#include "StOpt/sddp/SDDPLocalCut.h"
#include "StOpt/sddp/backwardForwardSDDP.h"
#include "test/c++/tools/simulators/MeanRevertingSimulator.h"
#include "test/c++/tools/sddp/OptimizeGasStorageSDDP.h"
#include "test/c++/tools/dp/OptimizeGasStorage.h"
#include "test/c++/tools/dp/DynamicProgrammingByRegression.h"
#ifdef USE_MPI
#include "test/c++/tools/dp/DynamicProgrammingByRegressionDist.h"
#endif

using namespace std;
using namespace Eigen ;
using namespace StOpt;


#if defined   __linux
#include <fenv.h>
#define enable_abort_on_floating_point_exception() feenableexcept(FE_DIVBYZERO | FE_INVALID)
#endif


/// For Clang < 3.7 (and above ?) to be compatible GCC 5.1 and above
namespace boost
{
namespace unit_test
{
namespace ut_detail
{
string normalize_test_case_name(const_string name)
{
    return (name[0] == '&' ? string(name.begin() + 1, name.size() - 1) : string(name.begin(), name.size()));
}
}
}
}

class ZeroFunction
{
public:
    ZeroFunction() {}
    double operator()(const int &, const ArrayXd &, const ArrayXd &) const
    {
        return 0. ;
    }
};


// low accuracy to have not too long test cases
double accuracyClose =  5.;


/// \brief Comparison SDDP , DP for gas storage
/// \param p_nbStorage   number of storage
/// \param p_nMeshSDDP   number of mesh for uncertainty in backward resolution,
/// \param p_nbsimulOpt  number of simulations in backward resolution.
/// \param p_nbSamSimul  number of sample used in simulation to derive new path used in the next backward resolution
/// \param p_world       MPI communicator
/// \param p_strAdd      to make difference for file with  different communicators
/// template the regressor
template< class LocalRegression, class LocalRegressionForSDDP>
void testStorageSDDP(const int &p_nbStorage, const int   &p_nMeshSDDP, const int   &p_nbsimulOpt, const int   &p_nbSamSimul
#ifdef USE_MPI
                     , const boost::mpi::communicator &p_world, const string &p_strAdd
#endif
                    )
{

    // storage
    /////////
    double maxLevelStorage  = 360000;
    double injectionRateStorage = 60000;
    double withdrawalRateStorage = 45000;
    double injectionCostStorage = 0.35;
    double withdrawalCostStorage = 0.35;

    double maturity = 1.;
    size_t nstep = 30;

    // Dynamic of the future
    //*********************

    // define a a time grid
    shared_ptr<OneDimRegularSpaceGrid> timeGrid(new OneDimRegularSpaceGrid(0., maturity / nstep, nstep));
    // future values
    shared_ptr<vector< double > > futValues(new vector<double>(nstep + 1));
    // periodicity factor
    int iPeriod = 52;
    for (size_t i = 0; i < nstep + 1; ++i)
        (*futValues)[i] = 50. + 20 * sin((M_PI * i * iPeriod) / nstep);
    // define the future curve
    shared_ptr<OneDimData<OneDimRegularSpaceGrid, double> > futureGrid(new OneDimData< OneDimRegularSpaceGrid, double> (timeGrid, futValues));
    // one dimensional factors
    int nDim = 1;
    VectorXd sigma = VectorXd::Constant(nDim, 0.94);
    VectorXd mr = VectorXd::Constant(nDim, 0.29);

    // value by DP
    //**************
    double valueOptim = 0;

    {
        // initial values
        ArrayXd initialStockDP = ArrayXd::Constant(1, maxLevelStorage);

        // grid
        //////
        int nGrid = 320;
        ArrayXd lowValues = ArrayXd::Constant(1, 0.);
        ArrayXd step = ArrayXd::Constant(1, maxLevelStorage / nGrid);
        ArrayXi nbStep = ArrayXi::Constant(1, nGrid);
        shared_ptr<RegularSpaceGrid> grid = make_shared<RegularSpaceGrid>(lowValues, step, nbStep);

        // no actualization
        double r  = 0 ;
        // a backward simulator
        ///////////////////////
        bool bForward = false;
        int nbSimulDP = p_nbsimulOpt;
        shared_ptr< MeanRevertingSimulator< OneDimData<OneDimRegularSpaceGrid, double> > > backSimulatorDP = make_shared<MeanRevertingSimulator< OneDimData<OneDimRegularSpaceGrid, double> > >(futureGrid, sigma, mr, r,  maturity, nstep, nbSimulDP, bForward);

        // optimizer
        ///////////
        shared_ptr< OptimizeGasStorage< MeanRevertingSimulator< OneDimData<OneDimRegularSpaceGrid, double> > > > storage = make_shared< OptimizeGasStorage< MeanRevertingSimulator< OneDimData<OneDimRegularSpaceGrid, double> > > > (injectionRateStorage, withdrawalRateStorage, injectionCostStorage, withdrawalCostStorage);
        // regressor
        ///////////
        int nMeshDP = p_nMeshSDDP ;
        ArrayXi nbMeshDP = ArrayXi::Constant(1, nMeshDP);
        shared_ptr< BaseRegression > regressor = make_shared<LocalRegression>(nbMeshDP);
        // final value
        function<double(const int &, const ArrayXd &, const ArrayXd &)>   vFunction = ZeroFunction();

        // initial regime
        int initialRegime = 0; // only one regime

        // Optimize
        ///////////
        string fileToDump = "CondExpGasStorageSDDP" ;
#ifdef USE_MPI
        fileToDump += p_strAdd;
#endif
        // link the simulations to the optimizer
        storage->setSimulator(backSimulatorDP);
#ifdef USE_MPI
        bool bOneFile = true;
        valueOptim =  DynamicProgrammingByRegressionDist(grid, storage, regressor, vFunction, initialStockDP, initialRegime, fileToDump, bOneFile, p_world);
#else
        valueOptim =  DynamicProgrammingByRegression(grid, storage, regressor, vFunction, initialStockDP, initialRegime, fileToDump);

#endif

#ifdef USE_MPI
        if (p_world.rank() == 0)
#endif
            cout << " Value by DP " << valueOptim << endl ;

    }
#ifdef USE_MPI
    p_world.barrier();
#endif

    // initial values
    ArrayXd initialStock = ArrayXd::Constant(p_nbStorage, maxLevelStorage);

    // no actualization
    double r  = 0 ;
    // a backward simulator
    bool bForward = false;
    int nbsimulOpt = p_nbsimulOpt;
    shared_ptr<MeanRevertingSimulator< OneDimData<OneDimRegularSpaceGrid, double> > > backSimulator = make_shared<MeanRevertingSimulator< OneDimData<OneDimRegularSpaceGrid, double> > >(futureGrid, sigma, mr, r, maturity, nstep, nbsimulOpt, bForward);

    // mesh for regressor
    int nMesh = p_nMeshSDDP;
    ArrayXi nbMesh = ArrayXi::Constant(nDim, nMesh);

    // a forward simulator
    int nbsimulSim = p_nbSamSimul;

    bForward = true;
    shared_ptr<MeanRevertingSimulator< OneDimData<OneDimRegularSpaceGrid, double> > > forSimulator = make_shared<MeanRevertingSimulator< OneDimData<OneDimRegularSpaceGrid, double> > >(futureGrid, sigma, mr, r, maturity, nstep, nbsimulSim, bForward);

    // value by SDDP
    //**************
    // SDDP one step optimizer
    shared_ptr< OptimizerSDDPBase> optimizer =   make_shared< OptimizeGasStorageSDDP< MeanRevertingSimulator< OneDimData<OneDimRegularSpaceGrid, double> > > > (maxLevelStorage, injectionRateStorage,
            withdrawalRateStorage, injectionCostStorage,
            withdrawalCostStorage, p_nbStorage,
            backSimulator, forSimulator);

    /// final cut
    ArrayXXd finalCut =  ArrayXXd::Zero(1 + p_nbStorage, 1);
    SDDPFinalCut finCut(finalCut);

    // optimisation dates
    ArrayXd dates = ArrayXd::LinSpaced(nstep + 1, 0., maturity);

    // names for archive
    string nameRegressor = "RegressorGasStorage" ;
    string nameCut = "CutGasStorage" ;
    string nameVisitedStates = "VisitedStateGasStorage" ;
#ifdef USE_MPI
    nameRegressor += p_strAdd;
    nameCut += p_strAdd;
    nameVisitedStates += p_strAdd ;
#endif

    // precision parameter
    int nIterMax = 100;
    double accuracy = accuracyClose / 100;
    int nstepIterations = 10; // take values between optimization n and optimization n+ nstepIterations for convergence criterion
    ostringstream stringStream;
    pair<double, double>  values = backwardForwardSDDP<LocalRegressionForSDDP>(optimizer, nbsimulOpt, initialStock, finCut, dates,  nbMesh,
                                   nameRegressor, nameCut, nameVisitedStates, nIterMax,
                                   accuracy, nstepIterations, stringStream
#ifdef USE_MPI
                                   , p_world
#endif
                                                                              );

#ifdef USE_MPI
    if (p_world.rank() == 0)
#endif
    {
        cout << stringStream.str() << endl ;
        cout << " Value Optim " <<  values.first << " and Simulation " << values.second << " Iteration " << nIterMax << endl ;

        BOOST_CHECK_CLOSE(valueOptim * p_nbStorage, values.first, accuracyClose);
    }
}

BOOST_AUTO_TEST_CASE(testSimpleStorageSDDP1D)
{
#ifdef USE_MPI
    boost::mpi::communicator world;
    world.barrier();
#endif
    int nbMeshLinear = 4 ;
    int nbsimulLinear = 4000;
    int nbSamSimulLinear = 4;
    testStorageSDDP<LocalLinearRegression, LocalLinearRegressionForSDDP>(1, nbMeshLinear, nbsimulLinear, nbSamSimulLinear
#ifdef USE_MPI
            , world, ""
#endif
                                                                        );
#ifdef USE_LONG_TEST
    int nbMeshConst = 30 ;
    int nbsimulConst  = 14000;
    int nbSamSimulConst = 30 ;
    testStorageSDDP<LocalConstRegression, LocalConstRegressionForSDDP>(1, nbMeshConst, nbsimulConst, nbSamSimulConst
#ifdef USE_MPI
            , world, ""
#endif
                                                                      );
#endif
}

#ifdef USE_MPI
BOOST_AUTO_TEST_CASE(testSimpleStorageSDDP1DCommunicator)
{
    boost::mpi::communicator world;
    if (world.size() > 1)
    {
        // create 2 communicators
        bool bSeparate = (2 * world.rank() < world.size());
        boost::mpi::communicator worldLoc = world.split(bSeparate ? 0 : 1);
        string strToAdd = (bSeparate ? "0" : "1");
        int nbMeshLinear = 4 ;
        int nbsimulLinear = 4000;
        int nbSamSimulLinear = 4;
        testStorageSDDP<LocalLinearRegression, LocalLinearRegressionForSDDP>(1, nbMeshLinear, nbsimulLinear, nbSamSimulLinear, worldLoc, strToAdd);
    }
}
#endif


#ifdef USE_LONG_TEST
BOOST_AUTO_TEST_CASE(testSimpleStorageSDDP2D)
{
#ifdef USE_MPI
    boost::mpi::communicator world;
    world.barrier();
#endif
    int nbMeshLinear = 4 ;
    int nbsimulLinear = 4000;
    int nbSamSimulLinear = 4;
    testStorageSDDP<LocalLinearRegression, LocalLinearRegressionForSDDP>(2, nbMeshLinear, nbsimulLinear, nbSamSimulLinear
#ifdef USE_MPI
            , world, ""
#endif
                                                                        );
    int nbMeshConst = 30 ;
    int nbsimulConst  = 14000;
    int nbSamSimulConst = 30 ;
    testStorageSDDP<LocalConstRegression, LocalConstRegressionForSDDP>(2, nbMeshConst, nbsimulConst, nbSamSimulConst
#ifdef USE_MPI
            , world, ""
#endif
                                                                      );
}


// BOOST_AUTO_TEST_CASE(testSimpleStorageSDDP5D)
// {
//     int nbMeshLinear = 4 ;
//     int nbsimulLinear = 2000;
//     int nbSamSimulLinear = 4;
//     testStorageSDDP<LocalLinearRegression, LocalLinearRegressionForSDDP>(5, nbMeshLinear, nbsimulLinear, nbSamSimulLinear);
//     int nbMeshConst = 30 ;
//     int nbsimulConst  = 14000;
//     int nbSamSimulConst = 30 ;
//     testStorageSDDP<LocalConstRegression, LocalConstRegressionForSDDP>(5, nbMeshConst, nbsimulConst, nbSamSimulConst);
// }

// BOOST_AUTO_TEST_CASE(testSimpleStorageSDDP10D)
// {
//     int nbMeshLinear = 4 ;
//     int nbsimulLinear = 2000;
//     int nbSamSimulLinear = 4;
//     testStorageSDDP<LocalLinearRegression, LocalLinearRegressionForSDDP>(10, nbMeshLinear, nbsimulLinear, nbSamSimulLinear);
//     int nbMeshConst = 30 ;
//     int nbsimulConst  = 14000;
//     int nbSamSimulConst = 30 ;
//     testStorageSDDP<LocalConstRegression, LocalConstRegressionForSDDP>(10, nbMeshConst, nbsimulConst, nbSamSimulConst);
// }
#endif

#ifdef USE_MPI
// (empty) Initialization function. Can't use testing tools here.
bool init_function()
{
    return true;
}

int main(int argc, char *argv[])
{
    boost::mpi::environment env(argc, argv);
    return ::boost::unit_test::unit_test_main(&init_function, argc, argv);
}
#endif
