// Copyright (C) 2016 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#include <memory>
#include "StOpt/regression/SparseRegression.h"
#include "StOpt/core/sparse/sparseGridNoBound.h"
#include "StOpt/core/utils/choleskiFunctionsVariants.h"

using namespace Eigen;
using namespace std;

namespace StOpt
{

SparseRegression::SparseRegression(const int &p_levelMax, const ArrayXd &p_weight, const int &p_degree, bool p_bNoRescale):
    BaseRegression(false), m_spGrid(new SparseSpaceGridNoBound(p_levelMax, p_weight, p_degree)), m_mesh(p_weight.size()),
    m_bNoRescale(p_bNoRescale)
{
}

SparseRegression::SparseRegression(const bool &p_bZeroDate,
                                   const ArrayXXd &p_particles,
                                   const int &p_levelMax, const ArrayXd &p_weight, const int &p_degree):
    BaseRegression(p_bZeroDate, p_particles, false), m_spGrid(new SparseSpaceGridNoBound(p_levelMax, p_weight, p_degree)), m_mesh(p_weight.size()), m_bNoRescale(false)
{
    if (!p_bZeroDate)
        createAndFactorize(p_levelMax, p_weight, p_degree);

}

SparseRegression::SparseRegression(const bool &p_bZeroDate, shared_ptr< SparseSpaceGridNoBound>  p_spGrid,  const vector< vector< double> > &p_mesh, const   ArrayXd &p_meanX,
                                   const   ArrayXd   &p_etypX, const   MatrixXd   &p_svdMatrix):
    BaseRegression(p_bZeroDate, p_meanX,  p_etypX,  p_svdMatrix, false), m_spGrid(p_spGrid), m_mesh(p_mesh), m_bNoRescale(false)
{
    // fill in function basis
    createBasisFunctionNoBound(m_spGrid->getLevelMax(), m_spGrid->getWeight(), m_spGrid->getDegree(), m_functionScal);

}

void SparseRegression::updateSimulations(const bool &p_bZeroDate, const Eigen::ArrayXXd &p_particles)
{
    BaseRegression::updateSimulationsBase(p_bZeroDate, p_particles);
    if (!p_bZeroDate)
        createAndFactorize(m_spGrid->getLevelMax(), m_spGrid->getWeight(), m_spGrid->getDegree());

}

void SparseRegression::createAndFactorize(const int &p_levelMax, const ArrayXd &p_weight, const int &p_degree)
{
    // fill in function basis
    createBasisFunctionNoBound(p_levelMax, p_weight, p_degree, m_functionScal);
    // transform particles to unit square
    m_partRescaled.resize(m_particles.rows(), m_particles.cols());
    if (m_bNoRescale)
        rescale(p_weight);
    else
        transformToUnitSquare(p_levelMax, p_weight);
    // calculate regression matrix
    MatrixXd uReg(m_spGrid->getNbPoints(), m_spGrid->getNbPoints());
    m_yReg.resize(m_spGrid->getNbPoints(), m_spGrid->getNbPoints());
    MatrixXd matReg(m_spGrid->getNbPoints(), m_spGrid->getNbPoints());
    fillInRegressionMatrix(matReg);
    // apriori not a singular matrix
    m_bSingular = false ;
    recurChol(matReg,  uReg, m_yReg, m_bSingular);
    if (m_bSingular)
        pseudoInverseFac(uReg, m_yReg, m_llt);
    else
    {
        // directly factorize the matrix
        m_yReg.resize(0, 0);
        m_llt.compute(matReg);
    }
}

void SparseRegression::rescaleAParticle(const Ref< const ArrayXd>  p_aParticle, Ref< ArrayXd>  p_aPartRescaled)const
{
    for (int id = 0; id < p_aParticle.size(); ++id)
    {
        double aleaInit = p_aParticle(id);
        int ipos = 1    ;
        while ((aleaInit > m_mesh[id][ipos]) && (ipos < (static_cast<int>(m_mesh[id].size()) - 1)))
            ipos += 1 ;
        p_aPartRescaled(id) = min(max((((aleaInit - m_mesh[id][ipos - 1]) / (m_mesh[id][ipos] - m_mesh[id][ipos - 1])) + (ipos - 1)) / (m_mesh[id].size() - 1), 0.), 1.) ;
        assert(p_aPartRescaled(id) >= -1e-9);
        assert(p_aPartRescaled(id) <= 1. + 1e-9);
    }
}

void SparseRegression::transformToUnitSquare(const int &p_levelMax,  const ArrayXd &p_weight)
{
    int nbSimul = m_particles.cols();
    // mesh
    for (int id = 0 ; id < p_weight.size(); ++id)
    {
        // calculate level max for that dimension
        int initLevel = static_cast<int>(p_levelMax / p_weight(id));
        // Below the size of the mesh is choosen equal to the nodal support size
        int nbMesh = pow(2., initLevel - 1); // nb of meshes on last level (dimension 1)
        // resize to store mesh position
        m_mesh[id].resize(nbMesh + 1);
        // now rescale the uncertainties in [0,1] with linear interpolation of the Cumulative distribution
        // such that all meshes have roughtly the same number of particules
        // use the meshing p above the very local meshing
        int nsimulPerStep = nbSimul / nbMesh;
        vector<double> vecToSort(nbSimul);
        for (int is = 0 ; is <  nbSimul ; ++is)
            vecToSort[is] = m_particles(id, is);
        vector<double>::iterator startD = vecToSort.begin();
        int iFirstStep = 0;
        int iLastStep = 1;
        for (int istep = 0 ; istep  < nbMesh ; ++istep)
        {
            nth_element(startD + iFirstStep, startD + iLastStep - 1, vecToSort.end());
            m_mesh[id][istep] = vecToSort[iLastStep];
            iFirstStep  = iLastStep;
            iLastStep   += nsimulPerStep;
        }
        nth_element(startD + iFirstStep, startD + nbSimul - 1, vecToSort.end());
        m_mesh[id][nbMesh] = vecToSort[nbSimul - 1];
    }
    // rescale
    for (int is = 0 ; is < nbSimul  ; ++is)
        rescaleAParticle(m_particles.col(is), m_partRescaled.col(is));

}

void SparseRegression::rescale(const ArrayXd &p_weight)
{
    int nbSimul = m_particles.cols();
    ArrayXd  vecToSort(nbSimul);
    // mesh
    for (int id = 0 ; id < p_weight.size(); ++id)
    {
        for (int is = 0 ; is <  nbSimul ; ++is)
            vecToSort(is) = m_particles(id, is);
        double  vMin = vecToSort.minCoeff();
        double vMax = vecToSort.maxCoeff();
        // resize to store mesh position
        m_mesh[id].resize(2);
        m_mesh[id][0] = vMin;
        m_mesh[id][1] = vMax;
    }
    // rescale
    for (int is = 0 ; is < nbSimul  ; ++is)
        rescaleAParticle(m_particles.col(is), m_partRescaled.col(is));

}

void SparseRegression::recursiveFillFunctionBasisWithSon(ArrayXc &p_levelCurrent,
        ArrayXui   &p_positionCurrent,
        const int &p_ipoint,
        ArrayXd &p_xMiddle,
        ArrayXd &p_dx,
        const ArrayXd &p_x,
        const unsigned short int &p_idimMin,
        ArrayXd &p_funcVal,
        const Array< array<int, 2 >, Dynamic, Dynamic >    &p_son,
        vector<double> &p_nonNullFunctionValues,
        vector<int>    &p_associatedFunctionNumber) const
{
    // add value
    p_nonNullFunctionValues.push_back(p_funcVal.prod());
    p_associatedFunctionNumber.push_back(p_ipoint);
    // iterate on dimension
    for (int idim = 0 ; idim < p_idimMin ;  ++idim)
    {
        double oldFuncVal = p_funcVal(idim);
        double oldDx = p_dx(idim);
        double dxModified = 0.5 * p_dx(idim) ;
        p_dx(idim) = dxModified;
        // child level
        p_levelCurrent(idim) += 1;
        if (p_x(idim) <= p_xMiddle(idim))
        {
            // child level
            int iLeft = p_son(p_ipoint, idim)[0];
            if (iLeft >= 0)
            {
                // go left
                p_xMiddle(idim) -= dxModified;
                // left
                p_positionCurrent(idim) *= 2;
                p_funcVal(idim) = m_functionScal[p_levelCurrent(idim) - 1][p_positionCurrent(idim)](p_x(idim));
                recursiveFillFunctionBasisWithSon(p_levelCurrent, p_positionCurrent, iLeft, p_xMiddle, p_dx, p_x, idim + 1, p_funcVal, p_son, p_nonNullFunctionValues, p_associatedFunctionNumber);
                p_positionCurrent(idim) /= 2;
                p_xMiddle(idim) += dxModified;
            }
        }
        else
        {
            // child level
            int iRight = p_son(p_ipoint, idim)[1];
            if (iRight >= 0)
            {
                // go right
                p_xMiddle(idim) += dxModified;
                p_positionCurrent(idim) = 2 * p_positionCurrent(idim) + 1;
                p_funcVal(idim) = m_functionScal[p_levelCurrent(idim) - 1][p_positionCurrent(idim)](p_x(idim));
                recursiveFillFunctionBasisWithSon(p_levelCurrent, p_positionCurrent, iRight, p_xMiddle, p_dx, p_x, idim + 1, p_funcVal, p_son, p_nonNullFunctionValues, p_associatedFunctionNumber);
                p_positionCurrent(idim) = (p_positionCurrent(idim) - 1) / 2;
                p_xMiddle(idim) -= dxModified;
            }
        }
        p_funcVal(idim) = oldFuncVal;
        p_dx(idim) = oldDx;
        p_levelCurrent(idim) -= 1;
    }
}


void SparseRegression::assessFuncBasis(const ArrayXd &p_particle, vector<double> &p_nonNullFunctionValues, vector<int>    &p_associatedFunctionNumber) const
{
    // to store function value
    ArrayXd  funcVal = ArrayXd::Constant(m_spGrid->getDimension(), 1);
    // level
    ArrayXc levelCurrent = ArrayXc::Constant(m_spGrid->getDimension(), 1);
    // position
    ArrayXui  positionCurrent = ArrayXui::Zero(m_spGrid->getDimension());
    // utilitary for terations
    ArrayXd dx = ArrayXd::Constant(m_spGrid->getDimension(), 0.5);
    ArrayXd  xMiddle = ArrayXd::Constant(m_spGrid->getDimension(), 0.5);
    p_nonNullFunctionValues.reserve(m_spGrid->getDataSetDepth());
    p_associatedFunctionNumber.reserve(m_spGrid->getDataSetDepth());
    // fill in non null function basis values
    recursiveFillFunctionBasisWithSon(levelCurrent,  positionCurrent, 0, xMiddle,  dx, p_particle, m_spGrid->getDimension(), funcVal, *m_spGrid->getSon(), p_nonNullFunctionValues,  p_associatedFunctionNumber);
}



void SparseRegression::fillInRegressionMatrix(MatrixXd &p_matReg)
{
    // parallel Matrix
#ifdef _OPENMP
    int nbThread = omp_get_max_threads();
#else
    int nbThread = 1 ;
#endif
    vector< MatrixXd> matrixThread(nbThread);
    for (int it = 0 ; it < nbThread; ++it)
        matrixThread[it].resize(p_matReg.rows(), p_matReg.cols());
    int it ;
#ifdef _OPENMP
    #pragma omp parallel for  private(it)
#endif
    for (it = 0 ; it < nbThread; ++it)
        matrixThread[it].setConstant(0.);

    int is ;
#ifdef _OPENMP
    #pragma omp parallel for  private(is)
#endif
    for (is = 0 ; is < m_particles.cols(); ++is)
    {
#ifdef _OPENMP
        int ithread = omp_get_thread_num();
#else
        int ithread = 0;
#endif
        // to store result
        vector<double> nonNullFunctionValues;
        vector<int>   associatedFunctionNumber;
        assessFuncBasis(m_partRescaled.col(is), nonNullFunctionValues, associatedFunctionNumber);

        // calculate  the matrix
        for (size_t ipx = 0; ipx < associatedFunctionNumber.size(); ++ipx)
        {
            int iposx =  associatedFunctionNumber[ipx];
            for (size_t ipy =  ipx ;  ipy <  associatedFunctionNumber.size(); ++ipy)
            {
                int iposy =  associatedFunctionNumber[ipy];
                matrixThread[ithread](iposx, iposy) += nonNullFunctionValues[ipx] * nonNullFunctionValues[ipy];
            }
        }
    }
    p_matReg = matrixThread[0];
    for (it = 1 ;  it < nbThread; ++it)
        p_matReg += matrixThread[it];
    //   only a part has been calculated, just symetrize
    for (int ipx = 0; ipx < p_matReg.rows(); ++ipx)
        for (int ipy = ipx + 1; ipy < p_matReg.rows(); ++ipy)
            p_matReg(ipx, ipy) += p_matReg(ipy, ipx);
    // just copy the low diagonal
    for (int ipx = 0; ipx < p_matReg.rows(); ++ipx)
        for (int ipy = 0; ipy < ipx; ++ipy)
            p_matReg(ipx, ipy) = p_matReg(ipy, ipx);

}

MatrixXd SparseRegression::fillInSecondMember(const MatrixXd &p_fToRegress) const
{
    MatrixXd secondMember(p_fToRegress.rows(), m_spGrid->getNbPoints());
    // parallel Matrix
#ifdef _OPENMP
    int nbThread = omp_get_max_threads();
#else
    int nbThread = 1 ;
#endif
    vector< MatrixXd> secMember(nbThread);
    for (int it = 0 ; it < nbThread; ++it)
        secMember[it].resize(secondMember.rows(), secondMember.cols());
    int it ;
#ifdef _OPENMP
    #pragma omp parallel for  private(it)
#endif
    for (it = 0 ; it < nbThread; ++it)
        secMember[it].setConstant(0.);

    int is ;
#ifdef _OPENMP
    #pragma omp parallel for  private(is)
#endif
    for (is = 0 ; is < m_particles.cols(); ++is)
    {
#ifdef _OPENMP
        int ithread = omp_get_thread_num();
#else
        int ithread = 0;
#endif
        // to store result
        vector<double> nonNullFunctionValues;
        vector<int>   associatedFunctionNumber;
        assessFuncBasis(m_partRescaled.col(is), nonNullFunctionValues, associatedFunctionNumber);

        for (size_t ip = 0 ; ip <  associatedFunctionNumber.size(); ++ip)
        {
            secMember[ithread].col(associatedFunctionNumber[ip]) += p_fToRegress.col(is) * nonNullFunctionValues[ip];
        }
    }
    secondMember =  secMember[0];
    for (it = 1 ;  it < nbThread; ++it)
        secondMember += secMember[it];
    return secondMember;
}


void SparseRegression::reconstruct(const ArrayXXd &p_basisCoeff, ArrayXXd   &p_fRegressed) const
{
    // reconstruct
    p_fRegressed.setConstant(0.);
    int is ;
#ifdef _OPENMP
    #pragma omp parallel for  private(is)
#endif
    for (is = 0 ; is < m_particles.cols(); ++is)
    {
        // to store result
        vector<double> nonNullFunctionValues;
        vector<int>   associatedFunctionNumber;
        assessFuncBasis(m_partRescaled.col(is), nonNullFunctionValues, associatedFunctionNumber);

        for (size_t ip = 0 ; ip <  associatedFunctionNumber.size(); ++ip)
            p_fRegressed.col(is) += p_basisCoeff.col(associatedFunctionNumber[ip]) * nonNullFunctionValues[ip];
    }
}

double  SparseRegression::reconstructOneParticle(const ArrayXd &p_basisCoeff, const ArrayXd   &p_aParticle) const
{
    // reconstruct return
    double ret  = 0;
    // to store result
    vector<double> nonNullFunctionValues;
    vector<int>   associatedFunctionNumber;
    assessFuncBasis(p_aParticle, nonNullFunctionValues, associatedFunctionNumber);

    for (size_t ip = 0 ; ip <  associatedFunctionNumber.size(); ++ip)
        ret += p_basisCoeff(associatedFunctionNumber[ip]) * nonNullFunctionValues[ip];
    return ret;
}

ArrayXd SparseRegression::getCoordBasisFunction(const ArrayXd &p_fToRegress) const
{
    if ((!m_bZeroDate) && (m_spGrid->getLevelMax() != 0))
    {
        Map<const MatrixXd>  fToRegress2D(p_fToRegress.data(), 1, p_fToRegress.size());
        MatrixXd secMember2D = fillInSecondMember(fToRegress2D);
        Map<const VectorXd > secMember(secMember2D.data(), secMember2D.size());
        if (m_bSingular)
        {
            // find pseudo inverse  yReg^T  M^{-2}  yReg secMember
            Eigen::VectorXd xInter = m_llt.solve(m_yReg * secMember)  ;
            return (m_yReg.transpose() * m_llt.solve(xInter)).array();

        }
        else
        {
            // Direct Choleski
            return m_llt.solve(secMember).array();
        }
    }
    else
    {
        ArrayXd retAverage(1);
        retAverage(0) = p_fToRegress.mean();
        return retAverage;
    }
}

ArrayXXd SparseRegression::getCoordBasisFunctionMultiple(const ArrayXXd &p_fToRegress) const
{
    if ((!m_bZeroDate) && (m_spGrid->getLevelMax() != 0))
    {
        MatrixXd secMember2D = fillInSecondMember(p_fToRegress).transpose();
        MatrixXd ret(secMember2D.rows(), secMember2D.cols()) ;
        if (m_bSingular)
        {
            for (int j = 0; j < secMember2D.cols(); ++j)
            {
                // find pseudo inverse  yReg^T  M^{-2}  yReg secMember
                Eigen::VectorXd xInter = m_llt.solve(m_yReg * secMember2D.col(j));
                ret.col(j) = m_yReg.transpose() * m_llt.solve(xInter);
            }
        }
        else
        {
            // direct inversion with Choleski
            for (int j = 0; j < secMember2D.cols(); ++j)
                ret.col(j) = m_llt.solve(secMember2D.col(j)) ;
        }
        return ret.transpose().array();
    }
    else
    {
        ArrayXXd retAverage(p_fToRegress.rows(), 1);
        for (int nsm = 0; nsm <  p_fToRegress.rows(); ++nsm)
            retAverage.row(nsm).setConstant(p_fToRegress.row(nsm).mean());
        return retAverage;
    }
}

double SparseRegression::reconstructionASim(const int &p_isim, const Eigen::ArrayXd   &p_basisCoefficients) const
{
    if ((m_bZeroDate) || (m_spGrid->getLevelMax() == 0))
        return p_basisCoefficients(0);
// to store result
    vector<double> nonNullFunctionValues;
    vector<int>   associatedFunctionNumber;
    assessFuncBasis(m_partRescaled.col(p_isim), nonNullFunctionValues, associatedFunctionNumber);
    double ret  = 0. ;
    for (size_t ip = 0 ; ip <  associatedFunctionNumber.size(); ++ip)
        ret += p_basisCoefficients(associatedFunctionNumber[ip]) * nonNullFunctionValues[ip];
    return ret;
}

ArrayXd SparseRegression::getAllSimulations(const ArrayXd &p_fToRegress) const
{
    if ((m_bZeroDate) || (m_spGrid->getLevelMax() == 0))
        return ArrayXd::Constant(p_fToRegress.size(), p_fToRegress.mean());
    ArrayXd basisCoefficients = getCoordBasisFunction(p_fToRegress);
    Map<const ArrayXXd> basisCoefficients2D(basisCoefficients.data(), 1, basisCoefficients.size());
    ArrayXXd regressed2D(1, p_fToRegress.size());
    reconstruct(basisCoefficients2D, regressed2D);
    return regressed2D.transpose().col(0);
}
ArrayXXd SparseRegression::getAllSimulationsMultiple(const ArrayXXd &p_fToRegress) const
{
    if ((m_bZeroDate) || (m_spGrid->getLevelMax() == 0))
    {
        ArrayXXd ret(p_fToRegress.rows(), p_fToRegress.cols());
        for (int ism = 0; ism < p_fToRegress.rows(); ++ism)
            ret.row(ism).setConstant(p_fToRegress.row(ism).mean());
        return ret;
    }
    ArrayXXd BasisCoefficients = getCoordBasisFunctionMultiple(p_fToRegress);
    ArrayXXd regressed(p_fToRegress.rows(), p_fToRegress.cols());
    reconstruct(BasisCoefficients, regressed);
    return regressed;
}

ArrayXd SparseRegression::reconstruction(const ArrayXd   &p_basisCoefficients) const
{
    if ((!m_bZeroDate) && (m_spGrid->getLevelMax() != 0))
    {
        ArrayXXd regressed2D(1, m_particles.cols());
        Map<const ArrayXXd>  BasisCoefficients(p_basisCoefficients.data(), 1, p_basisCoefficients.size());
        reconstruct(BasisCoefficients, regressed2D);
        return regressed2D.transpose().col(0);
    }
    else
    {
        ArrayXXd retValue(p_basisCoefficients.rows(), m_particles.cols());
        for (int nsm = 0; nsm < p_basisCoefficients.rows(); ++nsm)
            retValue.row(nsm).setConstant(p_basisCoefficients(nsm, 0));
        return retValue ;

    }
}
ArrayXXd SparseRegression::reconstructionMultiple(const ArrayXXd   &p_basisCoefficients) const
{
    if ((!m_bZeroDate) && (m_spGrid->getLevelMax() != 0))
    {
        ArrayXXd regressed(p_basisCoefficients.rows(), m_particles.cols());
        reconstruct(p_basisCoefficients, regressed);
        return  regressed ;
    }
    else
    {
        ArrayXXd retValue(p_basisCoefficients.rows(),  m_particles.cols());
        for (int nsm = 0; nsm < p_basisCoefficients.rows(); ++nsm)
            retValue.row(nsm).setConstant(p_basisCoefficients(nsm, 0));
        return retValue ;
    }
}

double SparseRegression::getValue(const  ArrayXd   &p_coordinates,	const ArrayXd   &p_coordBasisFunction) const
{
    if ((!m_bZeroDate) && (m_spGrid->getLevelMax() != 0))
    {
        // rotation
        VectorXd x = m_svdMatrix * ((p_coordinates - m_meanX) / m_etypX).matrix();
        // first rescale the particle
        ArrayXd aPartRescaled(p_coordinates.size());
        rescaleAParticle(x.array(), aPartRescaled);
        return reconstructOneParticle(p_coordBasisFunction, aPartRescaled);
    }
    else
    {
        return p_coordBasisFunction(0);
    }
}

double SparseRegression::getAValue(const ArrayXd &p_coordinates,  const ArrayXd &p_ptOfStock,
                                   const vector< shared_ptr<InterpolatorSpectral> > &p_interpFuncBasis) const
{
    Eigen::ArrayXd coordBasisFunction(p_interpFuncBasis.size());
    for (int i = 0; i < coordBasisFunction.size(); ++i)
    {
        coordBasisFunction(i) = p_interpFuncBasis[i]->apply(p_ptOfStock);
    }
    return  getValue(p_coordinates, coordBasisFunction);
}
}