// Copyright (C) 2016 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 localLinearMatrixCalculation(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() + 1;
    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 localLinearSecondMemberCalculation(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() + 1;
    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 localLinearSecondMemberCalculationOneCell(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() + 1;
    // 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;
}


void localLinearCholeski(ArrayXXd   &p_mat, ArrayXXd   &p_diag, Array<bool, Dynamic, 1> &p_bSingular)
{
    int nDim = p_diag.rows();
    int nbCell = p_diag.cols();
    int iCell;
    p_bSingular.setConstant(false);
    for (iCell = 0 ; iCell < nbCell ; iCell++)
    {
        for (int id = 0 ; id < nDim ; ++id)
        {
            double sum =  p_mat(id + id * nDim, iCell) ;
            for (int kd = 0; kd < id ; ++kd)
            {
                sum -=  p_mat(id + kd * nDim, iCell) * p_mat(id + kd * nDim, iCell) ;
            }
            if (sum <= tiny)
            {
                p_bSingular(iCell) = true;
                break;
            }
            p_diag(id, iCell) = sqrt(sum) ;
            for (int jd = id + 1; jd < nDim ; ++jd)
            {
                double sum =  p_mat(id + jd * nDim, iCell) ;
                for (int kd = 0; kd < id ; ++kd)
                {
                    sum -=  p_mat(id + kd * nDim, iCell) * p_mat(jd + kd * nDim, iCell) ;
                }
                p_mat(jd + id * nDim, iCell) = sum / p_diag(id, iCell) ;
            }
        }
    }
}

ArrayXd  localLinearCholeskiInversion(const ArrayXXd   &p_mat, const ArrayXXd   &p_diag,  const ArrayXd &p_secMember)
{
    int nDim = p_diag.rows();
    int nbCell = p_diag.cols();
    int iCell;
    ArrayXd solution(nDim * nbCell);
    for (iCell = 0 ; iCell < nbCell ; iCell++)
    {
        for (int id = 0 ; id < nDim; ++id)
        {
            int iloc = iCell * nDim + id;
            double sum = p_secMember(iloc);
            for (int  kd = id - 1; kd >= 0; kd--) sum -=  p_mat(id + kd * nDim, iCell) * solution(iCell * nDim + kd);
            solution(iloc) = sum / p_diag(id, iCell);
        }
        for (int id = nDim - 1; id >= 0; id--)
        {
            int iloc = iCell * nDim + id;
            double sum = solution(iloc) - p_mat.matrix().col(iCell).segment(id * (nDim + 1) + 1, nDim - id - 1).transpose() * solution.matrix().segment(id + 1 + iCell * nDim, nDim - id - 1);
            solution(iloc) = sum / p_diag(id, iCell);
        }
    }
    return solution;
}


ArrayXXd localLinearReconstruction(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() + 1 ;
    // 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  localLinearReconstructionASim(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() + 1 ;
    // 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  localLinearCholeskiInversionOneCell(const int &p_iCell, const ArrayXXd   &p_mat, const ArrayXXd   &p_diag,  const ArrayXd &p_secMember)
{
    int nDim = p_diag.rows();
    ArrayXd solution(nDim);
    for (int id = 0 ; id < nDim; ++id)
    {
        double sum = p_secMember(id);
        for (int  kd = id - 1; kd >= 0; kd--) sum -=  p_mat(id + kd * nDim, p_iCell) * solution(kd);
        solution(id) = sum / p_diag(id, p_iCell);
    }
    for (int id = nDim - 1; id >= 0; id--)
    {
        double sum = solution(id) - p_mat.matrix().col(p_iCell).segment(id * (nDim + 1) + 1, nDim - id - 1).transpose() * solution.matrix().segment(id + 1, nDim - id - 1);
        solution(id) = sum / p_diag(id, p_iCell);
    }
    return solution;
}



ArrayXd  localLinearReconstructionOnePoint(const ArrayXd &p_oneParticle,
        const vector< shared_ptr< ArrayXd > >   &p_mesh1D,
        const ArrayXXd   &p_foncBasisCoef)
{
    int nBase = p_oneParticle.size() + 1;
    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;
    }
    // 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  localLinearReconsOnePointSimStock(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() + 1;
    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;
    }
    // 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  localLinearReconstructionOnePointOneCell(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;
}

void  localLinearConvexHull(const ArrayXXd &p_particles,
                            const Array< array< double, 2>, Dynamic, Dynamic > &p_mesh,
                            const ArrayXXd   &p_foncBasisCoef,
                            const int &p_nbCell,
                            ArrayXi   &p_simToCell)
{
    int nbSimul = p_simToCell.size();
    int nBase =  p_particles.rows() + 1 ;
    // basis
    ArrayXd FBase(nBase);
    FBase(0) = 1;
    for (int is = 0; is < nbSimul ; ++is)
    {
        double solOpt = -1e10;
        // nest on all cells
        for (int icell = 0; icell < p_nbCell ; ++icell)
            // int icell = p_simToCell(is);
        {
            for (int id = 0 ; id < nBase - 1 ; id++)
            {
                double xPosMin = p_mesh(id, icell)[0] ;
                double xPosMax = p_mesh(id, icell)[1] ;
                FBase(id + 1) = (p_particles(id, is) - xPosMin) / (xPosMax - xPosMin);
            }
            int idec =  icell * nBase ;
            double solution = 0 ;
            for (int isecMem = 0; isecMem < p_foncBasisCoef.rows(); ++isecMem)
                for (int id = 0 ; id < nBase ; ++id)
                    solution += p_foncBasisCoef(isecMem, idec + id) * FBase(id);
            // convex hull
            if (solution > solOpt)
            {
                p_simToCell(is) = icell;
                solOpt = solution;
            }
        }
    }
}

}
