/**
 * @file
 * @copyright This code is licensed under the 3-clause BSD license.\n
 *            Copyright ETH Zurich, Laboratory for Physical Chemistry, Reiher Group.\n
 *            See LICENSE.txt for details.
 */

#include <Utils/DataStructures/OccupiedMolecularOrbitals.h>
#include <Utils/Scf/LcaoUtils/HFWaveFunctionOverlap.h>
#include <Eigen/Dense>

namespace Scine {
namespace Utils {

namespace LcaoUtils {

double HFWaveFunctionOverlap::calculateOrthonormalOverlap(const OccupiedMolecularOrbitals& c1,
                                                          const OccupiedMolecularOrbitals& c2) {
  // Convert to uhf to make comparison of RHF - UHF possible
  // TODO: make more efficient if no conversion is needed
  auto unrestricted1 = c1.toUnrestricted();
  auto unrestricted2 = c2.toUnrestricted();
  return unrestrictedOrthonormalOverlap(unrestricted1, unrestricted2);
}

double HFWaveFunctionOverlap::calculateNonOrthonormalOverlap(const OccupiedMolecularOrbitals& c1,
                                                             const OccupiedMolecularOrbitals& c2, const Eigen::MatrixXd& s) {
  // Convert to uhf to make comparison of RHF - UHF possible
  // TODO: make more efficient if no conversion is needed
  auto unrestricted1 = c1.toUnrestricted();
  auto unrestricted2 = c2.toUnrestricted();
  return unrestrictedNonOrthonormalOverlap(unrestricted1, unrestricted2, s);
}

double HFWaveFunctionOverlap::unrestrictedOrthonormalOverlap(const OccupiedMolecularOrbitals& c1,
                                                             const OccupiedMolecularOrbitals& c2) {
  double f1 = orthonormalContribution(c1.alphaMatrix(), c2.alphaMatrix());
  double f2 = orthonormalContribution(c1.betaMatrix(), c2.betaMatrix());
  // Return a product of determinants:
  // The determinant of a block diagonal matrix is the product of the determinants of the block
  return f1 * f2;
}

double HFWaveFunctionOverlap::unrestrictedNonOrthonormalOverlap(const OccupiedMolecularOrbitals& c1,
                                                                const OccupiedMolecularOrbitals& c2,
                                                                const Eigen::MatrixXd& s) {
  double f1 = nonOrthonormalContribution(c1.alphaMatrix(), c2.alphaMatrix(), s);
  double f2 = nonOrthonormalContribution(c1.betaMatrix(), c2.betaMatrix(), s);
  // Return a product of determinants:
  // The determinant of a block diagonal matrix is the product of the determinants of the block
  return f1 * f2;
}

double HFWaveFunctionOverlap::orthonormalContribution(const Eigen::MatrixXd& m1, const Eigen::MatrixXd& m2) {
  if (m1.cols() != m2.cols()) {
    throw std::runtime_error("Not possible to calculate overlap between systems with different number of electrons.");
  }
  return (m1.transpose() * m2).determinant();
}

double HFWaveFunctionOverlap::nonOrthonormalContribution(const Eigen::MatrixXd& m1, const Eigen::MatrixXd& m2,
                                                         const Eigen::MatrixXd& s) {
  if (m1.cols() != m2.cols()) {
    throw std::runtime_error("Not possible to calculate overlap between systems with different number of electrons.");
  }
  return (m1.transpose() * s * m2).determinant();
}

} // namespace LcaoUtils
} // namespace Utils
} // namespace Scine


/* ScaFES
 * Copyright (c) 2011-2015, ZIH, TU Dresden, Federal Republic of Germany.
 * For details, see the files COPYING and LICENSE in the base directory
 * of the package.
 */

/**
 *  @file EMFDTD.hpp
 *
 *  @brief Implementation of EMFDTD.
 */

#ifndef EMFDTD_HPP_
#define EMFDTD_HPP_

#include <fstream>
#include <string>
#include <vector>

#include <boost/property_tree/ini_parser.hpp>
#include <boost/tokenizer.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/serialization/vector.hpp>

#include <ScaFES.hpp>

#include "physical_constants.hpp"
#include "geometry/IndicatorEllipsoid.hpp"
#include "CPMLParams.hpp"
#include "stencils/CompleteCPMLCurlStencil.hpp"

#include "sources/ModulatedGaussianSource.hpp"
#include "conversion.hpp"

/*******************************************************************************
 ******************************************************************************/
/**
 * \class EMFDTDng
 *  @brief Class for discretized EM problem.
 *
 * \section EMFDTDng3D EMFDTD at ellipsoid.
 *
 * The EM problem is discretized using the FDTD method (Finite Differences
 * at the Time Domain).
 *
 * The workflow for running the simulation is as follows:
 *
 *      - Set up the computational domain within in the source code.
 *      - Recompile source code.
 *      - Set up EM parameters within the ini file.
 *      - Run simulation.
 *
 * Parameters of the PML (perfectly matched layers):
 *
 *      [PML]
        d=10
        m_a=1.
        m=3
        k_max=1.
        s_max=0.01222
        a_max=1.e-3

 * Location and parameters of the source.
 *
        [source]
        ; Location within computational domain
        location=25x25x25
        ; modulation frequency:
        f_mod = 0
        Gauss pulse: time step with maximal stimulation (=1)
        t0 = 40.0
        ; Width
        spread = 8.0

 * Location and target file of the sink.
 *
        [sink]
        ; Location within computational domain
        location=30x30x30
        ; Target file
        data_file = reference.txt

 * Domain initialization
 *
        [initialization]
        ; for ellipsoid: (px,py,pz,rx,ry,rz)
        parameters = 100 100 100 50 50 50
 */

/** data field access by symbolic name. */
enum class DFs: std::vector<std::string>::size_type {
    SIGMA=0,
    Ex, Ey, Ez,
    Hx, Hy, Hz,
    Dx, Dy, Dz,
    Ix, Iy, Iz,
    Psi_ExDy, Psi_ExDz, Psi_EyDx,
    Psi_EyDz, Psi_EzDx, Psi_EzDy,
    Psi_HxDy, Psi_HxDz, Psi_HyDx,
    Psi_HyDz, Psi_HzDx, Psi_HzDy
};

/** when to write which data field. */
typedef ScaFES::WriteHowOften W;
std::vector<W> initWritePolicy()
{
    std::vector<W> v;
    v.push_back(W::AT_START);
    v.push_back(W::LIKE_GIVEN_AT_CL); v.push_back(W::LIKE_GIVEN_AT_CL); v.push_back(W::LIKE_GIVEN_AT_CL);
    v.push_back(W::NEVER); v.push_back(W::NEVER); v.push_back(W::NEVER);
    v.push_back(W::NEVER); v.push_back(W::NEVER); v.push_back(W::NEVER);
    v.push_back(W::NEVER); v.push_back(W::NEVER); v.push_back(W::NEVER);
    v.push_back(W::NEVER); v.push_back(W::NEVER); v.push_back(W::NEVER);
    v.push_back(W::NEVER); v.push_back(W::NEVER); v.push_back(W::NEVER);
    v.push_back(W::NEVER); v.push_back(W::NEVER); v.push_back(W::NEVER);
    v.push_back(W::NEVER); v.push_back(W::NEVER); v.push_back(W::NEVER);
    return v;
}
static const std::vector<W> writePolicy = initWritePolicy();

/** data field names. */
std::vector<std::string> initDFNames()
{
    std::vector<std::string> v;
    v.push_back("SIGMA");
    v.push_back("Ex"); v.push_back("Ey"); v.push_back("Ez");
    v.push_back("Hx"); v.push_back("Hy"); v.push_back("Hz");
    v.push_back("Dx"); v.push_back("Dy"); v.push_back("Dz");
    v.push_back("Ix"); v.push_back("Iy"); v.push_back("Iz");
    v.push_back("Psi_ExDy"); v.push_back("Psi_ExDz"); v.push_back("Psi_EyDx");
    v.push_back("Psi_EYDz"); v.push_back("Psi_EzDx"); v.push_back("Psi_EzDy");
    v.push_back("Psi_HxDy"); v.push_back("Psi_HxDz"); v.push_back("Psi_HyDx");
    v.push_back("Psi_HyDz"); v.push_back("Psi_HzDx"); v.push_back("Psi_HzDy");
    return v;
}
static const std::vector<std::string> DFNames = initDFNames();

/** MPI halos. */
std::vector<int> initHaloWidths()
{
    std::vector<int> v;
    v.push_back(0);
    v.push_back(1); v.push_back(1); v.push_back(1);
    v.push_back(1); v.push_back(1); v.push_back(1);
    v.push_back(0); v.push_back(0); v.push_back(0);
    v.push_back(0); v.push_back(0); v.push_back(0);
    v.push_back(0); v.push_back(0); v.push_back(0);
    v.push_back(0); v.push_back(0); v.push_back(0);
    v.push_back(0); v.push_back(0); v.push_back(0);
    v.push_back(0); v.push_back(0); v.push_back(0);
    return v;
}
static const std::vector<int> HaloWidths = initHaloWidths();

template<typename CT, bool ForwardOnly>
class EMFDTDng: public ScaFES::Problem< EMFDTDng<CT, ForwardOnly>, CT, 3> {
public:
    static const size_t DIM=3;
private:
    typedef ScaFES::Problem< EMFDTDng<CT, ForwardOnly>, CT, 3> Super;
    typedef std::vector<std::string>::size_type size_type;
    typedef ScaFES::Ntuple<int, DIM> Index;
    typedef typename Index::value_type index_type;
    using PTree=boost::property_tree::ptree;

    // Stencils
    typedef CompleteCurlCPMLStencil<CT, Index, 0, 0, 1, 0, 1, 0> HxStencil;
    typedef CompleteCurlCPMLStencil<CT, Index, 1, 0, 0, 0, 0, 1> HyStencil;
    typedef CompleteCurlCPMLStencil<CT, Index, 0, 1, 0, 1, 0, 0> HzStencil;
    typedef CompleteCurlCPMLStencil<CT, Index, 0, -1, 0, 0, 0, -1> DxStencil;
    typedef CompleteCurlCPMLStencil<CT, Index, 0, 0, -1, -1, 0, 0> DyStencil;
    typedef CompleteCurlCPMLStencil<CT, Index, -1, 0, 0, 0, -1, 0> DzStencil;
public:

    EMFDTDng(
        ScaFES::Parameters const& cl,
        ScaFES::GridGlobal<DIM> const& gg,
        const PTree& ptree
        )
    : Super(cl, gg, true
        , DFNames // DF names
        , HaloWidths // MPI halo width
        , std::vector<bool>(DFNames.size(), false) // is known? --> false
        , std::vector<int>(DFNames.size(), 0) // border layers: 0 -> complete volume
        , std::vector<CT>(DFNames.size(), CT(0.)) // initial values for DFs
        , writePolicy
        , std::vector<bool>(DFNames.size(), false) // compute error? --> false
        , parse<CT>(ptree.get<std::string>("initialization.parameters"), " ") // params, vOld in initialize()
        )
    , ptree_(ptree)
    , cfln_(.5)
    , ds_(1.)
    , dt_((cfln_*ds_)/C_0)
    , cpmlp_(dt_,
        ptree_.get<std::size_t>("PML.d"),
        ptree_.get<CT>("PML.m_a"),
        ptree_.get<CT>("PML.m"),
        ptree_.get<CT>("PML.k_max"),
        ptree_.get<CT>("PML.s_max"),
        ptree_.get<CT>("PML.a_max")
        )
    , first_idx_(0, 0, 0)
    , last_idx_(cl.nNodes().at(0)-1, cl.nNodes().at(2)-1, cl.nNodes().at(2)-1)
    , loc_src_(vec2idx(parse<index_type>(ptree_.get<std::string>("source.location"))))
    , loc_sink_(vec2idx(parse<index_type>(ptree_.get<std::string>("sink.location"))))
    , reference_(load_reference(ptree_.get<std::string>("optimization.reference_data_file"), cl.nTimesteps()))
    , src_(dt_,ptree_)
    , hxstencil_(cpmlp_, Index(1, 0, 0), Index(0, -1, -1), last_idx_, cfln_)
    , hystencil_(cpmlp_, Index(0, 1, 0), Index(-1, 0, -1), last_idx_, cfln_)
    , hzstencil_(cpmlp_, Index(0, 0, 1), Index(-1, -1, 0), last_idx_, cfln_)
    , dxstencil_(cpmlp_, Index(0, 0, 1), Index(-1, 0, 0), last_idx_, -cfln_)
    , dystencil_(cpmlp_, Index(1, 0, 1), Index(0, -1, 0), last_idx_, -cfln_)
    , dzstencil_(cpmlp_, Index(1, 1, 0), Index(0, 0, -1), last_idx_, -cfln_)
    {
        boost::mpi::communicator world;

        if (world.rank()==0)
            std::cout<<"Source at "<<loc_src_
            <<"\nSink at "<<loc_sink_
            <<"\nInitialization parameters: "
            <<ptree.get<std::string>("initialization.parameters")
            <<std::endl;

        if (loc_src_==loc_sink_)
            throw std::runtime_error("Source and sink cannot be at the same position.");

        if ((loc_src_<first_idx_)||(loc_src_>last_idx_))
            throw std::runtime_error("Source not within simulation area.");
        if ((loc_sink_<first_idx_)||(loc_sink_>last_idx_))
            throw std::runtime_error("Sink not within simulation area.");

        if (cl.tau()!=1.0)
            throw std::runtime_error("Please choose \"--starttime\" and \"--endtime\" so that tau()=1.0.");
    }

    /**
     * Empty.
     */
    template<typename TT>
    void initInner(std::vector< ScaFES::DataField<TT, DIM> >& /*vNew*/,
        std::vector< ScaFES::DataField<TT, DIM> > const& /*vOld*/,
        Index const& /*idxNode*/,
        int const& /*timestep*/) { }

    /**
     * Empty.
     */
    template<typename TT>
    void initBorder(std::vector< ScaFES::DataField<TT, DIM> >& /*vNew*/,
            std::vector< ScaFES::DataField<TT, DIM> > const& /*vOld*/,
            Index const& /*idxNode*/,
            int const& /*timestep*/) { }

    /**
     * Initialize from p.
     */
    template <typename TT>
    void initInner
    (
     std::vector< ScaFES::DataField<TT, DIM> >& dfs,
     std::vector<TT> const& p,
     Index const& idx,
     int const& /*timestep*/
    )
    {
            typedef ScaFES::Ntuple<TT, DIM> TT3;
            TT3 fidx(TT(idx.elem(0)), TT(idx.elem(1)), TT(idx.elem(2)));

            const TT inside=3.546e7;
            const TT outside=0.;

            IndicatorEllipsoid<TT> e(p.at(0), p.at(1), p.at(2), p.at(3), p.at(4), p.at(5), inside, outside);

            const TT sigma=e(fidx);

            dfs[0](idx)=sigma;
    }

    /**
     * Wire to initInner.
     */
    template<typename TT>
    void initBorder(std::vector< ScaFES::DataField<TT, DIM> >& dfs,
            std::vector<TT> const& p,
            Index const& idx,
            int const& timestep)
    {
        this->initInner(dfs, p, idx, timestep);
    }

    /**
     * Compute H.
     */
    template<typename TT>
    void updateInner(std::vector< ScaFES::DataField<TT, DIM> >& n,
            std::vector< ScaFES::DataField<TT, DIM> > const& o,
            Index const& idx,
            int const& /*timestep*/)
    {
        // start by updating H from E
        n[sc(DFs::Hx)](idx)=hxstencil_(idx, o[sc(DFs::Hx)],
            o[sc(DFs::Ey)], o[sc(DFs::Psi_EyDz)], n[sc(DFs::Psi_EyDz)],
            o[sc(DFs::Ez)], o[sc(DFs::Psi_EzDy)], n[sc(DFs::Psi_EzDy)]);

        n[sc(DFs::Hy)](idx)=hystencil_(idx, o[sc(DFs::Hy)],
            o[sc(DFs::Ez)], o[sc(DFs::Psi_EzDx)], n[sc(DFs::Psi_EzDx)],
            o[sc(DFs::Ex)], o[sc(DFs::Psi_ExDz)], n[sc(DFs::Psi_ExDz)]);

        n[sc(DFs::Hz)](idx)=hzstencil_(idx, o[sc(DFs::Hz)],
            o[sc(DFs::Ex)], n[sc(DFs::Psi_ExDy)], n[sc(DFs::Psi_ExDy)],
            o[sc(DFs::Ey)], n[sc(DFs::Psi_EyDx)], n[sc(DFs::Psi_EyDx)]);
    }

    /**
     * Wire to updateInner.
     */
    template<typename TT>
    void updateBorder(std::vector< ScaFES::DataField<TT, DIM> >& vNew,
            std::vector< ScaFES::DataField<TT, DIM>>const& vOld,
            Index const& idx,
            int const& timestep)
    {
        updateInner(vNew, vOld, idx, timestep);
    }

    /**
     * Compute E.
     */
    template<typename TT>
    void updateInner2(std::vector< ScaFES::DataField<TT, DIM> >& n,
            std::vector< ScaFES::DataField<TT, DIM> > const& o,
            Index const& idx,
            int const& /*timestep*/)
    {
        // update D from H
        // n.b.:
        // 1) although comfortable, DO NOT USE "auto" for type deduction. ADOL-C yiels obscure and hard to debug compilation failures otherwise ("auto"-deduced rhs type is not assignable).
        // 2) although comfortable, NEVER violate the "compute -> use -> persist" order for (AD-)active variables. ADOL-C will simply not work if you persist a datum (i.e., write it to memory) and re-use it from there. The use of temporaries is a workaround.
        TT Dx=dxstencil_(idx, o[sc(DFs::Dx)],
            o[sc(DFs::Hz)], o[sc(DFs::Psi_HzDy)], n[sc(DFs::Psi_HzDy)],
            o[sc(DFs::Hy)], o[sc(DFs::Psi_HyDz)], n[sc(DFs::Psi_HyDz)]);

        TT Dy=dystencil_(idx, o[sc(DFs::Dy)],
            o[sc(DFs::Hx)], o[sc(DFs::Psi_HxDz)], n[sc(DFs::Psi_HxDz)],
            o[sc(DFs::Hz)], o[sc(DFs::Psi_HzDx)], n[sc(DFs::Psi_HzDx)]);

        TT Dz=dzstencil_(idx, o[sc(DFs::Dz)],
            o[sc(DFs::Hy)], o[sc(DFs::Psi_HyDx)], n[sc(DFs::Psi_HyDx)],
            o[sc(DFs::Hx)], o[sc(DFs::Psi_HxDy)], n[sc(DFs::Psi_HxDy)]);

        TT ga=1./(1.+(o[sc(DFs::SIGMA)](idx))*dt_/EPSILON_0);
        TT gb=(o[sc(DFs::SIGMA)](idx))*dt_/EPSILON_0;

        TT Ex=ga*(Dx-o[sc(DFs::Ix)](idx));
        TT Ey=ga*(Dy-o[sc(DFs::Iy)](idx));
        TT Ez=ga*(Dz-o[sc(DFs::Iz)](idx));

        n[sc(DFs::Ix)](idx)=o[sc(DFs::Ix)](idx)+(gb*Ex);
        n[sc(DFs::Iy)](idx)=o[sc(DFs::Iy)](idx)+(gb*Ey);
        n[sc(DFs::Iz)](idx)=o[sc(DFs::Iz)](idx)+(gb*Ez);

        // write temporaries to dependent fields, see above, item 2
        n[sc(DFs::Dx)](idx)=Dx;
        n[sc(DFs::Dy)](idx)=Dy;
        n[sc(DFs::Dz)](idx)=Dz;
        n[sc(DFs::Ex)](idx)=Ex;
        n[sc(DFs::Ey)](idx)=Ey;

        // inject source, use old time
        if (idx==loc_src_)
            n[sc(DFs::Ez)](idx)=src_((o[sc(DFs::Ez)]).time());
        else
            n[sc(DFs::Ez)](idx)=Ez;

        // collect sink
        if (idx==loc_sink_) {
            convert<CT, TT> c;
            CT val=c(n[sc(DFs::Ez)](idx));
            sink_.push_back(val);
        }
    }

    /**
     * Wire to updateInner2.
     */
    template<typename TT>
    void updateBorder2(std::vector< ScaFES::DataField<TT, DIM> >& vNew,
            std::vector< ScaFES::DataField<TT, DIM>>const& vOld,
            Index const& idx,
            int const& timestep)
    {
        updateInner2(vNew, vOld, idx, timestep);
    }

    /**
     * Empty.
     */
    template<typename TT>
    void evalInner(std::vector< ScaFES::DataField<TT, DIM> >& /*vNew*/,
        ScaFES::Ntuple<int, DIM> const& /*idxNode*/,
        int const& /*timestep*/) { }

    /**
     * Empty.
     */
    template<typename TT>
    void evalBorder(std::vector< ScaFES::DataField<TT, DIM> >& /*vNew*/,
            ScaFES::Ntuple<int, DIM> const& /*idxNode*/,
            int const& /*timestep*/) { }

    /**
     * Write sink data to file.
     */
    void writeSinkData()
    {
        boost::mpi::communicator world;

        if (world.size()!=1) {
            if (world.rank()==0) {
                if (sink_.size()==0)
                    world.recv(boost::mpi::any_source, 0, sink_);
            } else {
                if (sink_.size()!=0)
                    world.send(0, 0, sink_);
            }
        }

        if (world.rank()!=0)
            return;

        const std::string fname=ptree_.get<std::string>("sink.data_file");

        if (fname.size()==0)
            throw std::runtime_error("Sink file name empty.");

        std::cout<<"Writing sink data to file \""
            <<fname<<"\"..."
            <<std::flush;

        std::ofstream of(fname.c_str());
        of.setf(of.scientific);
        of.precision(20);

        for (auto it=sink_.begin(); it!=sink_.end(); ++it)
            of<<(*it)<<std::endl;

        std::cout<<"done."<<std::endl;
    }

private:
    /**
     * Layer structure for datafields with holes.
     */
    std::vector<int> layers(const PTree& p)
    {
        std::vector<int> ret(13, 0);
        const int d=p.get<int>("PML.d")+2;

        for (int i=0; i<12; ++i)
            ret.push_back(d);

        return ret;
    };

    /**
     * Parse into vector.
     * @param s
     * @param sep
     * @return
     */
    template <typename T>
    static std::vector<T> parse(const std::string& s, const char* sep="x")
    {
        boost::char_separator<char> csep(sep);
        boost::tokenizer<boost::char_separator<char> > tok(s, csep);

        std::vector<T> ret;

        for (auto ti=tok.begin(); ti!=tok.end(); ++ti) {
            try {
            ret.push_back(boost::lexical_cast<T>(*ti));
            }
            catch (...) {
                std::cout << "While parsing string \"" << s << "\":" << std::endl;
                throw;
            }
        }
        return ret;
    }

    /**
     * Make Index from vector.
     * @param v
     * @return
     */
    static Index vec2idx(const std::vector<typename Index::value_type>& v)
    {
        Index ret;

        for (unsigned long int i=0; i<Index::DIM; ++i)
            ret[i]=v.at(i);

        return ret;
    }

    /**
     * Cast "nice names" into accessible ints
     * @param df
     * @return
     */
    static size_type sc(const DFs & df)
    {
        return static_cast<size_type> (df);
    }

    /**
     * Load reference data from file.
     */
    static std::vector<CT> load_reference(const std::string& fname, const unsigned int& n)
    {
        boost::mpi::communicator world;
        std::vector<CT> ret;

        if (ForwardOnly==true) {
            if (world.rank()==0)
                std::cout<<"NOT reading reference from file. Calling TF evaluation routines WILL FAIL!"<<std::endl;
            return ret;
        }

        /** TODO: This statement is unreachable in case ForwardOnly=true. */
        if (world.rank()==0) {
            std::ifstream ifs(fname.c_str());

            if (!ifs.is_open())
                throw std::runtime_error("Failed opening file \""+fname+"\".");

            while ((ifs.good()) && (ret.size()<n)) {
                std::string line;
                std::getline(ifs, line);
                try {
                    ret.push_back(boost::lexical_cast<CT, std::string>(line));
                } catch (...) {
                    std::cout<<"While parsing reference file \"" << fname << "\": " <<std::endl;
                    throw;
                }
            }

            if (ret.size()<n)
                throw std::runtime_error("Reference file \""+fname+"\" contains too few items.");
        }

        if (world.size()>1)
            boost::mpi::broadcast(world, ret, 0);

        return ret;
    }

    const PTree ptree_;
    const CT cfln_, ds_, dt_;
    const CPMLParams<CT> cpmlp_;
    const Index first_idx_, last_idx_, loc_src_, loc_sink_;
    const std::vector<CT> reference_;
    const ModulatedGaussianSource<CT> src_;
    std::vector<CT> sink_;

    const HxStencil hxstencil_;
    const HyStencil hystencil_;
    const HzStencil hzstencil_;
    const DxStencil dxstencil_;
    const DyStencil dystencil_;
    const DzStencil dzstencil_;
};

#endif // EMFDTD_HPP_


#include "feature_estimator.h"

#include <Eigen/Dense>

namespace {

int calculate_medoid_index(std::vector<Eigen::Vector3d> &points) {
  std::vector<double> distances(points.size(), 0);
  for (auto i = 0u; i < points.size(); ++i) {
    for (auto j = 0u; j < points.size(); ++j) {
      const double distance = (points[i] - points[j]).eval().norm();
      distances[i] += distance;
    }
  }

  int min_idx = 0;
  double min_distance = std::numeric_limits<double>::max();
  for (auto i = 0u; i < distances.size(); ++i) {
    if (distances[i] < min_distance) {
      min_idx = i;
      min_distance = distances[i];
    }
  }

  return min_idx;
}

}  // namespace

namespace pcs {

FeatureEstimator::FeatureEstimator(std::shared_ptr<pcs::PointCloud> point_cloud,
                                   double voxel_size, int num_neighbors,
                                   int num_scales, int batch_size)
    : point_cloud_(point_cloud),
      voxel_size_(voxel_size),
      num_neighbors_(num_neighbors),
      num_scales_(num_scales),
      batch_size_(batch_size),
      tree_(nullptr),
      num_points_(0),
      has_colors_(false),
      pool_(std::thread::hardware_concurrency()) {
  if (point_cloud_ and not point_cloud_->empty()) {
    num_points_ = point_cloud_->points_.size();
    has_colors_ = point_cloud_->has_colors();
    tree_ = std::shared_ptr<KDTree<3>>(new KDTree<3>(
        point_cloud_->points_, std::thread::hardware_concurrency()));

    // Build points pyramid
    point_clouds_.push_back(point_cloud_);
    trees_.push_back(tree_);
    for (auto s = 1u; s < num_scales_; ++s) {
      const auto voxel_scale = std::pow(2, s);
      auto down_scaled_cloud =
          voxel_down_sample(*(point_clouds_[s - 1]), voxel_size_ * voxel_scale);
      point_clouds_.push_back(down_scaled_cloud);

      auto tree = std::shared_ptr<KDTree<3>>(new KDTree<3>(
          point_clouds_[s]->points_, std::thread::hardware_concurrency()));
      trees_.push_back(tree);
    }
  }
}

Eigen::VectorXd FeatureEstimator::get_features_for_point(
    std::size_t point_id) const {
  if (point_id >= num_points_) {
    std::string error_string("given point index { ");
    error_string +=
        std::to_string(point_id) + " } is bigger then point cloud size { ";
    error_string += std::to_string(num_points_) + " }";
    throw std::out_of_range(error_string);
  }

  const std::size_t num_features = feature_size();
  Eigen::VectorXd features(num_features);
  Eigen::Vector3d point = point_cloud_->points_[point_id];
  Eigen::Vector3d neig_color(0., 0., 0.);

  for (auto s = 0u; s < point_clouds_.size(); ++s) {
    const std::size_t scale_offset = s * features_per_scale;
    auto neighbors = trees_[s]->find_nns(point, num_neighbors_);
    if (s == 0 and has_colors_) {
      for (auto k = 0u; k < neighbors.size(); ++k) {
        neig_color += (point_cloud_->colors_[neighbors[k].first]);
      }
    }
    std::vector<Eigen::Vector3d> n_points;
    for (auto j = 0u; j < neighbors.size(); ++j) {
      n_points.push_back(point_clouds_[s]->points_[neighbors[j].first]);
    }

    Eigen::Vector3d medoid = n_points[calculate_medoid_index(n_points)];
    Eigen::Matrix3d cov = Eigen::Matrix3d::Zero();
    Eigen::MatrixXd data = Eigen::MatrixXd::Zero(num_neighbors_, 3);
    for (auto j = 0u; j < num_neighbors_; ++j) {
      data.row(j) = n_points[j].transpose();
    }

    Eigen::MatrixXd centered = data.rowwise() - data.colwise().mean();
    cov = (centered.transpose() * centered) / (num_neighbors_ - 1);

    Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> eig(cov);
    Eigen::Matrix3d eigenvectors = eig.eigenvectors();
    Eigen::Vector3d eigenvalues = eig.eigenvalues();
    eigenvalues.normalize();

    // Covariance features
    const double omnivariance =
        std::pow(eigenvalues(0) * eigenvalues(1) * eigenvalues(2), 1.0 / 3.0);

    features[scale_offset + 0] = omnivariance;
    double eigenentropy = 0.0;
    for (int k = 0; k < 3; ++k) {
      eigenentropy -= (eigenvalues(k) * std::log(eigenvalues(k)));
    }

    features[scale_offset + 1] = eigenentropy;

    const double anisotropy =
        (eigenvalues(2) - eigenvalues(0)) / eigenvalues(2);
    features[scale_offset + 2] = anisotropy;

    const double planarity = (eigenvalues(1) - eigenvalues(0)) / eigenvalues(2);
    features[scale_offset + 3] = planarity;

    const double linearity = (eigenvalues(2) - eigenvalues(1)) / eigenvalues(2);
    features[scale_offset + 4] = linearity;

    const double surface_variation = eigenvalues(0);
    features[scale_offset + 5] = surface_variation;

    const double scatter = eigenvalues(0) / eigenvalues(2);
    features[scale_offset + 6] = scatter;

    const double verticality = 1. - std::abs(eigenvectors.col(0)[2]);
    features[scale_offset + 7] = verticality;

    // Moments features
    double first_order_first_axis = 0;
    double first_order_second_axis = 0;
    double second_order_first_axis = 0;
    double second_order_second_axis = 0;
    for (auto &p : n_points) {
      Eigen::Vector3d diff = p - medoid;
      first_order_first_axis += diff.dot(eigenvectors.col(2));
      first_order_second_axis += diff.dot(eigenvectors.col(1));
      second_order_first_axis +=
          first_order_first_axis * first_order_first_axis;
      second_order_second_axis +=
          first_order_second_axis * first_order_second_axis;
    }
    features[scale_offset + 8] = first_order_first_axis;
    features[scale_offset + 9] = first_order_second_axis;
    features[scale_offset + 10] = second_order_first_axis;
    features[scale_offset + 11] = second_order_second_axis;

    // Height features
    double z_min = std::numeric_limits<double>::max();
    double z_max = -std::numeric_limits<double>::max();
    for (auto &p : n_points) {
      const double z = p.z();
      if (z < z_min) {
        z_min = z;
      }
      if (z > z_max) {
        z_max = z;
      }
    }

    const double vertical_range = z_max - z_min;
    features[scale_offset + 12] = vertical_range;

    const double height_below = point.z() - z_min;
    features[scale_offset + 13] = height_below;

    const double height_above = z_max - point.z();
    features[scale_offset + 14] = height_above;
  }

  // Color features calculated if initial cloud has colors
  if (has_colors_) {
    const std::size_t geometric_features_offset =
        num_scales_ * features_per_scale;
    Eigen::Vector3d main_color = point_cloud_->colors_[point_id];
    features[geometric_features_offset + 0] = main_color(0);
    features[geometric_features_offset + 1] = main_color(1);
    features[geometric_features_offset + 2] = main_color(2);

    neig_color /= num_neighbors_;
    features[geometric_features_offset + 3] = neig_color(0);
    features[geometric_features_offset + 4] = neig_color(1);
    features[geometric_features_offset + 5] = neig_color(2);
  }
  features =
      features.unaryExpr([](double v) { return std::isfinite(v) ? v : 0.0; });
  return features;
}

Eigen::MatrixXd FeatureEstimator::get_features_for_batch(
    std::size_t start_id, std::size_t end_id) const {
  if (end_id > point_cloud_->points_.size() or
      start_id >= point_cloud_->points_.size() or end_id <= start_id) {
    std::string error_string("invalid slice { ");
    error_string += std::to_string(start_id) + " : ";
    error_string += std::to_string(end_id) + " } for array with { ";
    error_string += std::to_string(point_cloud_->points_.size()) + " } points";
    throw std::out_of_range(error_string);
  }

  const std::size_t range = end_id - start_id;
  const std::size_t num_features = feature_size();
  Eigen::MatrixXd result(range, num_features);
  const double drange =
      static_cast<double>(range) / static_cast<double>(pool_.num_threads());
  const std::size_t num_points_per_thread =
      std::max<std::size_t>(1u, std::ceil(drange));
  const std::size_t max_threads = std::min(pool_.num_threads(), range);
  for (auto i = 0u; i < max_threads; ++i) {
    const auto s = start_id + i * num_points_per_thread;
    const auto e = std::min(end_id, s + num_points_per_thread);
    pool_.add_task([&result, s, e, start_id, this] {
      for (auto j = s; j < e; ++j) {
        result.row(j - start_id) = get_features_for_point(j);
      }
    });
  }
  pool_.wait();

  return result;
}

Eigen::MatrixXd FeatureEstimator::get_features_for_points(
    const std::vector<int> &idxs) const {
  const std::size_t range = idxs.size();
  const std::size_t num_features = feature_size();
  Eigen::MatrixXd result(range, num_features);
  const double drange =
      static_cast<double>(range) / static_cast<double>(pool_.num_threads());
  const std::size_t num_points_per_thread =
      std::max<std::size_t>(1u, std::ceil(drange));
  const std::size_t max_threads = std::min(pool_.num_threads(), range);
  std::vector<std::future<void>> futures;
  for (auto i = 0u; i < max_threads; ++i) {
    const auto s = i * num_points_per_thread;
    const auto e = std::min(range, s + num_points_per_thread);
    futures.push_back(pool_.add_task([&result, &idxs, s, e, this] {
      for (auto j = s; j < e; ++j) {
        result.row(j) = get_features_for_point(idxs[j]);
      }
    }));
  }
  for (auto &&x : futures) {
    x.get();
  }

  return result;
}

std::size_t FeatureEstimator::num_points() const { return num_points_; }

std::vector<unsigned int> FeatureEstimator::soft_voting_smoothing(
    const Eigen::MatrixXd &probabilities, std::size_t num_neighbors) const {
  std::vector<unsigned int> classes(probabilities.rows());

  if (probabilities.size() == 0) {
    return classes;
  }
  const std::size_t num_classes = probabilities.cols();

  const double drange = static_cast<double>(num_points_) /
                        static_cast<double>(pool_.num_threads());
  const std::size_t num_points_per_thread =
      std::max<std::size_t>(1u, std::ceil(drange));
  const std::size_t max_threads = std::min(pool_.num_threads(), num_points_);
  for (auto i = 0u; i < max_threads; ++i) {
    const auto s = i * num_points_per_thread;
    const auto e = std::min(num_points_, s + num_points_per_thread);
    pool_.add_task(
        [&classes, &probabilities, num_neighbors, num_classes, s, e, this] {
          for (auto point_id = s; point_id < e; ++point_id) {
            Eigen::Vector3d point = point_cloud_->points_[point_id];
            const Eigen::Vector3d main_color = point_cloud_->colors_[point_id];
            const auto neighbors = tree_->find_nns(point, num_neighbors);
            float weight_sum = 0.f;
            Eigen::VectorXd commite_result = Eigen::VectorXd::Zero(num_classes);
            for (const auto &neighbor : neighbors) {
              commite_result += probabilities.row(neighbor.first);
            }
            Eigen::VectorXd::Index index;
            commite_result.maxCoeff(&index);
            classes[point_id] = index;
          }
        });
  }
  pool_.wait();

  return classes;
}

std::size_t FeatureEstimator::num_batches() const {
  return static_cast<std::size_t>(
      std::ceil(static_cast<double>(num_points_) / batch_size_));
}

std::size_t FeatureEstimator::batch_size() const { return batch_size_; }

Eigen::MatrixXd FeatureEstimator::get_features_for_batch(
    std::size_t batch_id) const {
  const std::size_t start_id = batch_id * batch_size_;
  const std::size_t end_id =
      std::min<std::size_t>(start_id + batch_size_, num_points_);
  return get_features_for_batch(start_id, end_id);
}

std::size_t FeatureEstimator::feature_size() const {
  if (has_colors_) {
    return color_features + features_per_scale * num_scales_;
  } else {
    return features_per_scale * num_scales_;
  }
}

std::vector<unsigned int> FeatureEstimator::hard_voting_smoothing(
    const std::vector<unsigned int> &labels, std::size_t num_neighbors) const {
  std::vector<unsigned int> smooth_labels(labels.size(), 0);

  if (labels.empty()) {
    return smooth_labels;
  }

  const double drange = static_cast<double>(num_points_) /
                        static_cast<double>(pool_.num_threads());
  const std::size_t num_points_per_thread =
      std::max<std::size_t>(1u, std::ceil(drange));
  const std::size_t max_threads = std::min(pool_.num_threads(), num_points_);
  for (auto i = 0u; i < max_threads; ++i) {
    const auto s = i * num_points_per_thread;
    const auto e = std::min(num_points_, s + num_points_per_thread);
    pool_.add_task([this, s, e, num_neighbors, &smooth_labels, &labels] {
      for (auto point_id = s; point_id < e; ++point_id) {
        Eigen::Vector3d point = point_cloud_->points_[point_id];
        const auto neighbors = tree_->find_nns(point, num_neighbors);
        std::unordered_map<unsigned int, int> commite_result;
        for (const auto &neighbor : neighbors) {
          commite_result[labels[neighbor.first]]++;
        }
        unsigned int max_class = 0;
        int max_value = -1;
        for (auto &v : commite_result) {
          if (v.second > max_value) {
            max_class = v.first;
            max_value = v.second;
          }
        }
        smooth_labels[point_id] = max_class;
      }
    });
  }
  pool_.wait();

  return smooth_labels;
}

}  // namespace pcs


#include <limits>
#include <Eigen/Geometry>
#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <pybind11/numpy.h>

#include <igl/point_mesh_squared_distance.h>
#include <igl/AABB.h>

namespace py = pybind11;
using namespace Eigen;

typedef Eigen::Vector3d Vec3d;
typedef Eigen::Vector2d Vec2d;

std::tuple<py::array_t<int>, py::array_t<double>, py::array_t<int>, py::array_t<double>>
rayMeshIntersect(
        py::array_t<double> verts_in,
        py::array_t<int> tris_in,
        py::array_t<double> ray_pts_in,
        py::array_t<double> ray_dirs_in,
        double max_distance,
        double max_angle,
        bool allow_backface_hit)
{
    auto verts = verts_in.unchecked<2>();
    auto tris = tris_in.unchecked<2>();
    auto rayPts = ray_pts_in.unchecked<2>();
    auto rayDirs = ray_dirs_in.unchecked<2>();

    std::vector<int> hitTris;
    std::vector<Vec2d> hitUVs;
    std::vector<Vec3d> hitPoints;
    std::vector<int> hitRayIndices;

    double maxAngleCos = std::cos(max_angle);
    for(int iray = 0; iray < rayPts.shape(0); iray++) {
        Vec3d rayPos(rayPts(iray, 0), rayPts(iray, 1), rayPts(iray, 2));
        Vec3d rayDir(rayDirs(iray, 0), rayDirs(iray, 1), rayDirs(iray, 2));
        Vec3d rayDirNorm = rayDir.normalized();
        Vec3d bestHit;
        float bestU, bestV, bestAngleCos;
        int bestTri = -1;
        float best_t = std::numeric_limits<float>::infinity();
        for(int itri = 0; itri < tris.shape(0); itri++) {
            // put data into Vec3d for easy dot/cross operations
            auto i1 = tris(itri, 0); auto i2 = tris(itri, 1); auto i3 = tris(itri, 2);
            Vec3d v1(verts(i1, 0), verts(i1, 1), verts(i1, 2));
            Vec3d v2(verts(i2, 0), verts(i2, 1), verts(i2, 2));
            Vec3d v3(verts(i3, 0), verts(i3, 1), verts(i3, 2));
            // perform ray-triangle hit test
            Vec3d edge1 = v2 - v1;
            Vec3d edge2 = v3 - v1;
            Vec3d pvec = rayDir.cross(edge2);
            float det = edge1.dot(pvec);
            if(std::abs(det) < std::numeric_limits<double>::epsilon()) {
                continue;
            }
            float invDet = 1.0 / det;
            Vec3d tvec = rayPos - v1;
            float u = tvec.dot(pvec) * invDet;
            if(u < 0.0 || u > 1.0) {
                continue;
            }
            Vec3d qvec = tvec.cross(edge1);
            float v = rayDir.dot(qvec) * invDet;
            if(v < 0.0 || u + v > 1.0) {
                continue;
            }
            float t = edge2.dot(qvec) * invDet;
            Vec3d hitPoint = v1 + edge1*u + edge2*v;
            // check if nearest hit and if it is valid
            if(fabs(t) < fabs(best_t) && (hitPoint - rayPos).norm() < max_distance) {
                bestHit = hitPoint;
                bestU = u; bestV = v;
                bestTri = itri;
                best_t = t;
                Vec3d normal = edge1.normalized().cross(edge2.normalized()).normalized();
                if (allow_backface_hit) {
                    bestAngleCos = fmax(rayDirNorm.dot(normal), rayDirNorm.dot(normal * -1.f));
                }
                else {
                    bestAngleCos = rayDirNorm.dot(normal);
                }
            }
        }
        if(bestTri > -1) {
            // check ray-normal angle
            if(bestAngleCos < maxAngleCos) {
                continue;
            }
            hitUVs.emplace_back(bestU, bestV);
            hitPoints.push_back(bestHit);
            hitRayIndices.push_back(iray);
            hitTris.push_back(bestTri);
        }
    }

    return std::make_tuple(
            py::array_t<int>(hitTris.size(), hitTris.data()),
            py::array_t<double>({(unsigned long)hitUVs.size(), 2ul}, (double*)hitUVs.data()),
            py::array_t<int>(hitRayIndices.size(), hitRayIndices.data()),
            py::array_t<double>({(unsigned long)hitPoints.size(), 3ul}, (double*)hitPoints.data())
    );
}

std::tuple<VectorXi, MatrixX2d, VectorXi>
rayMeshIntersectFast(
        const MatrixXd& verts,
        const MatrixXi& tris,
        const MatrixXd& ray_pts,
        const MatrixXd& ray_dirs)
{
    igl::AABB<MatrixXd, 3> tree;
    tree.init(verts, tris);

    std::vector<std::pair<igl::Hit, int>> hits;
    for (int i = 0; i < ray_pts.rows(); i++) {
        igl::Hit hit;
        if (tree.intersect_ray(verts, tris, ray_pts.row(i), ray_dirs.row(i), hit)) {
            hits.emplace_back(hit, i);
        }
    }

    VectorXi tri_ixs(hits.size());
    MatrixX2d barys(hits.size(), 2);
    VectorXi ray_ixs(hits.size());

    for(int i = 0; i < hits.size(); i++) {
        ray_ixs(i) = hits[i].second;
        const igl::Hit & hit = hits[i].first;
        tri_ixs(i) = hit.id;
        barys.row(i) = Vec2d((double)hit.u, (double)hit.v);
    }

    return std::make_tuple(tri_ixs, barys, ray_ixs);
}

std::tuple<VectorXd, VectorXi, MatrixX3d, MatrixX2d>
closestPointOnMesh(const MatrixXd& P, const MatrixXd& V, const MatrixXi& tris)
{
    VectorXd sq_dists;
    MatrixX3d hit_pts;
    VectorXi tri_ixs;
    igl::point_mesh_squared_distance(P, V, tris, sq_dists, tri_ixs, hit_pts);

    // determine uv coordinates of those closest hits
    MatrixX2d hit_uv(hit_pts.rows(), 2);
    for (int i = 0; i < P.rows(); ++i) {
        Vector3d p = hit_pts.row(i);
        Vector3i tri = tris.row(tri_ixs(i));
        // setup local coordinate frame
        Vector3d e10 = V.row(tri(1)) - V.row(tri(0));
        Vector3d e20 = V.row(tri(2)) - V.row(tri(0));
        Vector3d n = e10.cross(e20); // normal
        Matrix3d F;
        F << e10, e20, n;
        // invert local coordinate frame to get
        // frame-local coordinates
        Vector3d v0 = hit_pts.row(i) - V.row(tri(0));
        Vector3d uvw = F.inverse() * v0;
        hit_uv(i, 0) = uvw(0);
        hit_uv(i, 1) = uvw(1);
    }

    return std::make_tuple(sq_dists, tri_ixs, hit_pts, hit_uv);
};

PYBIND11_PLUGIN(_intersections_ext) {
    using namespace pybind11::literals;

    py::module m("_intersections_ext");
    m.def("ray_mesh_intersect", &rayMeshIntersect,
            "verts"_a, "tris"_a, "ray_pts"_a, "ray_dirs"_a,
            "max_distance"_a = std::numeric_limits<double>::infinity(),
            "max_angle"_a = std::numeric_limits<double>::infinity(),
            "allow_backface_hit"_a = true);

    m.def("ray_mesh_intersect_fast", &rayMeshIntersectFast,
            "verts"_a, "tris"_a, "ray_pts"_a, "ray_dirs"_a);

    m.def("closest_points_on_mesh", &closestPointOnMesh,
          "points"_a, "vertices"_a, "triangles"_a);

    return m.ptr();
}


#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <vector>
#include <math.h>

#include "NRGclasses.hpp"
using namespace std;

//////////////
double CNRGarray::EniBlockj(int i, int j){

  // returns the energy of the i-th state in block j

  double aux=0.0;

  if ( (2*j-2)<BlockBegEnd.size() )
    {
      if ( (BlockBegEnd[2*j-2]+(i-1))<dEn.size() )
	{
	  aux=dEn[BlockBegEnd[2*j-2]+(i-1)];
	}
    }

  return(aux);

}
//////////////
// void CNRGarray::SetsizeBlock(){

//   int ii,i1;
//   for (ii=0;ii<BlockBegEnd.size();ii+=2)
//     {
//       sizeBlock.push_back(BlockBegEnd[ii+1]-BlockBegEnd[ii]+1);
//     }

// }
//////////////
void CNRGarray::PrintQNumbers(){

  int ii;
  cout << "Qnumbers : " << endl;
  for (ii=0;ii<QNumbers.size();ii++)
    cout << QNumbers[ii] << "  ";
  cout << endl;
}

//////////////
void CNRGarray::PrintAll(){

  int ii;
  void CNRGarray::PrintQNumbers();

  PrintQNumbers();

  cout << "dEn : " << endl;
  for (ii=0;ii<dEn.size();ii++)
    cout << dEn[ii] << " ";
  cout << endl;

  cout << "BlockBegEnd : " << endl;
  for (ii=0;ii<BlockBegEnd.size();ii+=2)
    cout << BlockBegEnd[ii] << " "<< BlockBegEnd[ii+1] << ", ";
  cout << endl;

  cout << "Size Blocks : " << endl;
  for (ii=0;ii<NumBlocks();ii++)
    cout << GetBlockSize(ii) << " ";
  cout << endl;
}
//////////////

int CNRGarray::NumBlocks(){

  return((int)(QNumbers.size()/NQNumbers));

}
//////////////
int CNRGarray::Nstates(){

  return(BlockBegEnd[NQNumbers*(NumBlocks()-1)+1]-BlockBegEnd[0]+1);

}

//////////////
void CNRGarray::SetDegen(){

  int nst,ii,jj;

  // Needs work

//   iDegen.push_back(1);
//   ii=1;
//   while ( ii<dEn.size() )
//     {
//       if (fabs(dEn[ii]-dEn[ii-1])<1.0E-10) nst=iDegen[ii-1]+1;
//       else nst=1;
//       iDegen.push_back(nst);
//       ii++;
//     }

}

void CNRGarray::FilterQNumbers(){

  int equal=0;
  int ii,i1;

  vector<double>::iterator qnums_iter,qnums_iter2,qnums_ii1,qnums_ii2;

  for (qnums_iter=QNumbers.begin(); qnums_iter<QNumbers.end();
       qnums_iter+=NQNumbers)
    {

      //      cout << "Iter1 = " << (*qnums_iter) << endl;
      ii=1;
      qnums_iter2=qnums_iter+ii*NQNumbers;
      while (qnums_iter2<QNumbers.end())
	{
	  //  cout << "Iter2 = " << (*qnums_iter2) << endl;
	  equal=1;
	  i1=0;
	  for (qnums_ii2=qnums_iter2;
	       qnums_ii2<qnums_iter2+NQNumbers;qnums_ii2++)
	    {
	      qnums_ii1=qnums_iter+i1;
	      if ( (*qnums_ii2==*qnums_ii1)&&(equal==1) ) equal=1;
	      else equal=0;
	      i1++;
	    }
	  if (equal==1)
	    {
	      QNumbers.erase(qnums_iter2,qnums_iter2+NQNumbers);
	    }
	  ii++;
	  qnums_iter2=qnums_iter+ii*NQNumbers;
	}
    }

}

//////////////
void CNRGarray::ClearAll(){

  QNumbers.clear();
  dEn.clear();
  dEigVec.clear();
  iDegen.clear();

  BlockBegEnd.clear();

}

//////////////
double CNRGarray::GetQNumber(int iblock, int whichqn){

  int iq=iblock*NQNumbers+whichqn;

  return(QNumbers[iq]);
}

//////////////
int CNRGarray::GetBlockLimit(int iblock, int whichlimit){

  return(BlockBegEnd[2*iblock+whichlimit]);

}

//////////////
int CNRGarray::GetBlockSize(int iblock){

  return(BlockBegEnd[2*iblock+1]-BlockBegEnd[2*iblock]+1);

}

////////////////////////////////////
// Class CNRGbasisarray functions //
////////////////////////////////////

//////////////
//void CNRGbasisarray::ClearVecsBasis(){
void CNRGbasisarray::ClearAll(){

  CNRGarray::ClearAll();

  iType.clear();
  StCameFrom.clear();

}
//////////////
void CNRGbasisarray::PrintAll(){

  CNRGarray::PrintAll();

  cout << "Type : " << endl;
  for (int ii=0;ii<iType.size();ii++)
    cout << iType[ii] << " ";
  cout << endl;

  cout << "StCameFrom : " << endl;
  for (int ii=0;ii<StCameFrom.size();ii++)
    cout << StCameFrom[ii] << " ";
  cout << endl;

}
///////////////////

////////////////////////////////////
// Class CNRGmatrix functions //
////////////////////////////////////

// vector<double> CNRGmatrix::GetFullMatrix(int Nst){

//   vector<double> Mat((Nst*Nst),0.0); //All zeros

//   int i1;
//   for (int ii=0;ii<vec.size();ii++)
// 	{
// 	  i1=ij2r(Nst,vec[ii].ist,vec[ii].jst);
// 	  if (i1<Mat.size()) Mat[i1]=vec[ii].val;
// 	  i1=ij2r(Nst,vec[ii].jst,vec[ii].ist);
// 	  if (i1<Mat.size()) Mat[i1]=vec[ii].val;
// 	}

//   return(Mat);

// }
//////////////
int CNRGmatrix::GetMatBlockLimit(int iblock1, int iblock2, int whichlimit){

  int ii=0;

  while ( ((iblock1!=MatBlockMap[ii])||(iblock2!=MatBlockMap[ii+1]))
	  &&(ii<MatBlockMap.size()) ) ii+=2;

  return(MatBlockBegEnd[2*ii+whichlimit]);

}

//////////////
int CNRGmatrix::GetMatBlockSize(int iblock1, int iblock2){

  int ii=0;

  while ( ((iblock1!=MatBlockMap[ii])||(iblock2!=MatBlockMap[ii+1]))
	  &&(ii<MatBlockMap.size()) ) ii+=2;

  return(MatBlockBegEnd[2*ii+1]-MatBlockBegEnd[2*ii]+1);

}

void CNRGmatrix::DiagBlock(int iblock,
			   vector<double> &eigvalues,
			   vector<double> &eigvectors){

  // LAPACK variables
  char JOBZ='N',UPLO='U';
  int Nst;
  int INFO=0,LDA,LWORK;

  // LDA=Nst
  // LWORK=3*Nst-1

  double *A;
  double *eigv;
  double *work;

  // Get Block info

  Nst=GetMatBlockSize(iblock,iblock);
//   cout << "Size block                 : " << GetBlockSize(iblock) << endl;
//   cout << "No of independent elements : " << GetMatrixBlockSize(iblock)
//        << endl;

  //Nst=2;
  LDA=Nst;
  LWORK=3*Nst-1;
  eigv = new double [Nst];
  work = new double [LWORK];
  A = new double[Nst*Nst];

  // I'll work on a routine to map upper triangular to r.
  // For now, let's use this:
  double **Aaux;
  Aaux = new double*[Nst];
  for (int ii=0;ii<Nst;ii++) Aaux[ii]=new double[Nst];

  for (int ii=0;ii<MatEl.size();ii++)
    {
//       // This takes a tremendous amount of CPU time! Must do better!
//       if ( (vec[ii].Bli==iblock)&&(vec[ii].Blj==iblock) )
// 	{
// 	  // Awful construcion. An invitation to segmentation faults.
// 	  Aaux[vec[ii].ibl][vec[ii].jbl]=vec[ii].val;
// 	  Aaux[vec[ii].jbl][vec[ii].ibl]=vec[ii].val;
// 	}
    }
  int i1=0;
  for (int ii=0;ii<Nst;ii++)
    {
      for (int jj=0;jj<Nst;jj++)
	{
	  A[i1]=Aaux[ii][jj];
	  i1++;
	}
    }
//   int i1=0;
//   int r1=0;
//   for (int i=GetBlockLimit(iblock,0);
//            i<GetBlockLimit(iblock,1);i++)
//     {
//       for (int j=i;
//            j<GetBlockLimit(iblock,1);j++)
// 	{
// 	  r1=ij2r(Nstates(),i,j);
//       Aaux[]=vec[r1].val;
// 	}
//       i1++;
//     }

//   A[0]=1.0;
//   A[1]=-sqrt(2.0);
//   A[2]=-sqrt(2.0);
//   A[3]=1.0;

  dsyev_(&JOBZ,&UPLO,&Nst,&A[0],&LDA,eigv,work,&LWORK,&INFO);

  for (int ii=0;ii<Nst;ii++)
    {
      eigvalues.push_back(eigv[ii]);
      eigvectors.push_back(eigv[ii]);
      cout << " Block  : " << iblock << endl;
      std::cout << "E = " << eigvalues[ii] << std::endl;
    }

  delete[] eigv;
  delete[] work;
  delete[] A;

  for (int ii=0;ii<Nst;ii++) delete[] Aaux[ii];
  delete[] Aaux;

}


#include "modprop/compo/AdditiveWrapper.hpp"
#include "modprop/neural/LinearLayer.hpp"
#include "modprop/neural/FullyConnectedNet.hpp"

#include "modprop/neural/HingeActivation.hpp"
#include "modprop/neural/SigmoidActivation.hpp"
#include "modprop/neural/NullActivation.hpp"
#include "modprop/neural/NetworkTypes.h"

#include "modprop/optim/SquaredLoss.hpp"
#include "modprop/optim/StochasticMeanCost.hpp"
#include "modprop/optim/ParameterL2Cost.hpp"

#include "optim/Optimizers.h"

#include <boost/random/uniform_real_distribution.hpp>
#include <boost/random/random_device.hpp>
#include <boost/random/mersenne_twister.hpp>

#include "modprop/utils/Derivatives.hpp"
#include "modprop/utils/Randomization.hpp"

#include <deque>
#include <cstdlib>
#include <iostream>

using namespace percepto;

// Comment the above and uncomment below to use Rectified Linear Units instead
typedef PerceptronNet TestNet;
// typedef ReLUNet TestNet;

unsigned int inputDim = 1;
unsigned int outputDim = 1;
unsigned int numHiddenLayers = 3;
unsigned int layerWidth = 20;
unsigned int batchSize = 20;

struct Regressor
{
	TerminalSource<VectorType> netInput;
	TestNet net;
	SquaredLoss<VectorType> loss;

	Regressor()
	: net( inputDim, outputDim, numHiddenLayers, layerWidth,
	       SigmoidActivation(), TestNet::OUTPUT_UNRECTIFIED )
	{
		net.SetSource( &netInput );
		loss.SetSource( &net.GetOutputSource() );
	}

	Regressor( const Regressor& other )
	: net( other.net )
	{
		net.SetSource( &netInput );
		loss.SetSource( &net.GetOutputSource() );
	}

	void Foreprop()
	{
		netInput.Foreprop();
	}

	double GetOutput()
	{
		return net.GetOutput()(0);
	}

	void Invalidate()
	{
		netInput.Invalidate();
	}
};

struct RegressionProblem
: public NaturalOptimizationProblem
{
	std::deque<Regressor> regressors;
	StochasticMeanCost<double> losses;
	ParameterL2Cost regularizer;
	AdditiveWrapper<double> objective;

	Parameters::Ptr params;

	RegressionProblem( Parameters::Ptr p,
	                   unsigned int batchSize,
	                   double l2Weight )
	{
		params = p;
		losses.SetBatchSize( batchSize );
		regularizer.SetWeight( l2Weight );
		regularizer.SetParameters( params );
		objective.SetSourceA( &losses );
		objective.SetSourceB( &regularizer );
	}

	virtual bool IsMinimization() const { return true; }

	virtual void Resample()
	{
		losses.Resample();
	}

	virtual double ComputeObjective()
	{
		Invalidate();
		Foreprop();
		return objective.GetOutput();
	}

	virtual VectorType ComputeGradient()
	{
		Invalidate();
		Foreprop();
		Backprop();
		return params->GetDerivs();
	}

	virtual VectorType ComputeNaturalGradient()
	{
		Invalidate();
		Foreprop();
		BackpropNatural();
		return params->GetDerivs();
	}

	virtual VectorType GetParameters() const
	{
		return params->GetParamsVec();
	}

	virtual void SetParameters( const VectorType& p )
	{
		params->SetParamsVec( p );
	}

	void Invalidate()
	{
		for( unsigned int i = 0; i < regressors.size(); i++ )
		{
			regressors[i].Invalidate();
		}
		regularizer.Invalidate();
		params->ResetAccumulators();
	}

	void Foreprop()
	{
		const std::vector<unsigned int>& inds = losses.GetActiveInds();
		BOOST_FOREACH( unsigned int i, inds )
		{
			regressors[i].Foreprop();
		}
		regularizer.Foreprop();
	}

	void ForepropAll()
	{
		for( unsigned int i = 0; i < regressors.size(); i++ )
		{
			regressors[i].Foreprop();
		}
		regularizer.Foreprop();
	}

	void Backprop()
	{
		objective.Backprop( MatrixType::Identity(1,1) );
	}

	void BackpropNatural()
	{
		const std::vector<unsigned int>& inds = losses.GetActiveInds();
		MatrixType dodw = MatrixType::Identity( 1, 1 ) / inds.size();
		BOOST_FOREACH( unsigned int i, inds )
		{
			regressors[i].loss.Backprop( dodw );
		}
	}
};

// The highly nonlinear function we will try to fit
double f( double x )
{
	return 8 * std::cos( x ) + 2.5 * x * sin( x ) + 2.8 * x;
}

// Generate a specified number of random xs and corresponding ys
void generate_data( std::vector<VectorType>& xs,
                    std::vector<VectorType>& ys,
                    unsigned int num_data )
{
	boost::random::mt19937 generator;
	boost::random::random_device rng;
	generator.seed( rng );
	boost::random::uniform_real_distribution<> xDist( -5.0, 5.0 );

	xs.clear();
	ys.clear();
	xs.reserve( num_data );
	ys.reserve( num_data );

	VectorType x(1,1);
	VectorType y(1,1);
	for( unsigned int i = 0; i < num_data; i++ )
	{
		x(0) = xDist( generator );
		y(0) = f( x(0) );
		xs.push_back( x );
		ys.push_back( y );
	}
}

int main( int argc, char** argv )
{
	unsigned int numTrain = 150;
	unsigned int numTest = 200;
	double l2Weight = 1E-3;

	std::vector<VectorType> xTest, yTest, xTrain, yTrain;
	generate_data( xTest, yTest, numTest );
	generate_data( xTrain, yTrain, numTrain );

	std::cout << "Initializing net..." << std::endl;
	std::cout << "Creating linear layers..." << std::endl;

	// // ReLU initialization
	Regressor reg;
	Parameters::Ptr params = reg.net.CreateParameters();

	// Randomize parameters
	VectorType p( params->ParamDim() );
	randomize_vector( p, -0.2, 0.2 );
	params->SetParamsVec( p );
	std::cout << "Initial net: " << std::endl << reg.net << std::endl;

	// Create the loss functions
	std::cout << "Generating losses..." << std::endl;
	RegressionProblem trainProblem( params, batchSize, l2Weight );
	RegressionProblem testProblem( params, batchSize, l2Weight );

	// NOTE If we don't reserve, the vector resizing and moving may
	// invalidate the references. Alternatively we can use a deque
	for( unsigned int i = 0; i < numTrain; i++ )
	{
		trainProblem.regressors.emplace_back( reg );
		trainProblem.regressors[i].netInput.SetOutput( xTrain[i] );
		trainProblem.regressors[i].loss.SetTarget( yTrain[i] );
		trainProblem.losses.AddSource( &trainProblem.regressors[i].loss );
	}

	for( unsigned int i = 0; i < numTest; i++ )
	{
		testProblem.regressors.emplace_back( reg );
		testProblem.regressors[i].netInput.SetOutput( xTest[i] );
		testProblem.regressors[i].loss.SetTarget( yTest[i] );
		testProblem.losses.AddSource( &testProblem.regressors[i].loss );
	}

	ModularOptimizer optimizer;

	// AdamSearchDirector::Ptr director = std::make_shared<AdamSearchDirector>();
	// GradientSearchDirector::Ptr director = std::make_shared<GradientSearchDirector>();
	NaturalSearchDirector::Ptr director = std::make_shared<NaturalSearchDirector>();
	optimizer.SetSearchDirector( director );

	BacktrackingSearchStepper::Ptr stepper = std::make_shared<BacktrackingSearchStepper>();
	stepper->SetInitialStep( 1E-1 );
	stepper->SetBacktrackingRatio( 0.5 );
	stepper->SetMaxBacktracks( 20 );
	stepper->SetImprovementRatio( 0.75 );

	// L1ConstrainedSearchStepper::Ptr stepper = std::make_shared<L1ConstrainedSearchStepper>();
	// stepper->SetMaxL1Norm( 1E-1 );
	// stepper->SetStepSize( 1E-2 );
	optimizer.SetSearchStepper( stepper );

	RuntimeTerminationChecker::Ptr runtimeChecker = std::make_shared<RuntimeTerminationChecker>();
	runtimeChecker->SetMaxRuntime( 20 );
	optimizer.AddTerminationChecker( runtimeChecker );

	// GradientTerminationChecker::Ptr gradientChecker = std::make_shared<GradientTerminationChecker>();
	// gradientChecker->SetMinGradientNorm( 1E-9 );
	// optimizer.AddTerminationChecker( gradientChecker );

	// IterationTerminationChecker::Ptr iterationChecker = std::make_shared<IterationTerminationChecker>();
	// iterationChecker->SetMaxIterations( 1e3 );
	// optimizer.AddTerminationChecker( iterationChecker );

	trainProblem.Invalidate();
	testProblem.Invalidate();
	trainProblem.ForepropAll();
	testProblem.ForepropAll();
	trainProblem.losses.ParentCost::Foreprop();
	testProblem.losses.ParentCost::Foreprop();
	std::cout << "initial train avg loss: " << trainProblem.losses.ParentCost::GetOutput() << std::endl;
	std::cout << "initial train max loss: " << trainProblem.losses.ParentCost::ComputeMax() << std::endl;
	std::cout << "initial test avg loss: " << testProblem.losses.ParentCost::GetOutput() << std::endl;
	std::cout << "initial test max loss: " << testProblem.losses.ParentCost::ComputeMax() << std::endl;

	std::cout << "Beginning optimization..." << std::endl;
	optimizer.ResetAll();
	trainProblem.Resample();
	OptimizationResults result = optimizer.Optimize( trainProblem );
	std::cout << "Terminated with condition: " << result.status << std::endl;
	std::cout << "Final net: " << std::endl << reg.net << std::endl;

	trainProblem.Invalidate();
	testProblem.Invalidate();
	trainProblem.ForepropAll();
	testProblem.ForepropAll();
	trainProblem.losses.ParentCost::Foreprop();
	testProblem.losses.ParentCost::Foreprop();
	std::cout << "train avg loss: " << trainProblem.losses.ParentCost::GetOutput() << std::endl;
	std::cout << "train max loss: " << trainProblem.losses.ParentCost::ComputeMax() << std::endl;
	std::cout << "test avg loss: " << testProblem.losses.ParentCost::GetOutput() << std::endl;
	std::cout << "test max loss: " << testProblem.losses.ParentCost::ComputeMax() << std::endl;

	for( unsigned int i = 0; i < numTest; i++ )
	{
		std::cout << "xtest: " << xTest[i] << " ytest: " << yTest[i]
		          << " regout: " << testProblem.regressors[i].GetOutput() << std::endl;
	}

	return 0;
}


#include <armadillo>
#include <emissions.hpp>
#include <ForwardBackward.hpp>
#include <cmath>
#include <json.hpp>
#include <iostream>
#include <vector>

using namespace arma;
using namespace std;
using json = nlohmann::json;

namespace hsmm {

    double gaussianlogpdf_(double x, double mu, double sigma) {
        double ret = ((x - mu)*(x - mu)) / (-2*sigma*sigma);
        ret = ret - log(sqrt(2 * M_PI) * sigma);
        return ret;
    }

    double gaussianpdf_(double x, double mu, double sigma) {
        return exp(gaussianlogpdf_(x, mu, sigma));
    }

    field<mat> fromMatToField(const mat& obs) {
        field<mat> ret(obs.n_cols);
        for(int i = 0; i < obs.n_cols; i++)
            ret(i) = obs.col(i);
        return ret;
    }

    mat fromFieldToMat(const field<mat>& obs) {
        int total_cols = 0;
        int total_rows = obs(0).n_rows;
        for(auto& m : obs) {
            total_cols += m.n_cols;
            assert(m.n_rows == total_rows);
        }
        mat ret(total_rows, total_cols);
        int idx = 0;
        for(auto& m: obs) {
            ret.cols(idx, m.n_cols - 1) = m;
            idx += m.n_cols;
        }
        return ret;
    }

    /**
     * Abstract emission implementation.
     */
    AbstractEmission::AbstractEmission(int nstates, int dimension) :
            nstates_(nstates), dimension_(dimension) {}

    int AbstractEmission::getNumberStates() const {
        return nstates_;
    }

    int AbstractEmission::getDimension() const {
        return dimension_;
    }

    cube AbstractEmission::likelihoodCube(int min_duration, int ndurations,
            const field<mat>& obs) const {
        return exp(loglikelihoodCube(min_duration, ndurations, obs));
    }

    // This should return a cube of dimensions (nstates, nobs, ndurations)
    // where the entry (i, j, k) is the log-likelihood of the observations
    // in the interval [j, min_duration + k - 1] being produced by state i.
    cube AbstractEmission::loglikelihoodCube(int min_duration, int ndurations,
            const field<mat>& obs) const {
        int nobs = obs.n_elem;
        cube pdf(getNumberStates(), nobs, ndurations);
        pdf.fill(-datum::inf);
        for(int i = 0; i < getNumberStates(); i++)
            for(int t = 0; t < nobs; t++)
                for(int d = 0; d < ndurations; d++) {
                    if (t + min_duration + d > nobs)
                        break;
                    int end_idx = t + min_duration + d - 1;
                    pdf(i, t, d) = loglikelihood(i, obs.rows(t, end_idx));
                }
        return pdf;
    }

    field<mat> AbstractEmission::sampleFromState(int state, int size) {
        return sampleFromState(state, size, rand_generator_);
    }

    json AbstractEmission::to_stream() const {
        cout << "Warning: serialization of emission parameters not implemented"
                << endl;
        return json::object();  // empty json object by default.
    }

    void AbstractEmission::from_stream(const json& emission_params) {
        cout << "Warning: not updating the emission parameters from the stream"
                << endl;
        return;
    }

    /*
     * AbstractEmissionOnlineSetting implementation
     */
    field<mat> AbstractEmissionOnlineSetting::sampleNextObsGivenPastObs(int state,
            int seg_dur, const field<mat>& past_obs) {
        return sampleNextObsGivenPastObs(state, seg_dur, past_obs,
                rand_generator_);
    }

    // Default implementation. Disregards any information from the last
    // segment.
    field<mat> AbstractEmissionOnlineSetting::sampleFirstSegmentObsGivenLastSegment(
            int curr_state, int curr_seg_dur, const field<mat> &last_segment,
            int last_state, std::mt19937 &rng) const {
        field<mat> empty_segment;
        return sampleNextObsGivenPastObs(curr_state, curr_seg_dur,
                empty_segment, rng);
    }

    field<mat> AbstractEmissionOnlineSetting::sampleFirstSegmentObsGivenLastSegment(
            int curr_state, int curr_seg_dur, const field<mat> &last_segment,
            int last_state) {
        return sampleFirstSegmentObsGivenLastSegment(curr_state, curr_seg_dur,
                last_segment, last_state, rand_generator_);
    }

    /**
     * AbstractEmissionIIDobs implementation
     */
    AbstractEmissionConditionalIIDobs::AbstractEmissionConditionalIIDobs(
            int nstates, int dimension) :
            AbstractEmission(nstates, dimension) {}

    double AbstractEmissionConditionalIIDobs::loglikelihood(int state,
            const field<mat>& obs) const {
        double ret = 0;
        int seg_dur = obs.n_elem;
        for(int i = 0; i < seg_dur; i++) {

            // If there are missing outputs just skip them.
            if (!obs(i).is_empty())
                ret += loglikelihoodIIDobs(state, seg_dur, i, obs(i));
        }
        return ret;
    }

    /**
     * DummyGaussianEmission implementation.
     */
    DummyGaussianEmission::DummyGaussianEmission(vec& means, vec& std_devs) :
            AbstractEmissionOnlineSetting(means.n_elem, 1), means_(means),
            std_devs_(std_devs) {
        assert(means_.n_elem == std_devs_.n_elem);
    }

    DummyGaussianEmission* DummyGaussianEmission::clone() const {
        return new DummyGaussianEmission(*this);
    }

    double DummyGaussianEmission::loglikelihood(int state,
            const field<mat>& obs) const {
        double ret = 0;
        int seg_dur = obs.n_elem;
        for(int i = 0; i < seg_dur; i++) {

            // If there are missing outputs just skip them.
            if (!obs(i).is_empty())
                ret += loglikelihoodIIDobs(state, seg_dur, i, obs(i));
        }
        return ret;
    }

    double DummyGaussianEmission::loglikelihoodIIDobs(int state, int seg_dur,
            int offset, const mat& single_obs) const {
        assert(single_obs.n_rows == 1 && single_obs.n_cols == 1);
        return gaussianlogpdf_(single_obs(0, 0), means_(state),
                std_devs_(state));
    }

    json DummyGaussianEmission::to_stream() const {
        vector<double> means = conv_to<vector<double>>::from(means_);
        vector<double> std_devs = conv_to<vector<double>>::from(std_devs_);
        json ret;
        ret["means"] = means;
        ret["std_devs"] = std_devs;
        return ret;
    }

    void DummyGaussianEmission::reestimate(int min_duration,
            const field<cube>& meta, const field<field<mat>>& mobs) {
        int nseq = mobs.n_elem;
        for(int i = 0; i < getNumberStates(); i++) {

            // Reestimating the mean.
            vector<double> num_mult;
            vector<double> num_obs;
            for(int s = 0; s < nseq; s++) {
                auto& obs = mobs(s);
                int nobs = obs.n_elem;
                const cube& eta = meta(s);
                int ndurations = eta.n_cols;
                for(int t = min_duration - 1; t < nobs; t++) {
                    for(int d = 0; d < ndurations; d++) {
                        int first_idx_seg = t - min_duration - d + 1;
                        if (first_idx_seg < 0)
                            break;

                        // Since the observations factorize given t, d and i.
                        for(int k = first_idx_seg; k <= t; k++) {
                            num_mult.push_back(eta(i, d, t));
                            num_obs.push_back(obs(k)(0,0));
                        }
                    }
                }
            }
            vec num_mult_v(num_mult);
            vec num_obs_v(num_obs);
            num_mult_v = num_mult_v - logsumexp(num_mult_v);
            num_mult_v = exp(num_mult_v);
            double new_mean = dot(num_mult_v, num_obs_v);

            // Reestimating the variance.
            vector<double> num_obs_var;
            for(int s = 0; s < nseq; s++) {
                const auto& obs = mobs(s);
                int nobs = obs.n_elem;
                const cube& eta = meta(s);
                int ndurations = eta.n_cols;
                for(int t = min_duration - 1; t < nobs; t++) {
                    for(int d = 0; d < ndurations; d++) {
                        int first_idx_seg = t - min_duration - d + 1;
                        if (first_idx_seg < 0)
                            break;

                        // Since the observations factorize given t, d and i.
                        for(int k = first_idx_seg; k <= t; k++) {
                            double diff = (obs(k)(0,0) - new_mean);
                            num_obs_var.push_back(diff * diff);
                        }
                    }
                }
            }
            vec num_obs_var_v(num_obs_var);
            double new_variance = dot(num_mult_v, num_obs_var_v);

            means_(i) = new_mean;
            std_devs_(i) = sqrt(new_variance);
        }
    }

    field<mat> DummyGaussianEmission::sampleFromState(int state,
            int size, mt19937 &rng) const {
        mat ret = randn<mat>(1, size) * std_devs_(state) + means_(state);
        return fromMatToField(ret);
    }

    field<mat> DummyGaussianEmission::sampleNextObsGivenPastObs(int state,
            int seg_dur, const field<mat>& past_obs, mt19937 &rng) const {
        assert(past_obs.n_elem < seg_dur);
        field<mat> ret(seg_dur);
        int idx = 0;
        for(int i = 0; i < past_obs.n_elem; i++, idx++)
            ret(i) = past_obs(i);
        for(int i = idx; i < seg_dur; i++)
            ret(i) = randn<mat>(1, 1) * std_devs_(state) + means_(state);
        return ret;
    }

    /**
     * DummyMultivariateGaussianEmission implementation.
     */
    DummyMultivariateGaussianEmission::DummyMultivariateGaussianEmission(
            mat& means, double std_dev_output_noise) :
            AbstractEmission(means.n_rows, means.n_cols), means_(means),
            std_dev_output_noise_(std_dev_output_noise) {}

    DummyMultivariateGaussianEmission* DummyMultivariateGaussianEmission::
            clone() const {
        return new DummyMultivariateGaussianEmission(*this);
    }

    double DummyMultivariateGaussianEmission::loglikelihood(int state,
                const field<mat>& obs) const {
        mat copy_obs = fromFieldToMat(obs);
        assert(copy_obs.n_rows == getDimension());
        int size = copy_obs.n_cols;
        for(int i = 0; i < getDimension(); i++)
            copy_obs.row(i) -= linspace<rowvec>(0.0, 1.0, size) +
                    means_(state, i);
        double ret = 0.0;
        for(int i = 0; i < getDimension(); i++)
            for(int j = 0; j < size; j++)
                ret += gaussianlogpdf_(copy_obs(i, j), 0,
                        std_dev_output_noise_);
        return ret;
    }

    void DummyMultivariateGaussianEmission::reestimate(int min_duration,
            const field<cube>& eta, const field<field<mat>>& mobs) {
        // TODO.
    }

    field<mat> DummyMultivariateGaussianEmission::sampleFromState(
            int state, int size, mt19937 &rng) const {
        mat ret = randn<mat>(getDimension(), size) * std_dev_output_noise_;
        for(int i = 0; i < getDimension(); i++)
            ret.row(i) += linspace<rowvec>(0.0, 1.0, size) + means_(state, i);
        return fromMatToField(ret);
    }

};


/* -*- mode: c++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

/*
 Copyright (C) 2015 Johannes Goettker-Schnetmann
 Copyright (C) 2015 Klaus Spanderen

 This file is part of QuantLib, a free-software/open-source library
 for financial quantitative analysts and developers - http://quantlib.org/

 QuantLib is free software: you can redistribute it and/or modify it
 under the terms of the QuantLib license.  You should have received a
 copy of the license along with this program; if not, please email
 <quantlib-dev@lists.sf.net>. The license is also available online at
 <http://quantlib.org/license.shtml>.

 This program is distributed in the hope that it will be useful, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 FOR A PARTICULAR PURPOSE.  See the license for more details.
*/

#include <ql/math/functional.hpp>
#include <ql/time/calendars/nullcalendar.hpp>
#include <ql/processes/hestonprocess.hpp>
#include <ql/processes/blackscholesprocess.hpp>
#include <ql/math/integrals/gausslobattointegral.hpp>
#include <ql/termstructures/volatility/equityfx/blackconstantvol.hpp>
#include <ql/experimental/finitedifferences/bsmrndcalculator.hpp>
#include <ql/experimental/finitedifferences/hestonrndcalculator.hpp>

#if defined(__GNUC__) && (((__GNUC__ == 4) && (__GNUC_MINOR__ >= 8)) || (__GNUC__ > 4))
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-local-typedefs"
#endif
#include <boost/bind.hpp>
#if defined(__GNUC__) && (((__GNUC__ == 4) && (__GNUC_MINOR__ >= 8)) || (__GNUC__ > 4))
#pragma GCC diagnostic pop
#endif

#include <complex>

namespace QuantLib {

namespace {
        struct HestonParams {
            Real v0, kappa, theta, sigma, rho;
        };

        HestonParams getHestonParams(
            const ext::shared_ptr<HestonProcess>& process) {
            const HestonParams p = { process->v0(),    process->kappa(),
                                     process->theta(), process->sigma(),
                                     process->rho() };
            return p;
        }

        std::complex<Real> gamma(const HestonParams& p, Real p_x) {
            return std::complex<Real>(p.kappa, p.rho*p.sigma*p_x);
        }

        std::complex<Real> omega(const HestonParams& p, Real p_x) {
           const std::complex<Real> g = gamma(p, p_x);
           return std::sqrt(g*g
                  + p.sigma*p.sigma*std::complex<Real>(p_x*p_x, -p_x));
        }

        class CpxPv_Helper
            : public std::unary_function<Real, Real > {
          public:
            CpxPv_Helper(const HestonParams& p, Real x, Time t)
              : p_(p), t_(t), x_(x),
                c_inf_(std::min(10.0, std::max(0.0001,
                      std::sqrt(1.0-square<Real>()(p_.rho))/p_.sigma))
                      *(p_.v0 + p_.kappa*p_.theta*t))  {}

            Real operator()(Real x) const {
                return std::real(transformPhi(x));
            }

            Real p0(Real p_x) const {
                if (p_x < QL_EPSILON) {
                    return 0.0;
                }

                const Real u_x = std::max(QL_EPSILON, -std::log(p_x)/c_inf_);
                return std::real(phi(u_x)
                        /((p_x*c_inf_)*std::complex<Real>(0.0, u_x)));
            }

          private:
            std::complex<Real> transformPhi(Real x) const {
                if (x < QL_EPSILON) {
                    return std::complex<Real>(0.0, 0.0);
                }

                const Real u_x = -std::log(x)/c_inf_;
                return phi(u_x)/(x*c_inf_);
            }

            std::complex<Real> phi(Real p_x) const {
                const Real sigma2 = p_.sigma*p_.sigma;
                const std::complex<Real> g = gamma(p_, p_x);
                const std::complex<Real> o = omega(p_, p_x);
                const std::complex<Real> gamma = (g-o)/(g+o);

                return 2.0*std::exp(std::complex<Real>(0.0, p_x*x_)
                        - p_.v0*std::complex<Real>(p_x*p_x, -p_x)
                          /(g+o*(1.0+std::exp(-o*t_))/(1.0-std::exp(-o*t_)))
                         +p_.kappa*p_.theta/sigma2*(
                           (g-o)*t_ - 2.0*std::log((1.0-gamma*std::exp(-o*t_))
                                                               /(1.0-gamma))));
            }

            const HestonParams& p_;
            const Time t_;
            const Real x_, c_inf_;
        };
    }

    HestonRNDCalculator::HestonRNDCalculator(
        const ext::shared_ptr<HestonProcess>& hestonProcess,
        Real integrationEps, Size maxIntegrationIterations)
    : hestonProcess_(hestonProcess),
      x0_(std::log(hestonProcess_->s0()->value())),
      integrationEps_(integrationEps),
      maxIntegrationIterations_(maxIntegrationIterations) { }

    Real HestonRNDCalculator::x_t(Real x, Time t) const {
        const DiscountFactor dr = hestonProcess_->riskFreeRate()->discount(t);
        const DiscountFactor dq = hestonProcess_->dividendYield()->discount(t);

        return x - x0_ + std::log(dr/dq);
    }

    Real HestonRNDCalculator::pdf(Real x, Time t) const {
        return GaussLobattoIntegral(
            maxIntegrationIterations_, 0.1*integrationEps_)(
            CpxPv_Helper(getHestonParams(hestonProcess_), x_t(x, t), t),
            0.0, 1.0)/M_TWOPI;
    }

    Real HestonRNDCalculator::cdf(Real x, Time t) const {
        return GaussLobattoIntegral(
            maxIntegrationIterations_, 0.1*integrationEps_)(
            boost::bind(&CpxPv_Helper::p0,
                CpxPv_Helper(getHestonParams(hestonProcess_), x_t(x,t),t),_1),
            0.0, 1.0)/M_TWOPI + 0.5;

    }
    Real HestonRNDCalculator::invcdf(Real p, Time t) const {
        const Real v0    = hestonProcess_->v0();
        const Real kappa = hestonProcess_->kappa();
        const Real theta = hestonProcess_->theta();

        const Volatility expVol
            = std::sqrt(theta + (v0-theta)*(1-std::exp(-kappa*t))/(t*kappa));

        const ext::shared_ptr<BlackScholesMertonProcess> bsmProcess(
            ext::make_shared<BlackScholesMertonProcess>(
                hestonProcess_->s0(),
                hestonProcess_->dividendYield(),
                hestonProcess_->riskFreeRate(),
                Handle<BlackVolTermStructure>(
                    ext::make_shared<BlackConstantVol>(
                            hestonProcess_->riskFreeRate()->referenceDate(),
                            NullCalendar(),
                            expVol,
                            hestonProcess_->riskFreeRate()->dayCounter()))));

        const Real guess = BSMRNDCalculator(bsmProcess).invcdf(p, t);

        return RiskNeutralDensityCalculator::InvCDFHelper(
            this, guess, 0.1*integrationEps_, maxIntegrationIterations_)
            .inverseCDF(p, t);
    }
}