/**
 * @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


#include "timer.hpp"
#include "read_map.hpp"
#include <jps_basis/data_utils.h>
#include <jps_planner/jps_planner/jps_planner.h>

#include <boost/geometry.hpp>
#include <boost/geometry/geometries/point_xy.hpp>
#include <boost/geometry/geometries/polygon.hpp>

using namespace JPS;

int main(int argc, char ** argv){
  if(argc != 2) {
    printf(ANSI_COLOR_RED "Input yaml required!\n" ANSI_COLOR_RESET);
    return -1;
  }

  // Read the map from yaml
  MapReader<Vec2i, Vec2f> reader(argv[1], true); // Map read from a given file
  if(!reader.exist()) {
    printf(ANSI_COLOR_RED "Cannot read input file [%s]!\n" ANSI_COLOR_RESET, argv[1]);
    return -1;
  }

  // store map in map_util
  std::shared_ptr<OccMapUtil> map_util = std::make_shared<OccMapUtil>();
  map_util->setMap(reader.origin(), reader.dim(), reader.data(), reader.resolution());

  const Vec2f start(reader.start(0), reader.start(1));
  const Vec2f goal(reader.goal(0), reader.goal(1));

  std::unique_ptr<JPSPlanner2D> planner_ptr(new JPSPlanner2D(false)); // Declare a planner
  planner_ptr->setMapUtil(map_util); // Set collision checking function
  planner_ptr->updateMap();

  Timer time_jps(true);
  bool valid_jps = planner_ptr->plan(start, goal, 1, true); // Plan from start to goal using JPS
  double dt_jps = time_jps.Elapsed().count();
  const auto path_jps = planner_ptr->getRawPath(); // Get the planned raw path from JPS
  printf("JPS Planner takes: %f ms\n", dt_jps);
  printf("JPS Path Distance: %f\n", total_distance2f(path_jps));

  Timer time_astar(true);
  bool valid_astar = planner_ptr->plan(start, goal, 1, false); // Plan from start to goal using A*
  double dt_astar = time_astar.Elapsed().count();
  const auto path_astar = planner_ptr->getRawPath(); // Get the planned raw path from A*
  printf("JPS Planner takes: %f ms\n", dt_jps);
  printf("AStar Planner takes: %f ms\n", dt_astar);
  printf("AStar Path Distance: %f\n", total_distance2f(path_astar));

  // Plot the result in svg image
  typedef boost::geometry::model::d2::point_xy<double> point_2d;
  // Declare a stream and an SVG mapper
  std::ofstream svg("output.svg");
  boost::geometry::svg_mapper<point_2d> mapper(svg, 1000, 1000);

  // Draw the canvas
  boost::geometry::model::polygon<point_2d> bound;
  const Vec2i dim = map_util->getDim();
  const Vec2f ori = map_util->getOrigin();
  const double res = map_util->getRes();

  const double origin_x = ori(0);
  const double origin_y = ori(1);
  const double range_x = dim(0) * res;
  const double range_y = dim(1) * res;
  std::vector<point_2d> points;
  points.push_back(point_2d(origin_x, origin_y));
  points.push_back(point_2d(origin_x, origin_y+range_y));
  points.push_back(point_2d(origin_x+range_x, origin_y+range_y));
  points.push_back(point_2d(origin_x+range_x, origin_y));
  points.push_back(point_2d(origin_x, origin_y));
  boost::geometry::assign_points(bound, points);
  boost::geometry::correct(bound);

  mapper.add(bound);
  mapper.map(bound, "fill-opacity:1.0;fill:rgb(255,255,255);stroke:rgb(0,0,0);stroke-width:2"); // White

  // Draw start and goal
  point_2d start_pt, goal_pt;
  boost::geometry::assign_values(start_pt, start(0), start(1));
  mapper.add(start_pt);
  mapper.map(start_pt, "fill-opacity:1.0;fill:rgb(255,0,0);", 10); // Red
  boost::geometry::assign_values(goal_pt, goal(0), goal(1));
  mapper.add(goal_pt);
  mapper.map(goal_pt, "fill-opacity:1.0;fill:rgb(255,0,0);", 10); // Red

  // Draw the obstacles
  for(int x = 0; x < dim(0); x ++) {
    for(int y = 0; y < dim(1); y ++) {
      if(!map_util->isFree(Vec2i(x, y))) {
        Vec2f pt = map_util->intToFloat(Vec2i(x, y));
        point_2d a;
        boost::geometry::assign_values(a, pt(0), pt(1));
        mapper.add(a);
        mapper.map(a, "fill-opacity:1.0;fill:rgb(0,0,0);", 1);
      }
    }
  }

  // Draw the path from JPS
  if(valid_jps) {
    vec_Vec2f path = path_jps;
    boost::geometry::model::linestring<point_2d> line;
    for(auto pt: path)
      line.push_back(point_2d(pt(0), pt(1)));
    mapper.add(line);
    mapper.map(line, "opacity:0.4;fill:none;stroke:rgb(212,0,0);stroke-width:5"); // Red
  }

  // Draw the path from A*
  if(valid_astar) {
    vec_Vec2f path = path_astar;
    boost::geometry::model::linestring<point_2d> line;
    for(auto pt: path)
      line.push_back(point_2d(pt(0), pt(1)));
    mapper.add(line);
    mapper.map(line, "opacity:0.4;fill:none;stroke:rgb(1,212,0);stroke-width:5"); // Green
  }

  // Write title at the lower right corner on canvas
  mapper.text(point_2d(origin_x + range_x - 6, origin_y+1.8), "test_planner_2d",
              "fill-opacity:1.0;fill:rgb(10,10,250);");

  mapper.text(point_2d(origin_x + range_x - 8, origin_y+1.2), "Green: ",
              "fill-opacity:1.0;fill:rgb(1,212,0);");
  mapper.text(point_2d(origin_x + range_x - 5.5, origin_y+1.2), "astar path",
              "fill-opacity:1.0;fill:rgb(0,0,0);");

  mapper.text(point_2d(origin_x + range_x - 8, origin_y+0.6), "Red: ",
              "fill-opacity:1.0;fill:rgb(212,0,0);");
  mapper.text(point_2d(origin_x + range_x - 5.5, origin_y+0.6), "jps path",
              "fill-opacity:1.0;fill:rgb(0,0,0);");

  return 0;
}


#ifndef __GIFS_ELECTRONIC_HPP
#define __GIFS_ELECTRONIC_HPP

#include <armadillo>

class Electronic{
public:
  Electronic(void) {}
  Electronic(const arma::cx_mat &cs) : amplitudes (cs) {reserve();}
  Electronic(const arma::cx_vec &c)  : amplitudes (c)  {reserve();}
  Electronic(arma::uword NStates, arma::uword NSets=1, arma::uword active=0) {
    reset(NStates, NSets, active);
  }

  void reset(arma::uword NStates, arma::uword NSets, arma::uword active){
    amplitudes.set_size(NStates, NSets);
    amplitudes.zeros();
    amplitudes.row(active).ones();
    reserve();
  }

  void set(const arma::cx_mat &cs) {amplitudes = cs;}
  void set(const arma::cx_vec &c) {amplitudes.col(0) = c;}

  const std::complex<double> operator()(arma::uword i, arma::uword j) const {return amplitudes(i,j);}
  const std::complex<double> operator()(arma::uword i) const {return amplitudes(i,0);}
  const arma::cx_vec operator()(void) const {return amplitudes.col(0);}

  const arma::cx_mat p(void) const { return amplitudes.col(0) * amplitudes.col(0).t();}
  const std::complex<double> p(arma::uword i, arma::uword j) const {return amplitudes(i,0) * std::conj(amplitudes(j,0));}

  // modification operator
  std::complex<double> & operator[](arma::uword i) {return amplitudes(i,0);}

  const arma::cx_mat get(void) const { return amplitudes; }

  void advance(const arma::cx_mat &H, double dt) {advance_exact(H, dt);}
  void advance_rk4(const arma::cx_mat &H, double dt);
  void advance_exact(const arma::cx_mat &H, double dt);

  static void phase_match(arma::mat &U, arma::vec &phases);
  static void phase_match(arma::cx_mat &U, arma::cx_vec &phases);

private:
  void reserve(void);

  arma::cx_mat amplitudes;
  arma::cx_mat k1, k2, k3, k4;
};

#endif


/** $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);


#include "ros/ros.h"
#include "drive_ros_msgs/VehicleEncoder.h"
#include "nav_msgs/Odometry.h"
#include "tf/tf.h"
#include "tf/transform_broadcaster.h"
#include "eigen3/Eigen/Dense"
#include "std_srvs/Trigger.h"
#include <rosbag/bag.h>
#include <rosbag/view.h>
#include <boost/foreach.hpp>

#include "drive_ros_localize_wheel_odometry/cov_elements.h"
#include "drive_ros_localize_wheel_odometry/moving_average.h"
#include "drive_ros_localize_wheel_odometry/save_odom_in_csv.h"

// matrix indices
static constexpr size_t X     = 0;
static constexpr size_t Y     = 1;
static constexpr size_t THETA = 2;

static constexpr size_t DELTA_R = 0;
static constexpr size_t DELTA_L = 1;

static constexpr size_t VEL_X = 0;
static constexpr size_t VEL_Y = 0;

static constexpr size_t VEL_R = 0;
static constexpr size_t VEL_L = 1;

const int encoder_ct = 4;                // number of encoder (assume we have 4 wheel encoder)
double theta = 0;                        // heading
double b_actual;                         // actual distance between wheels on an axis (effective track width)
double err_d;                            // error factor between left and right
double err_s;                            // scaling error of wheels
bool use_front;                          // whether to use front or rear axis encoder
bool broadcast_tf;                       // whether to broadcast tf
bool use_sensor_time = true;             // use sensor time or not
std::string static_frame_id;             // static frame id
std::string axis_frame;                  // axis being used
std::string input_topic;                 // input topic name
std::ofstream file_out_log;              // file output handle
bool use_bag;                            // debug mode

ros::ServiceServer reset_svr;            // reset service
ros::Publisher odom_pub;                 // publisher
nav_msgs::Odometry odom_out;             // published message
tf::TransformBroadcaster* br;            // broadcast transform

Eigen::Matrix<double, 3, 3> Sigma_p;     // covariance matrix of previous step
MovingAverage* theta_filter;             // moving average filter for theta
MovingAverage* delta_s_filter;           // moving average filter for delta_s
MovingAverage* v_filter;                 // moving average filter for v

std::vector<float> initial_cov;          // initial covariances

// reset
void reset()
{
  ROS_INFO_STREAM("Reset Wheel Odometry.");

  // new odometry should be zero
  odom_out = nav_msgs::Odometry();

  theta = 0;
  theta_filter->clear();
  delta_s_filter->clear();
  v_filter->clear();

  // set all not used odom covariances to -1
  for(int i=0; i<36; i++)
  {
    odom_out.twist.covariance[i] = -1;
    odom_out.pose.covariance[i] = -1;
  }

  // set initial covariances
  for(int i=0; i<9; i++)
  {
    Sigma_p(i/3, i%3) = initial_cov.at(i);
  }
  ROS_INFO_STREAM("Set Sigma_P: " << std::endl << Sigma_p);
}

// services to reset odometry
bool svrResetOdom(std_srvs::Trigger::Request  &req,
                  std_srvs::Trigger::Response &res)
{

  reset();
  res.message = "Reset wheel odometry.";
  return res.success = true;
}

// encoder callback
void encoderCallback(const drive_ros_msgs::VehicleEncoder::ConstPtr& msg)
{

  // choose front or rear
  uint8_t wheel_left = (use_front) ? msg->FRONT_WHEEL_LEFT : msg->REAR_WHEEL_LEFT;
  uint8_t wheel_right = (use_front) ? msg->FRONT_WHEEL_RIGHT : msg->REAR_WHEEL_RIGHT;

  // correct messages from error
  drive_ros_msgs::EncoderLinear left;
  double left_corr = err_s * 2/(err_d + 1);
  left.pos_rel      = msg->encoder[wheel_left].pos_rel      *     left_corr;
  left.pos_rel_var  = msg->encoder[wheel_left].pos_rel_var  * pow(left_corr, 2);
  left.pos_abs      = msg->encoder[wheel_left].pos_abs      *     left_corr;
  left.pos_abs_var  = msg->encoder[wheel_left].pos_abs_var  * pow(left_corr, 2);
  left.vel          = msg->encoder[wheel_left].vel          *     left_corr;
  left.vel_var      = msg->encoder[wheel_left].vel_var      * pow(left_corr, 2);

  drive_ros_msgs::EncoderLinear right;
  double right_corr = err_s * 2/(1/err_d + 1);
  right.pos_rel     = msg->encoder[wheel_right].pos_rel     *     right_corr;    ;
  right.pos_rel_var = msg->encoder[wheel_right].pos_rel_var * pow(right_corr, 2);;
  right.pos_abs     = msg->encoder[wheel_right].pos_abs     *     right_corr;    ;
  right.pos_abs_var = msg->encoder[wheel_right].pos_abs_var * pow(right_corr, 2);;
  right.vel         = msg->encoder[wheel_right].vel         *     right_corr;    ;
  right.vel_var     = msg->encoder[wheel_right].vel_var     * pow(right_corr, 2);;

  // set header
  odom_out.header.frame_id = static_frame_id;
  odom_out.header.stamp = (use_sensor_time) ? msg->header.stamp : ros::Time::now();
  odom_out.child_frame_id = axis_frame;

  // calculate mean
  double delta_s = (right.pos_rel + left.pos_rel) / 2.0;
  double vel     = (right.vel     + left.vel    ) / 2.0;
  double dtheta  = (right.pos_rel - left.pos_rel) / b_actual; // track width error already included

  delta_s = delta_s_filter->addAndGetCrrtAvg(delta_s);
  dtheta  = theta_filter->addAndGetCrrtAvg(dtheta);
  vel     = v_filter->addAndGetCrrtAvg(vel);

  // some intermediate variables
  double th = theta + dtheta/2;
  double costh = cos(static_cast<float>(th));
  double sinth = sin(static_cast<float>(th));

  // save position in odom_out
  odom_out.pose.pose.position.x += delta_s * costh;
  odom_out.pose.pose.position.y += delta_s * sinth;

  // integrate theta
  theta += dtheta;

  // save relative theta in odom_out
  tf::Quaternion q;
  q.setRPY(double(0), double(0), theta);
  tf::quaternionTFToMsg(q, odom_out.pose.pose.orientation);

  // create jacobians
  Eigen::Matrix<double, 3, 3> Jacobian_p;
  Jacobian_p.setIdentity();
  Jacobian_p(X, THETA) = - delta_s * sinth;
  Jacobian_p(Y, THETA) =   delta_s * costh;

  Eigen::Matrix<double, 3, 2> Jacobian_delta;
  Jacobian_delta(X, DELTA_R) = costh/2 - (delta_s * sinth)/(2 * b_actual);
  Jacobian_delta(X, DELTA_L) = costh/2 + (delta_s * sinth)/(2 * b_actual);
  Jacobian_delta(Y, DELTA_R) = sinth/2 + (delta_s * costh)/(2 * b_actual);
  Jacobian_delta(Y, DELTA_L) = sinth/2 - (delta_s * costh)/(2 * b_actual);
  Jacobian_delta(THETA, DELTA_R) =  1/b_actual;
  Jacobian_delta(THETA, DELTA_L) = -1/b_actual;

  Eigen::Matrix<double, 2, 2> Sigma_delta;
  Sigma_delta.setZero();

  Sigma_delta(DELTA_R, DELTA_R) = right.pos_rel_var;
  Sigma_delta(DELTA_L, DELTA_L) = left.pos_rel_var;

  // error propagation law
  Sigma_p = Jacobian_p     * Sigma_p     * Jacobian_p.transpose()       // covariances from previous state
          + Jacobian_delta * Sigma_delta * Jacobian_delta.transpose();  // covariances from new delta

  odom_out.pose.covariance[CovElem::lin_ang::linX_linX] = Sigma_p(X    , X    );
  odom_out.pose.covariance[CovElem::lin_ang::linX_linY] = Sigma_p(X    , Y    );
  odom_out.pose.covariance[CovElem::lin_ang::linX_angZ] = Sigma_p(X    , THETA);
  odom_out.pose.covariance[CovElem::lin_ang::linY_linX] = Sigma_p(Y    , X    );
  odom_out.pose.covariance[CovElem::lin_ang::linY_linY] = Sigma_p(Y    , Y    );
  odom_out.pose.covariance[CovElem::lin_ang::linY_angZ] = Sigma_p(Y    , THETA);
  odom_out.pose.covariance[CovElem::lin_ang::angZ_linX] = Sigma_p(THETA, X    );
  odom_out.pose.covariance[CovElem::lin_ang::angZ_linY] = Sigma_p(THETA, Y    );
  odom_out.pose.covariance[CovElem::lin_ang::angZ_angZ] = Sigma_p(THETA, THETA);

  // velocity is simply projected in driving direction
  odom_out.twist.twist.linear.x = vel * costh;
  odom_out.twist.twist.linear.y = vel * sinth;

  // compute Jacobian and covariances
  Eigen::Matrix<double, 2, 2> Jacobian_vel;
  Jacobian_vel.setZero();
  Jacobian_vel(VEL_X, VEL_R) = costh / 2.0;
  Jacobian_vel(VEL_X, VEL_L) = costh / 2.0;
  Jacobian_vel(VEL_Y, VEL_R) = sinth / 2.0;
  Jacobian_vel(VEL_Y, VEL_L) = sinth / 2.0;

  Eigen::Matrix<double, 2, 2> Sigma_vel;
  Sigma_vel.setZero();
  Sigma_vel(VEL_R, VEL_R) = right.vel_var;
  Sigma_vel(VEL_L, VEL_L) = left.vel_var;

  Eigen::Matrix<double, 2, 2> Sigma_v;
  Sigma_v = Jacobian_vel * Sigma_vel * Jacobian_vel.transpose();

  odom_out.twist.covariance[CovElem::lin_ang::linX_linX] = Sigma_v(VEL_X, VEL_X);
  odom_out.twist.covariance[CovElem::lin_ang::linX_linY] = Sigma_v(VEL_X, VEL_Y);
  odom_out.twist.covariance[CovElem::lin_ang::linY_linX] = Sigma_v(VEL_Y, VEL_X);
  odom_out.twist.covariance[CovElem::lin_ang::linY_linY] = Sigma_v(VEL_Y, VEL_Y);

  // send out odom
  odom_pub.publish(odom_out);

  // broadcast tf message
  if(broadcast_tf)
  {
    // tf to static frame
    geometry_msgs::TransformStamped tf_msg;
    tf::Transform trafo;
    tf_msg.header.stamp            = (use_sensor_time) ? msg->header.stamp : ros::Time::now();
    tf_msg.child_frame_id          = axis_frame;
    tf_msg.header.frame_id         = odom_out.header.frame_id;
    tf_msg.transform.translation.x = odom_out.pose.pose.position.x;
    tf_msg.transform.translation.y = odom_out.pose.pose.position.y;
    tf_msg.transform.translation.z = odom_out.pose.pose.position.z;
    tf_msg.transform.rotation.x    = odom_out.pose.pose.orientation.x;
    tf_msg.transform.rotation.y    = odom_out.pose.pose.orientation.y;
    tf_msg.transform.rotation.z    = odom_out.pose.pose.orientation.z;
    tf_msg.transform.rotation.w    = odom_out.pose.pose.orientation.w;
    br->sendTransform(tf_msg);
  }

  // save message to debug file
  if(use_bag){
    SaveOdomInCSV::writeMsg(odom_out, file_out_log);
  }
}

// read data from bag and feed it into the callback function
bool readFromBag(std::string bag_file_path)
{
  // open bag
  rosbag::Bag bag;
  bag.open(bag_file_path, rosbag::bagmode::Read);

  // topics to load
  std::vector<std::string> bag_topics;
  bag_topics.push_back(input_topic);

  // create bag view
  rosbag::View bag_view(bag, rosbag::TopicQuery(bag_topics));

  // loop over all messages
  BOOST_FOREACH(rosbag::MessageInstance const m, bag_view)
  {
    if(m.getTopic() == input_topic || ("/" + m.getTopic() == input_topic))
    {
      drive_ros_msgs::VehicleEncoder::ConstPtr enc = m.instantiate<drive_ros_msgs::VehicleEncoder>();
      if (enc != NULL){
        encoderCallback(enc);
      }
    }
  }

  // close bag
  bag.close();
  return true;
}

int main(int argc, char **argv)
{
  ros::init(argc, argv, "localize_wheel_odometry");
  ros::NodeHandle nh;
  ros::NodeHandle pnh("~");

  // get parameters
  b_actual = pnh.param<double>("b_actual", 0.22);
  ROS_INFO_STREAM("Loaded b_actual: " << b_actual);

  err_d = pnh.param<double>("err_d", 1);
  ROS_INFO_STREAM("Loaded err_d: " << err_d);

  err_s = pnh.param<double>("err_s", 1);
  ROS_INFO_STREAM("Loaded err_s: " << err_s);

  use_front = pnh.param<bool>("use_front", true);
  ROS_INFO_STREAM("Loaded use_front: " << use_front);

  static_frame_id = pnh.param<std::string>("static_frame_id", "");
  ROS_INFO_STREAM("Loaded static_frame: " << static_frame_id);

  input_topic = pnh.param<std::string>("input_topic", "");
  ROS_INFO_STREAM("Loaded input_topic: " << input_topic);

  broadcast_tf = pnh.param<bool>("broadcast_tf", true);
  ROS_INFO_STREAM("Loaded broadcast_tf: " << broadcast_tf);

  int filter_length = pnh.param<int>("filter_length", 10);
  ROS_INFO_STREAM("Loaded theta_filter_length: " << filter_length);

  axis_frame = pnh.param<std::string>("axis_frame", "");
  ROS_INFO_STREAM("Loaded axis_frame: " << axis_frame);

  use_bag = pnh.param<bool>("use_bag", false);
  ROS_INFO_STREAM("Loaded use_bag: " << use_bag);

  std::string bag_file = pnh.param<std::string>("bag_file", "");
  ROS_INFO_STREAM("Loaded bag_file: " << bag_file);

  std::string csv_out = pnh.param<std::string>("csv_out", "");
  ROS_INFO_STREAM("Loaded csv_out: " << csv_out);

  pnh.getParam("initial_cov", initial_cov);

  // initialize filter
  theta_filter   = new MovingAverage(filter_length);
  delta_s_filter = new MovingAverage(filter_length);
  v_filter = new MovingAverage(filter_length);

  // reset everything
  reset();

  // setup subscriber and publisher
  if(broadcast_tf)
  {
    br = new tf::TransformBroadcaster;
  }
  reset_svr = pnh.advertiseService("reset_odom", svrResetOdom);
  odom_pub = pnh.advertise<nav_msgs::Odometry>("odom_out", 100);

  // check if read from bag
  if(use_bag){

    SaveOdomInCSV::writeHeader(csv_out, file_out_log);

    // read data directly from a bag
    readFromBag(bag_file);
    ROS_INFO_STREAM("Finished Reading from Bag");

  }else{

    ros::Subscriber sub = pnh.subscribe(input_topic, 100, encoderCallback);
    // spin node normally
    while(ros::ok()){
      ros::spin();
    }
  }

  ROS_INFO_STREAM("Exit Node");
  pnh.shutdown();
  nh.shutdown();
  return 0;
}


//==============================================================================
//         Copyright 2003 - 2012   LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2012   LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_REDUCTION_FUNCTIONS_SUM_HPP_INCLUDED
#define BOOST_SIMD_REDUCTION_FUNCTIONS_SUM_HPP_INCLUDED
#include <boost/simd/include/functor.hpp>
#include <boost/dispatch/include/functor.hpp>
#include <boost/simd/operator/functions/plus.hpp>
#include <boost/simd/constant/constants/zero.hpp>

namespace boost { namespace simd { namespace tag
  {
   /*!
     @brief sum generic tag

     Represents the sum function in generic contexts.

     @par Models:
        Hierarchy
   **/
    struct sum_ : ext::reduction_<sum_, tag::plus_, tag::Zero>
    {
      /// @brief Parent hierarchy
      typedef ext::reduction_<sum_, tag::plus_, tag::Zero> parent;
      template<class... Args>
      static BOOST_FORCEINLINE BOOST_AUTO_DECLTYPE dispatch(Args&&... args)
      BOOST_AUTO_DECLTYPE_BODY( dispatching_sum_( ext::adl_helper(), static_cast<Args&&>(args)... ) )
    };
  }
  namespace ext
  {
   template<class Site, class... Ts>
   BOOST_FORCEINLINE generic_dispatcher<tag::sum_, Site> dispatching_sum_(adl_helper, boost::dispatch::meta::unknown_<Site>, boost::dispatch::meta::unknown_<Ts>...)
   {
     return generic_dispatcher<tag::sum_, Site>();
   }
   template<class... Args>
   struct impl_sum_;
  }
  /*!
    Returns the sum of the elements of the SIMD vector

    @par Semantic:

    For every parameter of type T0

    @code
    scalar<T0> r = sum(a0);
    @endcode

    is similar to:

    @code
    scalar<T0> r = Zero;
    for(result_type i = 0; i != meta::cardinal_of<T0>; ++i)
      r += a0[i];
    @endcode

    @param a0

    @return a value of the scalar type associated to the parameter
  **/
  BOOST_DISPATCH_FUNCTION_IMPLEMENTATION(tag::sum_, sum, 1)
  BOOST_DISPATCH_FUNCTION_IMPLEMENTATION(tag::sum_, sum, 2)

} }
#endif


//  Copyright John Maddock 2009.
//  Distributed under the Boost
//  Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#include <boost/integer/integer_mask.hpp> // must be the only #include!

int main()
{
   boost::high_bit_mask_t<20>::least l = boost::high_bit_mask_t<20>::high_bit;
   boost::high_bit_mask_t<12>::fast f = boost::high_bit_mask_t<12>::high_bit_fast;
   l += f + boost::high_bit_mask_t<12>::bit_position;
   (void)l;
   boost::low_bits_mask_t<20>::least l2 = boost::low_bits_mask_t<20>::sig_bits;
   boost::low_bits_mask_t<12>::fast f2 = boost::low_bits_mask_t<12>::sig_bits_fast;
   l2 += f2 + boost::low_bits_mask_t<12>::bit_count;
   (void)l2;
}


#ifndef SOPUHS_TESTS_HPP
#define SOPUHS_TESTS_HPP

#include <array>

#include <Eigen/StdVector>
#include <unsupported/Eigen/MatrixFunctions>

#include <sophus/average.hpp>
#include <sophus/interpolate.hpp>
#include <sophus/num_diff.hpp>
#include <sophus/test_macros.hpp>

#ifdef SOPHUS_CERES
#include <ceres/jet.h>
#endif

namespace Sophus {

template <class LieGroup_>
class LieGroupTests {
 public:
  using LieGroup = LieGroup_;
  using Scalar = typename LieGroup::Scalar;
  using Transformation = typename LieGroup::Transformation;
  using Tangent = typename LieGroup::Tangent;
  using Point = typename LieGroup::Point;
  using Line = typename LieGroup::Line;
  using Adjoint = typename LieGroup::Adjoint;
  static int constexpr N = LieGroup::N;
  static int constexpr DoF = LieGroup::DoF;
  static int constexpr num_parameters = LieGroup::num_parameters;

  LieGroupTests(
      std::vector<LieGroup, Eigen::aligned_allocator<LieGroup>> const&
          group_vec,
      std::vector<Tangent, Eigen::aligned_allocator<Tangent>> const&
          tangent_vec,
      std::vector<Point, Eigen::aligned_allocator<Point>> const& point_vec)
      : group_vec_(group_vec),
        tangent_vec_(tangent_vec),
        point_vec_(point_vec) {}

  bool adjointTest() {
    bool passed = true;
    for (size_t i = 0; i < group_vec_.size(); ++i) {
      Transformation T = group_vec_[i].matrix();
      Adjoint Ad = group_vec_[i].Adj();
      for (size_t j = 0; j < tangent_vec_.size(); ++j) {
        Tangent x = tangent_vec_[j];

        Transformation I;
        I.setIdentity();
        Tangent ad1 = Ad * x;
        Tangent ad2 = LieGroup::vee(T * LieGroup::hat(x) *
                                    group_vec_[i].inverse().matrix());
        SOPHUS_TEST_APPROX(passed, ad1, ad2, Scalar(10) * kSmallEps,
                           "Adjoint case %, %", i, j);
      }
    }
    return passed;
  }

  bool contructorAndAssignmentTest() {
    bool passed = true;
    for (LieGroup foo_T_bar : group_vec_) {
      LieGroup foo_T2_bar = foo_T_bar;
      SOPHUS_TEST_APPROX(passed, foo_T_bar.matrix(), foo_T2_bar.matrix(),
                         kSmallEps, "Copy constructor: %\nvs\n %",
                         transpose(foo_T_bar.matrix()),
                         transpose(foo_T2_bar.matrix()));
      LieGroup foo_T3_bar;
      foo_T3_bar = foo_T_bar;
      SOPHUS_TEST_APPROX(passed, foo_T_bar.matrix(), foo_T3_bar.matrix(),
                         kSmallEps, "Copy assignment: %\nvs\n %",
                         transpose(foo_T_bar.matrix()),
                         transpose(foo_T3_bar.matrix()));

      LieGroup foo_T4_bar(foo_T_bar.matrix());
      SOPHUS_TEST_APPROX(
          passed, foo_T_bar.matrix(), foo_T4_bar.matrix(), kSmallEps,
          "Constructor from homogeneous matrix: %\nvs\n %",
          transpose(foo_T_bar.matrix()), transpose(foo_T4_bar.matrix()));

      Eigen::Map<LieGroup> foo_Tmap_bar(foo_T_bar.data());
      LieGroup foo_T5_bar = foo_Tmap_bar;
      SOPHUS_TEST_APPROX(
          passed, foo_T_bar.matrix(), foo_T5_bar.matrix(), kSmallEps,
          "Assignment from Eigen::Map type: %\nvs\n %",
          transpose(foo_T_bar.matrix()), transpose(foo_T5_bar.matrix()));

      Eigen::Map<LieGroup const> foo_Tcmap_bar(foo_T_bar.data());
      LieGroup foo_T6_bar;
      foo_T6_bar = foo_Tcmap_bar;
      SOPHUS_TEST_APPROX(
          passed, foo_T_bar.matrix(), foo_T5_bar.matrix(), kSmallEps,
          "Assignment from Eigen::Map type: %\nvs\n %",
          transpose(foo_T_bar.matrix()), transpose(foo_T5_bar.matrix()));

      LieGroup I;
      Eigen::Map<LieGroup> foo_Tmap2_bar(I.data());
      foo_Tmap2_bar = foo_T_bar;
      SOPHUS_TEST_APPROX(passed, foo_Tmap2_bar.matrix(), foo_T_bar.matrix(),
                         kSmallEps, "Assignment to Eigen::Map type: %\nvs\n %",
                         transpose(foo_Tmap2_bar.matrix()),
                         transpose(foo_T_bar.matrix()));
    }
    return passed;
  }

  bool derivativeTest() {
    bool passed = true;

    LieGroup g;
    for (int i = 0; i < DoF; ++i) {
      Transformation Gi = g.Dxi_exp_x_matrix_at_0(i);
      Transformation Gi2 = curveNumDiff(
          [i](Scalar xi) -> Transformation {
            Tangent x;
            setToZero(x);
            setElementAt(x, xi, i);
            return LieGroup::exp(x).matrix();
          },
          Scalar(0));
      SOPHUS_TEST_APPROX(passed, Gi, Gi2, kSmallEpsSqrt,
                         "Dxi_exp_x_matrix_at_ case %", i);
    }

    return passed;
  }

  template <class G = LieGroup>
  enable_if_t<std::is_same<G, Sophus::SO2<Scalar>>::value ||
                  std::is_same<G, Sophus::SO3<Scalar>>::value ||
                  std::is_same<G, Sophus::SE2<Scalar>>::value ||
                  std::is_same<G, Sophus::SE3<Scalar>>::value,
              bool>
  additionalDerivativeTest() {
    bool passed = true;
    for (size_t j = 0; j < tangent_vec_.size(); ++j) {
      Tangent a = tangent_vec_[j];
      Eigen::Matrix<Scalar, num_parameters, DoF> J = LieGroup::Dx_exp_x(a);
      Eigen::Matrix<Scalar, num_parameters, DoF> J_num =
          vectorFieldNumDiff<Scalar, num_parameters, DoF>(
              [](Tangent const& x) -> Sophus::Vector<Scalar, num_parameters> {
                return LieGroup::exp(x).params();
              },
              a);

      SOPHUS_TEST_APPROX(passed, J, J_num, 3 * kSmallEpsSqrt,
                         "Dx_exp_x case: %", j);
    }

    Tangent o;
    setToZero(o);
    Eigen::Matrix<Scalar, num_parameters, DoF> J = LieGroup::Dx_exp_x_at_0();
    Eigen::Matrix<Scalar, num_parameters, DoF> J_num =
        vectorFieldNumDiff<Scalar, num_parameters, DoF>(
            [](Tangent const& x) -> Sophus::Vector<Scalar, num_parameters> {
              return LieGroup::exp(x).params();
            },
            o);
    SOPHUS_TEST_APPROX(passed, J, J_num, kSmallEpsSqrt, "Dx_exp_x_at_0");

    for (size_t i = 0; i < group_vec_.size(); ++i) {
      LieGroup T = group_vec_[i];

      Eigen::Matrix<Scalar, num_parameters, DoF> J = T.Dx_this_mul_exp_x_at_0();
      Eigen::Matrix<Scalar, num_parameters, DoF> J_num =
          vectorFieldNumDiff<Scalar, num_parameters, DoF>(
              [T](Tangent const& x) -> Sophus::Vector<Scalar, num_parameters> {
                return (T * LieGroup::exp(x)).params();
              },
              o);

      SOPHUS_TEST_APPROX(passed, J, J_num, kSmallEpsSqrt,
                         "Dx_this_mul_exp_x_at_0 case: %", i);
    }

    return passed;
  }

  template <class G = LieGroup>
  enable_if_t<!std::is_same<G, Sophus::SO2<Scalar>>::value &&
                  !std::is_same<G, Sophus::SO3<Scalar>>::value &&
                  !std::is_same<G, Sophus::SE2<Scalar>>::value &&
                  !std::is_same<G, Sophus::SE3<Scalar>>::value,
              bool>
  additionalDerivativeTest() {
    return true;
  }

  bool expLogTest() {
    bool passed = true;

    for (size_t i = 0; i < group_vec_.size(); ++i) {
      Transformation T1 = group_vec_[i].matrix();
      Transformation T2 = LieGroup::exp(group_vec_[i].log()).matrix();
      SOPHUS_TEST_APPROX(passed, T1, T2, kSmallEps, "G - exp(log(G)) case: %",
                         i);
    }
    return passed;
  }

  bool expMapTest() {
    bool passed = true;
    for (size_t i = 0; i < tangent_vec_.size(); ++i) {
      Tangent omega = tangent_vec_[i];
      Transformation exp_x = LieGroup::exp(omega).matrix();
      Transformation expmap_hat_x = (LieGroup::hat(omega)).exp();
      SOPHUS_TEST_APPROX(passed, exp_x, expmap_hat_x, Scalar(10) * kSmallEps,
                         "expmap(hat(x)) - exp(x) case: %", i);
    }
    return passed;
  }

  bool groupActionTest() {
    bool passed = true;

    for (size_t i = 0; i < group_vec_.size(); ++i) {
      for (size_t j = 0; j < point_vec_.size(); ++j) {
        Point const& p = point_vec_[j];
        Transformation T = group_vec_[i].matrix();
        Point point1 = group_vec_[i] * p;
        Point point2 = map(T, p);
        SOPHUS_TEST_APPROX(passed, point1, point2, kSmallEps,
                           "Transform point case: %", i);
      }
    }
    return passed;
  }

  bool lineActionTest() {
    bool passed = point_vec_.size() > 1;

    for (size_t i = 0; i < group_vec_.size(); ++i) {
      for (size_t j = 0; j + 1 < point_vec_.size(); ++j) {
        Point const& p1 = point_vec_[j];
        Point const& p2 = point_vec_[j + 1];
        Line l = Line::Through(p1, p2);
        Point p1_t = group_vec_[i] * p1;
        Point p2_t = group_vec_[i] * p2;
        Line l_t = group_vec_[i] * l;

        SOPHUS_TEST_APPROX(passed, l_t.squaredDistance(p1_t),
                           static_cast<Scalar>(0), kSmallEps,
                           "Transform line case (1st point) : %", i);
        SOPHUS_TEST_APPROX(passed, l_t.squaredDistance(p2_t),
                           static_cast<Scalar>(0), kSmallEps,
                           "Transform line case (2nd point) : %", i);
        SOPHUS_TEST_APPROX(passed, l_t.direction().squaredNorm(),
                           l.direction().squaredNorm(), kSmallEps,
                           "Transform line case (direction) : %", i);
      }
    }
    return passed;
  }

  bool lieBracketTest() {
    bool passed = true;
    for (size_t i = 0; i < tangent_vec_.size(); ++i) {
      for (size_t j = 0; j < tangent_vec_.size(); ++j) {
        Tangent tangent1 =
            LieGroup::lieBracket(tangent_vec_[i], tangent_vec_[j]);
        Transformation hati = LieGroup::hat(tangent_vec_[i]);
        Transformation hatj = LieGroup::hat(tangent_vec_[j]);

        Tangent tangent2 = LieGroup::vee(hati * hatj - hatj * hati);
        SOPHUS_TEST_APPROX(passed, tangent1, tangent2, kSmallEps,
                           "Lie Bracket case: %", i);
      }
    }
    return passed;
  }

  bool veeHatTest() {
    bool passed = true;
    for (size_t i = 0; i < tangent_vec_.size(); ++i) {
      SOPHUS_TEST_APPROX(passed, Tangent(tangent_vec_[i]),
                         LieGroup::vee(LieGroup::hat(tangent_vec_[i])),
                         kSmallEps, "Hat-vee case: %", i);
    }
    return passed;
  }

  bool newDeleteSmokeTest() {
    bool passed = true;
    LieGroup* raw_ptr = nullptr;
    raw_ptr = new LieGroup();
    SOPHUS_TEST_NEQ(passed, reinterpret_cast<std::uintptr_t>(raw_ptr), 0);
    delete raw_ptr;
    return passed;
  }

  bool interpolateAndMeanTest() {
    bool passed = true;
    using std::sqrt;
    Scalar const eps = Constants<Scalar>::epsilon();
    Scalar const sqrt_eps = sqrt(eps);
    // TODO: Improve accuracy of ``interpolate`` (and hence ``exp`` and ``log``)
    //       so that we can use more accurate bounds in these tests, i.e.
    //       ``eps`` instead of ``sqrt_eps``.

    for (LieGroup const& foo_T_bar : group_vec_) {
      for (LieGroup const& foo_T_baz : group_vec_) {
        // Test boundary conditions ``alpha=0`` and ``alpha=1``.
        LieGroup foo_T_quiz = interpolate(foo_T_bar, foo_T_baz, Scalar(0));
        SOPHUS_TEST_APPROX(passed, foo_T_quiz.matrix(), foo_T_bar.matrix(),
                           sqrt_eps);
        foo_T_quiz = interpolate(foo_T_bar, foo_T_baz, Scalar(1));
        SOPHUS_TEST_APPROX(passed, foo_T_quiz.matrix(), foo_T_baz.matrix(),
                           sqrt_eps);
      }
    }
    for (Scalar alpha :
         {Scalar(0.1), Scalar(0.5), Scalar(0.75), Scalar(0.99)}) {
      for (LieGroup const& foo_T_bar : group_vec_) {
        for (LieGroup const& foo_T_baz : group_vec_) {
          LieGroup foo_T_quiz = interpolate(foo_T_bar, foo_T_baz, alpha);
          // test left-invariance:
          //
          // dash_T_foo * interp(foo_T_bar, foo_T_baz)
          // == interp(dash_T_foo * foo_T_bar, dash_T_foo * foo_T_baz)

          if (interp_details::Traits<LieGroup>::hasShortestPathAmbiguity(
                  foo_T_bar.inverse() * foo_T_baz)) {
            // skip check since there is a shortest path ambiguity
            continue;
          }
          for (LieGroup const& dash_T_foo : group_vec_) {
            LieGroup dash_T_quiz = interpolate(dash_T_foo * foo_T_bar,
                                               dash_T_foo * foo_T_baz, alpha);
            SOPHUS_TEST_APPROX(passed, dash_T_quiz.matrix(),
                               (dash_T_foo * foo_T_quiz).matrix(), sqrt_eps);
          }
          // test inverse-invariance:
          //
          // interp(foo_T_bar, foo_T_baz).inverse()
          // == interp(foo_T_bar.inverse(), dash_T_foo.inverse())
          LieGroup quiz_T_foo =
              interpolate(foo_T_bar.inverse(), foo_T_baz.inverse(), alpha);
          SOPHUS_TEST_APPROX(passed, quiz_T_foo.inverse().matrix(),
                             foo_T_quiz.matrix(), sqrt_eps);
        }
      }

      for (LieGroup const& bar_T_foo : group_vec_) {
        for (LieGroup const& baz_T_foo : group_vec_) {
          LieGroup quiz_T_foo = interpolate(bar_T_foo, baz_T_foo, alpha);
          // test right-invariance:
          //
          // interp(bar_T_foo, bar_T_foo) * foo_T_dash
          // == interp(bar_T_foo * foo_T_dash, bar_T_foo * foo_T_dash)

          if (interp_details::Traits<LieGroup>::hasShortestPathAmbiguity(
                  bar_T_foo * baz_T_foo.inverse())) {
            // skip check since there is a shortest path ambiguity
            continue;
          }
          for (LieGroup const& foo_T_dash : group_vec_) {
            LieGroup quiz_T_dash = interpolate(bar_T_foo * foo_T_dash,
                                               baz_T_foo * foo_T_dash, alpha);
            SOPHUS_TEST_APPROX(passed, quiz_T_dash.matrix(),
                               (quiz_T_foo * foo_T_dash).matrix(), sqrt_eps);
          }
        }
      }
    }

    for (LieGroup const& foo_T_bar : group_vec_) {
      for (LieGroup const& foo_T_baz : group_vec_) {
        if (interp_details::Traits<LieGroup>::hasShortestPathAmbiguity(
                foo_T_bar.inverse() * foo_T_baz)) {
          // skip check since there is a shortest path ambiguity
          continue;
        }

        // test average({A, B}) == interp(A, B):
        LieGroup foo_T_quiz = interpolate(foo_T_bar, foo_T_baz, 0.5);
        optional<LieGroup> foo_T_iaverage = iterativeMean(
            std::array<LieGroup, 2>({{foo_T_bar, foo_T_baz}}), 20);
        optional<LieGroup> foo_T_average =
            average(std::array<LieGroup, 2>({{foo_T_bar, foo_T_baz}}));
        SOPHUS_TEST(passed, bool(foo_T_average),
                    "log(foo_T_bar): %\nlog(foo_T_baz): %",
                    transpose(foo_T_bar.log()), transpose(foo_T_baz.log()));
        if (foo_T_average) {
          SOPHUS_TEST_APPROX(
              passed, foo_T_quiz.matrix(), foo_T_average->matrix(), sqrt_eps,
              "log(foo_T_bar): %\nlog(foo_T_baz): %\n"
              "log(interp): %\nlog(average): %",
              transpose(foo_T_bar.log()), transpose(foo_T_baz.log()),
              transpose(foo_T_quiz.log()), transpose(foo_T_average->log()));
        }
        SOPHUS_TEST(passed, bool(foo_T_iaverage),
                    "log(foo_T_bar): %\nlog(foo_T_baz): %\n"
                    "log(interp): %\nlog(iaverage): %",
                    transpose(foo_T_bar.log()), transpose(foo_T_baz.log()),
                    transpose(foo_T_quiz.log()),
                    transpose(foo_T_iaverage->log()));
        if (foo_T_iaverage) {
          SOPHUS_TEST_APPROX(
              passed, foo_T_quiz.matrix(), foo_T_iaverage->matrix(), sqrt_eps,
              "log(foo_T_bar): %\nlog(foo_T_baz): %",
              transpose(foo_T_bar.log()), transpose(foo_T_baz.log()));
        }
      }
    }

    return passed;
  }

  bool testRandomSmoke() {
    bool passed = true;
    std::default_random_engine engine;
    for (int i = 0; i < 100; ++i) {
      LieGroup g = LieGroup::sampleUniform(engine);
      SOPHUS_TEST_EQUAL(passed, g.params(), g.params());
    }
    return passed;
  }

  template <class S = Scalar>
  enable_if_t<std::is_floating_point<S>::value, bool> doAllTestsPass() {
    return doesLargeTestSetPass();
  }

  template <class S = Scalar>
  enable_if_t<!std::is_floating_point<S>::value, bool> doAllTestsPass() {
    return doesSmallTestSetPass();
  }

 private:
  bool doesSmallTestSetPass() {
    bool passed = true;
    passed &= adjointTest();
    passed &= contructorAndAssignmentTest();
    passed &= expLogTest();
    passed &= groupActionTest();
    passed &= lineActionTest();
    passed &= lieBracketTest();
    passed &= veeHatTest();
    passed &= newDeleteSmokeTest();
    return passed;
  }

  bool doesLargeTestSetPass() {
    bool passed = true;
    passed &= doesSmallTestSetPass();
    passed &= additionalDerivativeTest();
    passed &= derivativeTest();
    passed &= expMapTest();
    passed &= interpolateAndMeanTest();
    passed &= testRandomSmoke();
    return passed;
  }

  Scalar const kSmallEps = Constants<Scalar>::epsilon();
  Scalar const kSmallEpsSqrt = Constants<Scalar>::epsilonSqrt();

  Eigen::Matrix<Scalar, N - 1, 1> map(
      Eigen::Matrix<Scalar, N, N> const& T,
      Eigen::Matrix<Scalar, N - 1, 1> const& p) {
    return T.template topLeftCorner<N - 1, N - 1>() * p +
           T.template topRightCorner<N - 1, 1>();
  }

  Eigen::Matrix<Scalar, N, 1> map(Eigen::Matrix<Scalar, N, N> const& T,
                                  Eigen::Matrix<Scalar, N, 1> const& p) {
    return T * p;
  }

  std::vector<LieGroup, Eigen::aligned_allocator<LieGroup>> group_vec_;
  std::vector<Tangent, Eigen::aligned_allocator<Tangent>> tangent_vec_;
  std::vector<Point, Eigen::aligned_allocator<Point>> point_vec_;
};

template <class Scalar>
std::vector<SE3<Scalar>, Eigen::aligned_allocator<SE3<Scalar>>> getTestSE3s() {
  Scalar const kPi = Constants<Scalar>::pi();
  std::vector<SE3<Scalar>, Eigen::aligned_allocator<SE3<Scalar>>> se3_vec;
  se3_vec.push_back(SE3<Scalar>(
      SO3<Scalar>::exp(Vector3<Scalar>(Scalar(0.2), Scalar(0.5), Scalar(0.0))),
      Vector3<Scalar>(Scalar(0), Scalar(0), Scalar(0))));
  se3_vec.push_back(SE3<Scalar>(
      SO3<Scalar>::exp(Vector3<Scalar>(Scalar(0.2), Scalar(0.5), Scalar(-1.0))),
      Vector3<Scalar>(Scalar(10), Scalar(0), Scalar(0))));
  se3_vec.push_back(SE3<Scalar>::trans(Scalar(0), Scalar(100), Scalar(5)));
  se3_vec.push_back(SE3<Scalar>::rotZ(Scalar(0.00001)));
  se3_vec.push_back(
      SE3<Scalar>::trans(Scalar(0), Scalar(-0.00000001), Scalar(0.0000000001)) *
      SE3<Scalar>::rotZ(Scalar(0.00001)));
  se3_vec.push_back(SE3<Scalar>::transX(Scalar(0.01)) *
                    SE3<Scalar>::rotZ(Scalar(0.00001)));
  se3_vec.push_back(SE3<Scalar>::trans(Scalar(4), Scalar(-5), Scalar(0)) *
                    SE3<Scalar>::rotX(kPi));
  se3_vec.push_back(
      SE3<Scalar>(SO3<Scalar>::exp(
                      Vector3<Scalar>(Scalar(0.2), Scalar(0.5), Scalar(0.0))),
                  Vector3<Scalar>(Scalar(0), Scalar(0), Scalar(0))) *
      SE3<Scalar>::rotX(kPi) *
      SE3<Scalar>(SO3<Scalar>::exp(Vector3<Scalar>(Scalar(-0.2), Scalar(-0.5),
                                                   Scalar(-0.0))),
                  Vector3<Scalar>(Scalar(0), Scalar(0), Scalar(0))));
  se3_vec.push_back(
      SE3<Scalar>(SO3<Scalar>::exp(
                      Vector3<Scalar>(Scalar(0.3), Scalar(0.5), Scalar(0.1))),
                  Vector3<Scalar>(Scalar(2), Scalar(0), Scalar(-7))) *
      SE3<Scalar>::rotX(kPi) *
      SE3<Scalar>(SO3<Scalar>::exp(Vector3<Scalar>(Scalar(-0.3), Scalar(-0.5),
                                                   Scalar(-0.1))),
                  Vector3<Scalar>(Scalar(0), Scalar(6), Scalar(0))));
  return se3_vec;
}

template <class T>
std::vector<SE2<T>, Eigen::aligned_allocator<SE2<T>>> getTestSE2s() {
  T const kPi = Constants<T>::pi();
  std::vector<SE2<T>, Eigen::aligned_allocator<SE2<T>>> se2_vec;
  se2_vec.push_back(SE2<T>());
  se2_vec.push_back(SE2<T>(SO2<T>(0.2), Vector2<T>(10, 0)));
  se2_vec.push_back(SE2<T>::transY(100));
  se2_vec.push_back(SE2<T>(SO2<T>(-1.), Vector2<T>(20, -1)));
  se2_vec.push_back(
      SE2<T>(SO2<T>(0.00001), Vector2<T>(-0.00000001, 0.0000000001)));
  se2_vec.push_back(SE2<T>(SO2<T>(0.3), Vector2<T>(2, 0)) * SE2<T>::rot(kPi) *
                    SE2<T>(SO2<T>(-0.3), Vector2<T>(0, 6)));
  return se2_vec;
}
}  // namespace Sophus
#endif  // TESTS_HPP


// Copyright (c) 2020 Graphcore Ltd. All rights reserved.

// Tool for applying various simple processing operations
// to EXR image files.
#include <cstdlib>
#include <vector>
#include <iostream>

#include <spdlog/spdlog.h>
#include <boost/program_options.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>

#include <light/src/exr/exr.hpp>

std::shared_ptr<spdlog::logger> logger() {
    static auto logger = spdlog::stdout_logger_mt("exrtool_logger");
    return spdlog::get("exrtool_logger");
}

boost::program_options::variables_map
parseOptions(int argc, char** argv) {

  namespace po = boost::program_options;
  po::options_description desc("Program for manipulating EXR files. Options");
  desc.add_options()
  ("help", "Show command help.")
  ("inputs,i", po::value<std::vector<std::string>>()->required()->multitoken(),
   "Input file(s)")
  ("output,o", po::value<std::string>()->required(), "Set the output file name.")
  ("op", po::value<std::string>()->default_value("error"),
    "Select processing operation from ['error', 'mean']."
    "\t\nerror: Compute the MSE and PSNR between two images (the first image is considered to be the noise free one)."
    "\t\nmean: Compute the mean of a list of images.");

  po::variables_map vm;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  if (vm.count("help")) {
    std::cout << desc << "\n";
    throw std::runtime_error("Show help");
  }

  po::notify(vm);

  auto numInputs = vm.at("inputs").as<std::vector<std::string>>().size();
  auto op = vm.at("op").as<std::string>();
  if (op == "error" && numInputs != 2) {
    logger()->error("'{}' operation requires exactly 2 inputs.", op);
    throw std::logic_error("Incorrect number of inputs for op.");
  }

  if (op != "error" && op != "mean") {
    logger()->error("Unknown operation '{}'.", op);
    throw std::logic_error("Unknown image operation");
  }

  return vm;
}

// Compute mean of all the images (e.g. use to combine
// parallel renders of same scene):
exr::Image meanImage(const std::vector<exr::Image>& images) {
  exr::Image mean(images.front().width, images.front().height);
  for (const auto& image : images) {
    for (auto& p : image.slices) {
      mean.slices[p.first] = std::vector<float>(p.second.size());
    }
  }

  logger()->info("Averaging {} images", images.size());
  const float scale = 1.f / images.size();
  for (const auto& image : images) {
    for (auto& p : image.slices) {
      const auto& name = p.first;
      auto& result = mean.slices.at(name);
      const auto& data = image.slices.at(name);
      for (auto i = 0u; i < data.size(); ++i) {
        result[i] += scale * data[i];
      }
    }
  }
  return mean;
}

struct Stats {
  double MSE;
  double PSNR;
};

// Compute difference image, mean squared error (MSE), and
// the peak signal to noise ratio (PSNR) between two input
// images. PSNR assumes the first image is the target/ground
// truth and so takes the max value from that image:
exr::Image errorImage(const std::vector<exr::Image>& images, Stats& errors) {
  auto err = images.front();
  auto& b = images.back();
  errors.MSE = 0.f;
  double max = 0.f;
  for (auto& p : err.slices) {
    const auto& name = p.first;
    auto& dataA = err.slices.at(name);
    auto& dataB = b.slices.at(name);
    if (dataA.size() != dataB.size()) {
      logger()->error("Image sizes do not match.");
      throw std::runtime_error("Incompatible image dimensions");
    }
    for (auto i = 0u; i < dataA.size(); ++i) {
      max = std::max(max, (double)dataA[i]);
      dataA[i] -= dataB[i];
      errors.MSE += dataA[i] * dataA[i];
    }
  }
  errors.MSE /= err.width * err.height;
  errors.PSNR = -10.0 * log10(errors.MSE / (max * max));
  return err;
}

int main(int argc, char** argv) {
  try {
    auto args = parseOptions(argc, argv);
    auto inputFiles = args.at("inputs").as<std::vector<std::string>>();
    auto outputFile = args.at("output").as<std::string>();
    auto op = args.at("op").as<std::string>();

    std::vector<exr::Image> images;
    images.reserve(inputFiles.size());
    for (const auto& fileName : inputFiles) {
      images.push_back(exr::read(fileName));
    }

    if (op == "error") {
      Stats errors;
      auto diff = errorImage(images, errors);
      logger()->info("Statistics: MSE {} PSNR {} dB", errors.MSE, errors.PSNR);
      exr::write(outputFile, diff);
    }

    if (op == "mean") {
      auto mean = meanImage(images);
      exr::write(outputFile, mean);
    }
  } catch (const std::exception& e) {
    logger()->error("{}", e.what());
    return EXIT_FAILURE;
  }

  return EXIT_SUCCESS;
}


/*
// Copyright 2018 University of Liege
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors:
// - Adrien Crovato
*/

//// Field variables writing
// Write field variable in external .dat and .pos files
//
// I/O:
// - outPath: path to write file
// - Minf: freestream Mach number
// - alpha: freestream angle of attack
// - fPan: (network of) field panels

#include <iostream>
#include <iomanip>
#include <fstream>
#include <Eigen/Dense>
#include "write_fv.h"

using namespace std;
using namespace Eigen;

void write_fv(string outPath, double alpha, double Minf, Field &fPan) {

    //// Write to .dat file
    cout << "Writing field variables file in 'fv.dat'... " << flush;
    string outPathDat = outPath + "fv.dat";
    ofstream fv;
    fv.open(outPathDat);

    // General information (header)
    fv << "Field variables file" << endl << endl;
    fv.width(25); fv << left << "Angle of attack: "; fv.width(20); fv << left << alpha*180/3.14159 << endl;
    fv.width(25); fv << left << "Mach: "; fv.width(20); fv << left << Minf << endl;
    fv.width(25); fv << left << "Field cells: "; fv.width(20); fv << left << fPan.nF << endl << endl;

    // Labeling
    fv.width(15); fv << right << "x";
    fv.width(15); fv << right << "y";
    fv.width(15); fv << right << "z";
    fv.width(15); fv << right << "U_x";
    fv.width(15); fv << right << "U_y";
    fv.width(15); fv << right << "U_z";
    fv.width(15); fv << right << "rho";
    fv.width(15); fv << right << "rho_X";
    fv.width(15); fv << right << "rho_Z";
    fv.width(15); fv << right << "Mach";
    fv.width(15); fv << right << "Sigma";
    fv.width(15); fv << right << "Epsilon" << endl;

    // Write field quantities
    for (int i = 0; i < fPan.nF; ++i) {
        fv.width(15); fv << right << fPan.CG(i,0);
        fv.width(15); fv << right << fPan.CG(i,1);
        fv.width(15); fv << right << fPan.CG(i,2);
        fv.width(15); fv << right << fPan.U(i,0);
        fv.width(15); fv << right << fPan.U(i,1);
        fv.width(15); fv << right << fPan.U(i,2);
        fv.width(15); fv << right << fPan.rho(i);
        fv.width(15); fv << right << fPan.dRho(i,0);
        fv.width(15); fv << right << fPan.dRho(i,2);
        fv.width(15); fv << right << fPan.M(i);
        fv.width(15); fv << right << fPan.sigma(i);
        fv.width(15); fv << right << fPan.epsilon(i) << endl;
    }

    // Close file
    fv.close();
    cout << "Done!" << endl;

    //// Write to gmsh .pos file
    cout << "Writing field variables file in 'M.pos'... " << flush;
    string outPathPos = outPath + "M.pos";
    fv.open(outPathPos);

    // Gmsh header
    fv << "View \"M\" {" << endl;

    // Mach number
    for (int f = 0; f < fPan.nF; ++f) {
        fv << "SH(";
        fv << fPan.vX(f,0) << "," << fPan.vY(f,0) << "," << fPan.vZ(f,0) << ',';
        fv << fPan.vX(f,1) << "," << fPan.vY(f,0) << "," << fPan.vZ(f,0) << ',';
        fv << fPan.vX(f,1) << "," << fPan.vY(f,1) << "," << fPan.vZ(f,0) << ',';
        fv << fPan.vX(f,0) << "," << fPan.vY(f,1) << "," << fPan.vZ(f,0) << ',';
        fv << fPan.vX(f,0) << "," << fPan.vY(f,0) << "," << fPan.vZ(f,1) << ',';
        fv << fPan.vX(f,1) << "," << fPan.vY(f,0) << "," << fPan.vZ(f,1) << ',';
        fv << fPan.vX(f,1) << "," << fPan.vY(f,1) << "," << fPan.vZ(f,1) << ',';
        fv << fPan.vX(f,0) << "," << fPan.vY(f,1) << "," << fPan.vZ(f,1) << "){";
        for (int i = 0; i < 7; ++i)
            fv << fPan.M(f) << ",";
        fv << fPan.M(f) << "};" << endl;
    }

    // Gmsh footer
    fv << "};" << endl << endl;

    // Close file
    fv.close();
    cout << "Done!" << endl << endl;
}