// Copyright (C) 2025 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#define BOOST_TEST_MODULE testLocalCutsGridAdapt1D
#define BOOST_TEST_DYN_LINK
#include <memory>
#include <iostream>
#include <fstream>
#include <boost/test/unit_test.hpp>
#include "StOpt/regression/LocalConstRegressionGeners.h"
#include "StOpt/regression/ContinuationCutsGridAdaptNonConcaveGeners.h"

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

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

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

    /// function to maximize on grid [-5,5]
    //   max  -x+ xi
    //   xi <= x

    // grid for storage
    double xMin = -5;
    double xMax = 5.;
    int nbMeshX = 10;

    // nb stock points
    int sizeForStock  = nbMeshX + 1;

    // grid
    shared_ptr< GridAdapt1D > grid1D = make_shared<GridAdapt1D>(xMin, xMax, nbMeshX);


    int nbSimul = 2;
    ArrayXXd  toRegressCut(sizeForStock, 2 * nbSimul);
    ArrayXXd x = ArrayXXd::Zero(1, nbSimul);
    for (int is  = 0;  is < nbSimul; ++is)
    {
        for (int i = 0; i < nbMeshX + 1; ++i)
        {
            int istock = i;
            toRegressCut(istock, is) =  2 * (xMin + i); // value of the VB
            toRegressCut(istock, is + nbSimul) = 1; // cut value is 1 in x
        }
    }
    // conditional expectation
    ArrayXi nbMesh(1);
    nbMesh(0) = 1;
    shared_ptr<LocalConstRegression> localRegressor = make_shared<LocalConstRegression>(false, x, nbMesh);

    // creation continuation value object
    ContinuationCutsGridAdaptNonConcave  contCut(grid1D, localRegressor,  toRegressCut.transpose());

    // points that can be reached : all
    ArrayXXd   hypStock(1, 2);
    hypStock(0, 0) = xMin;
    hypStock(0, 1) = xMax;


    // one simulation
    vector< pair <shared_ptr< StOpt::GridAdaptBase>, shared_ptr<ArrayXXd>  > >  meshsAndCuts = contCut.getCutsASim(hypStock, 0);

    cout << " meshsAndCuts.size()" << meshsAndCuts.size() << endl ;
    BOOST_CHECK_EQUAL(meshsAndCuts.size(), nbMeshX);

    for (const auto &aMeshAndCuts : meshsAndCuts)
    {
        // cout << " meshsAndCuts "  << aMeshAndCuts.second << endl ;
        // cout << " PT " << aMeshAndCuts.first->getXCoord(0) << " a " <<  aMeshAndCuts.first->getYCoord(0) << endl ;
        shared_ptr<GridAdapt1D> const pGrid1D = dynamic_pointer_cast<GridAdapt1D>(aMeshAndCuts.first);
        BOOST_CHECK_EQUAL((*aMeshAndCuts.second)(0, 0), pGrid1D->getXCoord(0));
        BOOST_CHECK_EQUAL((*aMeshAndCuts.second)(1, 0), 1.);
    }

}

/// the solution here is concave
/// calculating conditional ùmeshes while gathering should lead to a single grid
BOOST_AUTO_TEST_CASE(testLocalCutsAndGridAgregatingConcaveCase)
{
#if defined   __linux
    enable_abort_on_floating_point_exception();
#endif

    /// function to maximize on grid [-5,5]
    //   max  -x+ xi
    //   xi <= x

    // grid for storage
    double xMin = -5;
    double xMax = 5;
    int nbMeshX = 10;

    // grid
    shared_ptr< GridAdapt1D > grid1D = make_shared<GridAdapt1D>(xMin, xMax, nbMeshX);

    // refine in [-2,-1] and [0,2] but 2 times
    list< pair<shared_ptr< Mesh1D>, shared_ptr< vector< ArrayXi > > > >  meshes = grid1D->getMeshes();
    for (auto &mesh : meshes)
    {
        if (((mesh.first->getXL() >= -2) && (mesh.first->getXR() <= -1)) ||
                ((mesh.first->getXL() >= 0) && (mesh.first->getXR() <= 2)))
            grid1D->splitMesh(mesh);
    }

    BOOST_CHECK_EQUAL(grid1D->getNbPoints(), 14);
    meshes = grid1D->getMeshes();
    // in [0,2]
    for (auto &mesh : meshes)
    {
        if ((mesh.first->getXL() >= 0) && (mesh.first->getXR() <= 2))
            grid1D->splitMesh(mesh);
    }
    BOOST_CHECK_EQUAL(grid1D->getNbPoints(), 18);


    // get all points
    vector< ArrayXd > vecPoints = grid1D->getPoints();
    int nbSimul = 2;
    ArrayXXd  toRegressCut(vecPoints.size(), 2 * nbSimul);
    ArrayXXd x = ArrayXXd::Zero(1, nbSimul);
    for (int is  = 0;  is < nbSimul; ++is)
    {
        for (size_t i = 0  ; i < vecPoints.size();  ++i)
        {
            // f(x)= concave and increasing
            toRegressCut(i, is) = - pow(vecPoints[i][0] - 5, 2) + 100.  ; // value of the VB
            toRegressCut(i, is + nbSimul) =  -2 * vecPoints[i][0];
        }
    }

    // conditional expectation
    ArrayXi nbMesh(1);
    nbMesh(0) = 1;
    shared_ptr<LocalConstRegression> localRegressor = make_shared<LocalConstRegression>(false, x, nbMesh);

    // creation continuation value object
    ContinuationCutsGridAdaptNonConcave  contCut(grid1D, localRegressor,  toRegressCut.transpose());

    // first gather all points
    ArrayXXd   hypStock(1, 2);
    hypStock(0, 0) = xMin;
    hypStock(0, 1) = xMax;


    // one simulation
    vector< pair <shared_ptr< StOpt::GridAdaptBase>, shared_ptr<ArrayXXd>  > >  meshsAndCuts = contCut.getCutsConcGatherASim(hypStock, 0, false);

    BOOST_CHECK_EQUAL(meshsAndCuts.size(), 1);
    BOOST_CHECK_EQUAL(meshsAndCuts[0].first->getNbPoints(), grid1D->getNbPoints());


}


/// the solution here is NOT concave anymore
/// calculating conditional meshes while gathering should lead to some grids
BOOST_AUTO_TEST_CASE(testLocalCutsAndGridAgregatingNotConcaveCase)
{
#if defined   __linux
    enable_abort_on_floating_point_exception();
#endif

    /// function to maximize on grid [-5,5]
    //   max  -x + xi
    //   xi <= x

    // grid for storage
    double xMin =  0;
    double xMax = 5;
    int nbMeshX = 5;

    // grid
    shared_ptr< GridAdapt1D > grid1D = make_shared<GridAdapt1D>(xMin, xMax, nbMeshX);

    // refine [0,2] but 2 times
    list< pair<shared_ptr< Mesh1D>, shared_ptr< vector< ArrayXi > > > >  meshes = grid1D->getMeshes();
    for (auto &mesh : meshes)
    {
        if ((mesh.first->getXL() >= 0) && (mesh.first->getXR() <= 2))
            grid1D->splitMesh(mesh);
    }

    meshes = grid1D->getMeshes();
    // in [0,2] again
    for (auto &mesh : meshes)
    {
        if ((mesh.first->getXL() >= 0) && (mesh.first->getXR() <= 2))
            grid1D->splitMesh(mesh);
    }

    meshes = grid1D->getMeshes();
    // at last refine in [1,2]
    for (auto &mesh : meshes)
    {
        if ((mesh.first->getXL() >= 1) && (mesh.first->getXR() <= 2))
            grid1D->splitMesh(mesh);
    }

    meshes = grid1D->getMeshes();
    // at last refine in [3,5]
    for (auto &mesh : meshes)
    {
        if ((mesh.first->getXL() >= 3) && (mesh.first->getXR() <= 5))
            grid1D->splitMesh(mesh);
    }



    // get all points
    vector< ArrayXd > vecPoints = grid1D->getPoints();
    int nbSimul = 2;
    ArrayXXd  toRegressCut(vecPoints.size(), 2 * nbSimul);
    ArrayXXd x = ArrayXXd::Zero(1, nbSimul);
    double x1NCMin = 0.75;
    double x1NCMax = 1.75;
    double x2NCMin = 1.;
    double x2NCMax = 2;
    double x3NCMin = 3.5;
    double x3NCMax = 4.5;

    for (int is  = 0;  is < nbSimul; ++is)
    {
        for (size_t i = 0  ; i < vecPoints.size();  ++i)
        {
            // Concavity rupture
            if (((vecPoints[i][0] >= x1NCMin - 1e-7) && (vecPoints[i][0] <= x1NCMax + 1e-7))
                    || ((vecPoints[i][0] >= x2NCMin - 1e-7) && (vecPoints[i][0] <= x2NCMax + 1e-7))
                    || ((vecPoints[i][0] >= x3NCMin - 1e-7) && (vecPoints[i][0] <= x3NCMax + 1e-7)))
            {
                toRegressCut(i, is) = log(vecPoints[i][0]) ; // value of the VB
                toRegressCut(i, is + nbSimul) = 1. / vecPoints[i][0];
            }
            else
            {
                toRegressCut(i, is) = - 2 * pow(vecPoints[i][0] - xMax, 2) ; // value of the VB
                toRegressCut(i, is + nbSimul) = -4 * (vecPoints[i][0] - xMax);
            }
        }
    }

    // conditional expectation
    ArrayXi nbMesh(1);
    nbMesh(0) = 1;
    shared_ptr<LocalConstRegression> localRegressor = make_shared<LocalConstRegression>(false, x, nbMesh);

    // creation continuation value object
    ContinuationCutsGridAdaptNonConcave  contCut(grid1D, localRegressor,  toRegressCut.transpose());

    // first gather all points
    ArrayXXd   hypStock(1, 2);
    hypStock(0, 0) = xMin;
    hypStock(0, 1) = xMax;


    // one simulation
    vector< pair <shared_ptr< StOpt::GridAdaptBase>, shared_ptr<ArrayXXd>  > >  meshsAndCuts = contCut.getCutsConcGatherASim(hypStock, 0, false);

    cout << " NB GRIDS " << meshsAndCuts.size() << " Initial grid nb points " << grid1D->getNbPoints() <<  endl ;
    vector< ArrayXd >  points;
    points.reserve(grid1D->getNbPoints());
    for (const auto   &gridAndCuts : meshsAndCuts)
    {
        for (const auto &pt : gridAndCuts.first->getPoints())
        {
            bool bAdd = true;
            for (const auto &ptSt : points)
            {
                if (fabs(ptSt[0] - pt[0]) < tiny)
                {
                    bAdd = false;
                    break;
                }
            }
            if (bAdd)
            {
                points.push_back(pt);
            }
        }
    }
    // missing points
    vector< ArrayXd >  pointsG = grid1D->getPoints();
    for (const auto &pt : pointsG)
    {
        bool bPres = false;
        for (const auto &ptSt : points)
        {
            if (fabs(ptSt[0] - pt[0]) < tiny)
            {
                bPres = true;
                break;
            }
        }
        if (!bPres)
            cout << " MISSING POINTS IN GRID x :" << pt[0] << endl;
    }
    BOOST_CHECK_EQUAL(static_cast<int>(points.size()), grid1D->getNbPoints());
    // check that grids points are the same
    vector< ArrayXd >  pointsInit = grid1D->getPoints();
    for (const auto &ptInt : pointsInit)
    {
        bool bPres = false;
        for (const auto &pt : points)
        {
            if (fabs(ptInt[0] - pt[0]) < tiny)
            {
                bPres = true;
                break;
            }
        }
        BOOST_CHECK(bPres);
    }

    ofstream gridsFile("Grids.txt");
    for (const auto   &gridAndCuts : meshsAndCuts)
    {
        shared_ptr<GridAdapt1D> const pGrid1D = dynamic_pointer_cast<GridAdapt1D>(gridAndCuts.first);
        gridsFile <<  pGrid1D->getXMin() << "   " <<  pGrid1D->getXMax() - pGrid1D->getXMin() << endl;
    }
    gridsFile.close();
    ofstream initialMeshFile("GridInit.txt");
    list< pair<shared_ptr< Mesh1D>, shared_ptr< vector< ArrayXi > > > >  meshesAndPos = grid1D->getMeshes();
    for (const auto &mesh : meshesAndPos)
    {
        initialMeshFile << mesh.first->getXL() << "  " <<   mesh.first->getXR() -  mesh.first->getXL() << endl ;
        cout <<  mesh.first->getXL() << "  " <<   mesh.first->getXR() -  mesh.first->getXL();
        for (const auto &it : *mesh.second)
            cout << " ( " << it[0] << " ) ";
        cout << endl ;
    }
    initialMeshFile.close();
    ofstream nonConcaveZoneFile("NonConcaveZone.txt"); //not concave on [1,1.5] and [1.5,3]
    nonConcaveZoneFile << x1NCMin << " " << x1NCMax - x1NCMin << endl ;
    nonConcaveZoneFile << x2NCMin << " " << x2NCMax - x2NCMin << endl ;
    nonConcaveZoneFile << x3NCMin << " " << x3NCMax - x3NCMin << endl ;
    nonConcaveZoneFile.close();

    // get all non concave regions
    vector< shared_ptr< StOpt::GridAdaptBase> >  meshNotConcave = contCut.getMeshNotConcASim(hypStock, 0, false);
    ofstream nonConMeshes("MeshNotConv.txt");
    for (const auto &mesh : meshNotConcave)
    {
        shared_ptr<GridAdapt1D> const pGrid1D = dynamic_pointer_cast<GridAdapt1D>(mesh);
        nonConMeshes << pGrid1D->getXMin() << "  " <<   pGrid1D->getXMax() -  pGrid1D->getXMin() << endl ;
    }


}



// The function to approximate is defined on [0,1]
// It is defined as  f(x)=  g(x) with
// g(x)= log(1+x) for x < 1/3
// g(x)=  -log(4./3.) +2*log(1+x) for x in [1/3,2/3]
// g(x)=  -log(4/3.) -2 log(5/3) +4 log(1+x) for x >= 2/3.
// This function is concave by parts.

// calculate VB value and derivatives
Array2d  analVBAndDeriv(const double &x)
{
    Array2d ret ;

    double fx ;
    double dfx;

    if (x < 1. / 3)
    {
        fx = log(1 + x);
        dfx = 1. / (1. + x);
    }
    else if (x < 2 / 3.)
    {
        fx = -log(4. / 3) + 2 * log(1 + x);
        dfx = 2. / (1. + x);
    }
    else
    {
        fx = - log(4. / 3) - 2 * log(5. / 3.) + 4 * log(1 + x);
        dfx = 4. / (1 + x);
    }
    ret(0) = fx;
    ret(1) = dfx;
    return ret;
}

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

    // grid for storage
    double xMin =  0;
    double xMax = 1;
    int nbMeshX = 10;

    // grid
    shared_ptr< GridAdapt1D > grid1D = make_shared<GridAdapt1D>(xMin, xMax, nbMeshX);

    // refine in [1./6,5.6]
    list< pair<shared_ptr< Mesh1D>, shared_ptr< vector< ArrayXi > > > >  meshes = grid1D->getMeshes();
    for (auto &mesh : meshes)
    {
        if ((mesh.first->getXL() >= 1. / 6.) && (mesh.first->getXR() <= 5. / 6))
            grid1D->splitMesh(mesh);
    }

    meshes = grid1D->getMeshes();
    // in [1/4,5./12] \cup [7./12,3./4]
    for (auto &mesh : meshes)
    {
        if (((mesh.first->getXL() >= 1. / 4.) && (mesh.first->getXR() <= 5. / 12)) || ((mesh.first->getXL() >= 7. / 12.) && (mesh.first->getXR() <= 3. / 4)))
            grid1D->splitMesh(mesh);
    }

    meshes = grid1D->getMeshes();

    // gain refine in [1/4,5./12] \cup [7./12,3./4]
    for (auto &mesh : meshes)
    {
        if (((mesh.first->getXL() >= 1. / 4.) && (mesh.first->getXR() <= 5. / 12)) || ((mesh.first->getXL() >= 7. / 12.) && (mesh.first->getXR() <= 3. / 4)))
            grid1D->splitMesh(mesh);
    }


    // get all points
    vector< ArrayXd > vecPoints = grid1D->getPoints();
    int nbSimul = 2;
    // to store  vb and its derivatives
    ArrayXXd  toRegressCut(vecPoints.size(), 2 * nbSimul);
    // uncertainties set to 0
    ArrayXXd x = ArrayXXd::Zero(1, nbSimul);

    for (size_t i = 0  ; i < vecPoints.size();  ++i)
    {
        Array2d valAndDer = analVBAndDeriv(vecPoints[i][0]);
        for (int is  = 0;  is < nbSimul; ++is)
        {
            toRegressCut(i, is) = valAndDer(0); // value of the VB
            toRegressCut(i, is + nbSimul) = valAndDer(1) ; // x derivative
        }
    }

    // conditional expectation
    ArrayXi nbMesh(1);
    nbMesh(0) = 1;
    shared_ptr<LocalConstRegression> localRegressor = make_shared<LocalConstRegression>(false, x, nbMesh);

    // creation continuation value object
    ContinuationCutsGridAdaptNonConcave  contCut(grid1D, localRegressor,  toRegressCut.transpose());

    // first gather all points
    ArrayXXd   hypStock(1, 2);
    hypStock(0, 0) = xMin;
    hypStock(0, 1) = xMax;


    // one simulation
    vector< pair <shared_ptr< StOpt::GridAdaptBase>, shared_ptr<ArrayXXd>  > >  meshsAndCuts = contCut.getCutsConcGatherASim(hypStock, 0, false);

    cout << " NB GRIDS " << meshsAndCuts.size() << " Initial grid nb points " << grid1D->getNbPoints() <<  endl ;
    vector< ArrayXd >  points;
    points.reserve(grid1D->getNbPoints());
    for (const auto   &gridAndCuts : meshsAndCuts)
    {
        for (const auto &pt : gridAndCuts.first->getPoints())
        {
            bool bAdd = true;
            for (const auto &ptSt : points)
            {
                if (fabs(ptSt[0] - pt[0]) < tiny)
                {
                    bAdd = false;
                    break;
                }
            }
            if (bAdd)
            {
                points.push_back(pt);
            }
        }
    }
    // missing points
    vector< ArrayXd >  pointsG = grid1D->getPoints();
    for (const auto &pt : pointsG)
    {
        bool bPres = false;
        for (const auto &ptSt : points)
        {
            if (fabs(ptSt[0] - pt[0]) < tiny)
            {
                bPres = true;
                break;
            }
        }
        if (!bPres)
            cout << " MISSING POINTS IN GRID x :" << pt[0] << endl;
    }
    BOOST_CHECK_EQUAL(static_cast<int>(points.size()), grid1D->getNbPoints());

    // check that grids points are the same
    vector< ArrayXd >  pointsInit = grid1D->getPoints();
    for (const auto &ptInt : pointsInit)
    {
        bool bPres = false;
        for (const auto &pt : points)
        {
            if (fabs(ptInt[0] - pt[0]) < tiny)
            {
                bPres = true;
                break;
            }
        }
        BOOST_CHECK(bPres);
    }

    ofstream gridsFile("Grids1.txt");
    for (const auto   &gridAndCuts : meshsAndCuts)
    {
        shared_ptr<GridAdapt1D> const pGrid1D = dynamic_pointer_cast<GridAdapt1D>(gridAndCuts.first);
        gridsFile <<  pGrid1D->getXMin() << "   " <<  pGrid1D->getXMax() - pGrid1D->getXMin() << endl;
    }
    gridsFile.close();
    ofstream initialMeshFile("GridInit1.txt");
    list< pair<shared_ptr< Mesh1D>, shared_ptr< vector< ArrayXi > > > >  meshesAndPos = grid1D->getMeshes();
    for (const auto &mesh : meshesAndPos)
    {
        initialMeshFile << mesh.first->getXL() << "  " <<   mesh.first->getXR() -  mesh.first->getXL() << endl ;
        cout <<  mesh.first->getXL() << "  " <<   mesh.first->getXR() -  mesh.first->getXL();
        for (const auto &it : *mesh.second)
            cout << " ( " << it[0] << " ) ";
        cout << endl ;
    }
    initialMeshFile.close();

    // get all non concave regions
    vector< shared_ptr< StOpt::GridAdaptBase> >  meshNotConcave = contCut.getMeshNotConcASim(hypStock, 0, false);
    ofstream nonConMeshes("MeshNotConv1.txt");
    for (const auto &mesh : meshNotConcave)
    {
        shared_ptr<GridAdapt1D> const pGrid1D = dynamic_pointer_cast<GridAdapt1D>(mesh);
        nonConMeshes << pGrid1D->getXMin() << "  " <<   pGrid1D->getXMax() -  pGrid1D->getXMin() << endl;
    }
}
