// Copyright (C) 2021 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/LocalLinearDiscrLastDimRegression.h"
#include "StOpt/regression/meshCalculationLocalRegression.h"
#include "StOpt/regression/localLinearMatrixOperation.h"
#include "StOpt/regression/localLinearDiscrLastDimMatrixOperation.h"
#include "StOpt/core/grids/InterpolatorSpectral.h"

using namespace std;
using namespace Eigen;

namespace StOpt
{

LocalLinearDiscrLastDimRegression::LocalLinearDiscrLastDimRegression(const ArrayXi &p_nbMesh, bool  p_bRotationAndRecale): LocalDiscrLastDimRegression(p_nbMesh, p_bRotationAndRecale) {}

LocalLinearDiscrLastDimRegression::LocalLinearDiscrLastDimRegression(const bool &p_bZeroDate,
        const ArrayXXd  &p_particles,
        const ArrayXi &p_nbMesh,
        bool  p_bRotationAndRecale) : LocalDiscrLastDimRegression(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();
        // regression matrix
        m_matReg = localLinearDiscrLastDimMatrixCalculation(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) ;
    }
}

LocalLinearDiscrLastDimRegression:: LocalLinearDiscrLastDimRegression(const   LocalLinearDiscrLastDimRegression &p_object): LocalDiscrLastDimRegression(p_object), m_matReg(p_object.getMatReg()), m_diagReg(p_object.getDiagReg())

{
}

LocalLinearDiscrLastDimRegression::LocalLinearDiscrLastDimRegression(const bool &p_bZeroDate,
        const   ArrayXi &p_nbMesh,
        const   Array< std::array< double, 2>, Dynamic, Dynamic >   &p_mesh,
        const std::vector< std::shared_ptr< ArrayXd > > &p_mesh1D, const   ArrayXd &p_meanX,
        const   ArrayXd   &p_etypX, const   MatrixXd   &p_svdMatrix,
        const bool   &p_bRotationAndRecale) : LocalDiscrLastDimRegression(p_bZeroDate, p_nbMesh, p_mesh,  p_mesh1D, p_meanX,  p_etypX,  p_svdMatrix, p_bRotationAndRecale)
{
}


void LocalLinearDiscrLastDimRegression::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();
        }
        meshCalculationLocalRegressionDiscrLastDim(m_particles, m_nbMesh, m_simToCell, m_mesh, m_mesh1D);

        int nbCell = m_mesh.cols() ;
        int iMatrixSize =  m_particles.rows();
        // regression matrix
        m_matReg = localLinearDiscrLastDimMatrixCalculation(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 LocalLinearDiscrLastDimRegression::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 = localLinearDiscrLastDimSecondMemberCalculation(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 LocalLinearDiscrLastDimRegression::getCoordBasisFunctionMultiple(const ArrayXXd &p_fToRegress) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        ArrayXXd secMember = localLinearDiscrLastDimSecondMemberCalculation(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 LocalLinearDiscrLastDimRegression::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 localLinearDiscrLastDimReconstruction(m_particles, m_simToCell, m_mesh, BasisCoefficients).row(0);
    }
    else
        return ArrayXd::Constant(m_simToCell.size(), p_basisCoefficients(0));
}

ArrayXXd LocalLinearDiscrLastDimRegression::reconstructionMultiple(const ArrayXXd &p_basisCoefficients) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        return localLinearDiscrLastDimReconstruction(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 LocalLinearDiscrLastDimRegression::reconstructionASim(const int &p_isim, const ArrayXd   &p_basisCoefficients) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        return localLinearDiscrLastDimReconstructionASim(p_isim, m_particles, m_simToCell, m_mesh, p_basisCoefficients);
    }
    else
    {
        return p_basisCoefficients(0);
    }
}


ArrayXd LocalLinearDiscrLastDimRegression::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 = localLinearDiscrLastDimReconstruction(m_particles, m_simToCell, m_mesh, BasisCoefficients);
    return condEspectationValues.row(0);
}

ArrayXXd LocalLinearDiscrLastDimRegression::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 localLinearDiscrLastDimReconstruction(m_particles, m_simToCell, m_mesh, BasisCoefficients);
}



double LocalLinearDiscrLastDimRegression::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 localLinearDiscrLastDimReconstructionOnePoint(x.array(), m_mesh1D, coordBasisFunction2D)(0);
    }
    else
    {
        return p_coordBasisFunction(0);
    }
}

double LocalLinearDiscrLastDimRegression::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 localLinearDiscrLastDimReconsOnePointSimStock(x, p_ptOfStock, p_interpFuncBasis, m_mesh1D);
    }
    else
    {
        return p_interpFuncBasis[0]->apply(p_ptOfStock);
    }
}

ArrayXd LocalLinearDiscrLastDimRegression::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 = localLinearDiscrLastDimSecondMemberCalculationOneCell(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 LocalLinearDiscrLastDimRegression::getCoordBasisFunctionMultipleOneCell(const int &p_iCell, const ArrayXXd &p_fToRegress) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        ArrayXXd secMember = localLinearDiscrLastDimSecondMemberCalculationOneCell(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  LocalLinearDiscrLastDimRegression::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 localLinearDiscrLastDimReconstructionOnePointOneCell(x.array(), m_mesh.col(p_cell), p_foncBasisCoef);
    }
    else
    {
        return  p_foncBasisCoef.col(0);
    }
}

}
