// Copyright (C) 2016 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#include <vector>
#include <memory>
#include <iostream>
#include <Eigen/Dense>
#include "StOpt/regression/LocalSameSizeLinearRegression.h"
#include "StOpt/core/utils/constant.h"
#include "StOpt/core/grids/InterpolatorSpectral.h"
#include "StOpt/regression/localLinearMatrixOperation.h"

using namespace std;
using namespace Eigen;

namespace StOpt
{


LocalSameSizeLinearRegression::LocalSameSizeLinearRegression(const bool &p_bZeroDate,
        const ArrayXXd  &p_particles,
        const ArrayXd &p_lowValues,
        const ArrayXd &p_step,
        const  ArrayXi &p_nbStep) : LocalSameSizeRegression(p_bZeroDate, p_particles, p_lowValues, p_step, p_nbStep)
{
    if (!m_bZeroDate)
    {
        // regression matrix
        constructAndFactorize();
    }
}

LocalSameSizeLinearRegression:: LocalSameSizeLinearRegression(const   LocalSameSizeLinearRegression &p_object): LocalSameSizeRegression(p_object),
    m_matReg(p_object.getMatReg()), m_diagReg(p_object.getDiagReg()),
    m_bSingular(p_object.getBSingular())
{
    const  std::vector< std::shared_ptr< Eigen::MatrixXd > > &pseudoInverse = p_object.getPseudoInverse();
    m_pseudoInverse.resize(pseudoInverse.size());
    for (size_t i = 0 ; i < pseudoInverse.size(); ++i)
    {
        if (pseudoInverse[i])
            m_pseudoInverse[i] = make_shared<Eigen::MatrixXd>(*pseudoInverse[i]);
    }
}


void LocalSameSizeLinearRegression::constructAndFactorize()
{
    int nBase = m_particles.rows() + 1;
    int nbCell = m_nbMeshTotal;
    //to store fuction basis values
    ArrayXd FBase(nBase);

    // initialization
    m_matReg = ArrayXXd::Zero(nBase * nBase, nbCell);

    for (size_t i = 0; i <  m_simAndCell.size() ; ++i)
    {
        // simulation
        int isim = m_simAndCell[i](0);
        // cell number
        int icell = m_simAndCell[i](1) ;
        // calculate basis function values
        FBase(0) = 1;
        for (int id = 0 ; id < nBase - 1 ; id++)
        {

            int iposition = (m_particles(id, isim) - m_lowValues(id)) / m_step(id);
            double xPosMin =  m_lowValues(id) + iposition * m_step(id);
            FBase(id + 1) = (m_particles(id, isim) - xPosMin) / m_step(id);
            assert((FBase(id + 1) >= 0) && (FBase(id + 1) <= 1.)) ;
        }
        for (int k = 0 ; k < nBase ; ++k)
            for (int kk = k ; kk < nBase ; ++kk)
                m_matReg(k + kk * nBase, icell) += FBase(k) * FBase(kk);
    }

    // try to factorize by cholesky
    m_diagReg.resize(nBase, nbCell);
    m_bSingular.resize(nbCell);
    // try to factorize each matrix
    localLinearCholeski(m_matReg, m_diagReg, m_bSingular) ;
    // prepare pseudo factorizatiion
    m_pseudoInverse.resize(nbCell);
    for (int icell = 0 ; icell < nbCell; ++icell)
    {
        if (m_bSingular(icell))
        {
            // modified choleski to account for Singularity (Courrieu 2005 : Fast Computation of Moore-Penrose Inverse Matrices )
            // create the Choleski matrix
            MatrixXd factorized = MatrixXd::Zero(nBase, nBase);
            int iCol = -1;
            // iterate on columns
            for (int id = 0 ; id < nBase ; ++id)
            {
                iCol += 1;
                for (int jd = id ; jd < nBase ; ++jd)
                {
                    factorized(jd, iCol) = m_matReg(id + jd * nBase, icell);
                    for (int kd = 0; kd < iCol ; ++kd)
                        factorized(jd, iCol) -= factorized(id, kd) * factorized(jd, kd);
                }
                if (factorized(id, iCol) > tiny)
                {
                    factorized(id, iCol) = sqrt(factorized(id, iCol));
                    for (int jd = id + 1 ; jd < nBase ; ++jd)
                        factorized(jd, iCol) /= factorized(id, iCol);
                }
                else
                {
                    iCol -= 1 ;
                }
            }
            // construct small dimension matrix
            if (iCol >= 0)
            {
                MatrixXd facReduced = factorized.leftCols(iCol + 1);
                MatrixXd matRegInv = (facReduced.transpose() * facReduced).inverse();
                m_pseudoInverse[icell] = make_shared< MatrixXd >(nBase, nBase);
                *m_pseudoInverse[icell] = facReduced * matRegInv * matRegInv * facReduced.transpose();
            }

        }
    }
}

ArrayXd LocalSameSizeLinearRegression::secondMember(const ArrayXd &p_fToRegress) const
{
    int nBase = m_particles.rows() + 1;
    int nbCell = m_nbMeshTotal;
    //to store fuction basis values
    ArrayXd FBase(nBase);

    // initialization
    ArrayXd secMember = ArrayXd::Zero(nbCell * nBase);

    for (size_t i = 0; i <  m_simAndCell.size() ; ++i)
    {
        // simulation
        int isim = m_simAndCell[i](0);
        // cell number
        int icell = m_simAndCell[i](1) ;
        // calculate basis function values
        FBase(0) = 1;
        for (int id = 0 ; id < nBase - 1 ; id++)
        {

            int iposition = (m_particles(id, isim) - m_lowValues(id)) / m_step(id);
            double xPosMin = m_lowValues(id) + iposition * m_step(id);
            FBase(id + 1) = (m_particles(id, isim) - xPosMin) / m_step(id);
            assert((FBase(id + 1) >= 0) && (FBase(id + 1) <= 1.)) ;
        }
        int idec = icell * nBase;
        for (int id = 0 ; id < nBase ; ++id)
            secMember(idec + id) += p_fToRegress(isim) * FBase(id);
    }
    return secMember;
}

ArrayXd LocalSameSizeLinearRegression::inversion(const ArrayXd    &p_secMem) const
{
    int nBase = m_diagReg.rows();
    int nbCell = m_diagReg.cols();
    int iCell;
    ArrayXd solution(nBase * nbCell);
    for (iCell = 0 ; iCell < nbCell ; iCell++)
    {
        // if not singular
        if (m_bSingular(iCell))
        {
            if (m_pseudoInverse[iCell])
            {
                solution.matrix().segment(iCell * nBase, nBase) = (*m_pseudoInverse[iCell]) * p_secMem.matrix().segment(iCell * nBase, nBase);
            }
            else
            {
                // no particle on the cell
                for (int id = 0 ; id < nBase; ++id)
                    solution(iCell * nBase + id) = 0. ;
            }
        }
        else
        {
            for (int id = 0 ; id < nBase; ++id)
            {
                int iloc = iCell * nBase + id;
                double sum = p_secMem(iloc);
                for (int  kd = id - 1; kd >= 0; kd--) sum -=  m_matReg(id + kd * nBase, iCell) * solution(iCell * nBase + kd);
                solution(iloc) = sum / m_diagReg(id, iCell);
            }
            for (int id = nBase - 1; id >= 0; id--)
            {
                int iloc = iCell * nBase + id;
                double sum = solution(iloc) - m_matReg.matrix().col(iCell).segment(id * (nBase + 1) + 1, nBase - id - 1).transpose() * solution.matrix().segment(id + 1 + iCell * nBase, nBase - id - 1);
                solution(iloc) = sum / m_diagReg(id, iCell);
            }
        }
    }
    return solution;
}

ArrayXd  LocalSameSizeLinearRegression::reconstructionAllPoints(const ArrayXd &p_foncBasisCoef) const
{
    int nBase = m_particles.rows() + 1;
    int nbSimul = m_particles.cols();
    //to store fuction basis values
    ArrayXd FBase(nBase);

    // initialization
    ArrayXd solution = ArrayXd::Zero(nbSimul) ;

    for (size_t i = 0; i <  m_simAndCell.size() ; ++i)
    {
        // simulation
        int isim = m_simAndCell[i](0);
        // cell number
        int icell = m_simAndCell[i](1) ;
        // calculate basis function values
        FBase(0) = 1;
        for (int id = 0 ; id < nBase - 1 ; id++)
        {

            int iposition = (m_particles(id, isim) - m_lowValues(id)) / m_step(id);
            double xPosMin =  m_lowValues(id) + iposition * m_step(id);
            FBase(id + 1) = (m_particles(id, isim) - xPosMin) / m_step(id);
            assert((FBase(id + 1) >= 0) && (FBase(id + 1) <= 1.)) ;
        }
        int idec =  icell * nBase ;
        for (int id = 0 ; id < nBase ; ++id)
            solution(isim) += p_foncBasisCoef(idec + id) * FBase(id);
    }
    return solution;
}


double LocalSameSizeLinearRegression::reconstructionOnlyOnePoint(const ArrayXd   &p_coord, const int &p_cell, const ArrayXd &p_basisCoefficients) const
{
    int nBase =  p_coord.size() + 1 ;
    // initialization
    double  solution = 0 ;
    // basis
    ArrayXd FBase(nBase);
    // calculate basis function values
    FBase(0) = 1;
    for (int id = 0 ; id < nBase - 1 ; id++)
    {

        int iposition = (p_coord(id) - m_lowValues(id)) / m_step(id);
        double xPosMin =  m_lowValues(id) + iposition * m_step(id);
        FBase(id + 1) = (p_coord(id) - xPosMin) / m_step(id);
        assert((FBase(id + 1) >= 0) && (FBase(id + 1) <= 1.)) ;
    }
    int idec =  p_cell * nBase ;
    for (int id = 0 ; id < nBase ; ++id)
        solution += p_basisCoefficients(idec + id) * FBase(id);
    return solution;
}



void LocalSameSizeLinearRegression::updateSimulations(const bool &p_bZeroDate, const ArrayXXd  &p_particles)
{
    BaseRegression::updateSimulationsBase(p_bZeroDate, p_particles);
    if (!m_bZeroDate)
    {

        // update utilitary arrays
        fillInSimCell(m_particles);

        // construct matrix
        constructAndFactorize();
    }
}

ArrayXd LocalSameSizeLinearRegression::getCoordBasisFunction(const ArrayXd &p_fToRegress) const
{
    if (!m_bZeroDate)
    {
        ArrayXd secMem = secondMember(p_fToRegress);
        return inversion(secMem);
    }
    else
    {
        ArrayXd retAverage(1);
        retAverage(0) = p_fToRegress.mean();
        return retAverage;
    }
}

ArrayXXd LocalSameSizeLinearRegression::getCoordBasisFunctionMultiple(const ArrayXXd &p_fToRegress) const
{
    if (!m_bZeroDate)
    {
        ArrayXXd regFunc = ArrayXXd::Zero(p_fToRegress.rows(), m_diagReg.size());
        for (int nsm = 0; nsm < p_fToRegress.rows(); ++nsm)
        {
            ArrayXd secMem = secondMember(p_fToRegress.transpose().col(nsm));
            regFunc.row(nsm) = inversion(secMem).transpose();
        }
        return regFunc;
    }
    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 LocalSameSizeLinearRegression::reconstruction(const ArrayXd &p_basisCoefficients) const
{
    if (!m_bZeroDate)
    {
        return reconstructionAllPoints(p_basisCoefficients) ;
    }
    else
        return ArrayXd::Constant(m_simToCell.size(), p_basisCoefficients(0));
}

ArrayXXd LocalSameSizeLinearRegression::reconstructionMultiple(const ArrayXXd &p_basisCoefficients) const
{
    if (!m_bZeroDate)
    {
        ArrayXXd ret(p_basisCoefficients.rows(), m_particles.cols());
        for (int nsm = 0; nsm < p_basisCoefficients.rows(); ++nsm)
        {
            ret.row(nsm) = reconstructionAllPoints(p_basisCoefficients.row(nsm).transpose()).transpose();
        }
        return ret ;
    }
    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 LocalSameSizeLinearRegression::reconstructionASim(const int &p_isim, const ArrayXd   &p_basisCoefficients) const
{
    if (!m_bZeroDate)
    {
        if (m_simToCell(p_isim) < 0)
            return 0;
        else
        {
            return reconstructionOnlyOnePoint(m_particles.col(p_isim), m_simToCell(p_isim), p_basisCoefficients);
        }
    }
    else
    {
        return p_basisCoefficients(0);
    }
}


ArrayXd LocalSameSizeLinearRegression::getAllSimulations(const ArrayXd &p_fToRegress) const
{
    if (m_bZeroDate)
        return ArrayXd::Constant(p_fToRegress.size(), p_fToRegress.mean());
    ArrayXd secMem = secondMember(p_fToRegress);
    ArrayXd coefBasis = inversion(secMem);
    return reconstructionAllPoints(coefBasis);
}

ArrayXXd LocalSameSizeLinearRegression::getAllSimulationsMultiple(const ArrayXXd &p_fToRegress) const
{
    if (m_bZeroDate)
    {
        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;
    }
    // solution
    ArrayXXd regressedFunc(p_fToRegress.rows(), p_fToRegress.cols());

    for (int nsm = 0; nsm < p_fToRegress.rows(); ++nsm)
    {
        ArrayXd secMem = secondMember(p_fToRegress.transpose().col(nsm));
        ArrayXd coeffBasis = inversion(secMem);
        regressedFunc.row(nsm) = reconstructionAllPoints(coeffBasis).transpose();
    }
    return regressedFunc;
}

double LocalSameSizeLinearRegression::getValue(const ArrayXd &p_coordinates, const ArrayXd &p_coordBasisFunction) const
{
    if (!m_bZeroDate)
    {
        // point location
        int ipos = pointLocation(p_coordinates.array());
        if (ipos < 0)
            return 0. ;
        else
        {
            return reconstructionOnlyOnePoint(p_coordinates, ipos, p_coordBasisFunction);
        }
    }
    else
    {
        return p_coordBasisFunction(0);
    }
}

double LocalSameSizeLinearRegression::getAValue(const ArrayXd &p_coordinates,  const ArrayXd &p_ptOfStock,
        const vector< shared_ptr<InterpolatorSpectral> > &p_interpFuncBasis) const
{
    if (!m_bZeroDate)
    {
        // point location
        int ipos = pointLocation(p_coordinates);
        if (ipos < 0)
            return 0. ;
        else
        {
            int nBase =  m_particles.rows() + 1 ;
            // initialization
            double  solution = 0 ;
            // basis
            ArrayXd FBase(nBase);
            // calculate basis function values
            FBase(0) = 1;
            for (int id = 0 ; id < nBase - 1 ; id++)
            {

                int iposition = (p_coordinates(id) - m_lowValues(id)) / m_step(id);
                double xPosMin =  m_lowValues(id) + iposition * m_step(id);
                FBase(id + 1) = (p_coordinates(id) - xPosMin) / m_step(id);
            }
            int idec =  ipos * nBase ;
            for (int id = 0 ; id < nBase ; ++id)
                solution += p_interpFuncBasis[idec + id]->apply(p_ptOfStock) * FBase(id);
            return solution;
        }
    }
    else
    {
        return p_interpFuncBasis[0]->apply(p_ptOfStock);
    }
}
}