// Copyright (C) 2020 EDF
// Copyright (C) 2020 CSIRO Data61
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#include <vector>
#include <array>
#include <memory>
#include <iostream>
#include <Eigen/Dense>


using namespace std;
using namespace Eigen ;

/// \brief Calculate local sums in Langrené-Warin 2020
///  Only the useful one
/// \param  p_sX        sorted particules per dimension
/// \param  p_iSort     position sorted  : give back position in initial X array
/// \param  p_z         rectilinear points  (evaluation points)
/// \param  p_y         values to sum  (nbValues per sample, nb samples)
/// \param  p_nbPtZ     helper to store index in local sum
/// \param  p_nbPtZT    total number of local sum
/// \return Local summations of p_y values
ArrayXXd localSumCalcKDE(const vector< shared_ptr<ArrayXd> >    &p_sX,
                         const vector< shared_ptr<ArrayXi> >      &p_iSort,
                         const vector< shared_ptr<ArrayXd> >    &p_z,
                         const ArrayXXd &p_y,
                         const ArrayXi   &p_nbPtZ, const  int &p_nbPtZT)
{

    int nbSim = p_sX[0]->size();

    ArrayXXi idx(p_sX.size(), nbSim);
    // loop on dimension
    for (size_t id = 0; id < p_sX.size(); ++id)
    {
        int xidx = 0;
        int zidx = 0;
        // select only important
        while ((xidx < nbSim) && ((*p_sX[id])(xidx) <= (*p_z[id])(p_z[id]->size() - 1)))
        {
            if ((*p_sX[id])(xidx) <= (*p_z[id])(zidx))
            {
                idx(id, (*p_iSort[id])(xidx)) = zidx;
                xidx += 1;
            }
            else
            {
                zidx += 1;
            }
        }
        for (int ix = xidx; ix < nbSim; ++ix)
        {
            idx(id, (*p_iSort[id])(ix)) = -1;
        }
    }
    // return
    ArrayXXd localSum = ArrayXXd::Zero(p_y.rows(), p_nbPtZT);
    for (int is = 0; is < nbSim; ++is)
    {
        // index
        if (idx.col(is).minCoeff() >= 0)
        {
            int iindex = idx(0, is) ;
            for (size_t id = 1; id < p_sX.size() ; ++id)
                iindex += idx(id, is) * p_nbPtZ(id);
            for (int js = 0; js < p_y.rows(); ++js)
                localSum(js, iindex) += p_y(js, is);
        }
    }
    localSum /= nbSim;
    return localSum;
}

/// \brief Calculate   fast CDF members for regressions
/// \param p_delta     array of size the dimension with 1 (delta[i]=1) or -1 (delta[i]=-1)
////                   Used to sort output
/// \param p_sX        sorted particules per dimension
/// \param p_iSort     position sorted  : give back position in initial X array
/// \param p_z         rectilinear points sorted in increasing dimension
/// \param p_y         values to sum  (nbValues per sample, nb samples)
/// \param p_nbPtZ     helper to store index in local sum
/// \param p_nbPtZT    total number of local sum
/// \param p_nbPtZInv  helper to store index in local sum
ArrayXXd CDFFastForKDE(const ArrayXi &p_delta,
                       const vector< shared_ptr<ArrayXd> >    &p_sX,
                       const vector< shared_ptr<ArrayXi> >     &p_iSort,
                       const vector< shared_ptr<ArrayXd> >    &p_z,
                       const ArrayXXd &p_y,
                       const ArrayXi   &p_nbPtZ, const  int &p_nbPtZT, const ArrayXi   &p_nbPtZInv)
{

    // calculate local sum
    ArrayXXd localSum = localSumCalcKDE(p_sX, p_iSort, p_z,  p_y, p_nbPtZ, p_nbPtZT);

    // store index
    ArrayXi index = ArrayXi::Constant(p_z.size(), -1);
    // for return
    ArrayXXd retCDF = ArrayXXd::Zero(p_y.rows(), p_nbPtZT);
    // to store partial sum
    vector< shared_ptr< ArrayXXd>  >  sumOfLocal(p_z.size() + 1);
    int isumOfLocal = p_nbPtZT;
    sumOfLocal[0] = make_shared< ArrayXXd>(ArrayXXd::Zero(p_y.rows(), isumOfLocal));
    for (size_t id  = 1; id <= p_z.size(); ++id)
    {
        isumOfLocal /= p_z[id - 1]->size();
        sumOfLocal[id] = make_shared< ArrayXXd>(ArrayXXd::Zero(p_y.rows(), isumOfLocal));
    }
    (*sumOfLocal[0]) = localSum;


    // previous index
    ArrayXi indexPrev(p_z.size());
    // nest on number of Z points
    for (int iloc = 0; iloc  < p_nbPtZT; ++iloc)
    {
        indexPrev = index;
        // index calculation for CDF
        int ipt = iloc;
        for (size_t id = 0; id  < p_z.size(); ++id)
        {
            index(id) = static_cast<int>(ipt / p_nbPtZInv(id));
            ipt -= index(id) * p_nbPtZInv(id);
        }
        for (size_t id = 0; id < p_z.size(); ++id)
        {
            if (index(id) != indexPrev(id))
            {
                // update sum of local
                for (int ipt = 0; ipt < sumOfLocal[id + 1]->cols(); ++ipt)
                {
                    int iindexInSum = index(id) + ipt * p_z[id]->size();
                    // local index in sumOfLocal array
                    sumOfLocal[id + 1]->col(ipt) += sumOfLocal[id]->col(iindexInSum);
                }
                // update sum of local in other dimensions
                for (size_t idN = id + 1; idN <  p_z.size(); ++idN)
                {
                    // loop on points
                    for (int ipt = 0; ipt < sumOfLocal[idN + 1]->cols(); ++ipt)
                    {
                        sumOfLocal[idN + 1]->col(ipt) = sumOfLocal[idN]->col(ipt * p_z[idN]->size());
                    }
                }
                break;
            }
        }
        // position in return array according to p_delta
        int illoc = static_cast<int>((p_z[0]->size() - 1) * (1 - p_delta(0)) / 2) + p_delta(0) * index(0);
        for (size_t id = 1; id < p_z.size(); ++id)
            illoc += (static_cast<int>((p_z[id]->size() - 1) * (1 - p_delta(id)) / 2) + p_delta(id) * index(id)) * p_nbPtZ(id);
        // store result
        retCDF.col(illoc) = sumOfLocal[p_z.size()]->col(0);
    }
    return retCDF;
}


/// \brief Calculate all terms for kernel regressions
/// \param p_x  samples   (dimension, nbSim)
/// \param p_z  rectilinear points
/// \param p_h  window size per direction
/// \param p_y  value to regress  (nb function to regress, nb samples)
ArrayXXd  FastSumRegressionTerms(const ArrayXXd &p_x,
                                 const vector< shared_ptr<ArrayXd> >    &p_z,
                                 const ArrayXd &p_h,
                                 const ArrayXXd &p_y)
{
    int nbSim  = p_x.cols();
    array< vector< shared_ptr<ArrayXi> >, 2>  iSort;
    array< vector< shared_ptr<ArrayXd> >, 2>  sX;
    array< vector< shared_ptr<ArrayXd> >, 2>  sZ;
    sX[0].reserve(p_x.rows());
    sX[1].reserve(p_x.rows());
    iSort[0].reserve(p_x.rows());
    iSort[1].reserve(p_x.rows());
    sZ[0].reserve(p_x.rows());
    sZ[1].reserve(p_x.rows());
    for (int id = 0 ; id <  p_x.rows(); ++id)
    {
        vector<pair<double, int> > toSort(p_x.cols());
        for (int is = 0; is < p_x.cols(); ++is)
            toSort[is] = make_pair(p_x(id, is), is);
        sort(toSort.begin(), toSort.end());
        // increasing order
        sX[0].push_back(make_shared< ArrayXd>(nbSim));
        iSort[0].push_back(make_shared< ArrayXi>(nbSim));
        for (int is = 0; is < p_x.cols(); ++is)
        {
            (*iSort[0][id])(is) = toSort[is].second;
            (*sX[0][id])(is) = toSort[is].first;
        }
        sZ[0].push_back(p_z[id]);
        // decreasing order
        sX[1].push_back(make_shared< ArrayXd>(nbSim));
        iSort[1].push_back(make_shared< ArrayXi>(nbSim));
        for (int is = 0; is < p_x.cols(); ++is)
        {
            (*iSort[1][id])(is) = toSort[p_x.cols() - 1 - is].second;
            (*sX[1][id])(is) = -toSort[p_x.cols() - 1 - is].first;
        }
        sZ[1].push_back(make_shared< ArrayXd>(p_z[id]->size()));
        for (int is = 0; is < p_z[id]->size(); ++is)
            (*sZ[1][id])(is) = -(*p_z[id])(p_z[id]->size() - 1 - is);
    }

    // store nbpt per dimension before
    ArrayXi nbPtZ(p_z.size() + 1);
    nbPtZ(0) = 1;
    for (size_t id = 0; id < p_z.size(); ++id)
        nbPtZ(id + 1) = nbPtZ(id) * p_z[id]->size();
    // nb point
    int nbPtZT = nbPtZ(p_z.size());
    // helper
    ArrayXi nbPtZInv(p_z.size());
    nbPtZInv(p_z.size() - 1) = 1;
    for (int  id = p_z.size() - 2; id >= 0 ; --id)
        nbPtZInv(id) = nbPtZInv(id + 1) * p_z[id + 1]->size();

    // number of terms
    int nbterm  = (0x01 << p_z.size());

    ArrayXi delta(p_z.size());

    // array for calculations for each value of delta
    vector< shared_ptr<ArrayXi> >   iSortLoc(p_z.size());
    vector< shared_ptr<ArrayXd> >   sXLoc(p_z.size());
    vector< shared_ptr<ArrayXd> >   zLoc(p_z.size());

    // for return
    ArrayXXd retKDE = ArrayXXd::Zero(p_y.rows(), nbPtZT);

    // loop on number of terms
    for (int iterm = 0 ; iterm < nbterm; ++iterm)
    {
        unsigned int ires = iterm ;
        for (int id = p_z.size() - 1 ; id >= 0  ; --id)
        {
            unsigned int idec = (ires >> id) ;
            delta(id) = -1 + 2 * idec;
            ires -= (idec << id);
        }
        for (size_t id = 0; id <  p_z.size(); ++id)
        {
            if (delta(id) == 1)
            {
                iSortLoc[id] = iSort[0][id];
                sXLoc[id] = sX[0][id];
                zLoc[id] = sZ[0][id];
            }
            else
            {
                iSortLoc[id] = iSort[1][id];
                sXLoc[id] = sX[1][id];
                zLoc[id] = sZ[1][id];
            }
        }
        // weight
        ArrayXd weight =  ArrayXd::Zero(nbPtZT);
        for (int ip  = 0; ip < nbPtZT  ; ++ip)
        {
            // index
            int ipt = ip;
            for (int  id = p_z.size() - 1; id  >= 0; --id)
            {
                int index = static_cast<int>(ipt / nbPtZ(id));
                ipt -= index * nbPtZ(id);
                weight(ip) -= delta(id) * (*p_z[id])(index) / p_h(id);
            }
        }
        weight = weight.exp();
        // calculate y value
        ArrayXd y  = ArrayXd::Zero(nbSim);
        for (int  ip = 0; ip < nbSim; ++ip)
            for (size_t id = 0; id < p_z.size(); ++id)
                y(ip) += p_x(id, ip) * delta(id) / p_h(id);
        y = y.exp();

        // developp for all terms of p_y
        ArrayXXd yDev = ArrayXXd::Zero(p_y.rows(), nbSim);
        for (int irow = 0; irow < p_y.rows(); ++irow)
            yDev.row(irow) = y.transpose() * p_y.row(irow);


        /// KDE terms
        ArrayXXd KDETerms = CDFFastForKDE(delta, sXLoc, iSortLoc, zLoc,  yDev, nbPtZ, nbPtZT, nbPtZInv);

        // Local KDE
        for (int icol = 0; icol < retKDE.cols() ; ++icol)
            retKDE.col(icol) += weight(icol) * KDETerms.col(icol);
    }

    return retKDE / (pow(2, p_z.size()) * p_h.prod());
}
