// Copyright (C) 2018 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#include <Eigen/SVD>
#include <Eigen/Cholesky>
#include "StOpt/regression/LaplacianLinearKernelRegression.h"
#include "StOpt/regression/nDDominanceKernel.h"

using namespace std ;
using namespace Eigen ;

namespace StOpt
{

LaplacianLinearKernelRegression::LaplacianLinearKernelRegression(const bool &p_bZeroDate,
        const ArrayXXd  &p_particles,
        const ArrayXd   &p_h):
    BaseRegression(p_bZeroDate, p_particles, false), m_h(p_h), m_tree(p_particles)
{
}

LaplacianLinearKernelRegression::LaplacianLinearKernelRegression(const ArrayXd   &p_h):
    BaseRegression(false), m_h(p_h)
{
}

LaplacianLinearKernelRegression::LaplacianLinearKernelRegression(const bool &p_bZeroDate,
        const ArrayXXd  &p_particles):
    BaseRegression(p_bZeroDate, p_particles, false), m_tree(p_particles)
{
}


void LaplacianLinearKernelRegression::updateSimulations(const bool &p_bZeroDate, const ArrayXXd &p_particles)
{
    BaseRegression::updateSimulationsBase(p_bZeroDate, p_particles);
    m_tree = KDTree(p_particles);

}

ArrayXXd LaplacianLinearKernelRegression::regressFunction(const ArrayXXd &p_fToRegress) const
{
    // dimension
    int nD = m_particles.rows();
    // number of functions to calculate for regressions
    int nbFuncReg = (nD + 1) * (nD + 2) / 2;
    int nbFuncSecMem = (nD + 1) * p_fToRegress.rows()   ;
    // creation of the 2^d terms
    int nbSum = pow(2, nD);
    vector< shared_ptr<ArrayXXd> > vecToAdd(nbSum);
    // calculate exp values
    Eigen::ArrayXi iCoord(nD) ;
    for (int i = 0; i < nbSum; ++i)
    {
        int ires = i;
        for (int id = nD - 1 ; id >= 0  ; --id)
        {
            unsigned int idec = (ires >> id) ;
            iCoord(id) = -(2 * idec - 1);
            ires -= (idec << id);
        }
        vecToAdd[i] = make_shared<ArrayXXd>(nbFuncReg + nbFuncSecMem, m_particles.cols());
        for (int is = 0; is < m_particles.cols(); ++is)
        {
            double ssum = 0;
            for (int id = 0; id < nD; ++id)
                ssum += iCoord(id) * m_particles(id, is) / m_h(id);
            double expSum =  exp(ssum);
            int iloc = 0;
            // lower triangular matrix
            (*vecToAdd[i])(iloc++, is) = expSum;
            for (int id = 0; id < nD; ++id)
            {
                (*vecToAdd[i])(iloc++, is) = expSum * m_particles(id, is);
                for (int idd = 0; idd <= id; ++idd)
                    (*vecToAdd[i])(iloc++, is) = expSum * m_particles(id, is) * m_particles(idd, is)  ;
            }
            for (int ifunc = 0; ifunc < p_fToRegress.rows(); ++ifunc)
            {
                (*vecToAdd[i])(iloc++, is) =  expSum * p_fToRegress(ifunc, is);
                for (int id = 0; id < nD ; ++id)
                    (*vecToAdd[i])(iloc++, is) =  expSum * p_fToRegress(ifunc, is) * m_particles(id, is) ;
            }
        }
    }

    vector< shared_ptr<ArrayXXd> > fDomin(nbSum);


    // kernel resolution
    nDDominanceKernel(m_particles, vecToAdd, fDomin);


    // reconstruction of the 2^d terms for each matrix term
    ArrayXd forMatrix(nbFuncReg);
    ArrayXd secMem(nbFuncSecMem);
    MatrixXd  matA(1 + nD, 1 + nD);
    VectorXd  vecB(1 + nD);
    ArrayXXd ret(p_fToRegress.rows(), p_fToRegress.cols());
    for (int is = 0; is < m_particles.cols(); ++is)
    {
        // calculate matrix Coeff
        int iPosLoc = 0;
        forMatrix(iPosLoc) = 1.;
        for (int iSum = 0; iSum < nbSum; ++iSum)
        {
            forMatrix(iPosLoc) += (*fDomin[iSum])(0, is) * (*vecToAdd[nbSum - 1 - iSum])(0, is);
        }
        iPosLoc += 1;
        for (int id = 0; id < nD; ++id)
        {
            forMatrix(iPosLoc) = m_particles(id, is);
            for (int iSum = 0; iSum < nbSum; ++iSum)
                forMatrix(iPosLoc) += (*fDomin[iSum])(iPosLoc, is) * (*vecToAdd[nbSum - 1 - iSum])(0, is);
            iPosLoc += 1;
            for (int idd = 0; idd <= id; ++idd)
            {
                forMatrix(iPosLoc) = m_particles(id, is) * m_particles(idd, is)  ;
                for (int iSum = 0; iSum < nbSum; ++iSum)
                    forMatrix(iPosLoc) += (*fDomin[iSum])(iPosLoc, is) * (*vecToAdd[nbSum - 1 - iSum])(0, is);
                iPosLoc += 1;
            }
        }
        // create second member
        int iSecMem = 0 ;
        for (int ifunc = 0 ; ifunc < p_fToRegress.rows(); ++ifunc)
        {
            secMem(iSecMem)  = p_fToRegress(ifunc, is);
            for (int iSum = 0; iSum < nbSum; ++iSum)
            {
                secMem(iSecMem) += (*fDomin[iSum])(nbFuncReg + iSecMem, is) * (*vecToAdd[nbSum - 1 - iSum])(0, is);
            }
            iSecMem += 1;
            for (int id = 0; id < nD; ++id)
            {
                secMem(iSecMem)  = p_fToRegress(ifunc, is) * m_particles(id, is) ;
                for (int iSum = 0; iSum < nbSum; ++iSum)
                {
                    secMem(iSecMem) += (*fDomin[iSum])(nbFuncReg + iSecMem, is) * (*vecToAdd[nbSum - 1 - iSum])(0, is);
                }
                iSecMem += 1;
            }
        }
        // create regression matrix
        int iloc = 0;
        for (int id = 0; id <= nD; ++id)
            for (int idd = 0; idd <= id; ++idd)
                matA(id, idd) = forMatrix(iloc++);
        for (int id = 0; id <= nD; ++id)
            for (int idd = id + 1; idd <= nD; ++idd)
                matA(id, idd) =  matA(idd, id);
        // second member and inverse
        int iloc1 = 0;
        // inverse
        LLT<MatrixXd>  lltA(matA);
        for (int ifunc = 0 ; ifunc < p_fToRegress.rows(); ++ifunc)
        {
            for (int id = 0; id <= nD; ++id)
                vecB(id) = secMem(iloc1++);
            VectorXd coeff = lltA.solve(vecB);
            ret(ifunc, is) = coeff(0);
            for (int id  = 0; id < nD; ++id)
                ret(ifunc, is)  += coeff(id + 1) * m_particles(id, is);
        }
    }
    return ret;
}



ArrayXd  LaplacianLinearKernelRegression::getCoordBasisFunction(const ArrayXd &p_fToRegress) const
{
    if (!BaseRegression::m_bZeroDate)
    {
        const ArrayXXd  fToRegress = Map<const ArrayXXd>(p_fToRegress.data(), 1, p_fToRegress.size()) ;
        ArrayXXd  regressed = regressFunction(fToRegress);
        ArrayXd toReturn = Map<ArrayXd>(regressed.data(), regressed.size());
        return toReturn;
    }
    else
    {
        ArrayXd retAverage(1);
        retAverage(0) = p_fToRegress.mean();
        return retAverage;
    }
}

ArrayXXd  LaplacianLinearKernelRegression::getCoordBasisFunctionMultiple(const ArrayXXd &p_fToRegress) const
{
    if (!BaseRegression::m_bZeroDate)
    {
        return regressFunction(p_fToRegress);
    }
    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;
    }
}

ArrayXd  LaplacianLinearKernelRegression::getAllSimulations(const ArrayXd &p_fToRegress) const
{
    if (!BaseRegression::m_bZeroDate)
    {
        const ArrayXXd  fToRegress = Map<const ArrayXXd>(p_fToRegress.data(), 1, p_fToRegress.size()) ;
        ArrayXXd  regressed = regressFunction(fToRegress);
        ArrayXd toReturn = Map<ArrayXd>(regressed.data(), regressed.size());
        return toReturn;
    }
    else
    {
        return ArrayXd::Constant(p_fToRegress.size(), p_fToRegress.mean());
    }
}

ArrayXXd  LaplacianLinearKernelRegression::getAllSimulationsMultiple(const ArrayXXd &p_fToRegress) const
{
    if (!BaseRegression::m_bZeroDate)
    {
        return regressFunction(p_fToRegress);
    }
    else
    {
        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;
    }
}


ArrayXd LaplacianLinearKernelRegression::reconstruction(const ArrayXd   &p_basisCoefficients) const
{
    if (!BaseRegression::m_bZeroDate)
        return p_basisCoefficients; // basis function are here regressed values !
    else
    {
        return ArrayXd::Constant(m_particles.cols(), p_basisCoefficients(0));
    }
}


ArrayXXd LaplacianLinearKernelRegression::reconstructionMultiple(const ArrayXXd   &p_basisCoefficients) const
{
    if (!BaseRegression::m_bZeroDate)
    {
        return p_basisCoefficients; // basis function are here regressed values !
    }
    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  LaplacianLinearKernelRegression::reconstructionASim(const int &p_isim, const ArrayXd   &p_basisCoefficients) const
{
    double ret ;
    if (!BaseRegression::m_bZeroDate)
    {
        ret = p_basisCoefficients(p_isim);
    }
    else
    {
        ret = p_basisCoefficients(0);
    }
    return ret ;
}

double LaplacianLinearKernelRegression::getValue(const ArrayXd   &p_coordinates,
        const ArrayXd   &p_coordBasisFunction)  const
{
    double ret  ;
    if (!BaseRegression::m_bZeroDate)
    {
        // Use KDTree to find nearest point close to a given one and regress the associated
        // regressed value
        return p_coordBasisFunction(m_tree.nearestIndex(p_coordinates));
    }
    else
        ret =  p_coordBasisFunction(0);
    return ret ;
}


double LaplacianLinearKernelRegression::getAValue(const ArrayXd &p_coordinates,  const ArrayXd &p_ptOfStock,
        const vector< shared_ptr<InterpolatorSpectral> > &p_interpFuncBasis) const
{
    if (!BaseRegression::m_bZeroDate)
    {
        return p_interpFuncBasis[m_tree.nearestIndex(p_coordinates)]->apply(p_ptOfStock);
    }
    else
    {
        return p_interpFuncBasis[0]->apply(p_ptOfStock);
    }
}
}
