// 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/LocalLinearRegression.h"
#include "StOpt/regression/meshCalculationLocalRegression.h"
#include "StOpt/regression/localLinearMatrixOperation.h"
#include "StOpt/core/grids/InterpolatorSpectral.h"

using namespace std;
using namespace Eigen;

namespace StOpt
{

LocalLinearRegression::LocalLinearRegression(const ArrayXi &p_nbMesh, bool  p_bRotationAndRecale): LocalRegression(p_nbMesh, p_bRotationAndRecale) {}

LocalLinearRegression::LocalLinearRegression(const bool &p_bZeroDate,
        const ArrayXXd  &p_particles,
        const ArrayXi &p_nbMesh,
        bool  p_bRotationAndRecale) : LocalRegression(p_bZeroDate, p_particles, p_nbMesh, p_bRotationAndRecale)
{
    if ((!m_bZeroDate) && (p_nbMesh.size() != 0))
    {
        int nbCell = m_mesh.cols() ;
        int iMatrixSize =  m_particles.rows() + 1;
        // regression matrix
        m_matReg = localLinearMatrixCalculation(m_particles, m_simToCell, m_mesh);
        // factorize
        m_diagReg.resize(iMatrixSize, nbCell);
        // utilitary
        Array<bool, Dynamic, 1> bSingular(nbCell);
        localLinearCholeski(m_matReg, m_diagReg, bSingular) ;
    }
}

LocalLinearRegression:: LocalLinearRegression(const   LocalLinearRegression &p_object): LocalRegression(p_object), m_matReg(p_object.getMatReg()), m_diagReg(p_object.getDiagReg())

{}


void LocalLinearRegression::updateSimulations(const bool &p_bZeroDate, const ArrayXXd  &p_particles)
{
    BaseRegression::updateSimulationsBase(p_bZeroDate, p_particles);
    m_simToCell.resize(p_particles.cols());
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        if (p_particles.rows() != m_nbMesh.size())
        {
            cout << " Dimension nd  of particles of size (nd, nbSimu) is " << p_particles.rows();
            cout << " and   should be equal to the size of the array describing the mesh refinement " << m_nbMesh.transpose() << endl ;
            abort();
        }
        meshCalculationLocalRegression(m_particles, m_nbMesh, m_simToCell, m_mesh, m_mesh1D);

        int nbCell = m_mesh.cols() ;
        int iMatrixSize =  m_particles.rows() + 1;
        // regression matrix
        m_matReg = localLinearMatrixCalculation(m_particles, m_simToCell, m_mesh);
        // factorize
        m_diagReg.resize(iMatrixSize, nbCell);
        // utilitary
        Array<bool, Dynamic, 1> bSingular(nbCell);
        localLinearCholeski(m_matReg, m_diagReg, bSingular) ;
    }
    else
    {
        m_simToCell.setConstant(0);
    }
}

ArrayXd LocalLinearRegression::getCoordBasisFunction(const ArrayXd &p_fToRegress) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        Map<const ArrayXXd>  fToRegress2D(p_fToRegress.data(), 1, p_fToRegress.size());
        ArrayXXd secMember2D = localLinearSecondMemberCalculation(m_particles, m_simToCell, m_mesh, fToRegress2D);
        Map<const ArrayXd > secMember(secMember2D.data(), secMember2D.size());
        return localLinearCholeskiInversion(m_matReg, m_diagReg, secMember);
    }
    else
    {
        ArrayXd retAverage(1);
        retAverage(0) = p_fToRegress.mean();
        return retAverage;
    }
}

ArrayXXd LocalLinearRegression::getCoordBasisFunctionMultiple(const ArrayXXd &p_fToRegress) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        ArrayXXd secMember = localLinearSecondMemberCalculation(m_particles, m_simToCell, m_mesh, p_fToRegress);
        ArrayXXd regFunc(p_fToRegress.rows(), secMember.cols());
        for (int nsm = 0; nsm <  p_fToRegress.rows(); ++nsm)
        {
            ArrayXd secMemberLoc(secMember.cols());
            secMemberLoc = secMember.row(nsm);
            regFunc.row(nsm) = localLinearCholeskiInversion(m_matReg, m_diagReg, secMemberLoc).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 LocalLinearRegression::reconstruction(const ArrayXd &p_basisCoefficients) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        Map<const ArrayXXd> BasisCoefficients(p_basisCoefficients.data(), 1, p_basisCoefficients.size());
        return localLinearReconstruction(m_particles, m_simToCell, m_mesh, BasisCoefficients).row(0);
    }
    else
        return ArrayXd::Constant(m_simToCell.size(), p_basisCoefficients(0));
}

ArrayXXd LocalLinearRegression::reconstructionMultiple(const ArrayXXd &p_basisCoefficients) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        return localLinearReconstruction(m_particles, m_simToCell, m_mesh, p_basisCoefficients);
    }
    else
    {
        ArrayXXd retValue(p_basisCoefficients.rows(), m_simToCell.size());
        for (int nsm = 0; nsm < p_basisCoefficients.rows(); ++nsm)
            retValue.row(nsm).setConstant(p_basisCoefficients(nsm, 0));
        return retValue ;
    }
}

double LocalLinearRegression::reconstructionASim(const int &p_isim, const ArrayXd   &p_basisCoefficients) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        return localLinearReconstructionASim(p_isim, m_particles, m_simToCell, m_mesh, p_basisCoefficients);
    }
    else
    {
        return p_basisCoefficients(0);
    }
}


ArrayXd LocalLinearRegression::getAllSimulations(const ArrayXd &p_fToRegress) const
{
    Map<const ArrayXXd>  fToRegress2D(p_fToRegress.data(), 1, p_fToRegress.size());
    ArrayXXd BasisCoefficients = getCoordBasisFunctionMultiple(fToRegress2D);
    if ((m_bZeroDate) || (m_nbMesh.size() == 0))
    {
        return  ArrayXd::Constant(p_fToRegress.size(), BasisCoefficients(0, 0));
    }
    ArrayXXd  condEspectationValues = localLinearReconstruction(m_particles, m_simToCell, m_mesh, BasisCoefficients);
    return condEspectationValues.row(0);
}

ArrayXXd LocalLinearRegression::getAllSimulationsMultiple(const ArrayXXd &p_fToRegress) const
{
    ArrayXXd BasisCoefficients = getCoordBasisFunctionMultiple(p_fToRegress);
    if ((m_bZeroDate) || (m_nbMesh.size() == 0))
    {
        ArrayXXd ret(p_fToRegress.rows(), p_fToRegress.cols());
        for (int ism = 0; ism < p_fToRegress.rows(); ++ism)
            ret.row(ism).setConstant(BasisCoefficients(ism, 0));
        return ret;
    }
    return localLinearReconstruction(m_particles, m_simToCell, m_mesh, BasisCoefficients);
}

ArrayXd LocalLinearRegression::getAllSimulationsConvex(const ArrayXd &p_fToRegress, const int &p_nbIterMax)
{
    if ((m_bZeroDate) || (m_nbMesh.size() == 0))
        return ArrayXd::Constant(p_fToRegress.size(), p_fToRegress.mean());

    // number of cells
    int nbCell = m_mesh.cols() ;
    // utilitary
    Array<bool, Dynamic, 1> bSingular(nbCell);
    // eigen map to avoid copy
    Map<const ArrayXXd>  fToRegress2D(p_fToRegress.data(), 1, p_fToRegress.size());
    // error
    double error = 1e6;
    // iteration number
    int iterN = 0 ;
    // to store old position
    ArrayXi simToCellOld = m_simToCell;
    // second member
    ArrayXXd BasisCoefficients = getCoordBasisFunctionMultiple(fToRegress2D);
    // iterate on the regressions
    while ((error > 0) && (iterN++ < p_nbIterMax))
    {

        localLinearConvexHull(m_particles, m_mesh, BasisCoefficients, nbCell,   m_simToCell);
        // update error
        error = (simToCellOld - m_simToCell).abs().sum();
        if (error > 0)
        {
            // update matrix
            m_matReg = localLinearMatrixCalculation(m_particles, m_simToCell, m_mesh);
            // utilitary
            localLinearCholeski(m_matReg, m_diagReg, bSingular) ;
            // second member update
            BasisCoefficients = getCoordBasisFunctionMultiple(fToRegress2D);
        }
        // keep position
        simToCellOld = m_simToCell;
    }
    // final reconstrucion
    ArrayXXd  condEspectationValues = localLinearReconstruction(m_particles, m_simToCell, m_mesh, BasisCoefficients);
    return condEspectationValues.row(0);
}


double LocalLinearRegression::getValue(const ArrayXd &p_coordinates, const ArrayXd &p_coordBasisFunction) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        // rotation
        VectorXd x = m_svdMatrix * ((p_coordinates - m_meanX) / m_etypX).matrix();
        Map<const ArrayXXd> coordBasisFunction2D(p_coordBasisFunction.data(), 1, p_coordBasisFunction.size());
        return localLinearReconstructionOnePoint(x.array(), m_mesh1D, coordBasisFunction2D)(0);
    }
    else
    {
        return p_coordBasisFunction(0);
    }
}

double LocalLinearRegression::getAValue(const ArrayXd &p_coordinates,  const ArrayXd &p_ptOfStock,
                                        const vector< shared_ptr<InterpolatorSpectral> > &p_interpFuncBasis) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        // rotation
        VectorXd x = m_svdMatrix * ((p_coordinates - m_meanX) / m_etypX).matrix();
        return localLinearReconsOnePointSimStock(x, p_ptOfStock, p_interpFuncBasis, m_mesh1D);
    }
    else
    {
        return p_interpFuncBasis[0]->apply(p_ptOfStock);
    }
}

ArrayXd LocalLinearRegression::getCoordBasisFunctionOneCell(const int &p_iCell, const ArrayXd &p_fToRegress) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        Map<const ArrayXXd>  fToRegress2D(p_fToRegress.data(), 1, p_fToRegress.size());
        ArrayXXd secMember2D = localLinearSecondMemberCalculationOneCell(m_particles, *(m_simulBelongingToCell[p_iCell]), m_mesh.col(p_iCell), fToRegress2D);
        Map<const ArrayXd > secMember(secMember2D.data(), secMember2D.size());
        return localLinearCholeskiInversionOneCell(p_iCell, m_matReg, m_diagReg, secMember);
    }
    else
    {
        ArrayXd retAverage(1);
        retAverage(0) = p_fToRegress.mean();
        return retAverage;
    }
}

ArrayXXd LocalLinearRegression::getCoordBasisFunctionMultipleOneCell(const int &p_iCell, const ArrayXXd &p_fToRegress) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        ArrayXXd secMember = localLinearSecondMemberCalculationOneCell(m_particles, *(m_simulBelongingToCell[p_iCell]), m_mesh.col(p_iCell), p_fToRegress);
        // normalization
        secMember /= m_particles.cols();
        ArrayXXd regFunc(p_fToRegress.rows(), secMember.cols());
        for (int nsm = 0; nsm <  p_fToRegress.rows(); ++nsm)
        {
            ArrayXd secMemberLoc(secMember.cols());
            secMemberLoc = secMember.row(nsm);
            regFunc.row(nsm) = localLinearCholeskiInversionOneCell(p_iCell, m_matReg, m_diagReg, secMemberLoc).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  LocalLinearRegression::getValuesOneCell(const ArrayXd &p_oneParticle, const int &p_cell, const ArrayXXd   &p_foncBasisCoef) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        // rotation
        VectorXd x = p_oneParticle.matrix();
        x = ((x.array() - m_meanX) / m_etypX).matrix();
        x = m_svdMatrix * x;
        return localLinearReconstructionOnePointOneCell(x.array(), m_mesh.col(p_cell), p_foncBasisCoef);
    }
    else
    {
        return  p_foncBasisCoef.col(0);
    }
}

}
