// Copyright (C) 2020 EDF
// Copyright (C) 2020 CSIRO Data61
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#include <Eigen/SVD>
#include <Eigen/Cholesky>
#include "StOpt/core/utils/constant.h"
#include "StOpt/core/grids/RegularSpaceGrid.h"
#include "StOpt/core/grids/LinearInterpolator.h"
#include "StOpt/regression/LaplacianGridKernelRegression.h"
#include "StOpt/regression/fastLaplacianKDE.h"

using namespace std ;
using namespace Eigen ;

namespace StOpt
{

LaplacianGridKernelRegression::LaplacianGridKernelRegression(const bool &p_bZeroDate,
        const ArrayXXd  &p_particles,
        const ArrayXd   &p_h,
        const double   &p_coefNbGridPoint,
        const bool &p_bLinear
                                                            ):
    BaseRegression(p_bZeroDate, p_particles, false), m_h(p_h), m_coefNbGridPoint(p_coefNbGridPoint), m_z(p_h.size()), m_bLinear(p_bLinear)
{
    rectiConst();
}

LaplacianGridKernelRegression::LaplacianGridKernelRegression(const bool &p_bZeroDate,
        const ArrayXd   &p_h,
        const double   &p_coefNbGridPoint,
        const bool &p_bLinear,
        const std::vector< std::shared_ptr<Eigen::ArrayXd> >   &p_z):
    BaseRegression(p_bZeroDate, false), m_h(p_h), m_coefNbGridPoint(p_coefNbGridPoint), m_z(p_z), m_bLinear(p_bLinear)
{
}

LaplacianGridKernelRegression::LaplacianGridKernelRegression(const ArrayXd   &p_h,
        const double   &p_coefNbGridPoint,
        const bool &p_bLinear):
    BaseRegression(false), m_h(p_h), m_coefNbGridPoint(p_coefNbGridPoint), m_z(p_h.size()), m_bLinear(p_bLinear)
{
}


void LaplacianGridKernelRegression::updateSimulations(const bool &p_bZeroDate, const ArrayXXd &p_particles)
{
    BaseRegression::updateSimulationsBase(p_bZeroDate, p_particles);

    // calculate the grid associated
    rectiConst();
}


void   LaplacianGridKernelRegression::rectiConst()
{
    // number of points per direction
    int nbPt = static_cast<int>(pow(m_particles.cols() * m_coefNbGridPoint, 1. / m_particles.rows()));
    for (int id = 0; id < m_particles.rows() ; ++id)
    {
        ArrayXd zDim = ArrayXd::LinSpaced(nbPt, m_particles.row(id).minCoeff() + tiny, m_particles.row(id).maxCoeff() - tiny);
        m_z[id] = make_shared<ArrayXd>(zDim);
    }
    m_nbPtZ  = static_cast<int>(pow(nbPt, m_particles.rows()));
}

ArrayXXd  LaplacianGridKernelRegression::kernelCoeffCalcConst(const ArrayXXd &p_fToRegress) const
{
    // number of functions to calculate for regressions
    int nbFuncReg = 1;
    int nbFuncSecMem =  p_fToRegress.rows()   ;

    // store all weights for kernel
    ArrayXXd y(nbFuncReg + nbFuncSecMem, p_fToRegress.cols());
    for (int is = 0; is <  p_fToRegress.cols(); ++is)
    {
        int iloc = 0;
        // lower triangular matrix
        y(iloc++, is) = 1.;
        for (int ifunc = 0; ifunc < p_fToRegress.rows(); ++ifunc)
        {
            y(iloc++, is) = p_fToRegress(ifunc, is);
        }
    }

    // kernel coefficients on grid
    ArrayXXd coeffOnZ = FastSumRegressionTerms(m_particles,  m_z,  m_h,  y);

    // now calculate the regressed values on the grid
    ArrayXXd regOnZ(p_fToRegress.rows(), m_nbPtZ);

    for (int ipt = 0 ; ipt < m_nbPtZ; ++ipt)
    {
        for (int ifunc = 0 ; ifunc < p_fToRegress.rows(); ++ifunc)
        {
            regOnZ(ifunc, ipt) = coeffOnZ(ifunc + 1, ipt) / coeffOnZ(0, ipt);
        }
    }
    return regOnZ;
}

ArrayXXd  LaplacianGridKernelRegression::kernelCoeffCalcLinear(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()   ;

    // store all weights for kernel
    ArrayXXd y(nbFuncReg + nbFuncSecMem, p_fToRegress.cols());
    for (int is = 0; is <  p_fToRegress.cols(); ++is)
    {
        int iloc = 0;
        // lower triangular matrix
        y(iloc++, is) = 1.;
        for (int id = 0; id < nD; ++id)
        {
            y(iloc++, is) = m_particles(id, is);
            for (int idd = 0; idd <= id; ++idd)
                y(iloc++, is) = m_particles(id, is) * m_particles(idd, is);
        }
        for (int ifunc = 0; ifunc < p_fToRegress.rows(); ++ifunc)
        {
            y(iloc++, is) = p_fToRegress(ifunc, is);
            for (int id = 0; id < nD ; ++id)
                y(iloc++, is) = p_fToRegress(ifunc, is) * m_particles(id, is);
        }
    }

    // kernel coefficients on grid
    ArrayXXd coeffOnZ = FastSumRegressionTerms(m_particles,  m_z,  m_h,  y);

    // now calculate the regressed values on the grid
    ArrayXXd regOnZ(p_fToRegress.rows(), m_nbPtZ);

    // for regressions
    MatrixXd  matA(1 + nD, 1 + nD);
    VectorXd  vecB(1 + nD);

    ArrayXi coord =  ArrayXi::Zero(nD);

    for (int ipt = 0 ; ipt < m_nbPtZ; ++ipt)
    {
        // create regression matrix
        int iloc = 0;
        for (int id = 0; id <= nD; ++id)
            for (int idd = 0; idd <= id; ++idd)
                matA(id, idd) = coeffOnZ(iloc++, ipt);
        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
        // inverse
        LLT<MatrixXd>  lltA(matA);
        for (int ifunc = 0 ; ifunc < p_fToRegress.rows() ; ++ifunc)
        {
            for (int id = 0; id <= nD; ++id)
                vecB(id) = coeffOnZ(iloc++, ipt);
            VectorXd coeff = lltA.solve(vecB);
            regOnZ(ifunc, ipt) = coeff(0);
            for (int id  = 0; id < nD; ++id)
                regOnZ(ifunc, ipt)  += coeff(id + 1) * (*m_z[id])(coord(id));
        }
        // update coordinates
        for (int id = 0; id < nD; ++id)
        {
            if (coord(id) < m_z[id]->size() - 1)
            {
                coord(id) += 1;
                break;
            }
            else
            {
                coord(id) = 0;
            }
        }
    }
    return regOnZ;
}


ArrayXXd  LaplacianGridKernelRegression::interpolateAtSample(const ArrayXXd &p_regOnGrid) const
{

    int nD = m_z.size() ;
    // create a grid
    Eigen::ArrayXd lowValues(nD);
    for (int i = 0; i < nD; ++i)
        lowValues(i) = (*m_z[i])[0];
    Eigen::ArrayXd step(nD);
    for (int i = 0; i < nD; ++i)
        step(i) = ((*m_z[i])[m_z[i]->size() - 1] - (*m_z[i])[0]) / (m_z[i]->size() - 1);
    Eigen::ArrayXi  nbStep(nD);
    for (int i = 0; i < nD; ++i)
        nbStep(i) = m_z[i]->size() - 1 ;

    // regular
    RegularSpaceGrid regGrid(lowValues, step, nbStep);

    // for return
    ArrayXXd regressRet(p_regOnGrid.rows(), m_particles.cols());

    // transpose once for all
    ArrayXXd regTrans = p_regOnGrid.transpose();

    // now interpolate
    for (int isim = 0; isim < m_particles.cols(); ++isim)
    {
        Eigen::ArrayXd point = m_particles.col(isim);

        // create the interpolator
        LinearInterpolator  regLin(&regGrid, point);
        for (int j = 0; j < p_regOnGrid.rows(); ++j)
        {
            regressRet(j, isim) = regLin.apply(regTrans.col(j));
        }
    }
    return regressRet;
}


double   LaplacianGridKernelRegression::interpolateAtPoint(const ArrayXd   &p_coordinates,  const ArrayXd &p_regOnGrid) const
{

    int nD = m_z.size();
    // create a grid
    Eigen::ArrayXd lowValues(nD);
    for (int i = 0; i < nD; ++i)
        lowValues(i) = (*m_z[i])[0];
    Eigen::ArrayXd step(nD);
    for (int i = 0; i < nD; ++i)
        step(i) = ((*m_z[i])[m_z[i]->size() - 1] - (*m_z[i])[0]) / (m_z[i]->size() - 1);
    Eigen::ArrayXi  nbStep(nD);
    for (int i = 0; i < nD; ++i)
        nbStep(i) = m_z[i]->size() - 1 ;

    // regular
    RegularSpaceGrid regGrid(lowValues, step, nbStep);

    // create the interpolator
    LinearInterpolator  regLin(&regGrid, p_coordinates);

    return regLin.apply(p_regOnGrid);
}


ArrayXXd LaplacianGridKernelRegression::regressFunction(const ArrayXXd &p_fToRegress) const
{

    // calculate the coefficients on the grid
    ArrayXXd  coefOnGrid = ((m_bLinear == true) ? kernelCoeffCalcLinear(p_fToRegress) :  kernelCoeffCalcConst(p_fToRegress));
    // interpolate at sample points
    return interpolateAtSample(coefOnGrid);
}



ArrayXd  LaplacianGridKernelRegression::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 = ((m_bLinear) ? kernelCoeffCalcLinear(fToRegress) : kernelCoeffCalcConst(fToRegress));
        ArrayXd toReturn = Map<ArrayXd>(regressed.data(), regressed.size());
        return toReturn;
    }
    else
    {
        ArrayXd retAverage(1);
        retAverage(0) = p_fToRegress.mean();
        return retAverage;
    }
}

ArrayXXd  LaplacianGridKernelRegression::getCoordBasisFunctionMultiple(const ArrayXXd &p_fToRegress) const
{
    if (!BaseRegression::m_bZeroDate)
    {
        if (m_bLinear)
        {
            return kernelCoeffCalcLinear(p_fToRegress);
        }
        else
        {
            return kernelCoeffCalcConst(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  LaplacianGridKernelRegression::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  LaplacianGridKernelRegression::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 LaplacianGridKernelRegression::reconstruction(const ArrayXd   &p_basisCoefficients) const
{
    if (!BaseRegression::m_bZeroDate)
    {
        const ArrayXXd basisCoefficients  = Map<const ArrayXXd>(p_basisCoefficients.data(), 1, p_basisCoefficients.size()) ;
        // basis function are here regressed values but on the grid, so interpolate
        ArrayXXd retD = interpolateAtSample(basisCoefficients);
        return  Map<ArrayXd>(retD.data(), retD.size());
    }
    else
    {
        return ArrayXd::Constant(m_particles.cols(), p_basisCoefficients(0));
    }
}


ArrayXXd LaplacianGridKernelRegression::reconstructionMultiple(const ArrayXXd   &p_basisCoefficients) const
{
    if (!BaseRegression::m_bZeroDate)
    {
        // basis function are here regressed values but on the grid
        return interpolateAtSample(p_basisCoefficients);
    }
    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  LaplacianGridKernelRegression::reconstructionASim(const int &p_isim, const ArrayXd   &p_basisCoefficients) const
{
    double ret ;
    if (!BaseRegression::m_bZeroDate)
    {
        ret =  interpolateAtPoint(m_particles.col(p_isim), p_basisCoefficients);
    }
    else
    {
        ret = p_basisCoefficients(0);
    }
    return ret ;
}

double LaplacianGridKernelRegression::getValue(const ArrayXd   &p_coordinates,
        const ArrayXd   &p_coordBasisFunction)  const
{
    double ret  ;
    if (!BaseRegression::m_bZeroDate)
    {
        return  interpolateAtPoint(p_coordinates, p_coordBasisFunction);
    }
    else
        ret =  p_coordBasisFunction(0);
    return ret ;
}


double LaplacianGridKernelRegression::getAValue(const ArrayXd &p_coordinates,  const ArrayXd &p_ptOfStock,
        const vector< shared_ptr<InterpolatorSpectral> > &p_interpFuncBasis) const
{
    if (!BaseRegression::m_bZeroDate)
    {
        // coordinates in Z grid
        ArrayXi coordMin(p_coordinates.size());
        for (int id = 0; id <  p_coordinates.size(); ++id)
        {
            if (p_coordinates(id) > (*m_z[id])(m_z[id]->size() - 1))
            {
                coordMin(id) = m_z[id]->size() - 2;
            }
            else
            {
                coordMin(id) =  m_z[id]->size() - 2;
                while ((*m_z[id])(coordMin(id)) > p_coordinates(id))
                {
                    coordMin(id) -= 1;
                    if (coordMin(id) == 0)
                        break;
                }
            }
        }
        // weight
        Eigen::ArrayXd weightPerDim(p_coordinates.size());
        for (int id = 0; id < p_coordinates.size(); ++id)
        {
            // weights have to be positive : so force in case is nearly 0 (rounding error)
            // they have to be below 1
            weightPerDim(id) = std::min(std::max(0., (p_coordinates(id) - (*m_z[id])(coordMin(id))) / ((*m_z[id])(coordMin(id) + 1) - (*m_z[id])(coordMin(id)))), 1.);
        }
        int nbWeigth = (0x01 << p_coordinates.size());
        double retInterp = 0.;
        ArrayXi iCoord(p_coordinates.size());
        // iterate on all vertex of the hypercube
        for (int j = 0 ; j < nbWeigth ; ++j)
        {
            unsigned int ires = j ;
            double weightLocal  = 1. ;
            for (int id = p_coordinates.size() - 1 ; id >= 0  ; --id)
            {
                unsigned int idec = (ires >> id) ;
                iCoord(id) = coordMin(id) + idec;
                weightLocal *= (1 - weightPerDim(id)) * (1 - idec) + weightPerDim(id) * idec ;
                ires -= (idec << id);
            }
            // global coordinate in grid
            int iCoorGlob = iCoord(0);
            int idec = m_z[0]->size();
            for (int id = 1; id < p_coordinates.size(); ++id)
            {
                iCoorGlob += iCoord(id) * idec ;
                idec *= m_z[id]->size();
            }
            retInterp += weightLocal * p_interpFuncBasis[iCoorGlob]->apply(p_ptOfStock);
        }
        return retInterp;
    }
    else
    {
        return p_interpFuncBasis[0]->apply(p_ptOfStock);
    }
}
}