// Copyright (C) 2019 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#define BOOST_TEST_MODULE testTreeExpCOnd
#define BOOST_TEST_DYN_LINK
#include <memory>
#include <vector>
#include <array>
#include <boost/test/unit_test.hpp>
#include "geners/BinaryFileArchive.hh"
#include "geners/Record.hh"
#include "geners/Reference.hh"
#include "StOpt/tree/Tree.h"
#include "StOpt/tree/TreeGeners.h"
#include "StOpt/tree/ContinuationValueTreeGeners.h"
#include "StOpt/tree/ContinuationCutsTreeGeners.h"
#include "StOpt/core/grids/RegularSpaceGridGeners.h"
#include "StOpt/core/utils/constant.h"
#include "test/c++/tools/simulators/TrinomialTreeOUSimulator.h"

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

double accuracyEqual = 1e-5;
double accuracyNearlyEqual = 0.1;


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



void testTreeSer()
{

    // mean reverting
    double mr = 0.3;
    // volatility
    double sig = 0.6;

    // step
    double dt = 1. / 100;
    // simulation dates
    ArrayXd dates =  dt * ArrayXd::LinSpaced(14, 0., 1.);

    // simulaton dates
    TrinomialTreeOUSimulator tree(mr, sig, dates);

    int iFirst = 10;
    int iLast = 13;

    // nodes at dates 5
    ArrayXXd points = tree.getPoints(iFirst);

    // nodes at last date
    ArrayXXd pointsNext = tree.getPoints(iLast);

    // probabilities
    ArrayXXd  proba = tree.getProbability(iFirst, iLast);

    // function to regress
    ArrayXd toTreeress =  1 + pointsNext.row(0).transpose();

    // connection matrix and proba
    pair< vector< vector< array<int, 2> > >, vector< double >  >   connectAndProba =  tree.calConnected(proba);

    Tree treeEsp(connectAndProba.second, connectAndProba.first);

    // single member
    ArrayXd valTree = treeEsp.expCond(toTreeress);


    // two members
    ArrayXXd toTreeress2D(2, pointsNext.cols());
    toTreeress2D.row(0) = 1 + pointsNext.row(0);
    toTreeress2D.row(1) =  1 + pointsNext.row(0);


    // single member
    ArrayXXd valTree2D = treeEsp.expCondMultiple(toTreeress2D);

    for (int i = 0; i < points.cols(); ++i)
    {
        BOOST_CHECK_CLOSE(points(0, i) + 1, valTree2D(0, i), accuracyNearlyEqual);
        BOOST_CHECK_CLOSE(points(0, i) + 1, valTree2D(1, i), accuracyNearlyEqual);
    }


    // serialization
    {
        BinaryFileArchive ar("archiveT1", "w");
        ar << Record(treeEsp, "Tree", "Top") ;
    }
    {
        BinaryFileArchive ar("archiveT1", "r");

        // Deserialization
        Tree treeEspLoc;
        Reference< Tree> (ar, "Tree", "Top").restore(0, &treeEspLoc);

        // single member
        ArrayXd valTreeLoc = treeEspLoc.expCond(toTreeress);

        for (int i = 0; i < points.cols(); ++i)
            BOOST_CHECK_CLOSE(points(0, i) + 1, valTreeLoc(i), accuracyNearlyEqual);

    }

}



BOOST_AUTO_TEST_CASE(testTree)
{
#if defined   __linux
    enable_abort_on_floating_point_exception();
#endif
    testTreeSer();
}


/// \param  p_nbStock number of stock poin
void testContinationTreeSer(const int   &p_nbStock)
{

    // mean reverting
    double mr = 0.3;
    // volatility
    double sig = 0.6;

    // step
    double dt = 1. / 100;
    // simulation dates
    ArrayXd dates =  dt * ArrayXd::LinSpaced(15, 0., 1.);

    // simulaton dates
    TrinomialTreeOUSimulator tree(mr, sig, dates);

    int iFirst = 10;
    int iLast = 14;

    // nodes at dates 5
    ArrayXXd points = tree.getPoints(iFirst);

    // nodes at last date
    ArrayXXd pointsNext = tree.getPoints(iLast);

    // probabilities
    ArrayXXd  proba = tree.getProbability(iFirst, iLast);

    // connection matrix
    pair< vector< vector< array<int, 2> > >, vector< double >  >   connectAndProba =  tree.calConnected(proba);

    shared_ptr<Tree> treeEsp  = make_shared<Tree>(connectAndProba.second, connectAndProba.first);


    // two members
    ArrayXXd toTreeress(p_nbStock, pointsNext.cols());
    for (int i = 0; i < p_nbStock; ++i)
    {
        toTreeress.row(i) = static_cast<double>(i) + pointsNext.row(0);
    }

    // single member
    ArrayXXd valTree = treeEsp->expCondMultiple(toTreeress);


    // grid for stock
    Eigen::ArrayXd lowValues(1), step(1);
    lowValues(0) = 0. ;
    step(0) = 1;
    Eigen::ArrayXi  nbStep(1);
    nbStep(0) = p_nbStock - 1;
    // grid
    shared_ptr< RegularSpaceGrid > regular = make_shared<RegularSpaceGrid>(lowValues, step, nbStep);

    // creation continuation value object
    ContinuationValueTree  continuation(regular, treeEsp,  toTreeress.transpose());

    ArrayXd ptStock(1) ;
    ptStock(0) = p_nbStock / 2;

    ArrayXd treeByContinuation = continuation.getValueAtNodes(ptStock);

    for (int is  = 0;  is <  points.cols(); ++is)
        BOOST_CHECK_CLOSE(treeByContinuation(is), valTree(p_nbStock / 2, is), accuracyEqual);

    {
        // default non compression
        BinaryFileArchive ar("archiveT2", "w");
        ar << Record(continuation, "FirstContinuation", "Top") ;
    }
    {
        // read archive
        BinaryFileArchive ar("archiveT2", "r");
        ContinuationValueTree contRead;
        Reference< ContinuationValueTree >(ar, "FirstContinuation", "Top").restore(0, &contRead);
        ArrayXd ptStock(1) ;
        ptStock(0) = p_nbStock / 2;
        for (int is  = 0;  is < points.cols(); ++is)
            BOOST_CHECK_CLOSE(contRead.getValueAtANode(is, ptStock), valTree(p_nbStock / 2, is), accuracyEqual);
    }
}

BOOST_AUTO_TEST_CASE(testContinuationTree)
{
#if defined   __linux
    enable_abort_on_floating_point_exception();
#endif

    int nbStock = 4 ;
    testContinationTreeSer(nbStock);
}
/// \param  p_nbStock number of stock poin
void testContinationCutTreeSer(const int   &p_nbStock)
{

    // mean reverting
    double mr = 0.3;
    // volatility
    double sig = 0.6;

    // step
    double dt = 1. / 100;
    // simulation dates
    ArrayXd dates =  dt * ArrayXd::LinSpaced(15, 0., 1.);

    // simulaton dates
    TrinomialTreeOUSimulator tree(mr, sig, dates);

    int iFirst = 10;
    int iLast = 14;

    // nodes at dates 5
    ArrayXXd points = tree.getPoints(iFirst);

    // nodes at last date
    ArrayXXd pointsNext = tree.getPoints(iLast);

    // probabilities
    ArrayXXd  proba = tree.getProbability(iFirst, iLast);

    // connection matrix
    pair< vector< vector< array<int, 2> > >, vector< double >  >  connectAndProba =  tree.calConnected(proba);

    shared_ptr<Tree> treeEsp  = make_shared<Tree>(connectAndProba.second, connectAndProba.first);

    // two members
    ArrayXXd toTreeress(p_nbStock, pointsNext.cols());
    ArrayXXd toTreeressCut(p_nbStock, pointsNext.cols() * 2);
    for (int i = 0; i < p_nbStock; ++i)
    {
        toTreeress.row(i) = static_cast<double>(i + 1) + pointsNext.row(0);
        toTreeressCut.row(i).head(pointsNext.cols()) = static_cast<double>(i + 1) + pointsNext.row(0);
        toTreeressCut.row(i).tail(pointsNext.cols()).setConstant(1);
    }

    // single member
    ArrayXXd valTree = treeEsp->expCondMultiple(toTreeress);

    // grid for stock
    Eigen::ArrayXd lowValues(1), step(1);
    lowValues(0) = 0. ;
    step(0) = 1;
    Eigen::ArrayXi  nbStep(1);
    nbStep(0) = p_nbStock - 1;
    // grid
    shared_ptr< RegularSpaceGrid > regular = make_shared<RegularSpaceGrid>(lowValues, step, nbStep);

    // creation continuation value object
    ContinuationCutsTree  contCut(regular, treeEsp,  toTreeressCut.transpose());

    // calculate all cuts
    ArrayXXd  hyperStock(1, 2) ;
    hyperStock(0, 0) = lowValues(0);
    hyperStock(0, 1) = lowValues(0) + nbStep(0) * step(0);

    ArrayXXd treeressedCuts = contCut.getCutsAllNodes(hyperStock);

    // test only values
    shared_ptr<GridIterator> iterGrid = regular->getGridIterator();
    while (iterGrid->isValid())
    {
        for (int is  = 0;  is < points.cols(); ++is)
        {
            // coordinates
            ArrayXd pointCoord = iterGrid->getCoordinate();
            // reconstruct value from cuts coefficients
            double valReconst = treeressedCuts(is, iterGrid->getCount()) + treeressedCuts(is + points.cols(), iterGrid->getCount()) * pointCoord(0);
            BOOST_CHECK_CLOSE(valReconst, valTree(iterGrid->getCount(), is), accuracyEqual);
        }
        iterGrid->next();
    }

    {
        // default non compression
        BinaryFileArchive ar("archiveT3", "w");
        ar << Record(contCut, "FirstContinuation", "Top") ;

    }
    {
        // read archive
        BinaryFileArchive ar("archiveT3", "r");
        ContinuationCutsTree contCutRead;
        Reference< ContinuationCutsTree >(ar, "FirstContinuation", "Top").restore(0, &contCutRead);

        // test only values
        shared_ptr<GridIterator> iterGrid = regular->getGridIterator();
        while (iterGrid->isValid())
        {
            for (int is  = 0;  is < points.cols(); ++is)
            {
                // coordinates
                ArrayXd pointCoord = iterGrid->getCoordinate();

                // cut from a sim
                ArrayXXd treeCuts = contCutRead.getCutsANode(hyperStock, is);

                // reconstruct value from cuts coefficients
                double valReconst = treeCuts(0, iterGrid->getCount()) + treeCuts(1, iterGrid->getCount()) * pointCoord(0);
                BOOST_CHECK_CLOSE(valReconst, valTree(iterGrid->getCount(), is), accuracyEqual);
            }
            iterGrid->next();
        }
    }
}

BOOST_AUTO_TEST_CASE(testContinuationCutTree)
{
#if defined   __linux
    enable_abort_on_floating_point_exception();
#endif

    int nbStock = 4 ;
    testContinationCutTreeSer(nbStock);
}
