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

using namespace std;
using namespace Eigen;

namespace StOpt
{

// Only upper part is filled in
ArrayXXd localLinearDiscrLastDimMatrixCalculation(const ArrayXXd &p_particles,
        const ArrayXi &p_simToCell,
        const Array< array< double, 2>, Dynamic, Dynamic > &p_mesh)
{
    int nbSimul =  p_simToCell.size();
    int nBase = p_particles.rows() ;
    int nbCell = p_mesh.cols();
    //to store fuction basis values
    ArrayXd FBase(nBase);

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

    for (int is = 0; is < nbSimul ; ++is)
    {
        // cell number
        int ncell = p_simToCell(is) ;
        // calculate basis function values
        FBase(0) = 1;
        for (int id = 0 ; id < nBase - 1 ; id++)
        {
            double xPosMin =  p_mesh(id, ncell)[0];
            double xPosMax =  p_mesh(id, ncell)[1] ;
            FBase(id + 1) = (p_particles(id, is) - xPosMin) / (xPosMax - xPosMin);
        }
        for (int k = 0 ; k < nBase ; ++k)
            for (int kk = k ; kk < nBase ; ++kk)
                matReg(k + kk * nBase, ncell) += FBase(k) * FBase(kk);
    }

    // normalization
    matReg /= nbSimul ;
    return matReg;

}

ArrayXXd localLinearDiscrLastDimSecondMemberCalculation(const ArrayXXd &p_particles,
        const ArrayXi &p_simToCell,
        const Array< array< double, 2>, Dynamic, Dynamic > &p_mesh,
        const ArrayXXd &p_fToRegress)
{
    int nbSimul = p_simToCell.size();
    int nBase = p_particles.rows() ;
    int nbCell =  p_mesh.cols();
    // number of function to regress
    int iSecMem = p_fToRegress.rows();

    ArrayXd  FBase(nBase);
    ArrayXXd secMember = ArrayXXd::Zero(p_fToRegress.rows(), nbCell * nBase);
    for (int is = 0; is < nbSimul ; ++is)
    {
        int nCell = p_simToCell(is);
        FBase(0) = 1;
        for (int id = 0 ; id < nBase - 1; id++)
        {
            double xPosMin = p_mesh(id, nCell)[0] ;
            double xPosMax = p_mesh(id, nCell)[1] ;
            FBase(id + 1) = (p_particles(id, is) - xPosMin) / (xPosMax - xPosMin);
        }
        int idec = nCell * nBase;
        // nest on second members
        for (int nsm = 0 ; nsm < iSecMem ; ++nsm)
        {
            double xtemp = p_fToRegress(nsm, is) ;
            // second member of the regression problem
            for (int id = 0 ; id < nBase ; ++id)
                secMember(nsm, idec + id) += xtemp * FBase(id);
        }
    }
    // normalization
    secMember /= nbSimul;
    return secMember;
}

ArrayXXd localLinearDiscrLastDimSecondMemberCalculationOneCell(const ArrayXXd &p_particles,
        const std::vector<int>   &p_SimulBelongingToCell,
        const Eigen::Ref<const Eigen::Array< array<double, 2 >, Eigen::Dynamic, 1 > >    &p_mesh,
        const ArrayXXd &p_fToRegress)
{
    int nbSimul =  p_SimulBelongingToCell.size();
    int nBase = p_particles.rows() ;
    // number of function to regress
    int iSecMem = p_fToRegress.rows();

    ArrayXd  FBase(nBase);
    ArrayXXd secMember = ArrayXXd::Zero(p_fToRegress.rows(), nBase);
    for (int is = 0; is < nbSimul ; ++is)
    {
        FBase(0) = 1;
        for (int id = 0 ; id < nBase - 1; id++)
        {
            double xPosMin = p_mesh(id)[0] ;
            double xPosMax = p_mesh(id)[1] ;
            FBase(id + 1) = (p_particles(id,  p_SimulBelongingToCell[is]) - xPosMin) / (xPosMax - xPosMin);
        }
        // nest on second members
        for (int nsm = 0 ; nsm < iSecMem ; ++nsm)
        {
            double xtemp = p_fToRegress(nsm, is) ;
            // second member of the regression problem
            for (int id = 0 ; id < nBase ; ++id)
                secMember(nsm, id) += xtemp * FBase(id);
        }
    }
    // no normalization
    return secMember;
}


ArrayXXd localLinearDiscrLastDimReconstruction(const ArrayXXd &p_particles, const ArrayXi &p_simToCell,
        const Array< array< double, 2>, Dynamic, Dynamic > &p_mesh,
        const ArrayXXd   &p_foncBasisCoef)
{
    int nbSimul = p_simToCell.size();
    int nBase =  p_particles.rows();
    // basis
    ArrayXd FBase(nBase);
    // initialization
    ArrayXXd solution = ArrayXXd::Zero(p_foncBasisCoef.rows(), nbSimul) ;

    for (int is = 0; is < nbSimul ; ++is)
    {
        int nCell = p_simToCell(is) ;

        FBase(0) = 1;
        for (int id = 0 ; id < nBase - 1 ; id++)
        {
            double xPosMin = p_mesh(id, nCell)[0] ;
            double xPosMax = p_mesh(id, nCell)[1] ;
            FBase(id + 1) = (p_particles(id, is) - xPosMin) / (xPosMax - xPosMin);
        }

        int idec =  nCell * nBase ;
        for (int isecMem = 0; isecMem < p_foncBasisCoef.rows(); ++isecMem)
            for (int id = 0 ; id < nBase ; ++id)
                solution(isecMem, is) += p_foncBasisCoef(isecMem, idec + id) * FBase(id);
    }
    return solution;
}

double  localLinearDiscrLastDimReconstructionASim(const int &p_isim, const ArrayXXd &p_particles, const ArrayXi &p_simToCell,
        const Array< array< double, 2>, Dynamic, Dynamic > &p_mesh,
        const ArrayXd   &p_foncBasisCoef)
{
    int nBase =  p_particles.rows()  ;
    // basis
    ArrayXd FBase(nBase);
    // initialization
    double  solution = 0 ;
    int nCell = p_simToCell(p_isim) ;
    FBase(0) = 1;
    for (int id = 0 ; id < nBase - 1 ; id++)
    {
        double xPosMin = p_mesh(id, nCell)[0] ;
        double xPosMax = p_mesh(id, nCell)[1] ;
        FBase(id + 1) = (p_particles(id, p_isim) - xPosMin) / (xPosMax - xPosMin);
    }
    int idec =  nCell * nBase ;
    for (int id = 0 ; id < nBase ; ++id)
        solution += p_foncBasisCoef(idec + id) * FBase(id);
    return solution;
}


ArrayXd  localLinearDiscrLastDimReconstructionOnePoint(const ArrayXd &p_oneParticle,
        const vector< shared_ptr< ArrayXd > >   &p_mesh1D,
        const ArrayXXd   &p_foncBasisCoef)
{
    int nBase = p_oneParticle.size();
    ArrayXd  FBase(nBase);
    // Values of the functon basis and position of the  particle in the mesh
    FBase(0) = 1;
    int iCell = 0 ;
    int idecCell = 1;
    for (int id = 0 ; id < nBase - 1 ; id++)
    {
        int iMesh = 1 ;
        while ((p_oneParticle(id) > (*p_mesh1D[id])(iMesh)) && (iMesh < p_mesh1D[id]->size() - 1)) iMesh++;
        double xPosMin = (*p_mesh1D[id])(iMesh - 1) ;
        double xPosMax = (*p_mesh1D[id])(iMesh) ;
        FBase(id + 1) = (p_oneParticle(id) - xPosMin) / (xPosMax - xPosMin);
        iCell += (iMesh - 1) * idecCell;
        idecCell *= p_mesh1D[id]->size() - 1;
    }
    // position last dimension
    int iMesh = 1 ;
    while ((p_oneParticle(nBase - 1) > (*p_mesh1D[nBase - 1])(iMesh)) && (iMesh < p_mesh1D[nBase - 1]->size() - 1)) iMesh++;
    iCell += (iMesh - 1) * idecCell;

    // reconstruction
    int idec = nBase * iCell ;
    ArrayXd solution = ArrayXd::Zero(p_foncBasisCoef.rows());
    for (int isecMem = 0; isecMem < solution.size(); ++isecMem)
        for (int id = 0 ; id < nBase ; ++id)
            solution(isecMem) += p_foncBasisCoef(isecMem, idec + id) * FBase(id);
    return solution;
}

double  localLinearDiscrLastDimReconsOnePointSimStock(const ArrayXd &p_oneParticle, const ArrayXd &p_stock,
        const std::vector< std::shared_ptr<InterpolatorSpectral> > &p_interpBaseFunc,
        const std::vector< std::shared_ptr< ArrayXd >  > &p_mesh1D)
{
    int nBase = p_oneParticle.size();
    ArrayXd  FBase(nBase);
    // Values of the functon basis and position of the  particle in the mesh
    FBase(0) = 1;
    int iCell = 0 ;
    int idecCell = 1;
    for (int id = 0 ; id < nBase - 1 ; id++)
    {
        int iMesh = 1 ;
        while ((p_oneParticle(id) > (*p_mesh1D[id])(iMesh)) && (iMesh < p_mesh1D[id]->size() - 1)) iMesh++;
        double xPosMin = (*p_mesh1D[id])(iMesh - 1) ;
        double xPosMax = (*p_mesh1D[id])(iMesh) ;
        FBase(id + 1) = (p_oneParticle(id) - xPosMin) / (xPosMax - xPosMin);
        iCell += (iMesh - 1) * idecCell;
        idecCell *= p_mesh1D[id]->size() - 1;
    }
    // position last dimension
    int iMesh = 1 ;
    while ((p_oneParticle(nBase - 1) > (*p_mesh1D[nBase - 1])(iMesh)) && (iMesh < p_mesh1D[nBase - 1]->size() - 1)) iMesh++;
    iCell += (iMesh - 1) * idecCell;
    // reconstruction
    int idec = nBase * iCell ;
    double solution = 0;
    for (int id = 0 ; id < nBase ; ++id)
        solution +=  p_interpBaseFunc[idec + id]->apply(p_stock) * FBase(id);
    return solution;
}

ArrayXd  localLinearDiscrLastDimReconstructionOnePointOneCell(const ArrayXd &p_oneParticle,
        const Array< array< double, 2>, Dynamic, 1 > &p_mesh,
        const ArrayXXd   &p_foncBasisCoef)
{
    int nBase =  p_foncBasisCoef.cols();
    // basis
    ArrayXd FBase(nBase);
    // initialization
    ArrayXd solution = ArrayXd::Zero(p_foncBasisCoef.rows()) ;
    FBase(0) = 1;
    for (int id = 0 ; id < nBase - 1 ; id++)
    {
        double xPosMin = p_mesh(id)[0] ;
        double xPosMax = p_mesh(id)[1] ;
        FBase(id + 1) = (p_oneParticle(id) - xPosMin) / (xPosMax - xPosMin);
    }
    for (int isecMem = 0; isecMem < p_foncBasisCoef.rows(); ++isecMem)
        for (int id = 0 ; id < nBase ; ++id)
            solution(isecMem) += p_foncBasisCoef(isecMem, id) * FBase(id);
    return solution;
}

}
