/**
 * @file   Optimizer.cpp
 * @brief  Implementations of Optimizer class for pose & map data optimization.
 * @author Charlie Li
 * @date   2019.09.17
 */

#include "Optimizer.hpp"

#include <map>
#include <memory>
#include <set>
#include <vector>

#include <g2o/core/sparse_optimizer.h>
#include <g2o/core/block_solver.h>
//#include <g2o/core/factory.h>
#include <g2o/core/optimization_algorithm_levenberg.h>
#include <g2o/core/robust_kernel_impl.h> // g2o::RobustKernelHuber, ...
#include <g2o/core/solver.h>
//#include <g2o/solvers/cholmod/linear_solver_cholmod.h> // for global BA
#include <g2o/solvers/csparse/linear_solver_csparse.h> // for global BA
#include <g2o/solvers/dense/linear_solver_dense.h> // for pose optimization
#include <g2o/solvers/eigen/linear_solver_eigen.h> // for local BA
#include <g2o/types/sba/types_six_dof_expmap.h> // g2o::EdgeSE3ProjectXYZ, ...
#include <opencv2/core.hpp>
#include <opencv2/core/eigen.hpp> // cv::cv2eigen()
#include <Eigen/Core>
#include <Eigen/Geometry> // Eigen::Quaternion
#include "Config.hpp"
#include "Frame.hpp"
#include "KeyFrame.hpp"
#include "Map.hpp"
#include "MapPoint.hpp"
#include "Tracker.hpp"

namespace SLAM_demo {

using std::map;
using std::make_shared;
using std::set;
using std::shared_ptr;
using std::vector;
using cv::Mat;

const int Optimizer::TH_MIN_NUM_MAPPOINT = 5;
const float Optimizer::TH_MAX_CHI2_FACTOR = 9.0f;

Optimizer::Optimizer(const std::shared_ptr<Map>& pMap) : mpMap(pMap) {}

int Optimizer::globalBundleAdjustment(unsigned nKFs, int nIter,
                                      bool bRobust) const
{
    // set definition of solver
    g2o::BlockSolver_6_3::LinearSolverType* pLinearSolver;
    pLinearSolver =
        new g2o::LinearSolverCSparse<g2o::BlockSolver_6_3::PoseMatrixType>();
    g2o::BlockSolver_6_3* pBlkSolver;
    pBlkSolver = new g2o::BlockSolver_6_3(pLinearSolver);
    g2o::OptimizationAlgorithmLevenberg* pSolverLM;
    pSolverLM = new g2o::OptimizationAlgorithmLevenberg(pBlkSolver);
    // configure optimizer
    g2o::SparseOptimizer optimizer;
    optimizer.setVerbose(false);
    optimizer.setAlgorithm(pSolverLM);

    // add vertices: poses
    vector<shared_ptr<KeyFrame>> vpKFs = mpMap->getLastNKFs(nKFs);
    unsigned idxKFMax = 0;
    for (const auto& pKF : vpKFs) {
        g2o::VertexSE3Expmap* pVSE3 = new g2o::VertexSE3Expmap();
        Mat Tcw = pKF->mPose.getPose();
        unsigned idxKF = pKF->index();
        pVSE3->setEstimate(cvMat2SE3Quat(Tcw));
        pVSE3->setId(idxKF);
        bool bFixed = idxKF == 0;
        pVSE3->setFixed(bFixed);
        optimizer.addVertex(pVSE3);
        if (idxKF > idxKFMax) {
            idxKFMax = idxKF;
        }
    }
    // add vertices: map points, and add edges for each map point
    vector<shared_ptr<MapPoint>> vpMPts = mpMap->getAllMPts();
    int nMPts = vpMPts.size();
    // skip optimization if map points are not enough
    if (nMPts < TH_MIN_NUM_MAPPOINT) {
        return 0;
    }
    vector<bool> vbMPtOptimized(nMPts, true);
    vector<vector<g2o::EdgeSE3ProjectXYZ*>> vvpEdges(nMPts);
    for (int i = 0; i < nMPts; ++i) {
        const shared_ptr<MapPoint>& pMPt = vpMPts[i];
        g2o::VertexSBAPointXYZ* pVPt = new g2o::VertexSBAPointXYZ();
        pVPt->setEstimate(cvMat2Vector3d(pMPt->X3D()));
        pVPt->setId(idxKFMax + 1 + i);
        pVPt->setMarginalized(true); // why?? (to decrease the size of Hessian?)
        optimizer.addVertex(pVPt);

        // add edges
        vector<shared_ptr<KeyFrame>> vpKFsMPt = (nKFs == 0) ?
            pMPt->getRelatedKFs() : vpKFs;
        int nKFsMpt = vpKFsMPt.size();
        vvpEdges[i].resize(nKFsMpt, nullptr);
        bool bHasEdge = false; // check whether the vertec has edges
        for (int j = 0; j < nKFsMpt; ++j) {
            auto& pKF = vpKFsMPt[j];
            // check whether the map point is observed by the target keyframe
            if (!pMPt->isObservedBy(pKF)) {
                continue;
            }
            // form an edge
            bHasEdge = true;
            cv::KeyPoint kpt = pMPt->keypoint(pKF);
            Eigen::Matrix<double, 2, 1> obs;
            obs << kpt.pt.x, kpt.pt.y;
            g2o::EdgeSE3ProjectXYZ* pEdge = new g2o::EdgeSE3ProjectXYZ();
            // vertex 0: map point
            pEdge->setVertex(0, dynamic_cast<g2o::OptimizableGraph::Vertex*>(
                                 pVPt));
            // vertex 1: pose
            pEdge->setVertex(1, dynamic_cast<g2o::OptimizableGraph::Vertex*>(
                                 optimizer.vertex(pKF->index())));
            pEdge->setMeasurement(obs);
            // set element in information matrix (value = 1 / sigma^2)
            float sigma = std::pow(Config::scaleFactor(), kpt.octave);
            float invSigma2 = 1.0f / (sigma * sigma);
            pEdge->setInformation(Eigen::Matrix2d::Identity() * invSigma2);
            // set robust kernel
            if (bRobust) {
                g2o::RobustKernelHuber* pRK = new g2o::RobustKernelHuber();
                pEdge->setRobustKernel(pRK);
                // rho(x) = x^2 if |x| < delta else 2*delta*|x| - delta^2
                pRK->setDelta(sigma);
            }
            // set cam intrinsics
            pEdge->fx = Config::fx();
            pEdge->fy = Config::fy();
            pEdge->cx = Config::cx();
            pEdge->cy = Config::cy();
            // record all the edges
            vvpEdges[i][j] = pEdge;
            // only optimize outliers at the last iteration
            if (pMPt->isOutlier()) {
                pEdge->setLevel(1);
            }
            optimizer.addEdge(pEdge);
        }
        if (!bHasEdge) {
            optimizer.removeVertex(pVPt);
            vbMPtOptimized[i] = false;
        }
    }

    // optimize
    int nIt = 2;
    for (int it = 0; it < nIt; ++it) {
        optimizer.initializeOptimization(0);
        optimizer.optimize(nIter);
        // exclude outliers
        for (int i = 0; i < nMPts; ++i) {
            auto& pMPt = vpMPts[i];
            if (vbMPtOptimized[i]) {
                vector<shared_ptr<KeyFrame>> vpKFsMPt = (nKFs == 0) ?
                    pMPt->getRelatedKFs() : vpKFs;
                int nKFsMpt = vpKFsMPt.size();
                for (int j = 0; j < nKFsMpt; ++j) {
                    auto& pEdge = vvpEdges[i][j];
                    if (!pEdge) {
                        continue;
                    }
                    // optimize all edges for the last iteration
                    // exclude outliers for other iterations
                    if (it == nIt - 2) {
                        pEdge->setLevel(0);
                    } else {
                        float chi2 = pEdge->chi2();
                        cv::KeyPoint kpt = pMPt->keypoint(vpKFsMPt[j]);
                        float sigma = std::pow(Config::scaleFactor(),
                                               kpt.octave);
                        if (chi2 > sigma*sigma ||
                            !pEdge->isDepthPositive()) {
                            pEdge->setLevel(1);
                        } else {
                            pEdge->setLevel(0);
                        }
                    }
                }
            }
        }
    }

    // update results back to the keyframes/map
    // pose update via vpKFs
    for (const auto& pKF : vpKFs) {
        g2o::VertexSE3Expmap* pVSE3 = dynamic_cast<g2o::VertexSE3Expmap*>(
            optimizer.vertex(pKF->index()));
        g2o::SE3Quat T = pVSE3->estimate();
        pKF->mPose.setPose(SE3Quat2cvMat(T));
    }
    // map point data update via vpMPts
    // count the number of map point outliers after optimization
    int nInliers = nMPts;
    for (int i = 0; i < nMPts; ++i) {
        if (vbMPtOptimized[i]) {
            shared_ptr<MapPoint>& pMPt = vpMPts[i];
            g2o::VertexSBAPointXYZ* pVPt = dynamic_cast<
                g2o::VertexSBAPointXYZ*>(optimizer.vertex(idxKFMax + 1 + i));
            Eigen::Vector3d X = pVPt->estimate();
            pMPt->setX3D(Vector3d2cvMat(X));

            // set outlier status for all map points
            pMPt->setOutlier(false);
            vector<shared_ptr<KeyFrame>> vpKFsMPt = (nKFs == 0) ?
                pMPt->getRelatedKFs() : vpKFs;
            int nKFsMpt = vpKFsMPt.size();
            for (int j = 0; j < nKFsMpt; ++j) {
                auto& pEdge = vvpEdges[i][j];
                if (!pEdge) {
                    continue;
                }
                float chi2 = pEdge->chi2();
                cv::KeyPoint kpt = pMPt->keypoint(vpKFsMPt[j]);
                float sigma = std::pow(Config::scaleFactor(), kpt.octave);
                if (chi2 > sigma*sigma * TH_MAX_CHI2_FACTOR ||
                    !pEdge->isDepthPositive()) {
                    pMPt->setOutlier(true);
                    --nInliers;
                    break;
                }
            }
        }
    }

    // clean up resources
    optimizer.clear();
    optimizer.clearParameters();

    return nInliers;
}

//int Optimizer::frameBundleAdjustment(unsigned nFrames, int nIter,
//                                     bool bRobust) const
//{
//    // set definition of solver
//    g2o::BlockSolver_6_3::LinearSolverType* pLinearSolver;
//    pLinearSolver =
//        new g2o::LinearSolverCSparse<g2o::BlockSolver_6_3::PoseMatrixType>();
//    g2o::BlockSolver_6_3* pBlkSolver;
//    pBlkSolver = new g2o::BlockSolver_6_3(pLinearSolver);
//    g2o::OptimizationAlgorithmLevenberg* pSolverLM;
//    pSolverLM = new g2o::OptimizationAlgorithmLevenberg(pBlkSolver);
//    // configure optimizer
//    g2o::SparseOptimizer optimizer;
//    optimizer.setVerbose(false);
//    optimizer.setAlgorithm(pSolverLM);
//
//    // add vertices: poses
//    vector<shared_ptr<Frame>> vpKFs = mpMap->getLastNKFs(nFrames);
//    shared_ptr<Frame> pFrameCur = nullptr;
//
//    unsigned idxFMax = 0;
//    for (const auto& pFrame : vpFrames) {
//        g2o::VertexSE3Expmap* pVSE3 = new g2o::VertexSE3Expmap();
//        Mat Tcw = pFrame->mPose.getPose();
//        unsigned idxF = pFrame->index();
//        pVSE3->setEstimate(cvMat2SE3Quat(Tcw));
//        pVSE3->setId(idxF);
//        bool bFixed = idxF != System::nCurrentFrame;
//        if (!bFixed) {
//            pFrameCur = pFrame;
//        }
//        pVSE3->setFixed(bFixed);
//        optimizer.addVertex(pVSE3);
//        if (idxF > idxFMax) {
//            idxFMax = idxF;
//        }
//    }
//
//    // no target frame to be optimized
//    if (!pFrameCur) {
//        return 0;
//    }
//
//    // add vertices: map points, and add edges for each map point
//    vector<shared_ptr<MapPoint>> vpMPts = pFrameCur->getpMPtsObserved();
//    int nMPts = vpMPts.size();
//    // skip optimization if map points are not enough
//    if (nMPts < TH_MIN_NUM_MAPPOINT) {
//        return 0;
//    }
//    vector<vector<g2o::EdgeSE3ProjectXYZ*>> vvpEdges(nMPts);
//    for (int i = 0; i < nMPts; ++i) {
//        const shared_ptr<MapPoint>& pMPt = vpMPts[i];
//        g2o::VertexSBAPointXYZ* pVPt = new g2o::VertexSBAPointXYZ();
//        pVPt->setEstimate(cvMat2Vector3d(pMPt->X3D()));
//        pVPt->setId(idxFMax + 1 + i);
//        pVPt->setMarginalized(true);
//        optimizer.addVertex(pVPt);
//
//        // add edges
//        vvpEdges[i].resize(nFrames, nullptr);
//        for (unsigned j = 0; j < nFrames; ++j) {
//            auto& pFrame = vpFrames[j];
//            // check whether the map point is observed by the target frame
//            if (!pMPt->isObservedBy(pFrame)) {
//                continue;
//            }
//            cv::KeyPoint kpt = pMPt->keypoint(pFrame);
//            Eigen::Matrix<double, 2, 1> obs;
//            obs << kpt.pt.x, kpt.pt.y;
//            g2o::EdgeSE3ProjectXYZ* pEdge = new g2o::EdgeSE3ProjectXYZ();
//            // vertex 0: map point
//            pEdge->setVertex(0, dynamic_cast<g2o::OptimizableGraph::Vertex*>(
//                                 pVPt));
//            // vertex 1: pose
//            pEdge->setVertex(1, dynamic_cast<g2o::OptimizableGraph::Vertex*>(
//                                 optimizer.vertex(pFrame->index())));
//            pEdge->setMeasurement(obs);
//            // set element in information matrix (value = 1 / sigma^2)
//            float sigma = std::pow(Config::scaleFactor(), kpt.octave);
//            float invSigma2 = 1.0f / (sigma * sigma);
//            pEdge->setInformation(Eigen::Matrix2d::Identity() * invSigma2);
//            // set robust kernel
//            if (bRobust) {
//                g2o::RobustKernelHuber* pRK = new g2o::RobustKernelHuber();
//                pEdge->setRobustKernel(pRK);
//                // rho(x) = x^2 if |x| < delta else 2*delta*|x| - delta^2
//                pRK->setDelta(sigma);
//            }
//            // set cam intrinsics
//            pEdge->fx = Config::fx();
//            pEdge->fy = Config::fy();
//            pEdge->cx = Config::cx();
//            pEdge->cy = Config::cy();
//            // only optimize outliers at the last iteration
//            if (pMPt->isOutlier()) {
//                pEdge->setLevel(1);
//            }
//            // record all the edges
//            vvpEdges[i][j] = pEdge;
//            optimizer.addEdge(pEdge);
//        }
//    }
//
//    // optimize
//    //optimizer.initializeOptimization();
//    //optimizer.optimize(nIter);
//    int nIt = 3;
//    for (int it = 0; it < nIt; ++it) {
//        optimizer.initializeOptimization(0);
//        optimizer.optimize(nIter);
//        // exclude outliers
//        for (int i = 0; i < nMPts; ++i) {
//            auto& pMPt = vpMPts[i];
//            for (unsigned j = 0; j < nFrames; ++j) {
//                auto& pEdge = vvpEdges[i][j];
//                if (!pEdge) {
//                    continue;
//                }
//                // optimize all edges for the last iteration
//                // exclude outliers for other iterations
//                if (it == nIt - 2) {
//                    pEdge->setLevel(0);
//                } else {
//                    float chi2 = pEdge->chi2();
//                    cv::KeyPoint kpt = pMPt->keypoint(vpFrames[j]);
//                    float sigma = std::pow(Config::scaleFactor(),
//                                           kpt.octave);
//                    if (chi2 > sigma*sigma ||
//                        !pEdge->isDepthPositive()) {
//                        pEdge->setLevel(1);
//                    } else {
//                        pEdge->setLevel(0);
//                    }
//                }
//            }
//        }
//    }
//
//    // update results back to the frames/map
//    // pose update via vpFrames
//    g2o::VertexSE3Expmap* pVSE3 = dynamic_cast<g2o::VertexSE3Expmap*>(
//        optimizer.vertex(System::nCurrentFrame));
//    g2o::SE3Quat T = pVSE3->estimate();
//    pFrameCur->mPose.setPose(SE3Quat2cvMat(T));
//    // map point data update via vpMPts
//    // count the number of map point outliers after optimization
//    int nInliers = nMPts;
//    for (int i = 0; i < nMPts; ++i) {
//        shared_ptr<MapPoint>& pMPt = vpMPts[i];
//        g2o::VertexSBAPointXYZ* pVPt = dynamic_cast<g2o::VertexSBAPointXYZ*>(
//            optimizer.vertex(idxFMax + 1 + i));
//        Eigen::Vector3d X = pVPt->estimate();
//        pMPt->setX3D(Vector3d2cvMat(X));
//
//        // set outlier status for all map points
//        pMPt->setOutlier(false);
//        for (unsigned j = 0; j < nFrames; ++j) {
//            auto& pEdge = vvpEdges[i][j];
//            if (!pEdge) {
//                continue;
//            }
//            float chi2 = pEdge->chi2();
//            cv::KeyPoint kpt = pMPt->keypoint(vpFrames[j]);
//            float sigma = std::pow(Config::scaleFactor(), kpt.octave);
//            if (chi2 > sigma*sigma * TH_MAX_CHI2_FACTOR ||
//                !pEdge->isDepthPositive()) {
//                pMPt->setOutlier(true);
//                --nInliers;
//                break;
//            }
//        }
//    }
//
//    // clean up resources
//    optimizer.clear();
//    optimizer.clearParameters();
//
//    return nInliers;
//}

int Optimizer::poseOptimization(const std::shared_ptr<Frame>& pFrame) const
{
    // count the number of map point outliers after optimization
    int nInliers = 0;
    // set definition of solver
    g2o::BlockSolver_6_3::LinearSolverType* pLinearSolver;
    pLinearSolver =
        new g2o::LinearSolverDense<g2o::BlockSolver_6_3::PoseMatrixType>();
    g2o::BlockSolver_6_3* pBlkSolver;
    pBlkSolver = new g2o::BlockSolver_6_3(pLinearSolver);
    g2o::OptimizationAlgorithmLevenberg* pSolverLM;
    pSolverLM = new g2o::OptimizationAlgorithmLevenberg(pBlkSolver);
    // configure optimizer
    g2o::SparseOptimizer optimizer;
    optimizer.setVerbose(false);
    optimizer.setAlgorithm(pSolverLM);

    // add vertices: the pose to be optimized
    g2o::VertexSE3Expmap* pVSE3 = new g2o::VertexSE3Expmap();
    Mat Tcw = pFrame->mPose.getPose();
    pVSE3->setEstimate(cvMat2SE3Quat(Tcw));
    pVSE3->setId(0);
    pVSE3->setFixed(false);
    optimizer.addVertex(pVSE3);

    // add edges: unary edge with map point data as measurement
    map<int, shared_ptr<MapPoint>> mpMPts = pFrame->getMPtsMap();
    int nMPts = mpMPts.size();
    // skip optimization if map points are not enough
    if (nMPts < TH_MIN_NUM_MAPPOINT) {
        return 0;
    }
    vector<g2o::EdgeSE3ProjectXYZOnlyPose*> vpEdges;
    vpEdges.reserve(nMPts);
    for (const auto& pair : mpMPts) {
        const shared_ptr<MapPoint>& pMPt = pair.second;
        // is it necessary to check whether the frame is valid?
        cv::KeyPoint kpt = pFrame->keypoint(pair.first);
        Eigen::Matrix<double, 2, 1> obs;
        obs << kpt.pt.x, kpt.pt.y;
        g2o::EdgeSE3ProjectXYZOnlyPose* pEdge =
            new g2o::EdgeSE3ProjectXYZOnlyPose();
        // vertex 0: pose
        pEdge->setVertex(0, dynamic_cast<g2o::OptimizableGraph::Vertex*>(
                             pVSE3));
        pEdge->setMeasurement(obs);
        // set element in information matrix (value = 1 / sigma^2)
        float sigma = std::pow(Config::scaleFactor(), kpt.octave);
        float invSigma2 = 1.0f / (sigma * sigma);
        pEdge->setInformation(Eigen::Matrix2d::Identity() * invSigma2);
        // set robust kernel
        g2o::RobustKernelHuber* pRK = new g2o::RobustKernelHuber();
        pEdge->setRobustKernel(pRK);
        // rho(x) = x^2 if |x| < delta else 2*delta*|x| - delta^2
        pRK->setDelta(sigma);
        // set map point position
        pEdge->Xw = cvMat2Vector3d(pMPt->X3D());
        // set cam intrinsics
        pEdge->fx = Config::fx();
        pEdge->fy = Config::fy();
        pEdge->cx = Config::cx();
        pEdge->cy = Config::cy();
        vpEdges.push_back(pEdge);
        // only optimize outliers at the last iteration
        //if (pMPt->isOutlier()) {
        //    pEdge->setLevel(1);
        //}
        optimizer.addEdge(pEdge);
    }

    // optimize (multi-pass?)
    int nIt = 4;
    for (int it = 0; it < nIt; ++it) {
        //pVSE3->setEstimate(cvMat2SE3Quat(Tcw));
        optimizer.initializeOptimization(0); // only optimize level 0 edges
        optimizer.optimize(10);
        // exclude outliers
        auto cit = mpMPts.cbegin();
        for (int i = 0; i < nMPts; ++i, ++cit) {
            assert(cit != mpMPts.cend());
            auto& pEdge = vpEdges[i];
            // optimize all edges for the last iteration
            // exclude outliers for other iterations
            if (it == nIt - 2) {
                pEdge->setLevel(0);
            } else {
                float chi2 = pEdge->chi2();
                cv::KeyPoint kpt = pFrame->keypoint(cit->first);
                float sigma = std::pow(Config::scaleFactor(), kpt.octave);
                if (chi2 > sigma*sigma || !pEdge->isDepthPositive()) {
                    pEdge->setLevel(1);
                } else {
                    pEdge->setLevel(0);
                }
            }
        }
    }

    // update optimized pose and map point outlier status
    auto cit = mpMPts.cbegin();
    for (int i = 0; i < nMPts ; ++i, ++cit) {
        assert(cit != mpMPts.cend());
        auto& pMPt = cit->second;
        auto& pEdge = vpEdges[i];
        float chi2 = pEdge->chi2();
        cv::KeyPoint kpt = pFrame->keypoint(cit->first);
        float sigma = std::pow(Config::scaleFactor(), kpt.octave);
        if (chi2 > sigma*sigma * TH_MAX_CHI2_FACTOR ||
            !pEdge->isDepthPositive()) {
            pMPt->setOutlier(true);
        } else {
            pMPt->setOutlier(false);
            nInliers++;
        }
    }
    g2o::SE3Quat T = pVSE3->estimate();
    pFrame->mPose.setPose(SE3Quat2cvMat(T));

    // clean up resources
    optimizer.clear();
    optimizer.clearParameters();

    return nInliers;
}

void Optimizer::localBundleAdjustment(const std::shared_ptr<KeyFrame>& pKFin,
                                      int nIter, bool bRobust) const
{
    // set definition of solver
    g2o::BlockSolver_6_3::LinearSolverType* pLinearSolver;
    pLinearSolver =
        new g2o::LinearSolverEigen<g2o::BlockSolver_6_3::PoseMatrixType>();
    g2o::BlockSolver_6_3* pBlkSolver;
    pBlkSolver = new g2o::BlockSolver_6_3(pLinearSolver);
    g2o::OptimizationAlgorithmLevenberg* pSolverLM;
    pSolverLM = new g2o::OptimizationAlgorithmLevenberg(pBlkSolver);
    // configure optimizer
    g2o::SparseOptimizer optimizer;
    optimizer.setVerbose(false);
    optimizer.setAlgorithm(pSolverLM);

    // add vertices: poses that can be optimized
    vector<shared_ptr<KeyFrame>> vpConnectedKFs = pKFin->getConnectedKFs();
    set<shared_ptr<KeyFrame>> spKFs;
    spKFs.insert(pKFin);
    spKFs.insert(vpConnectedKFs.cbegin(), vpConnectedKFs.cend());
    set<shared_ptr<MapPoint>> spMPts; // temp container for all local map points
    unsigned idxKFMax = 0;
    for (const auto& pKF : spKFs) {
        g2o::VertexSE3Expmap* pVSE3 = new g2o::VertexSE3Expmap();
        Mat Tcw = pKF->mPose.getPose();
        unsigned idxKF = pKF->index();
        pVSE3->setEstimate(cvMat2SE3Quat(Tcw));
        pVSE3->setId(idxKF);
        bool bFixed = pKF->index() == 0;
        pVSE3->setFixed(bFixed);
        optimizer.addVertex(pVSE3);
        if (idxKF > idxKFMax) {
            idxKFMax = idxKF;
        }
        // get local map points
        vector<shared_ptr<MapPoint>> vpMPtsKF = pKF->mappoints();
        for (const auto& pMPt : vpMPtsKF) {
            if (pMPt) {
                spMPts.insert(pMPt);
            }
        }
    }
    // add vertices: fixed poses
    vector<shared_ptr<KeyFrame>> vpWeakKFs = pKFin->getWeakKFs();
    set<shared_ptr<KeyFrame>> spFixedKFs;
    spFixedKFs.insert(vpWeakKFs.cbegin(), vpWeakKFs.cend());
    for (const auto& pMPt: spMPts) {
        vector<shared_ptr<KeyFrame>> vpKFRelated = pMPt->getRelatedKFs();
        for (const auto& pKF : vpKFRelated) {
            // add poses that is not in the local map to the fixed pose set
            if (spKFs.find(pKF) == spKFs.end()) {
                spFixedKFs.insert(pKF);
            }
        }
    }
    for (const auto& pKF : spFixedKFs) {
        g2o::VertexSE3Expmap* pVSE3 = new g2o::VertexSE3Expmap();
        Mat Tcw = pKF->mPose.getPose();
        unsigned idxKF = pKF->index();
        pVSE3->setEstimate(cvMat2SE3Quat(Tcw));
        pVSE3->setId(idxKF);
        pVSE3->setFixed(true);
        optimizer.addVertex(pVSE3);
        if (idxKF > idxKFMax) {
            idxKFMax = idxKF;
        }
    }

    // add vertices: map points, and add edges for each map point
    int nMPts = spMPts.size();
    // skip optimization if map points are not enough
    if (nMPts < TH_MIN_NUM_MAPPOINT) {
        return;
    }
    // use std::vector for indexing on each map point
    vector<shared_ptr<MapPoint>> vpMPts(spMPts.cbegin(), spMPts.cend());
    vector<bool> vbMPtOptimized(nMPts, true);
    vector<vector<g2o::EdgeSE3ProjectXYZ*>> vvpEdges(nMPts);
    for (int i = 0; i < nMPts; ++i) {
        const shared_ptr<MapPoint>& pMPt = vpMPts[i];
        g2o::VertexSBAPointXYZ* pVPt = new g2o::VertexSBAPointXYZ();
        pVPt->setEstimate(cvMat2Vector3d(pMPt->X3D()));
        pVPt->setId(idxKFMax + 1 + i);
        pVPt->setMarginalized(true); // why?? (to decrease the size of Hessian?)
        optimizer.addVertex(pVPt);

        // add edges
        vector<shared_ptr<KeyFrame>> vpKFsMPt = pMPt->getRelatedKFs();
        int nKFsMpt = vpKFsMPt.size();
        vvpEdges[i].resize(nKFsMpt, nullptr);
        bool bHasEdge = false; // check whether the vertec has edges
        for (int j = 0; j < nKFsMpt; ++j) {
            auto& pKF = vpKFsMPt[j];
            // check whether the map point is observed by the target keyframe
            if (!pMPt->isObservedBy(pKF)) {
                continue;
            }
            // form an edge
            bHasEdge = true;
            cv::KeyPoint kpt = pMPt->keypoint(pKF);
            Eigen::Matrix<double, 2, 1> obs;
            obs << kpt.pt.x, kpt.pt.y;
            g2o::EdgeSE3ProjectXYZ* pEdge = new g2o::EdgeSE3ProjectXYZ();
            // vertex 0: map point
            pEdge->setVertex(0, dynamic_cast<g2o::OptimizableGraph::Vertex*>(
                                 pVPt));
            // vertex 1: pose
            pEdge->setVertex(1, dynamic_cast<g2o::OptimizableGraph::Vertex*>(
                                 optimizer.vertex(pKF->index())));
            pEdge->setMeasurement(obs);
            // set element in information matrix (value = 1 / sigma^2)
            float sigma = std::pow(Config::scaleFactor(), kpt.octave);
            float invSigma2 = 1.0f / (sigma * sigma);
            pEdge->setInformation(Eigen::Matrix2d::Identity() * invSigma2);
            // set robust kernel
            if (bRobust) {
                g2o::RobustKernelHuber* pRK = new g2o::RobustKernelHuber();
                pEdge->setRobustKernel(pRK);
                // rho(x) = x^2 if |x| < delta else 2*delta*|x| - delta^2
                pRK->setDelta(sigma);
            }
            // set cam intrinsics
            pEdge->fx = Config::fx();
            pEdge->fy = Config::fy();
            pEdge->cx = Config::cx();
            pEdge->cy = Config::cy();
            // record all the edges
            vvpEdges[i][j] = pEdge;
            // only optimize outliers at the last iteration
            //if (pMPt->isOutlier()) {
            //    pEdge->setLevel(1);
            //}
            optimizer.addEdge(pEdge);
        }
        if (!bHasEdge) {
            optimizer.removeVertex(pVPt);
            vbMPtOptimized[i] = false;
        }
    }

    // optimize
    int nIt = 4;
    for (int it = 0; it < nIt; ++it) {
        optimizer.initializeOptimization(0);
        optimizer.optimize(nIter);
        // exclude outliers
        for (int i = 0; i < nMPts; ++i) {
            auto& pMPt = vpMPts[i];
            if (vbMPtOptimized[i]) {
                vector<shared_ptr<KeyFrame>> vpKFsMPt = pMPt->getRelatedKFs();
                int nKFsMpt = vpKFsMPt.size();
                for (int j = 0; j < nKFsMpt; ++j) {
                    auto& pEdge = vvpEdges[i][j];
                    if (!pEdge) {
                        continue;
                    }
                    // optimize all edges for the last iteration
                    // exclude outliers for other iterations
                    if (it == nIt - 2) {
                        pEdge->setLevel(0);
                    } else {
                        float chi2 = pEdge->chi2();
                        cv::KeyPoint kpt = pMPt->keypoint(vpKFsMPt[j]);
                        float sigma = std::pow(Config::scaleFactor(),
                                               kpt.octave);
                        if (chi2 > sigma*sigma ||
                            !pEdge->isDepthPositive()) {
                            pEdge->setLevel(1);
                        } else {
                            pEdge->setLevel(0);
                        }
                    }
                }
            }
        }
    }

    // update results back to the keyframes/map
    // pose update via spKFs
    for (const auto& pKF : spKFs) {
        g2o::VertexSE3Expmap* pVSE3 = dynamic_cast<g2o::VertexSE3Expmap*>(
            optimizer.vertex(pKF->index()));
        g2o::SE3Quat T = pVSE3->estimate();
        pKF->mPose.setPose(SE3Quat2cvMat(T));
    }
    // map point data update via vpMPts
    // count the number of map point outliers after optimization
    for (int i = 0; i < nMPts; ++i) {
        if (vbMPtOptimized[i]) {
            shared_ptr<MapPoint>& pMPt = vpMPts[i];
            g2o::VertexSBAPointXYZ* pVPt = dynamic_cast<
                g2o::VertexSBAPointXYZ*>(optimizer.vertex(idxKFMax + 1 + i));
            Eigen::Vector3d X = pVPt->estimate();
            pMPt->setX3D(Vector3d2cvMat(X));

            // set outlier status for all map points
            pMPt->setOutlier(false);
            vector<shared_ptr<KeyFrame>> vpKFsMPt = pMPt->getRelatedKFs();
            int nKFsMpt = vpKFsMPt.size();
            for (int j = 0; j < nKFsMpt; ++j) {
                auto& pEdge = vvpEdges[i][j];
                if (!pEdge) {
                    continue;
                }
                float chi2 = pEdge->chi2();
                cv::KeyPoint kpt = pMPt->keypoint(vpKFsMPt[j]);
                float sigma = std::pow(Config::scaleFactor(), kpt.octave);
                if (chi2 > sigma*sigma * TH_MAX_CHI2_FACTOR ||
                    !pEdge->isDepthPositive()) {
                    pMPt->setOutlier(true);
                    break;
                }
            }
        }
    }

    // clean up resources
    optimizer.clear();
    optimizer.clearParameters();
}

g2o::SE3Quat Optimizer::cvMat2SE3Quat(const cv::Mat& Tcw) const
{
    // pose data must be double for g2o to use!!!
    Eigen::Matrix<double, 3, 3> R;
    Eigen::Matrix<double, 3, 1> t;
    cv::cv2eigen(Tcw.colRange(0, 3).rowRange(0, 3), R); // float -> double
    cv::cv2eigen(Tcw.col(3), t);
    return g2o::SE3Quat(R, t);
}

cv::Mat Optimizer::SE3Quat2cvMat(const g2o::SE3Quat& T) const
{
    Mat Tcw(3, 4, CV_32FC1);
    // Eigen::Quaternion -> Eigen::Matrix3 for R
    Eigen::Matrix<float, 3, 3> R =
        T.rotation().toRotationMatrix().cast<float>();
    Eigen::Matrix<float, 3, 1> t = T.translation().cast<float>();
    cv::eigen2cv(R, Tcw.colRange(0, 3).rowRange(0, 3));
    cv::eigen2cv(t, Tcw.col(3));
    return Tcw;
}

Eigen::Vector3d Optimizer::cvMat2Vector3d(const cv::Mat& X3D) const
{
    Eigen::Matrix<double, 3, 1> X;
    cv::cv2eigen(X3D, X);
    return X;
}

cv::Mat Optimizer::Vector3d2cvMat(const Eigen::Vector3d& X) const
{
    Mat X3D(3, 1, CV_32FC1);
    Eigen::Vector3f Xf = X.cast<float>();
    cv::eigen2cv(Xf, X3D);
    return X3D;
}

} // namespace SLAM_demo


/** $Id: WeightCalculator.cxx 178081 2019-12-18 19:13:26Z jvansanten $
 * @file
 * @author Jakob van Santen <vansanten@wisc.edu>
 *
 * $Revision: 178081 $
 * $Date: 2019-12-18 12:13:26 -0700 (Wed, 18 Dec 2019) $
 */

#include <MuonGun/WeightCalculator.h>
#include <MuonGun/Generator.h>
#include <MuonGun/SamplingSurface.h>
#include <MuonGun/Cylinder.h>
#include <MuonGun/Flux.h>
#include <MuonGun/RadialDistribution.h>
#include <MuonGun/EnergyDistribution.h>
#include <MuonGun/I3MuonGun.h>
#include <MuonGun/Track.h>
#include <boost/foreach.hpp>

#include <icetray/I3Module.h>
#include <dataclasses/physics/I3MCTreeUtils.h>
#include <dataclasses/I3Double.h>
#include <simclasses/I3MMCTrack.h>
#include <phys-services/I3Calculator.h>
#include <boost/make_shared.hpp>
#include <boost/numeric/ublas/vector.hpp>

namespace I3MuonGun {

double
WeightCalculator::GetWeight(const I3Particle &axis, const BundleConfiguration &bundlespec) const
{
	std::pair<double, double> steps = surface_->GetIntersection(axis.GetPos(), axis.GetDir());
	// This shower axis doesn't intersect the sampling surface. Bail.
	if (!std::isfinite(steps.first))
		return 0.;

	double h = GetDepth(axis.GetPos().GetZ() + steps.first*axis.GetDir().GetZ());
	double coszen = cos(axis.GetDir().GetZenith());
	unsigned m = unsigned(bundlespec.size());

	double rate = flux_->GetLog(h, coszen, m) - generator_->GetLogGeneratedEvents(axis, bundlespec);

	BOOST_FOREACH(const BundleEntry &track, bundlespec){
		double logprob = energy_->GetLog(h, coszen, m, track.radius,
		    EnergyDistribution::log_value(std::log(track.energy)));
		if (!std::isfinite(logprob)){
                    log_warn("Log Energy weight of a least one muon is -inf, weight will be 0!");
                }
		rate += logprob;
        }
	// assert(std::isfinite(std::exp(rate)));
	return std::exp(rate);
}

// Possibly throw-away utility function: "track" muons to a fixed surface using the
// same method as WeightCalculatorModule
std::vector<I3Particle>
GetMuonsAtSurface(I3FramePtr frame, I3Surfaces::SurfaceConstPtr surface)
{
	std::vector<I3Particle> final_states;

	I3MCTreeConstPtr mctree = frame->Get<I3MCTreeConstPtr>();
	I3MMCTrackListConstPtr mmctracks = frame->Get<I3MMCTrackListConstPtr>("MMCTrackList");
	if (!mctree)
		log_fatal("I3MCTree missing!");
	if (!mmctracks)
		log_fatal("I3MMCTrackList missing!");
	BOOST_FOREACH(const Track &track, Track::Harvest(*mctree, *mmctracks)) {
		std::pair<double, double> steps =
		    surface->GetIntersection(track.GetPos(), track.GetDir());
		double energy = track.GetEnergy(steps.first);
		if (steps.first >= 0 && energy > 0) {
			final_states.push_back(track);
			I3Particle &p = final_states.back();
			p.SetEnergy(energy);
			p.SetPos(track.GetPos(steps.first));
			p.SetTime(track.GetTime(steps.first));
			p.SetLength(track.GetLength()-steps.first);
		}
	}

	return final_states;
}

namespace {

namespace ublas = boost::numeric::ublas;
typedef ublas::bounded_vector<double, 3> vector;

vector
make_vector(double x, double y, double z)
{
	vector v(3);
	v[0] = x; v[1] = y; v[2] = z;

	return v;
}

vector
make_vector(const I3Direction &dir)
{
	return make_vector(dir.GetX(), dir.GetY(), dir.GetZ());
}

inline vector
subtract(const I3Position &p1, const I3Position &p2)
{
	return make_vector(p1.GetX()-p2.GetX(), p1.GetY()-p2.GetY(), p1.GetZ()-p2.GetZ());
}

inline double
GetRadius(const I3Particle &axis, const I3Position &pos)
{
        vector r = subtract(pos,axis.GetPos());
	double l = ublas::inner_prod(make_vector(axis.GetDir()), r);

	return sqrt(std::max(0., ublas::inner_prod(r, r) - l*l));
}

}

MuonBundleConverter::MuonBundleConverter(size_t maxMultiplicity, SamplingSurfaceConstPtr surface)
    : maxMultiplicity_(maxMultiplicity),
    surface_(surface ? surface : boost::make_shared<Cylinder>(1600, 800))
{}

I3TableRowDescriptionPtr
MuonBundleConverter::CreateDescription(const I3MCTree&)
{
	I3TableRowDescriptionPtr desc(new I3TableRowDescription());

	desc->AddField<uint32_t>("multiplicity", "", "Number of muons in the bundle");
	desc->AddField<float>("depth", "km", "Vertical depth of intersection with the sampling surface");
	desc->AddField<float>("cos_theta", "", "Cosine of the shower zenith angle");
	desc->AddField<float>("energy", "GeV", "Muon energy at sampling surface",
	    maxMultiplicity_);
	desc->AddField<float>("radius", "m", "Perpendicular distance from of track "
	    "from the bundle axis at the sampling surface", maxMultiplicity_);

	return desc;
}

size_t
MuonBundleConverter::FillRows(const I3MCTree &mctree, I3TableRowPtr rows)
{
	I3MMCTrackListConstPtr mmctracks = currentFrame_->Get<I3MMCTrackListConstPtr>("MMCTrackList");
	if (!mmctracks)
		log_fatal("I3MMCTrackList missing!");

	const I3MCTree::const_iterator primary = mctree.begin();
	std::pair<double, double> primary_steps =
	    surface_->GetIntersection(primary->GetPos(), primary->GetDir());
	if (primary_steps.first > 0) {
		rows->Set<float>("depth", float(GetDepth(primary->GetPos().GetZ() + primary_steps.first*primary->GetDir().GetZ())));
		rows->Set<float>("cos_theta", float(cos(primary->GetDir().GetZenith())));
	}

	uint32_t m = 0;
	float *energies = rows->GetPointer<float>("energy");
	float *radii = rows->GetPointer<float>("radius");

	log_trace("%zu total tracks", Track::Harvest(mctree, *mmctracks).size());

	BOOST_FOREACH(const Track &track, Track::Harvest(mctree, *mmctracks)) {
		// MuonGun bundles are attached directly to the primary
		if (mctree.depth(track) > 1)
			continue;
		std::pair<double, double> steps =
		    surface_->GetIntersection(track.GetPos(), track.GetDir());
		float energy = float(track.GetEnergy(steps.first));
		log_trace("energy after %f m: %.1e", steps.first, energy);
		if (energy > 0) {
			if (m < maxMultiplicity_) {
				energies[m] = energy;
				radii[m] = float(GetRadius(*primary, track.GetPos(steps.first)));
			}
			m++;
		}
	}

	rows->Set("multiplicity", m);

	return 1;
}

/**
 * @brief Interface between WeightCalculator and IceTray
 *
 * WeightCalculatorModule handles the details of extracting energies and
 * radial offsets of muons from an I3MCTree and MMCTrackList.
 */
class WeightCalculatorModule : public I3Module, protected WeightCalculator {
public:
	WeightCalculatorModule(const I3Context &ctx) : I3Module(ctx)
	{
		AddOutBox("OutBox");
		AddParameter("Model", "Muon flux model for which to calculate a weight", boost::shared_ptr<BundleModel>());
		AddParameter("Generator", "Generation spectrum for the bundles to be weighted", generator_);
	}

	void Configure()
	{
		boost::shared_ptr<BundleModel> model;
		GetParameter("Model", model);
		GetParameter("Generator", generator_);

		if (!model)
			log_fatal("No flux model configured!");
		flux_ = model->flux;
		radius_ = model->radius;
		energy_ = model->energy;

		if (!generator_)
			log_fatal("No generator configured!");

		surface_ = generator_->GetInjectionSurface();
		if (!surface_)
			log_fatal("No surface configured!");
	}

	void DAQ(I3FramePtr frame)
	{
		// First, harvest the muons in the bundle at their points of injection, storing
		// everything that's necessary to estimate the energy lost up to an arbitrary point
		I3MCTreeConstPtr mctree = frame->Get<I3MCTreeConstPtr>();
		I3MMCTrackListConstPtr mmctracks = frame->Get<I3MMCTrackListConstPtr>("MMCTrackList");
		if (!mctree)
			log_fatal("I3MCTree missing!");
		// if (!mmctracks)
		// 	log_fatal("I3MMCTrackList missing!");

		const I3MCTree::const_iterator primary = mctree->begin();
		std::pair<double, double> steps =
		    surface_->GetIntersection(primary->GetPos(), primary->GetDir());
		BundleConfiguration bundlespec;

		if (mmctracks) {
			std::list<Track> tracks = Track::Harvest(*mctree, *mmctracks);
			BOOST_FOREACH(const Track &track, tracks) {
				// Omit secondary muons
				boost::optional<I3Particle> parent = mctree->parent(track);
				if (parent && parent->GetType() == I3Particle::NuclInt)
					continue;
				bundlespec.push_back(BundleEntry(
				    GetRadius(*primary, track.GetPos(steps.first)), track.GetEnergy(steps.first)));
			}
		} else {
			// log_warn("No MMCTrackList found in the frame! Assuming that everything starts on the sampling surface...");
			BOOST_FOREACH(const I3Particle &track, std::make_pair(mctree->begin(), mctree->end())) {
				if (track.GetType() == I3Particle::MuMinus || track.GetType() == I3Particle::MuPlus)
					bundlespec.push_back(BundleEntry(
					    GetRadius(*primary, track.GetPos()), track.GetEnergy()));
			}
		}

		frame->Put(GetName(), boost::make_shared<I3Double>(GetWeight(*primary, bundlespec)));
		PushFrame(frame);
	}

	void Finish();
};

// Out-of-line virtual method definition to force the vtable into this translation unit
void WeightCalculatorModule::Finish() {}

}

I3_MODULE(I3MuonGun::WeightCalculatorModule);


/*
 * Copyright (c) 2015-2016, Princeton University, Johannes M Dieterich, Emily A Carter
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation and/or
 * other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its contributors may
 * be used to endorse or promote products derived from this software without specific
 * prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
 * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#ifndef WANGTETER_HPP
#define	WANGTETER_HPP

#include <armadillo>
#include <memory>
#include "KEDF.hpp"
#include "FourierGrid.hpp"
#include "HelperFunctions.hpp"
using namespace std;
using namespace arma;

template<class GridType>
class WangTeter: public KEDF<GridType> {

public:

    WangTeter(GridType* example, const double alpha, const double beta, cube* keKernel)
    : _alpha(alpha), _beta(beta), _keKernel(keKernel){
    }

    ~WangTeter() {
    }

    string getMethodDescription() const {
        return "Wang-Teter KEDF";
    }

    vector<string> getCitations() const {

        vector<string> citations(0);
        citations.push_back("L.-W. Wang and M. P. Teter, Phys. Rev. B 45, 13196 (1992).");

        return citations;
    }

    vector<string> getWorkingEquations() const {

        vector<string> equations(0);
        equations.push_back("");

        //   WTEnergy = cTF * SUM(rhoR_SI**alpha * FFT(FFT(rhoR_SI**beta) * keKernel(:,:,:,1)))*/

        return equations;
    }

    double calcEnergy(const GridType& grid) const {

        unique_ptr<GridType> workGridB = grid.duplicate();

        workGridB->powGrid(_beta);

        cx_cube* recB = workGridB->getReciprocalGrid();

        const uword recElems = recB->n_elem;

        #pragma omp parallel for default(none) shared(recB)
        for (uword x = 0; x < recElems; ++x) {
            recB->at(x) *= _keKernel->at(x);
        }
        workGridB->completeReciprocal(recB);

        cube* realB = workGridB->getRealGrid();

        const uword elems = realB->n_elem;

        const cube* dens = grid.tryReadRealGrid();
        if (!dens) {
            // unfortunately, we need to get a copy of the grid
            unique_ptr<FourierGrid> workGrid = grid.createFourierDuplicate();
            dens = workGrid->readRealGrid();
        }

        #pragma omp parallel for default(none) shared(realB,dens)
        for (uword x = 0; x < elems; ++x) {
            realB->at(x) *= pow(dens->at(x), _alpha);
        }
        workGridB->complete(realB);

        const double eWT = _CTF * workGridB->integrate();

        return eWT;
    }

    double calcPotential(const GridType& grid, GridType& potential) const {

        const cube* dens = grid.tryReadRealGrid();
        if (!dens) {
            // unfortunately, we need to get a copy of the grid
            unique_ptr<GridType> workGrid = grid.duplicate();
            dens = workGrid->readRealGrid();
        }

        unique_ptr<GridType> workGridA = grid.duplicate();
        cube* densityA = workGridA->getRealGrid();

        const uword elems = densityA->n_elem;
        const size_t nSlices = densityA->n_slices;
        const size_t nRows = densityA->n_rows;
        const size_t nCols = densityA->n_cols;

        auto densityB = MemoryFunctions::allocateScratch(nRows, nCols, nSlices);
        cube* poten = potential.getRealGrid();

        #pragma omp parallel for default(none) shared(densityA,densityB,poten)
        for (uword x = 0; x < elems; ++x) {
            const double rho = densityA->at(x);
            const double rhoPBM = pow(rho, (_beta - 1));
            densityB->at(x) = rhoPBM; // density B contains rho**(beta-1)
            poten->at(x) = rhoPBM*rho; // poten contains rho**(beta)
            densityA->at(x) = pow(rho, _alpha); // density A contains rho**(alpha)
        }
        potential.complete(poten);
        workGridA->complete(densityA);

        cx_cube* recPot = potential.getReciprocalGrid();

        const uword recElems = recPot->n_elem;

        #pragma omp parallel for default(none) shared(recPot)
        for (uword x = 0; x < recElems; ++x) {
            recPot->at(x) *= _keKernel->at(x);
        }
        potential.completeReciprocal(recPot);

        potential.multiplyElementwise(workGridA.get());

        // get the energy
        const double eWT = _CTF * potential.integrate();

        // transform what is currently in potential into the first part of the actual potential
        cube* realPot = potential.getRealGrid();

        const double preAl = _CTF*_alpha;
        #pragma omp parallel for default(none) shared(realPot,dens)
        for (uword x = 0; x < elems; ++x) {
            realPot->at(x) *= preAl / dens->at(x);
        }

        // now take what is in densityA and FFT it
        cx_cube* recDensA = workGridA->getReciprocalGrid();

        #pragma omp parallel for default(none) shared(recDensA)
        for (size_t x = 0; x < recElems; ++x) {
            recDensA->at(x) *= _keKernel->at(x);
        }
        workGridA->completeReciprocal(recDensA);

        densityA = workGridA->getRealGrid();

        const double preBe = _CTF*_beta;
        #pragma omp parallel for default(none) shared(realPot,densityA,densityB)
        for (size_t x = 0; x < elems; ++x) {
            const double pot = densityA->at(x) * preBe * densityB->at(x);
            realPot->at(x) += pot;
        }
        potential.complete(realPot);

        return eWT;
    }

    unique_ptr<StressTensor> calcStress(const GridType& grid) const {
        throw runtime_error("not yet implemented");
    }

private:

    const double _CTF = 2.87123400018819;
    const double _alpha;
    const double _beta;
    const cube * _keKernel;
};

class WangTeterKernel {
public:
    WangTeterKernel(const double alpha, const double beta, const double rho0, const double lambdaTF, const double muVW, const double ft = (5.0/3.0));
    ~WangTeterKernel();

    void fillWTKernelReciprocal(cube* kernel, const cube* gNorms);

private:
    double _alpha;
    double _beta;
    double _rho0;
    double _ft;
    double _lambda;
    double _mu;
    double _coeff;
    double _tkF;
};

#ifdef LIBKEDF_OCL
#include "WangTeterOCL.hpp"
#endif

#endif	/* WANGTETER_HPP */


#ifndef USE_CUDA

#include <Spirit_Defines.h>
#include <engine/Hamiltonian_Heisenberg_Neighbours.hpp>
#include <engine/Vectormath.hpp>
#include <engine/Neighbours.hpp>
#include <data/Spin_System.hpp>
#include <utility/Constants.hpp>

#include <Eigen/Dense>

using namespace Data;
using namespace Utility;
using Utility::Constants::mu_B;
using Utility::Constants::mu_0;
using Utility::Constants::Pi;
using Engine::Vectormath::check_atom_type;
using Engine::Vectormath::idx_from_pair;

namespace Engine
{
    Hamiltonian_Heisenberg_Neighbours::Hamiltonian_Heisenberg_Neighbours(
        scalarfield mu_s,
        scalar external_field_magnitude, Vector3 external_field_normal,
        intfield anisotropy_indices, scalarfield anisotropy_magnitudes, vectorfield anisotropy_normals,
        scalarfield exchange_magnitudes,
        scalarfield dmi_magnitudes, int dm_chirality,
        scalar ddi_radius,
        std::shared_ptr<Data::Geometry> geometry,
        intfield boundary_conditions
    ) :
        Hamiltonian(boundary_conditions),
        geometry(geometry),
        mu_s(mu_s),
        external_field_magnitude(external_field_magnitude * mu_B), external_field_normal(external_field_normal),
        anisotropy_indices(anisotropy_indices), anisotropy_magnitudes(anisotropy_magnitudes), anisotropy_normals(anisotropy_normals),
        exchange_magnitudes(exchange_magnitudes),
        dmi_magnitudes(dmi_magnitudes),
        ddi_radius(ddi_radius)
    {
        // Generate Exchange neighbours
        exchange_neighbours = Neighbours::Get_Neighbours_in_Shells(*geometry, exchange_magnitudes.size());

        // Generate DMI neighbours and normals
        dmi_neighbours = Neighbours::Get_Neighbours_in_Shells(*geometry, dmi_magnitudes.size());
        for (unsigned int ineigh = 0; ineigh < dmi_neighbours.size(); ++ineigh)
        {
            dmi_normals.push_back(Neighbours::DMI_Normal_from_Pair(*geometry, dmi_neighbours[ineigh], dm_chirality));
        }

        // Generate DDI neighbours, magnitudes and normals
        this->ddi_neighbours = Engine::Neighbours::Get_Neighbours_in_Radius(*this->geometry, ddi_radius);
        scalar magnitude;
        Vector3 normal;
        for (unsigned int i=0; i<ddi_neighbours.size(); ++i)
        {
            Engine::Neighbours::DDI_from_Pair(*this->geometry, ddi_neighbours[i], magnitude, normal);
            this->ddi_magnitudes.push_back(magnitude);
            this->ddi_normals.push_back(normal);
        }

        this->Update_Energy_Contributions();
    }

    void Hamiltonian_Heisenberg_Neighbours::Update_N_Neighbour_Shells(int n_shells_exchange, int n_shells_dmi)
    {
        if (this->exchange_magnitudes.size() != n_shells_exchange)
        {
            this->exchange_magnitudes = scalarfield(n_shells_exchange);
            // Re-calculate exchange neighbour list
        }
        if (this->dmi_magnitudes.size() != n_shells_dmi)
        {
            this->dmi_magnitudes = scalarfield(n_shells_dmi);
            // Re-calculate dmi neighbour list
        }
    }

    void Hamiltonian_Heisenberg_Neighbours::Update_Energy_Contributions()
    {
        this->energy_contributions_per_spin = std::vector<std::pair<std::string, scalarfield>>(0);

        // External field
        if (this->external_field_magnitude > 0)
        {
            this->energy_contributions_per_spin.push_back({"Zeeman", scalarfield(0)});
            this->idx_zeeman = this->energy_contributions_per_spin.size()-1;
        }
        else this->idx_zeeman = -1;
        // Anisotropy
        if (this->anisotropy_indices.size() > 0)
        {
            this->energy_contributions_per_spin.push_back({"Anisotropy", scalarfield(0) });
            this->idx_anisotropy = this->energy_contributions_per_spin.size()-1;
        }
        else this->idx_anisotropy = -1;
        // Exchange
        if (this->exchange_neighbours.size() > 0)
        {
            this->energy_contributions_per_spin.push_back({"Exchange", scalarfield(0) });
            this->idx_exchange = this->energy_contributions_per_spin.size()-1;
        }
        else this->idx_exchange = -1;
        // DMI
        if (this->dmi_neighbours.size() > 0)
        {
            this->energy_contributions_per_spin.push_back({"DMI", scalarfield(0) });
            this->idx_dmi = this->energy_contributions_per_spin.size()-1;
        }
        else this->idx_dmi = -1;
        // Dipole-Dipole
        if (this->ddi_neighbours.size() > 0)
        {
            this->energy_contributions_per_spin.push_back({"DD", scalarfield(0) });
            this->idx_ddi = this->energy_contributions_per_spin.size()-1;
        }
        else this->idx_ddi = -1;
    }

    void Hamiltonian_Heisenberg_Neighbours::Energy_Contributions_per_Spin(const vectorfield & spins, std::vector<std::pair<std::string, scalarfield>> & contributions)
    {
        if (contributions.size() != this->energy_contributions_per_spin.size())
        {
            contributions = this->energy_contributions_per_spin;
        }

        int nos = spins.size();
        for (auto& pair : energy_contributions_per_spin)
        {
            // Allocate if not already allocated
            if (pair.second.size() != nos) pair.second = scalarfield(nos, 0);
            // Otherwise set to zero
            else for (auto& pair : energy_contributions_per_spin) Vectormath::fill(pair.second, 0);
        }

        // External field
        if (this->idx_zeeman >=0 )     E_Zeeman(spins, energy_contributions_per_spin[idx_zeeman].second);
        // Anisotropy
        if (this->idx_anisotropy >=0 ) E_Anisotropy(spins, energy_contributions_per_spin[idx_anisotropy].second);

        // Exchange
        if (this->idx_exchange >=0 )   E_Exchange(spins,energy_contributions_per_spin[idx_exchange].second);
        // DMI
        if (this->idx_dmi >=0 )        E_DMI(spins, energy_contributions_per_spin[idx_dmi].second);
        // DDI
        if (this->idx_ddi >=0 )        E_DDI(spins, energy_contributions_per_spin[idx_ddi].second);
    }

    void Hamiltonian_Heisenberg_Neighbours::E_Zeeman(const vectorfield & spins, scalarfield & Energy)
    {
        const int N = geometry->n_cell_atoms;

        #pragma omp parallel for
        for (int icell = 0; icell < geometry->n_cells_total; ++icell)
        {
            for (int ibasis = 0; ibasis < N; ++ibasis)
            {
                int ispin = icell*N + ibasis;
                if (check_atom_type(this->geometry->atom_types[ispin]))
                    Energy[ispin] -= this->mu_s[ibasis] * this->external_field_magnitude * this->external_field_normal.dot(spins[ispin]);
            }
        }
    }

    void Hamiltonian_Heisenberg_Neighbours::E_Anisotropy(const vectorfield & spins, scalarfield & Energy)
    {
        const int N = geometry->n_cell_atoms;

        #pragma omp parallel for
        for (int icell = 0; icell < geometry->n_cells_total; ++icell)
        {
            for (int iani = 0; iani < anisotropy_indices.size(); ++iani)
            {
                int ispin = icell*N + anisotropy_indices[iani];
                if (check_atom_type(this->geometry->atom_types[ispin]))
                    Energy[ispin] -= this->anisotropy_magnitudes[iani] * std::pow(anisotropy_normals[iani].dot(spins[ispin]), 2.0);
            }
        }
    }

    void Hamiltonian_Heisenberg_Neighbours::E_Exchange(const vectorfield & spins, scalarfield & Energy)
    {
        #pragma omp parallel for
        for (unsigned int ispin = 0; ispin < spins.size(); ++ispin)
        {
            // auto translations = Vectormath::translations_from_idx(geometry->n_cells, geometry->n_cell_atoms, ispin);
            for (unsigned int ineigh = 0; ineigh < exchange_neighbours.size(); ++ineigh)
            {
                int jspin = idx_from_pair(ispin, boundary_conditions, geometry->n_cells, geometry->n_cell_atoms, geometry->atom_types, exchange_neighbours[ineigh]);
                if ( jspin >= 0 )
                {
                    auto& ishell = exchange_neighbours[ineigh].idx_shell;
                    Energy[ispin] -= 0.5 * exchange_magnitudes[ishell] * spins[ispin].dot(spins[jspin]);
                }
            }
        }
    }

    void Hamiltonian_Heisenberg_Neighbours::E_DMI(const vectorfield & spins, scalarfield & Energy)
    {
        #pragma omp parallel for
        for (unsigned int ispin = 0; ispin < spins.size(); ++ispin)
        {
            for (unsigned int ineigh = 0; ineigh < dmi_neighbours.size(); ++ineigh)
            {
                int jspin = idx_from_pair(ispin, boundary_conditions, geometry->n_cells, geometry->n_cell_atoms, geometry->atom_types, dmi_neighbours[ineigh]);
                if ( jspin >= 0 )
                {
                    auto& ishell = dmi_neighbours[ineigh].idx_shell;
                    Energy[ispin] -= 0.5 * dmi_magnitudes[ishell] * dmi_normals[ineigh].dot(spins[ispin].cross(spins[jspin]));
                }
            }
        }
    }

    void Hamiltonian_Heisenberg_Neighbours::E_DDI(const vectorfield & spins, scalarfield & Energy)
    {
        // The translations are in angstr�m, so the |r|[m] becomes |r|[m]*10^-10
        const scalar mult = mu_0 * std::pow(mu_B, 2) / ( 4*Pi * 1e-30 );

        scalar result = 0.0;
        const int N = geometry->n_cell_atoms;

        #pragma omp parallel for
        for (int icell = 0; icell < geometry->n_cells_total; ++icell)
        {
            for (int ibasis = 0; ibasis < N; ++ibasis)
            {
                int ispin = icell * N + ibasis;
                for (unsigned int ineigh = 0; ineigh < ddi_neighbours.size(); ++ineigh)
                {
                    if (ddi_magnitudes[ineigh] > 0.0)
                    {
                        int jbasis = ddi_neighbours[ineigh].j;
                        int jspin = idx_from_pair(ispin, boundary_conditions, geometry->n_cells, geometry->n_cell_atoms, geometry->atom_types, ddi_neighbours[ineigh]);
                        if (jspin >= 0)
                        {
                            Energy[ispin] -= 0.5 * this->mu_s[ibasis] * this->mu_s[jbasis] * mult / std::pow(ddi_magnitudes[ineigh], 3.0) *
                                (3 * spins[jspin].dot(ddi_normals[ineigh]) * spins[ispin].dot(ddi_normals[ineigh]) - spins[ispin].dot(spins[jspin]));
                            Energy[jspin] -= 0.5 * this->mu_s[ibasis] * this->mu_s[jbasis] * mult / std::pow(ddi_magnitudes[ineigh], 3.0) *
                                (3 * spins[jspin].dot(ddi_normals[ineigh]) * spins[ispin].dot(ddi_normals[ineigh]) - spins[ispin].dot(spins[jspin]));
                        }
                    }
                }
            }
        }
    }// end DipoleDipole

    void Hamiltonian_Heisenberg_Neighbours::Gradient(const vectorfield & spins, vectorfield & gradient)
    {
        // Set to zero
        Vectormath::fill(gradient, {0,0,0});

        // External field
        Gradient_Zeeman(gradient);

        // Anisotropy
        Gradient_Anisotropy(spins, gradient);

        // Exchange
        this->Gradient_Exchange(spins, gradient);
        // DMI
        this->Gradient_DMI(spins, gradient);
        // DD
        this->Gradient_DDI(spins, gradient);
    }

    void Hamiltonian_Heisenberg_Neighbours::Gradient_Zeeman(vectorfield & gradient)
    {
        const int N = geometry->n_cell_atoms;

        #pragma omp parallel for
        for (int icell = 0; icell < geometry->n_cells_total; ++icell)
        {
            for (int ibasis = 0; ibasis < N; ++ibasis)
            {
                int ispin = icell*N + ibasis;
                if (check_atom_type(this->geometry->atom_types[ispin]))
                    gradient[ispin] -= this->mu_s[ibasis] * this->external_field_magnitude * this->external_field_normal;
            }
        }
    }

    void Hamiltonian_Heisenberg_Neighbours::Gradient_Anisotropy(const vectorfield & spins, vectorfield & gradient)
    {
        const int N = geometry->n_cell_atoms;

        #pragma omp parallel for
        for (int icell = 0; icell < geometry->n_cells_total; ++icell)
        {
            for (int iani = 0; iani < anisotropy_indices.size(); ++iani)
            {
                int ispin = icell*N + anisotropy_indices[iani];
                if (check_atom_type(this->geometry->atom_types[ispin]))
                    gradient[ispin] -= 2.0 * this->anisotropy_magnitudes[iani] * this->anisotropy_normals[iani] * anisotropy_normals[iani].dot(spins[ispin]);
            }
        }
    }

    void Hamiltonian_Heisenberg_Neighbours::Gradient_Exchange(const vectorfield & spins, vectorfield & gradient)
    {
        #pragma omp parallel for
        for (unsigned int ispin = 0; ispin < spins.size(); ++ispin)
        {
            for (unsigned int ineigh = 0; ineigh < exchange_neighbours.size(); ++ineigh)
            {
                int jspin = idx_from_pair(ispin, boundary_conditions, geometry->n_cells, geometry->n_cell_atoms, geometry->atom_types, exchange_neighbours[ineigh]);
                if ( jspin >= 0 )
                {
                    auto& ishell = exchange_neighbours[ineigh].idx_shell;
                    gradient[ispin] -= exchange_magnitudes[ishell] * spins[jspin];
                }
            }
        }
    }

    void Hamiltonian_Heisenberg_Neighbours::Gradient_DMI(const vectorfield & spins, vectorfield & gradient)
    {
        #pragma omp parallel for
        for (unsigned int ispin = 0; ispin < spins.size(); ++ispin)
        {
            auto translations = Vectormath::translations_from_idx(geometry->n_cells, geometry->n_cell_atoms, ispin);
            for (unsigned int ineigh = 0; ineigh < dmi_neighbours.size(); ++ineigh)
            {
                int jspin = idx_from_pair(ispin, boundary_conditions, geometry->n_cells, geometry->n_cell_atoms, geometry->atom_types, dmi_neighbours[ineigh]);
                if ( jspin >= 0 )
                {
                    auto& ishell = dmi_neighbours[ineigh].idx_shell;
                    gradient[ispin] -= dmi_magnitudes[ishell] * spins[jspin].cross(dmi_normals[ineigh]);
                }
            }
        }
    }

    void Hamiltonian_Heisenberg_Neighbours::Gradient_DDI(const vectorfield & spins, vectorfield & gradient)
    {
        // The translations are in angstr�m, so the |r|[m] becomes |r|[m]*10^-10
        const scalar mult = mu_0 * std::pow(mu_B, 2) / ( 4*Pi * 1e-30 );

        const int N = geometry->n_cell_atoms;

        #pragma omp parallel for
        for (int icell = 0; icell < geometry->n_cells_total; ++icell)
        {
            for (int ibasis = 0; ibasis < N; ++ibasis)
            {
                int ispin = icell * N + ibasis;
                for (unsigned int ineigh = 0; ineigh < ddi_neighbours.size(); ++ineigh)
                {
                    if (ddi_magnitudes[ineigh] > 0.0)
                    {
                        int jbasis = ddi_neighbours[ineigh].j;
                        int jspin = idx_from_pair(ispin, boundary_conditions, geometry->n_cells, geometry->n_cell_atoms, geometry->atom_types, ddi_neighbours[ineigh]);
                        if (jspin >= 0)
                        {
                            scalar skalar_contrib = mult / std::pow(ddi_magnitudes[ineigh], 3.0);
                            gradient[ispin] -= this->mu_s[jbasis] * skalar_contrib * (3 * ddi_normals[ineigh] * spins[jspin].dot(ddi_normals[ineigh]) - spins[jspin]);
                            gradient[jspin] -= this->mu_s[ibasis] * skalar_contrib * (3 * ddi_normals[ineigh] * spins[ispin].dot(ddi_normals[ineigh]) - spins[ispin]);
                        }
                    }
                }
            }
        }
    }//end Field_DipoleDipole

    void Hamiltonian_Heisenberg_Neighbours::Hessian(const vectorfield & spins, MatrixX & hessian)
    {
        int nos = spins.size();

        // Set to zero
        // for (auto& h : hessian) h = 0;
        hessian.setZero();

        // Single Spin elements
        for (int alpha = 0; alpha < 3; ++alpha)
        {
            for (unsigned int i = 0; i < anisotropy_indices.size(); ++i)
            {
                int idx = anisotropy_indices[i];
                // scalar x = -2.0*this->anisotropy_magnitudes[i] * std::pow(this->anisotropy_normals[i][alpha], 2);
                hessian(3*idx + alpha, 3*idx + alpha) += -2.0*this->anisotropy_magnitudes[i]*std::pow(this->anisotropy_normals[i][alpha],2);
            }
        }

        // std::cerr << "calculated hessian" << std::endl;

        // Spin Pair elements
        // Exchange
        for (unsigned int ispin = 0; ispin < spins.size(); ++ispin)
        {
            auto translations = Vectormath::translations_from_idx(geometry->n_cells, geometry->n_cell_atoms, ispin);
            for (unsigned int ineigh = 0; ineigh < this->exchange_neighbours.size(); ++ineigh)
            {
                for (int alpha = 0; alpha < 3; ++alpha)
                {
                    //int idx_i = 3 * exchange_neighbours[i_pair][0] + alpha;
                    //int idx_j = 3 * exchange_neighbours[i_pair][1] + alpha;
                    int jspin = Vectormath::idx_from_translations(geometry->n_cells, geometry->n_cell_atoms, translations, exchange_neighbours[ineigh].translations);
                    int ishell = exchange_neighbours[ineigh].idx_shell;
                    hessian(ispin, jspin) += -exchange_magnitudes[ineigh];
                    hessian(jspin, ispin) += -exchange_magnitudes[ineigh];
                }
            }
        }
        // DMI
        for (unsigned int ispin = 0; ispin < spins.size(); ++ispin)
        {
            auto translations = Vectormath::translations_from_idx(geometry->n_cells, geometry->n_cell_atoms, ispin);
            for (unsigned int ineigh = 0; ineigh < this->dmi_neighbours.size(); ++ineigh)
            {
                for (int alpha = 0; alpha < 3; ++alpha)
                {
                    for (int beta = 0; beta < 3; ++beta)
                    {
                        int idx_i = 3 * dmi_neighbours[ineigh].i + alpha;
                        int idx_j = 3 * dmi_neighbours[ineigh].j + beta;
                        if ((alpha == 0 && beta == 1))
                        {
                            hessian(idx_i, idx_j) +=
                                -dmi_magnitudes[ineigh] * dmi_normals[ineigh][2];
                            hessian(idx_j, idx_i) +=
                                -dmi_magnitudes[ineigh] * dmi_normals[ineigh][2];
                        }
                        else if ((alpha == 1 && beta == 0))
                        {
                            hessian(idx_i, idx_j) +=
                                dmi_magnitudes[ineigh] * dmi_normals[ineigh][2];
                            hessian(idx_j, idx_i) +=
                                dmi_magnitudes[ineigh] * dmi_normals[ineigh][2];
                        }
                        else if ((alpha == 0 && beta == 2))
                        {
                            hessian(idx_i, idx_j) +=
                                dmi_magnitudes[ineigh] * dmi_normals[ineigh][1];
                            hessian(idx_j, idx_i) +=
                                dmi_magnitudes[ineigh] * dmi_normals[ineigh][1];
                        }
                        else if ((alpha == 2 && beta == 0))
                        {
                            hessian(idx_i, idx_j) +=
                                -dmi_magnitudes[ineigh] * dmi_normals[ineigh][1];
                            hessian(idx_j, idx_i) +=
                                -dmi_magnitudes[ineigh] * dmi_normals[ineigh][1];
                        }
                        else if ((alpha == 1 && beta == 2))
                        {
                            hessian(idx_i, idx_j) +=
                                -dmi_magnitudes[ineigh] * dmi_normals[ineigh][0];
                            hessian(idx_j, idx_i) +=
                                -dmi_magnitudes[ineigh] * dmi_normals[ineigh][0];
                        }
                        else if ((alpha == 2 && beta == 1))
                        {
                            hessian(idx_i, idx_j) +=
                                dmi_magnitudes[ineigh] * dmi_normals[ineigh][0];
                            hessian(idx_j, idx_i) +=
                                dmi_magnitudes[ineigh] * dmi_normals[ineigh][0];
                        }
                    }
                }
            }
        }
        //// Dipole-Dipole
        //for (unsigned int i_pair = 0; i_pair < this->ddi_neighbours[i_periodicity].size(); ++i_pair)
        //{
        //	// indices
        //	int idx_1 = ddi_neighbours[i_periodicity][i_pair][0];
        //	int idx_2 = ddi_neighbours[i_periodicity][i_pair][1];
        //	// prefactor
        //	scalar prefactor = 0.0536814951168
        //		* this->mu_s[idx_1] * this->mu_s[idx_2]
        //		/ std::pow(ddi_magnitude[i_periodicity][i_pair], 3);
        //	// components
        //	for (int alpha = 0; alpha < 3; ++alpha)
        //	{
        //		for (int beta = 0; beta < 3; ++beta)
        //		{
        //			int idx_h = idx_1 + alpha*nos + 3 * nos*(idx_2 + beta*nos);
        //			if (alpha == beta)
        //				hessian[idx_h] += prefactor;
        //			hessian[idx_h] += -3.0*prefactor*DD_normal[i_periodicity][i_pair][alpha] * DD_normal[i_periodicity][i_pair][beta];
        //		}
        //	}
        //}
    }

    // Hamiltonian name as string
    static const std::string name = "Heisenberg (Neighbours)";
    const std::string& Hamiltonian_Heisenberg_Neighbours::Name() { return name; }
}

#endif


#include "resource.hpp"

#include "instance.hpp"
#include "traits.hpp"

#include <algorithm>
#include <boost/container/small_vector.hpp>
#include <cassert> // for assert
#include <cmath>   // for pow
#include <iterator>
#include <memory>
#include <pstl/glue_algorithm_defs.h>
#include <stdexcept>

// TODO factor polynomial stuff out into util
double
apply_polynomial(const polynomial & poly, double x)
{
	double sum = 0;
	for (auto term : poly) {
		sum += std::get<0>(term) * std::pow(x, std::get<1>(term));
	}
	return sum;
}

polynomial
add_poly(const polynomial & lhs, const polynomial & rhs)
{
	polynomial result(lhs.begin(), lhs.end());
	result.insert(result.end(), rhs.begin(), rhs.end());
	std::sort(result.begin(), result.end(),
	          [](const poly_term & sort_lhs, const poly_term & sort_rhs) {
		          return sort_lhs.second < sort_rhs.second;
	          });

	size_t len = result.size();
	for (size_t i = 0; (len > 0) && (i < (len - 1)); ++i) {
		// TODO float-equal comparison.
		if (double_eq(result[i].second, result[i + 1].second)) {
			// Equal exponents. Combine, mark for swap-and-delete
			result[i].first += result[i + 1].first;
			result[i + 1].first = 0;
			i++;
		}
	}

	auto new_end =
	    std::remove_if(result.begin(), result.end(), [](const poly_term & t) {
		    return double_eq(t.first, 0.0);
	    });
	result.erase(new_end, result.end());

	return result;
}
// TODO test add_poly

Resource::Resource(unsigned int id)
    : rid(id), availability(0), overshoot_costs({})
{}

Resource
Resource::clone() const
{
	Resource cloned(this->rid);
	cloned.set_availability(Availability(this->availability));
	cloned.set_investment_costs(this->investment_costs);
	cloned.set_overshoot_costs(FlexCost(this->overshoot_costs));

	return cloned;
}

void
Resource::set_availability(Availability && availability_in)
{
	this->availability = std::move(availability_in);
}

const Availability &
Resource::get_availability() const
{
	return this->availability;
}

unsigned int
Resource::get_rid()
{
	return this->rid;
}

void
Resource::set_id(unsigned int id)
{
	this->rid = id;
}

void
Resource::set_investment_costs(polynomial costs)
{
	this->investment_costs = costs;
}

void
Resource::set_overshoot_costs(FlexCost && costs)
{
	this->overshoot_costs = std::move(costs);
}

const polynomial &
Resource::get_investment_costs() const
{
	return this->investment_costs;
}

const polynomial &
Resource::get_overshoot_costs() const
{
	return this->overshoot_costs.get_base();
}

const FlexCost &
Resource::get_flex_overshoot() const
{
	return this->overshoot_costs;
}

const polynomial &
Resource::get_overshoot_costs(unsigned int pos) const
{
	return this->overshoot_costs.get_at(pos);
}

bool
Resource::is_overshoot_flat() const
{
	return this->overshoot_costs.is_flat();
}

FlexCost::FlexCost(polynomial base_in) : base(std::move(base_in)) {}

bool
FlexCost::is_flat() const
{
	return this->points.empty();
}

const polynomial &
FlexCost::get_base() const noexcept
{
	return this->base;
}

void
FlexCost::set_flexible(
    std::vector<std::pair<unsigned int, polynomial>> && new_points)
{
	if (this->base.empty()) {
		this->points = std::move(new_points);
	} else {
		this->points.reserve(new_points.size());
		this->points.clear();

		std::transform(new_points.begin(), new_points.end(),
		               std::back_inserter(this->points),
		               [&](const std::pair<unsigned int, polynomial> point) {
			               return std::pair<unsigned int, polynomial>(
			                   point.first, add_poly(this->base, point.second));
		               });
	}
}

const polynomial &
FlexCost::get_at(unsigned int pos) const noexcept
{
	if (this->points.empty()) {
		return this->base;
	} else {

		/* This comparator inverses comparison s.t. we can just use
		 * lower_bound on the reversed points to find the correct point. */
		struct Comp
		{
			bool
			operator()(const std::pair<unsigned int, polynomial> & lhs,
			           unsigned int rhs) const noexcept
			{
				return lhs.first > rhs;
			}
			bool
			operator()(unsigned int lhs,
			           const std::pair<unsigned int, polynomial> & rhs) const noexcept
			{
				return lhs > rhs.first;
			}
		};

		auto it = std::lower_bound(this->points.rbegin(), this->points.rend(), pos,
		                           Comp{});
		if (it == this->points.rend()) {
			return this->base;
		} else {
			return it->second;
		}
	}
}
// TODO test get_at

std::vector<std::pair<unsigned int, polynomial>>::const_iterator
FlexCost::begin() const
{
	return this->points.begin();
}
std::vector<std::pair<unsigned int, polynomial>>::const_iterator
FlexCost::end() const
{
	return this->points.end();
}

Availability::Availability(double start_amount) : points(1, {0, start_amount})
{}

double
Availability::get_at(unsigned int pos) const noexcept
{
	if (this->points.empty()) {
		return 0;
	} else {

		/* This comparator inverses comparison s.t. we can just use
		 * lower_bound on the reversed points to find the correct point. */
		struct Comp
		{
			bool
			operator()(const std::pair<unsigned int, double> & lhs,
			           unsigned int rhs) const noexcept
			{
				return lhs.first > rhs;
			}
			bool
			operator()(unsigned int lhs,
			           const std::pair<unsigned int, double> & rhs) const noexcept
			{
				return lhs > rhs.first;
			}
		};

		auto it = std::lower_bound(this->points.rbegin(), this->points.rend(), pos,
		                           Comp{});
		if (it == this->points.rend()) {
			return 0.0;
		} else {
			return it->second;
		}
	}
}

void
Availability::set(std::vector<std::pair<unsigned int, double>> && new_points)
{
	this->points = std::move(new_points);
}

std::vector<std::pair<unsigned int, double>>::const_iterator
Availability::begin() const
{
	return this->points.begin();
}

std::vector<std::pair<unsigned int, double>>::const_iterator
Availability::end() const
{
	return this->points.end();
}

double
Availability::get_flat_available() const
{
	assert(this->points.size() == 1);
	return this->points[0].second;
}

Resources::Resources() {}

Resources::Resources(double _usage) : usage{_usage} {}

Resources::Resources(const Instance * in, const ResVec & u)
    : instance(in), usage(u)
{
	if (in == NULL && u.size() > 1) {
		throw std::invalid_argument("Instance in is NULL");
	} else if (u.size() > 1 &&
	           !in->get_traits().has_flag(Traits::FLAT_AVAILABILITY)) {
		throw TraitUnfulfilledError("FLAT_AVAILABILITY required!");
	}
}

Resources::Resources(const Instance * in, const std::vector<double> & u)
    : instance(in), usage(u.begin(), u.end())
{
	if (in == NULL && u.size() > 1) {
		throw std::invalid_argument("Instance in is NULL");
	} else if (u.size() > 1 &&
	           !in->get_traits().has_flag(Traits::FLAT_AVAILABILITY)) {
		throw TraitUnfulfilledError("FLAT_AVAILABILITY required!");
	}
}

Resources::Resources(const Instance * in, ResVec && u)
    : instance(in), usage(std::move(u))
{
	if (in == NULL && u.size() > 1) {
		throw std::invalid_argument("Instance in is NULL");
	} else if (u.size() > 1 &&
	           !in->get_traits().has_flag(Traits::FLAT_AVAILABILITY)) {
		throw TraitUnfulfilledError("FLAT_AVAILABILITY required!");
	}
}

Resources::Resources(const Instance * in)
    : instance(in), usage(in->resource_count(), 0.0)
{}

const ResVec &
Resources::getUsage() const
{
	return usage;
}

ResVec &
Resources::getUsage()
{
	return usage;
}

Resources
Resources::operator+(const Resources & other) const
{
	Resources res(*this);
	res += other;
	return res;
}

Resources
Resources::operator-(const Resources & other) const
{
	Resources res(*this);
	res -= other;
	return res;
}

Resources
Resources::operator*(const Resources & other) const
{
	Resources res(*this);
	res *= other;
	return res;
}

Resources
Resources::operator/(const Resources & other) const
{
	Resources res(*this);
	res /= other;
	return res;
}

void
Resources::operator+=(const Resources & other)
{
	if (instance == NULL) {
		instance = other.instance;
	}
	if (usage.size() < other.usage.size()) {
		usage.resize(other.usage.size());
	}
	for (size_t i = 0; i < other.usage.size(); i++) {
		usage[i] += other.usage[i];
	}
}

void
Resources::operator*=(const Resources & other)
{
	if (instance == NULL) {
		instance = other.instance;
	}
	if (usage.size() < other.usage.size()) {
		usage.resize(other.usage.size());
	}
	for (size_t i = 0; i < other.usage.size(); i++) {
		usage[i] *= other.usage[i];
	}
}

void
Resources::operator/=(const Resources & other)
{
	if (instance == NULL) {
		instance = other.instance;
	}
	if (usage.size() < other.usage.size()) {
		usage.resize(other.usage.size());
	}
	for (size_t i = 0; i < other.usage.size(); i++) {
		usage[i] /= other.usage[i];
	}
}

void
Resources::operator-=(const Resources & other)
{
	if (instance == NULL) {
		instance = other.instance;
	}
	if (usage.size() < other.usage.size()) {
		usage.resize(other.usage.size());
	}
	for (size_t i = 0; i < other.usage.size(); i++) {
		usage[i] -= other.usage[i];
	}
}

bool
Resources::operator<(const Resources & other) const
{
	return this->getCosts() < other.getCosts();
}

bool
Resources::operator>(const Resources & other) const
{
	return this->getCosts() > other.getCosts();
}

bool
Resources::operator<=(const Resources & other) const
{
	return this->getCosts() <= other.getCosts();
}

bool
Resources::operator>=(const Resources & other) const
{
	return this->getCosts() >= other.getCosts();
}

bool
Resources::operator!=(const Resources & other) const
{
	return this->usage != other.usage;
}

bool
Resources::operator==(const Resources & other) const
{
	return this->usage == other.usage;
}

double
Resources::getCosts() const
{
	if (usage.size() == 0) {
		return 0;
	} else if (usage.size() == 1) {
		return usage[0];
	} else if (!cached) {
		cache = instance->calculate_costs(usage);
		cached = true;
	}
	return cache;
}


//=======================================================================
// Copyright 2015 Tsinghua University.
// Authors: Fuan Pu (Pu.Fuan@gmail.com)
//
// Dung's abstract argumentation framework
//=======================================================================

#ifndef DUNG_PREFERRED_REASONER_HPP
#define DUNG_PREFERRED_REASONER_HPP

#include <iostream>
#include <string>
#include <utility>
#include <algorithm>

//boost
#include <boost/graph/graph_traits.hpp>
#include <boost/graph/adjacency_list.hpp>

#include "config/config.hpp"
#include "ArgumentProperty.hpp"
#include "AttackProperty.hpp"
#include "bitmatrix/bitvector.hpp"
#include "bitmatrix/bitmatrix.hpp"
#include "DungAF.hpp"
#include "Reasoner.hpp"

namespace argumatrix{

using namespace std;

/**
* This reasoner for Dung theories performs inference on the preferred extension.
* Computes the (unique) grounded extension, i.e., the least fixpoint of
* the characteristic function.
*/
class PreferredReasoner : public Reasoner {
public:
	PreferredReasoner(const DungAF& daf, streambuf* osbuff = std::cout.rdbuf())
		: Reasoner(daf, osbuff) {
		// m_argNum = m_daf.getNumberOfArguments();
		// m_BmAtkMtx = m_daf.getAttackMatrix();
		m_attackedBy = m_BmAtkMtx.transpose();
	}

	/**
	 * Compute all extensions given a specific semantics. Each extension is a set
	 * of arguments. Here, we use a bitvector to represent an extension. The results
	 * are stored in m_extensions. When a new extension bv is computed, it can be added
	 * to m_extensions by m_extensions.push_back(bv).
	 * @return no return. The results are stored in m_extensions.
	 */
	void computeExtensions();

	/**
	 * Get attackers of arguments in _bv, R^-(S)
	 * R^-(S) = {x|x attacks some argument in S}.
	 * R^-(S_bv) = D^T*S_bv
	 * @param _bv the bitvector of a set of arguments S.
	 * @return a set of arguments in bitvector form.
	 */
	bitvector getAttackers(const bitvector& _bv);

private:
	enum SELECT_TYPE { TYPE_A, TYPE_B };
	enum LABELS { BLANK = 0, _IN_, _OUT_, MUST_OUT, UNDEC, Label_Num };

	/**
	 * Before computing the preferred extension, we will preprocess all arguments
	 * by assigning initial labels to some arguments for reducing search space.
	 * The preprocessing operations contain (1) setting all self-attacking arguments
	 * with OUT label; (2) setting all unattacked arguments with IN label;
	 * (3) setting all arguments attacked by IN-label arguments with OUT label;
	 * Lastly, all arguments labeled IN or OUT will be removed from BLANK.
	 */
	void preprocessing(bitvector& _blank, bitvector& _in, bitvector& _out, bitvector& _undec);

	void findPreferredExtensions(bitvector& _blank, bitvector& _in,
			bitvector& _out, bitvector& _undec, bitvector& _must_out);

	bool lookAhead(bitvector& _blank_new, bitvector& _tmp_must_out);

	pair<size_type, SELECT_TYPE> selectArgument(const bitvector& _blank, const bitvector& _in,
		const bitvector& _out, const bitvector& _undec, const bitvector& _must_out);

private:
	/**
	 * The transpose of the attack matrix, it provides an effective way to
	 * access all attacked arguments of a given argument. The attacked arguments
	 * of the argument with index i is m_attackedBy[i].
	 */
	bitmatrix m_attackedBy;

};  // class GroundedReasoner

void PreferredReasoner::computeExtensions()
{
	m_extensions.clear();

	// We using five bitvector to encode the labellings of all arguments.
	// Each argument is merely labeled with one of the following five labellings.
	bitvector _blank = bitvector::UniversalSet(m_argNum);  // At beginning, all arguments are labeled BLANK
	bitvector _out = bitvector::EmptySet(m_argNum);   // Empty set
	bitvector _in = bitvector::EmptySet(m_argNum);
	bitvector _must_out = bitvector::EmptySet(m_argNum);
	bitvector _undec = bitvector::EmptySet(m_argNum);

	// Preprocessing
	preprocessing(_blank, _in, _out, _undec);

	findPreferredExtensions(_blank, _in, _out, _undec, _must_out);
}

void PreferredReasoner::preprocessing(bitvector& _blank, bitvector& _in, bitvector& _out, bitvector& _undec)
{
	// setting all self-attacking arguments with OUT label
	// _out |= m_attackMatrix.diag();
	_undec |= getSelfAttackingArguments();

	// getting all unattacked arguments by characteristic function
	_in |= characteristic(bitvector::EmptySet(m_argNum));

	// setting all arguments, attacked by IN-label arguments, with OUT label;
	_out |= getAttacked(_in);

	// Removing IN and OUT from BLANK
	// _blank ^= (_in|_out);
	_blank -= _in;
	_blank -= _out;
	_blank -= _undec;
}

void PreferredReasoner::findPreferredExtensions(bitvector& _blank,
		bitvector& _in, bitvector& _out, bitvector& _undec, bitvector& _must_out)
{
	size_type i;
	SELECT_TYPE _s_type;

	boost::tie(i, _s_type) = selectArgument(_blank, _in, _out, _undec, _must_out);

	while (i != bitvector::npos)
	{
		bitvector _blank_new(_blank);
		bitvector _out_new(_out);
		bitvector _in_new(_in);
		bitvector _must_out_new(_must_out);
		bitvector _undec_new(_undec);

		_in_new[i] = true; _blank_new[i] = false;  // label i with IN

		_out_new |= m_attackedBy[i];
		_blank_new -= m_attackedBy[i];
		_undec_new -= m_attackedBy[i];
		_must_out_new -= m_attackedBy[i];

		// get argument i's attackers which are labeled by BLANK or UNDEC
		bitvector _tmp_must_out( m_BmAtkMtx[i]&(_blank_new|_undec_new) );
		_must_out_new |= _tmp_must_out;
		_blank_new -= _tmp_must_out;
		_undec_new -= _tmp_must_out;

		//bool forwardCheck = lookAhead(_blank_new, _tmp_must_out);
		bool forwardCheck = lookAhead(_blank_new, _must_out_new);
		if ( forwardCheck )
		{
			findPreferredExtensions(_blank_new, _in_new, _out_new, _undec_new, _must_out_new);
		}

		if (_s_type == TYPE_B)
		{
			_undec[i] = true;  _blank[i] = false;
			if ( !lookAhead(_blank, _must_out) )
				return;
		}
		else   // TYPE_A
		{
			if ( !forwardCheck )
				return;

			_blank = _blank_new;
			_in = _in_new;
			_out = _out_new;
			_must_out = _must_out_new;
			_undec = _undec_new;
		}

		boost::tie(i, _s_type) = selectArgument(_blank, _in, _out, _undec, _must_out);
	}

	// if there are no argument labeled MUST_OUT, then inserts the argument
	// set labeled IN into extensions
	if ( _must_out.is_emptyset() )
	{
		std::set<bitvector>::iterator sa_itr;

		for (sa_itr = m_extensions.begin(); sa_itr != m_extensions.end(); sa_itr++)
		{
			if (_in.is_subset_of( *sa_itr ))
			{
				//cout << "duplicate: " << _in << "is subset of" << *sa_itr << endl;
				return;
			}
		}

		m_extensions.insert( _in );
	}
}

bool PreferredReasoner::lookAhead(bitvector& _blank_new, bitvector& _tmp_must_out)
{
	for (size_type i = _tmp_must_out.find_first(); i != bitvector::npos; i = _tmp_must_out.find_next(i))
	{
		if ( !m_BmAtkMtx[i].intersects(_blank_new) )
			return false;
	}

	return true;
}

pair<size_type, PreferredReasoner::SELECT_TYPE> PreferredReasoner::selectArgument(
	const bitvector& _blank, const bitvector& _in, const bitvector& _out,
	const bitvector& _undec, const bitvector& _must_out)
{
	size_type i = bitvector::npos;
	for (size_type j = _blank.find_first(); j != bitvector::npos; j = _blank.find_next(j))
	{
		//bool _m_in = true;
		if( m_BmAtkMtx[j].is_subset_of(_out|_must_out) )
		//if ( !m_attackMatrix[j].intersects(_blank|_undec) )
		{
			return make_pair(j, TYPE_A);
		}

		if ( i == bitvector::npos) 	{
			i = j;
		} else if (m_attackedBy[j].count()+m_BmAtkMtx[j].count() >
			m_attackedBy[i].count()+m_BmAtkMtx[i].count()) {
			i = j;
		}
	}

	return make_pair(i, TYPE_B);
}

bitvector PreferredReasoner::getAttackers(const bitvector& _bv)
{
	assert( _bv.size() == m_argNum );

	return m_attackedBy * _bv;
}

} // namespace argumatrix

#endif  //DUNG_PREFERRED_REASONER_HPP


/**
 * @file gan_impl.hpp
 * @author Kris Singh
 * @author Shikhar Jaiswal
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license. You should have received a copy of the
 * 3-clause BSD license along with mlpack. If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */
#ifndef MLPACK_METHODS_ANN_GAN_GAN_IMPL_HPP
#define MLPACK_METHODS_ANN_GAN_GAN_IMPL_HPP

#include "gan.hpp"

#include <mlpack/core.hpp>

#include <mlpack/methods/ann/ffn.hpp>
#include <mlpack/methods/ann/init_rules/network_init.hpp>
#include <mlpack/methods/ann/visitor/output_parameter_visitor.hpp>
#include <mlpack/methods/ann/activation_functions/softplus_function.hpp>
#include <boost/serialization/variant.hpp>

namespace mlpack {
namespace ann /** Artifical Neural Network.  */ {
template<
  typename Model,
  typename InitializationRuleType,
  typename Noise,
  typename PolicyType
>
GAN<Model, InitializationRuleType, Noise, PolicyType>::GAN(
    arma::mat& predictors,
    Model generator,
    Model discriminator,
    InitializationRuleType& initializeRule,
    Noise& noiseFunction,
    const size_t noiseDim,
    const size_t batchSize,
    const size_t generatorUpdateStep,
    const size_t preTrainSize,
    const double multiplier,
    const double clippingParameter,
    const double lambda):
    generator(std::move(generator)),
    discriminator(std::move(discriminator)),
    initializeRule(initializeRule),
    noiseFunction(noiseFunction),
    noiseDim(noiseDim),
    batchSize(batchSize),
    generatorUpdateStep(generatorUpdateStep),
    preTrainSize(preTrainSize),
    multiplier(multiplier),
    clippingParameter(clippingParameter),
    lambda(lambda),
    reset(false)
{
  // Insert IdentityLayer for joining the Generator and Discriminator.
  this->discriminator.network.insert(
      this->discriminator.network.begin(),
      new IdentityLayer<>());

  counter = 0;
  currentBatch = 0;

  this->discriminator.deterministic = this->generator.deterministic = true;

  this->predictors.set_size(predictors.n_rows, predictors.n_cols + batchSize);
  this->predictors.cols(0, predictors.n_cols - 1) = predictors;
  this->discriminator.predictors = arma::mat(this->predictors.memptr(),
      this->predictors.n_rows, this->predictors.n_cols, false, false);

  responses.ones(1, predictors.n_cols + batchSize);
  responses.cols(predictors.n_cols,
      predictors.n_cols + batchSize - 1) = arma::zeros(1, batchSize);
  this->discriminator.responses = arma::mat(this->responses.memptr(),
      this->responses.n_rows, this->responses.n_cols, false, false);

  numFunctions = predictors.n_cols;

  noise.set_size(noiseDim, batchSize);

  this->generator.predictors.set_size(noiseDim, batchSize);
  this->generator.responses.set_size(predictors.n_rows, batchSize);
}

template<
  typename Model,
  typename InitializationRuleType,
  typename Noise,
  typename PolicyType
>
GAN<Model, InitializationRuleType, Noise, PolicyType>::GAN(
    const GAN& network):
    predictors(network.predictors),
    responses(network.responses),
    generator(network.generator),
    discriminator(network.discriminator),
    initializeRule(network.initializeRule),
    noiseFunction(network.noiseFunction),
    noiseDim(network.noiseDim),
    batchSize(network.batchSize),
    generatorUpdateStep(network.generatorUpdateStep),
    preTrainSize(network.preTrainSize),
    multiplier(network.multiplier),
    clippingParameter(network.clippingParameter),
    lambda(network.lambda),
    reset(network.reset),
    counter(network.counter),
    currentBatch(network.currentBatch),
    parameter(network.parameter),
    numFunctions(network.numFunctions),
    noise(network.noise)
{
  /* Nothing to do here */
}

template<
  typename Model,
  typename InitializationRuleType,
  typename Noise,
  typename PolicyType
>
GAN<Model, InitializationRuleType, Noise, PolicyType>::GAN(
    GAN&& network):
    predictors(std::move(network.predictors)),
    responses(std::move(network.responses)),
    generator(std::move(network.generator)),
    discriminator(std::move(network.discriminator)),
    initializeRule(std::move(network.initializeRule)),
    noiseFunction(std::move(network.noiseFunction)),
    noiseDim(network.noiseDim),
    batchSize(network.batchSize),
    generatorUpdateStep(network.generatorUpdateStep),
    preTrainSize(network.preTrainSize),
    multiplier(network.multiplier),
    clippingParameter(network.clippingParameter),
    lambda(network.lambda),
    reset(network.reset),
    counter(network.counter),
    currentBatch(network.currentBatch),
    parameter(std::move(network.parameter)),
    numFunctions(network.numFunctions),
    noise(std::move(network.noise))
{
  /* Nothing to do here */
}

template<
  typename Model,
  typename InitializationRuleType,
  typename Noise,
  typename PolicyType
>
void GAN<Model, InitializationRuleType, Noise, PolicyType>::Reset()
{
  size_t genWeights = 0;
  size_t discWeights = 0;

  NetworkInitialization<InitializationRuleType> networkInit(initializeRule);

  for (size_t i = 0; i < generator.network.size(); ++i)
  {
    genWeights += boost::apply_visitor(weightSizeVisitor, generator.network[i]);
  }

  for (size_t i = 0; i < discriminator.network.size(); ++i)
  {
    discWeights += boost::apply_visitor(weightSizeVisitor,
        discriminator.network[i]);
  }

  parameter.set_size(genWeights + discWeights, 1);
  generator.Parameters() = arma::mat(parameter.memptr(), genWeights, 1, false,
      false);
  discriminator.Parameters() = arma::mat(parameter.memptr() + genWeights,
      discWeights, 1, false, false);

  // Initialize the parameters generator
  networkInit.Initialize(generator.network, parameter);
  // Initialize the parameters discriminator
  networkInit.Initialize(discriminator.network, parameter, genWeights);

  reset = true;
}

template<
  typename Model,
  typename InitializationRuleType,
  typename Noise,
  typename PolicyType
>
template<typename OptimizerType>
double GAN<Model, InitializationRuleType, Noise, PolicyType>::Train(
    OptimizerType& Optimizer)
{
  if (!reset)
    Reset();
  return Optimizer.Optimize(*this, parameter);
}

template<
  typename Model,
  typename InitializationRuleType,
  typename Noise,
  typename PolicyType
>
template<typename Policy>
typename std::enable_if<std::is_same<Policy, StandardGAN>::value ||
                        std::is_same<Policy, DCGAN>::value, double>::type
GAN<Model, InitializationRuleType, Noise, PolicyType>::Evaluate(
    const arma::mat& /* parameters */,
    const size_t i,
    const size_t /* batchSize */)
{
  if (!reset)
    Reset();

  currentInput = arma::mat(predictors.memptr() + (i * predictors.n_rows),
      predictors.n_rows, batchSize, false, false);
  currentTarget = arma::mat(responses.memptr() + i, 1, batchSize, false,
      false);

  discriminator.Forward(std::move(currentInput));
  double res = discriminator.outputLayer.Forward(
      std::move(boost::apply_visitor(
      outputParameterVisitor,
      discriminator.network.back())), std::move(currentTarget));

  noise.imbue( [&]() { return noiseFunction();} );
  generator.Forward(std::move(noise));

  predictors.cols(numFunctions, numFunctions + batchSize - 1) =
      boost::apply_visitor(outputParameterVisitor, generator.network.back());
  discriminator.Forward(std::move(predictors.cols(numFunctions,
      numFunctions + batchSize - 1)));
  responses.cols(numFunctions, numFunctions + batchSize - 1) =
      arma::zeros(1, batchSize);

  currentTarget = arma::mat(responses.memptr() + numFunctions,
      1, batchSize, false, false);
  res += discriminator.outputLayer.Forward(
      std::move(boost::apply_visitor(
      outputParameterVisitor,
      discriminator.network.back())), std::move(currentTarget));

  return res;
}

template<
  typename Model,
  typename InitializationRuleType,
  typename Noise,
  typename PolicyType
>
template<typename GradType, typename Policy>
typename std::enable_if<std::is_same<Policy, StandardGAN>::value ||
                        std::is_same<Policy, DCGAN>::value, double>::type
GAN<Model, InitializationRuleType, Noise, PolicyType>::
EvaluateWithGradient(const arma::mat& /* parameters */,
                     const size_t i,
                     GradType& gradient,
                     const size_t /* batchSize */)
{
  if (!reset)
    Reset();

  if (gradient.is_empty())
  {
    if (parameter.is_empty())
      Reset();
    gradient = arma::zeros<arma::mat>(parameter.n_elem, 1);
  }
  else
    gradient.zeros();

  if (noiseGradientDiscriminator.is_empty())
  {
    noiseGradientDiscriminator = arma::zeros<arma::mat>(
        gradientDiscriminator.n_elem, 1);
  }
  else
  {
    noiseGradientDiscriminator.zeros();
  }

  gradientGenerator = arma::mat(gradient.memptr(),
      generator.Parameters().n_elem, 1, false, false);

  gradientDiscriminator = arma::mat(gradient.memptr() +
      gradientGenerator.n_elem,
      discriminator.Parameters().n_elem, 1, false, false);

  // Get the gradients of the Discriminator.
  double res = discriminator.EvaluateWithGradient(discriminator.parameter,
      i, gradientDiscriminator, batchSize);

  noise.imbue( [&]() { return noiseFunction();} );
  generator.Forward(std::move(noise));
  predictors.cols(numFunctions, numFunctions + batchSize - 1) =
      boost::apply_visitor(outputParameterVisitor, generator.network.back());
  responses.cols(numFunctions, numFunctions + batchSize - 1) =
      arma::zeros(1, batchSize);

  // Get the gradients of the Generator.
  res += discriminator.EvaluateWithGradient(discriminator.parameter,
      numFunctions, noiseGradientDiscriminator, batchSize);
  gradientDiscriminator += noiseGradientDiscriminator;

  if (currentBatch % generatorUpdateStep == 0 && preTrainSize == 0)
  {
    // Minimize -log(D(G(noise))).
    // Pass the error from Discriminator to Generator.
    responses.cols(numFunctions, numFunctions + batchSize - 1) =
        arma::ones(1, batchSize);
    discriminator.Gradient(discriminator.parameter, numFunctions,
        noiseGradientDiscriminator, batchSize);
    generator.error = boost::apply_visitor(deltaVisitor,
        discriminator.network[1]);

    generator.Predictors() = noise;
    generator.ResetGradients(gradientGenerator);
    generator.Gradient(generator.parameter, 0, gradientGenerator, batchSize);

    gradientGenerator *= multiplier;
  }

  counter++;
  currentBatch++;

  // Revert the counter to zero, if the total dataset get's covered.
  if (counter * batchSize >= numFunctions)
  {
    counter = 0;
  }

  if (preTrainSize > 0)
  {
    preTrainSize--;
  }

  return res;
}

template<
  typename Model,
  typename InitializationRuleType,
  typename Noise,
  typename PolicyType
>
template<typename Policy>
typename std::enable_if<std::is_same<Policy, StandardGAN>::value ||
                        std::is_same<Policy, DCGAN>::value, void>::type
GAN<Model, InitializationRuleType, Noise, PolicyType>::
Gradient(const arma::mat& parameters,
         const size_t i,
         arma::mat& gradient,
         const size_t batchSize)
{
  this->EvaluateWithGradient(parameters, i, gradient, batchSize);
}

template<
  typename Model,
  typename InitializationRuleType,
  typename Noise,
  typename PolicyType
>
void GAN<Model, InitializationRuleType, Noise, PolicyType>::Shuffle()
{
  const arma::uvec ordering = arma::shuffle(arma::linspace<arma::uvec>(0,
      numFunctions - 1, numFunctions));
  predictors.cols(0, numFunctions - 1) = predictors.cols(ordering);
}

template<
  typename Model,
  typename InitializationRuleType,
  typename Noise,
  typename PolicyType
>
void GAN<Model, InitializationRuleType, Noise, PolicyType>::Forward(
    arma::mat&& input)
{
  if (!reset)
    Reset();

  generator.Forward(std::move(input));
  ganOutput = boost::apply_visitor(
      outputParameterVisitor,
      generator.network.back());

  discriminator.Forward(std::move(ganOutput));
}

template<
  typename Model,
  typename InitializationRuleType,
  typename Noise,
  typename PolicyType
>
void GAN<Model, InitializationRuleType, Noise, PolicyType>::
Predict(arma::mat&& input, arma::mat& output)
{
  if (!reset)
    Reset();

  Forward(std::move(input));

  output = boost::apply_visitor(outputParameterVisitor,
      discriminator.network.back());
}

template<
  typename Model,
  typename InitializationRuleType,
  typename Noise,
  typename PolicyType
>
template<typename Archive>
void GAN<Model, InitializationRuleType, Noise, PolicyType>::
serialize(Archive& ar, const unsigned int /* version */)
{
  ar & BOOST_SERIALIZATION_NVP(parameter);
  ar & BOOST_SERIALIZATION_NVP(generator);
  ar & BOOST_SERIALIZATION_NVP(discriminator);
  ar & BOOST_SERIALIZATION_NVP(noiseFunction);
}

} // namespace ann
} // namespace mlpack
# endif