#define _LIBCPP_ENABLE_THREAD_SAFETY_ANNOTATIONS
#include "recognition_performance.h"

#include "icp.h"
#include "keypoint_transform.h"

#include <boost/histogram/ostream.hpp>
#include <fstream>
#include <gsl/gsl>
#include <iterator>
#include <opencv2/core/hal/interface.h>
#include <opencv2/core/mat.hpp>
#include <opencv2/core/types.hpp>
#include <opencv2/features2d.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/rgbd/depth.hpp>
#include <sens_loc/analysis/distance.h>
#include <sens_loc/analysis/match.h>
#include <sens_loc/analysis/recognition_performance.h>
#include <sens_loc/camera_models/pinhole.h>
#include <sens_loc/camera_models/projection.h>
#include <sens_loc/io/histogram.h>
#include <sens_loc/io/image.h>
#include <sens_loc/io/intrinsics.h>
#include <sens_loc/io/pose.h>
#include <sens_loc/math/coordinate.h>
#include <sens_loc/math/image.h>
#include <sens_loc/math/pointcloud.h>
#include <sens_loc/plot/backprojection.h>
#include <sens_loc/util/console.h>
#include <sens_loc/util/thread_analysis.h>
#include <util/batch_visitor.h>
#include <util/statistic_visitor.h>

using namespace std;
using namespace sens_loc;
using namespace gsl;
using namespace cv;

namespace {

/// Capsulate all required data for back-and-forth projection as well
/// as precision-recall computation.
struct reprojection_data {
    vector<cv::KeyPoint> keypoints;
    cv::Mat              descriptors;
    math::image<ushort>  depth_image;
    math::pose_t         absolute_pose;

    reprojection_data(string_view feature_path,
                      string_view depth_path,
                      string_view pose_path) noexcept(false) {
        const FileStorage fs = io::open_feature_file(string(feature_path));
        keypoints            = io::load_keypoints(fs);
        descriptors          = io::load_descriptors(fs);

        optional<math::image<ushort>> d_img =
            io::load_image<ushort>(string(depth_path), IMREAD_UNCHANGED);
        if (!d_img) {
            ostringstream oss;
            oss << "Could not load depth image from " << depth_path << "!";
            throw runtime_error{oss.str()};
        }
        depth_image = move(*d_img);

        ifstream               pose_file{string(pose_path)};
        optional<math::pose_t> pose = io::load_pose(pose_file);
        if (!pose) {
            ostringstream oss;
            oss << "Could not load pose from " << pose_path << "!";
            throw runtime_error{oss.str()};
        }
        absolute_pose = move(*pose);
    }
};

size_t mask_backprojection(const math::image<uchar>& mask,
                           math::imagepoints_t&      points) noexcept {
    size_t counter = 0UL;
    for (math::pixel_coord<float>& p : points) {
        if (p.u() < 0.0F || p.v() < 0.0F ||
            p.u() > narrow_cast<float>(mask.w()) ||
            p.v() > narrow_cast<float>(mask.h()) || mask.at(p) == 0U) {
            p = {-1.0F, -1.0F};
            ++counter;
        }
    }
    return counter;
}

struct recognition_data {
    recognition_data() = default;

    void insert_recognition(span<float>                         distances,
                            const analysis::element_categories& classification,
                            size_t masked_points) noexcept {
        lock_guard l{_mutex};

        _selected_elements_distance.insert(_selected_elements_distance.end(),
                                           begin(distances), end(distances));
        _stats.account(classification);
        _totally_masked += narrow<int>(masked_points);
    }

    tuple<vector<float>, analysis::recognition_statistic, int64_t>
    extract() noexcept {
        lock_guard l{_mutex};
        tuple<vector<float>, analysis::recognition_statistic, int64_t> t{
            move(_selected_elements_distance), move(_stats), _totally_masked};

        _selected_elements_distance = vector<float>();
        _stats                      = analysis::recognition_statistic();
        _totally_masked             = 0L;
        return t;
    }

  private:
    mutex                                     _mutex;
    vector<float> _selected_elements_distance GUARDED_BY(_mutex);
    analysis::recognition_statistic _stats    GUARDED_BY(_mutex);
    int64_t _totally_masked                   GUARDED_BY(_mutex) = 0L;
};

template <template <typename> typename Model = sens_loc::camera_models::pinhole,
          typename Real                      = float>
class prec_recall_analysis {
  public:
    prec_recall_analysis(
        string_view                                      feature_file_pattern,
        const apps::recognition_analysis_input&          input,
        const apps::recognition_analysis_output_options& output_options,
        recognition_data&                                accumulated_data,
        const apps::backproject_config&                  backproject_config)
        : _feature_file_pattern{feature_file_pattern}
        , _input{input}
        , _output_options{output_options}
        , _matcher{cv::BFMatcher::create(_input.matching_norm,
                                         /*crosscheck=*/true)}
        , _mask{nullopt}
        , _accumulated_data{accumulated_data}
        , _backprojection_config{backproject_config} {
        Expects(!_feature_file_pattern.empty());
        Expects(!_input.depth_image_pattern.empty());
        Expects(!_input.pose_file_pattern.empty());
        Expects(!_input.intrinsic_file.empty());

        ifstream intrinsic{string(_input.intrinsic_file)};

        auto maybe_intrinsic =
            io::camera<Real, Model>::load_intrinsic(intrinsic);
        if (!maybe_intrinsic.has_value()) {
            stringstream ss;
            ss << "Intrinsic file " << _input.intrinsic_file
               << " could not be loaded!";
            throw std::invalid_argument{ss.str()};
        }
        _intrinsic = *maybe_intrinsic;

        if (_input.mask_file) {
            _mask = io::load_image<uchar>(string(*_input.mask_file),
                                          IMREAD_GRAYSCALE);
            if (!_mask) {
                auto s = synced();
                cerr << util::err{} << "Could not load mask file '"
                     << *_input.mask_file << "'!\n"
                     << "Continuing without masking.\n";
            }
            // The mask file is loaded. It is necessary to check, that the
            // dimensions match the loaded intrinsic. If not, masking is
            // disabled.
            else {
                if (_mask->w() != _intrinsic.w() ||
                    _mask->h() != _intrinsic.h()) {
                    auto s = synced();
                    cerr << util::err{}
                         << "Dimensions of the mask mismatch the dimensions of "
                            "the intrinsic!\n"
                         << "Continuing without masking masking.\n"
                         << "Mask {" << _mask->w() << ", " << _mask->h()
                         << "} vs Intrinsic {" << _intrinsic.w() << ", "
                         << _intrinsic.h() << "}\n";
                    // Resetting the mask to none!
                    _mask = nullopt;
                }
            }
        }

        if constexpr (is_same_v<decltype(_intrinsic),
                                camera_models::pinhole<Real>>) {
            _icp = rgbd::FastICPOdometry::create(
                apps::cv_camera_matrix(_intrinsic));
        }
    }

    void operator()(int idx,
                    // NOLINTNEXTLINE(performance-unnecessary-value-param)
                    optional<vector<KeyPoint>> keypoints,
                    // NOLINTNEXTLINE(performance-unnecessary-value-param)
                    optional<Mat> descriptors) noexcept try {
        Expects(!keypoints);
        Expects(!descriptors);

        const int previous_idx = idx - 1;

        using namespace math;
        using namespace apps;

        reprojection_data prev{
            fmt::format(_feature_file_pattern, previous_idx),
            fmt::format(_input.depth_image_pattern, previous_idx),
            fmt::format(_input.pose_file_pattern, previous_idx)};
        reprojection_data curr{fmt::format(_feature_file_pattern, idx),
                               fmt::format(_input.depth_image_pattern, idx),
                               fmt::format(_input.pose_file_pattern, idx)};

        if (prev.keypoints.empty() || curr.keypoints.empty())
            return;

        // == Calculate relative pose between the two frames.
        pose_t rel_pose = relative_pose(prev.absolute_pose, curr.absolute_pose);

        // Refine that pose with an ICP if possible.
        if (_icp) {
            auto [icp_pose, icp_success] =
                refine_pose(*_icp, prev.depth_image, curr.depth_image,
                            _input.unit_factor, rel_pose);
            if (icp_success) {
                rel_pose = icp_pose;
            } else {
                auto s = synced();
                cerr << util::warn{} << "No ICP result for index " << idx
                     << "! Using loaded pose.\n";
            }
        }

        // == get keypoints as world points
        pointcloud_t prev_points = keypoints_to_pointcloud(
            prev.keypoints, prev.depth_image, _intrinsic, _input.unit_factor);

        imagepoints_t prev_in_img =
            project_to_other_camera(rel_pose, prev_points, _intrinsic);

        // Some backprojected points might land outside of the viewport of
        // the camera model, but still on valid pixel coordinates (inside the
        // image). This can be detected with using a mask.
        // Such points are set to {-1, -1} to be recognizable invalid.
        const size_t masked_points =
            _mask ? mask_backprojection(*_mask, prev_in_img) : 0UL;
        Expects(prev_in_img.size() == prev.keypoints.size());

        // == Match the keypoints with cross-checking.
        vector<DMatch> matches;
        // QueryDescriptors: first argument
        // TrainDescriptors: second argument
        _matcher->match(curr.descriptors, prev.descriptors, matches);

        using analysis::element_categories;
        using camera_models::keypoint_to_coords;
        const imagepoints_t curr_keypoints = keypoint_to_coords(curr.keypoints);

        const element_categories classification(
            curr_keypoints, prev_in_img, matches,
            _input.keypoint_distance_threshold);
        // True positives keypoints from this and previous frame in this
        // frames coordinate system.
        auto [t_p_t, t_p_o] = analysis::gather_correspondences(
            classification.true_positives, curr_keypoints, prev_in_img);

        if (_output_options.backproject_pattern) {
            string orig_f_name =
                fmt::format(*_output_options.original_files, idx);
            optional<math::image<uchar>> orig_img =
                io::load_as_8bit_gray(orig_f_name);

            if (orig_img) {
                // Plot false positives with orange distance line
                auto [f_p_t, f_p_o] = analysis::gather_correspondences(
                    classification.false_positives, curr_keypoints,
                    prev_in_img);
                Mat out_img = plot::backprojection_correspondence(
                     orig_img->data(), f_p_t, f_p_o,  // Base image and keypoints
                    _backprojection_config.false_positive.color,
                    _backprojection_config.false_positive.strength);

                // Plot false negatives with purple line
                auto [f_n_t, f_n_o] = analysis::gather_correspondences(
                    classification.false_negatives, curr_keypoints,
                    prev_in_img);
                out_img = plot::backprojection_correspondence(
                    out_img, f_n_t, f_n_o,  // Base image and keypoints
                    _backprojection_config.false_negative.color,
                    _backprojection_config.false_negative.strength);

                // Plot true positives with green line.
                out_img = plot::backprojection_correspondence(
                    out_img, t_p_t, t_p_o,  // Keypoints
                    _backprojection_config.true_positive.color,
                    _backprojection_config.true_positive.strength);

                const bool write_success = imwrite(
                    fmt::format(*_output_options.backproject_pattern, idx),
                    out_img);

                if (!write_success) {
                    auto s = synced();
                    cerr << util::err{}
                         << "Could not write backprojection correspondence to "
                            "disk for: "
                         << idx << "!\n";
                }
            } else {
                auto s = synced();
                cerr << util::warn{} << "Original file for index '" << idx
                     << "' ('" << orig_f_name
                     << "') could not be loaded for backprojection plot!\n";
            }
        }

        auto distances = pointwise_distance(t_p_t, t_p_o);
        _accumulated_data.insert_recognition(distances, classification,
                                             masked_points);
    } catch (const exception& e) {
        auto s = synced();
        cerr << util::err{} << "Could not analyze data for " << idx << "!\n"
             << e.what() << "\n";
        return;
    }

    size_t postprocess() {
        auto [distances, classification, masked_point_count] =
            _accumulated_data.extract();

        if (classification.total_elements() == 0L)
            return 0UL;

        sort(begin(distances), end(distances));
        const auto         dist_bins = 20;
        analysis::distance distance_stat{distances, dist_bins,
                                         "backprojection error pixels"};

        if (_output_options.stat_file) {
            cv::FileStorage recognize_out{*_output_options.stat_file,
                                          cv::FileStorage::WRITE |
                                              cv::FileStorage::FORMAT_YAML};
            write(recognize_out, "classification", classification);
            write(recognize_out, "masked_points",
                  narrow<int>(masked_point_count));
            recognize_out.release();
        } else {
            using namespace boost;
            cout << "# relevant:   " << distance_stat.count() << "\n"
                 << "min:          " << distance_stat.min() << "\n"
                 << "max:          " << distance_stat.max() << "\n"
                 << "median:       " << distance_stat.median() << "\n"
                 << "mean:         " << distance_stat.mean() << "\n"
                 << "Variance:     " << distance_stat.variance() << "\n"
                 << "StdDev:       " << distance_stat.stddev() << "\n"
                 << "Skewness:     " << distance_stat.skewness() << "\n"
                 << "Precision     " << classification.precision() << "\n"
                 << "Recall:       " << classification.recall() << "\n"
                 << "Sensitivity:  " << classification.sensitivity() << "\n"
                 << "Specificity:  " << classification.specificity() << "\n"
                 << "Fallout:      " << classification.fallout() << "\n"
                 << "Rand-Index:   " << classification.rand_index() << "\n"
                 << "Youden-Index: " << classification.youden_index() << "\n"
                 << "Min-TruePos:  "
                 << accumulators::min(
                        classification.true_positive_distribution().stat)
                 << "\n"
                 << "Avg-TruePos:  "
                 << accumulators::mean(
                        classification.true_positive_distribution().stat)
                 << "\n"
                 << "Min-Relevant: "
                 << accumulators::min(
                        classification.relevant_element_distribution().stat)
                 << "\n"
                 << "Avg-Relevant: "
                 << accumulators::mean(
                        classification.relevant_element_distribution().stat)
                 << "\n"
                 << "Avg-FalsePos: "
                 << accumulators::mean(
                        classification.false_positive_distribution().stat)
                 << "\n"
                 << "Masked pts:   " << masked_point_count << "\n";
        }

        if (_output_options.backprojection_selected_histo) {
            ofstream gnuplot_data{
                *_output_options.backprojection_selected_histo};
            gnuplot_data << sens_loc::io::to_gnuplot(distance_stat.histogram())
                         << endl;
        } else {
            cout << distance_stat.histogram() << "\n";
        }

        classification.make_histogram();

        if (_output_options.relevant_histo) {
            ofstream gnuplot_data{*_output_options.relevant_histo};
            gnuplot_data
                << sens_loc::io::to_gnuplot(
                       classification.relevant_element_distribution().histo)
                << endl;
        } else {
            cout << classification.relevant_element_distribution().histo
                 << "\n";
        }

        if (_output_options.true_positive_histo) {
            ofstream gnuplot_data{*_output_options.true_positive_histo};
            gnuplot_data
                << sens_loc::io::to_gnuplot(
                       classification.true_positive_distribution().histo)
                << endl;
        } else {
            cout << classification.true_positive_distribution().histo << "\n";
        }

        if (_output_options.false_positive_histo) {
            ofstream gnuplot_data{*_output_options.false_positive_histo};
            gnuplot_data
                << sens_loc::io::to_gnuplot(
                       classification.false_positive_distribution().histo)
                << endl;
        } else {
            cout << classification.false_positive_distribution().histo << "\n";
        }

        if (_output_options.true_positive_distance_histo) {
            ofstream gnuplot_data{
                *_output_options.true_positive_distance_histo};
            gnuplot_data
                << sens_loc::io::to_gnuplot(
                       classification.true_positive_distance().histogram())
                << endl;
        } else {
            cout << classification.true_positive_distance().histogram() << "\n";
        }

        if (_output_options.false_positive_distance_histo) {
            ofstream gnuplot_data{
                *_output_options.false_positive_distance_histo};
            gnuplot_data
                << sens_loc::io::to_gnuplot(
                       classification.false_positive_distance().histogram())
                << endl;
        } else {
            cout << classification.false_positive_distance().histogram()
                 << "\n";
        }

        return classification.total_elements();
    }

  private:
    std::string_view                                 _feature_file_pattern;
    const apps::recognition_analysis_input&          _input;
    const apps::recognition_analysis_output_options& _output_options;

    Model<Real>                  _intrinsic;
    Ptr<rgbd::Odometry>          _icp;
    Ptr<BFMatcher>               _matcher;
    optional<math::image<uchar>> _mask;
    recognition_data&            _accumulated_data;

    const apps::backproject_config& _backprojection_config;
};
}  // namespace

namespace sens_loc::apps {
int analyze_recognition_performance(
    util::processing_input                     in,
    const recognition_analysis_input&          required_data,
    const recognition_analysis_output_options& output_options,
    const backproject_config&                  backproject_config) {
    Expects(in.start < in.end &&
            "Recognition Performance calculation requires at least two images");

    // The code will load all data directly and does not rely on loading through
    // the statistics code.
    using visitor =
        statistic_visitor<prec_recall_analysis<>, required_data::none>;

    recognition_data accumulator;
    // The odd-looking double arguments comes from the genericity of the
    // statistic-visitation. The first argument goes to \c statistic_visitor
    // and the second one to \c prec_recall_analysis
    auto analysis_v =
        visitor{in.input_pattern, in.input_pattern, required_data,
                output_options,   accumulator,      backproject_config};

    // Consecutive images are matched and analysed, therefore the first
    // index must be skipped.
    auto   f          = parallel_visitation(in.start + 1, in.end, analysis_v);
    size_t n_elements = f.postprocess();

    return n_elements > 0L ? 0 : 1;
}

}  // namespace sens_loc::apps


#include <iostream>

#include <boost/timer.hpp>

#include <vector>
#include <cmath>
#include <cstring>

#include "NRGclasses.hpp"
#include "NRGfunctions.hpp"
#include "SpecFuncClass.hpp"
//#include "CondSIAM.hpp"
#include "NRGOpMatRules.hpp"

#ifndef pi
#define pi 3.141592653589793238462643383279502884197169
#endif

int main (int argc, char* argv[]){

  CNRGCodeHandler ThisCode;

  CNRGbasisarray* AcutN;
  CNRGbasisarray AbasisN;

  CNRGbasisarray SingleSite;

  CNRGmatrix** OpArrayN;

//   vector<double> ParamsTemp;
  double betabar=0.727;
  double DM,TM;
  // 0 - betabar; 1 - TM; 2 - DM ?
  double Temp; // In units of the bandwidth (fixed!)
               // Temp/DN = TempBar = 1/betabar

  char CNsites[8];
  char CNmat[8];
  char arqname[32];
  char ext[32];

  int NshellMax=3;
  int NFermiOps=2;

  // Check time
  boost::timer MyTime;
  double time_elapsed;

  int UseCFS=0;
  int Nw=1; // No of omegas in each shell

  int c;
  opterr = 0;

  char cmNtemp[8],cbetabar[8];

  while ((c = getopt (argc, argv, "h b: n:")) != -1)
    switch (c){
    case 'b':
      strcpy(cbetabar,optarg);
      betabar=atof(cbetabar);
      break;
    case 'n':
      strcpy(cmNtemp,optarg);
      Nw = atoi(cmNtemp);
      break;
    case 'h':
      cout << "Usage: " << argv[0] << " (-b betabar -n NtempsEachN)" << endl;
      cout << " Default values: betabar=0.727 ; NtempsEachN=1 " << endl;
      // 	cout << " Additional options: " << endl;
      // 	cout << " -C : uses CFS method " << endl;
      exit(0);
      break;
    case '?':
      cout << "Usage: " << argv[0] << " -b betabar -n NtempsEachN" << endl;
      cout << " Default values: betabar=0.727 ; NtempsEachN=1 " << endl;
      exit(0);
      break;
    default:
      break;
    }
  // end CommandLine

  // Read code parameters, including z_twist
  ThisCode.ReadGenPars(true);

  // Set Temp
  cout << " betabar = " << betabar
       << " NomegasEachShell = " << Nw << " (if =1 then use standard interpolation) "
       << endl;

  cout << "Lambda    = " << ThisCode.Lambda << endl;
  cout << "Nsites0   = " << ThisCode.Nsites0 << endl;
  cout << "Nsitesmax = " << ThisCode.Nsitesmax << endl;
  cout << "NFermiOps = " << ThisCode.NopsSaved << endl;
  cout << "Ext arq   = " << ThisCode.SaveArraysFileName << endl;
  cout << "Symmetry  = " << ThisCode.SymNo << endl;
  cout << "Model     = " << ThisCode.ModelNo << endl;
  if ( (ThisCode.ModelNo != 0)  ){
    cout << " CalcSIAM: Only conductance in the Anderson model supported at this point." << endl;
    cout << " CalcSIAM: Proceed at your own risk ... " << endl;
  }
  if (ThisCode.totalS){cout << " SU(2) symmetry detected. " << endl;}
  cout << "Oliveira z = " << ThisCode.chain.z_twist << endl;

  ThisCode.ReadParams((char *)"input_nrg.dat",1); // Needs Gamma

  NshellMax=ThisCode.Nsitesmax-1;
  NFermiOps= ThisCode.NopsSaved;

  // TotS QNs in ThisCode
  ThisCode.SetTotS();
  // Setting SingleSite (depends only on SymNo)
  ThisCode.SetSingleSite(&SingleSite);

  // Allocate matrices: Tricky

  OpArrayN=new CNRGmatrix* [NFermiOps];
  for (int iop=0;iop<NFermiOps;iop++){
    OpArrayN[iop]=new CNRGmatrix [NshellMax+1];
  }

  AcutN = new CNRGbasisarray [NshellMax+1];

  // SIAM Conductance calculation

  // Read Acut files and calculates the conductance

  //for (int Nshell=NshellMax;Nshell>=ThisCode.Nsites0;Nshell--){
  for (int Nshell=ThisCode.Nsites0;Nshell<=NshellMax;Nshell++){

    cout << "CondSIAM: Reading files for Nshell = " << Nshell << endl;

    AcutN[Nshell].ClearAll();
    AcutN[Nshell].Nshell=Nshell;

    sprintf(CNsites,"%d",Nshell);
    //strcpy(ext,"test_N");
    strcpy(ext,ThisCode.SaveArraysFileName);
    strcat(ext,"_N");
    strcat(ext,CNsites);
    strcat(ext,".bin");

    // Read Acut
    strcpy(arqname,"Acut_");
    strcat(arqname,ext);
    AcutN[Nshell].ReadBin(arqname);

    // Read Operators
    for (int iop=0;iop<NFermiOps;iop++){
      sprintf(CNmat,"%d",iop); // Get all SavedMatrices
      strcpy(arqname,"Mat");
      strcat(arqname,CNmat);
      strcat(arqname,"_");
      strcat(arqname,ext);
      OpArrayN[iop][Nshell].ReadBin(arqname);
    }
    // end read operators

  }
  // end loop in Nshell

  // Define spectral function

  CSpecFunction spec1;

  spec1.Lambda=ThisCode.Lambda;
  spec1.z_twist=ThisCode.chain.z_twist;
  spec1.NshellMax=ThisCode.Nsitesmax;
  spec1.NshellMin=ThisCode.Nsites0;
  spec1.AcutN=AcutN;

  //spec1.Temp=Temp;
  //spec1.Mtemp=Mtemp;
  spec1.Betabar=betabar;

  //spec1.TwindowFac=twindow;
  //spec1.dBroadTemp=broadtemp;

  spec1.NonDiagGF=false;
  spec1.Op1N=OpArrayN[0];

  spec1.CalcCondSIAM_ManyTs(Nw);

  // Need to multiply it by  Gamma

  // File Name
  strcpy(spec1.Name,"GT");

  // Will add OmegaRhow, z-value and .dat
  spec1.SaveOmegaRhow();

// Need to multiply it by Gamma...

// de-allocate matrices. Watch out.
  delete[] AcutN;
  for (int iop=NFermiOps-1;iop>=0;iop--){
    delete[] OpArrayN[iop];
  }
  delete[] OpArrayN;

}
// end MAIN


/*    Copyright (c) 2010-2019, Delft University of Technology
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 */
#define BOOST_TEST_MAIN

#include <string>
#include <thread>

#include <boost/test/unit_test.hpp>
#include <boost/random/uniform_01.hpp>

#include "Tudat/Astrodynamics/BasicAstrodynamics/orbitalElementConversions.h"
#include "Tudat/Astrodynamics/BasicAstrodynamics/physicalConstants.h"
#include "Tudat/Basics/testMacros.h"

#include "Tudat/Astrodynamics/Gravitation/thirdBodyPerturbation.h"
#include "Tudat/Astrodynamics/Gravitation/mutualSphericalHarmonicGravityModel.h"
#include "Tudat/External/SpiceInterface/spiceInterface.h"
#include "Tudat/InputOutput/basicInputOutput.h"
#include "Tudat/SimulationSetup/PropagationSetup/createAccelerationModels.h"
#include "Tudat/SimulationSetup/EnvironmentSetup/createBodies.h"
#include "Tudat/SimulationSetup/EnvironmentSetup/createGravityField.h"
#include "Tudat/SimulationSetup/EnvironmentSetup/defaultBodies.h"

namespace tudat
{

namespace unit_tests
{

using namespace tudat::spice_interface;
using namespace tudat::simulation_setup;
using namespace tudat::gravitation;
using namespace tudat::basic_astrodynamics;
using namespace tudat::simulation_setup;

//! Generate (dummy_ spherical harmonic coefficients.
std::pair< Eigen::MatrixXd, Eigen::MatrixXd > generateCosineSineCoefficients(
        const int maximumDegree, const int maximumOrder, const int bodyIndex )
{
    Eigen::MatrixXd cosineCoefficients = Eigen::MatrixXd::Zero( maximumDegree + 1, maximumOrder + 1 );
    Eigen::MatrixXd sineCoefficients = Eigen::MatrixXd::Zero( maximumDegree + 1, maximumOrder + 1 );

    cosineCoefficients( 0, 0 ) = 1.0;

    basic_mathematics::GlobalRandomNumberGeneratorType randomNumberGenerator(
                static_cast< unsigned int >( bodyIndex ) );
    boost::uniform_01< boost::mt19937> distribution( randomNumberGenerator );

    for( int i = 1; i < maximumDegree + 1; i++ )
    {
        for( int j = 0; ( j < maximumOrder + 1 ) && ( j <= i ); j++ )
        {
            cosineCoefficients( i, j ) = ( ( distribution( ) > 0.5 ) ? ( 1.0 ): ( -1.0 ) ) * distribution( ) * 1.0E-2;
            if( j > 0 )
            {
                sineCoefficients( i, j ) =  ( ( distribution( ) > 0.5 ) ? ( 1.0 ): ( -1.0 ) ) * distribution( ) * 1.0E-2;
            }

        }
    }

    return std::make_pair( cosineCoefficients, sineCoefficients );
}

//! Generate gravity field object (with severely exagerated magnitude).
std::shared_ptr< tudat::simulation_setup::GravityFieldSettings > getDummyJovianSystemGravityField(
        const std::string& bodyName )
{
    std::shared_ptr< GravityFieldSettings > gravityFieldSettings;

    std::vector< double > randomNumberSettings;
    randomNumberSettings.push_back( 0.0 );
    randomNumberSettings.push_back( 1.0E-4 );

    std::pair< Eigen::MatrixXd, Eigen::MatrixXd > coefficients;

    if( bodyName == "Jupiter" )
    {
        coefficients = generateCosineSineCoefficients( 10, 10, 0 );

        gravityFieldSettings = std::make_shared< SphericalHarmonicsGravityFieldSettings >
                ( getBodyGravitationalParameter( "Jupiter" ), getAverageRadius( "Jupiter" ),
                  coefficients.first, coefficients.second, "IAU_Jupiter" );
    }
    else if( bodyName == "Io" )
    {
        coefficients = generateCosineSineCoefficients( 10, 10, 1 );

        gravityFieldSettings = std::make_shared< SphericalHarmonicsGravityFieldSettings >
                ( 5.959916033410404E012, 200.0 * getAverageRadius( "Io" ),
                  coefficients.first, coefficients.second, "IAU_Io" );
    }
    else if( bodyName == "Europa" )
    {
        coefficients = generateCosineSineCoefficients( 10, 10, 2 );

        gravityFieldSettings = std::make_shared< SphericalHarmonicsGravityFieldSettings >
                ( 3.202738774922892E12, 200.0 * getAverageRadius( "Europa" ),
                  coefficients.first, coefficients.second, "IAU_Europa" );
    }

    return gravityFieldSettings;

}

BOOST_AUTO_TEST_SUITE( test_mutual_spherical_harmonic_gravity )

//! Test mutual spherical harmonic acceleration against manually combined spherical harmonic accelerations.
//! Note that the size of the Galilean moons, as well as the magnitude of the spherical harmonic coefficients has
//! been exaggerated to perform a more robust test (i.e ensure that typical errors in implementation are well above
//! numerical errors).
BOOST_AUTO_TEST_CASE( testMutualSphericalHarmonicGravity )
{
    // Load spice kernels.
    spice_interface::loadStandardSpiceKernels( );

    // Create list of bodies to create.
    std::vector< std::string > bodyNames;
    bodyNames.push_back( "Jupiter" );
    bodyNames.push_back( "Io" );
    bodyNames.push_back( "Europa" );
    bodyNames.push_back( "Sun" );

    // Specify initial time
    double initialTime = 1.0E7;
    double finalTime = 1.2E7;

    // Get body settings.
    std::map< std::string, std::shared_ptr< BodySettings > > bodySettings =
            getDefaultBodySettings( bodyNames, initialTime, finalTime );
    bodySettings[ "Jupiter" ]->gravityFieldSettings = getDummyJovianSystemGravityField( "Jupiter" );
    bodySettings[ "Io" ]->gravityFieldSettings = getDummyJovianSystemGravityField( "Io" );
    bodySettings[ "Europa" ]->gravityFieldSettings = getDummyJovianSystemGravityField( "Europa" );

    bodySettings[ "Jupiter" ]->ephemerisSettings = std::make_shared< KeplerEphemerisSettings >(
                ( Eigen::Vector6d( ) << 778.57E9, 0.0489, 1.3 / 60.0, 0.0, 0.0, 0.0 ).finished( ), 0.0,
                  getBodyGravitationalParameter( "Sun" ), "Sun", "ECLIPJ2000" );
    bodySettings[ "Io" ]->ephemerisSettings = std::make_shared< KeplerEphemerisSettings >(
                ( Eigen::Vector6d( ) << 421.8E6, 0.004, 0.04 / 60.0, 0.0, 0.0, 0.0 ).finished( ), 0.0,
                  getBodyGravitationalParameter( "Jupiter" ), "Sun", "ECLIPJ2000" );
    bodySettings[ "Europa" ]->ephemerisSettings = std::make_shared< KeplerEphemerisSettings >(
                ( Eigen::Vector6d( ) << 671.1E6, 0.009, 0.47 / 60.0, 0.0, 0.0, 0.0 ).finished( ), 0.0,
                  getBodyGravitationalParameter( "Jupiter" ), "Sun", "ECLIPJ2000" );

    // Create bodies needed in simulation
    NamedBodyMap bodyMap = createBodies( bodySettings );
    setGlobalFrameBodyEphemerides( bodyMap, "SSB", "ECLIPJ2000" );

    // Set current state and rotation of bodies.
    double currentTime = 1.1E7;
    bodyMap[ "Jupiter" ]->setCurrentRotationToLocalFrameFromEphemeris( currentTime );
    bodyMap[ "Jupiter" ]->setStateFromEphemeris( currentTime );
    bodyMap[ "Io" ]->setCurrentRotationToLocalFrameFromEphemeris( currentTime );
    bodyMap[ "Io" ]->setStateFromEphemeris( currentTime );
    bodyMap[ "Europa" ]->setCurrentRotationToLocalFrameFromEphemeris( currentTime );
    bodyMap[ "Europa" ]->setStateFromEphemeris( currentTime );

    // Retrieve gravity fields.
    std::shared_ptr< SphericalHarmonicsGravityField > jupiterGravityField = std::dynamic_pointer_cast< SphericalHarmonicsGravityField >(
                ( bodyMap.at( "Jupiter" ) )->getGravityFieldModel( ) );
    std::shared_ptr< SphericalHarmonicsGravityField > ioGravityField = std::dynamic_pointer_cast< SphericalHarmonicsGravityField >(
                ( bodyMap.at( "Io" ) )->getGravityFieldModel( ) );
    std::shared_ptr< SphericalHarmonicsGravityField > europaGravityField = std::dynamic_pointer_cast< SphericalHarmonicsGravityField >(
                ( bodyMap.at( "Europa" ) )->getGravityFieldModel( ) );

    // Create central gravity acceleration (mu = Io + Jupiter)
    std::shared_ptr< AccelerationSettings > centralGravitySettings = std::make_shared< AccelerationSettings >( central_gravity );
    std::shared_ptr< CentralGravitationalAccelerationModel3d > centralGravity =
            std::dynamic_pointer_cast< CentralGravitationalAccelerationModel3d >(
                createAccelerationModel( bodyMap.at( "Io" ), bodyMap.at( "Jupiter" ), centralGravitySettings, "Io", "Jupiter",
                                         bodyMap.at( "Jupiter" ), "Jupiter" ) );

    // Calculate central gravity acceleration.
    centralGravity->updateMembers( );
    Eigen::Vector3d centralGravityAcceleration = centralGravity->getAcceleration( );

    // Create spherical harmonic gravity of Jupiter on Io, Jupiter-fixed (mu = Io + Jupiter)
    std::shared_ptr< AccelerationSettings > sphericalHarmonicGravityOnIoFromJupiterSettings =
            std::make_shared< SphericalHarmonicAccelerationSettings >( 7, 7 );
    std::shared_ptr< SphericalHarmonicsGravitationalAccelerationModel > sphericalHarmonicGravityOnIoFromJupiter =
            std::dynamic_pointer_cast< SphericalHarmonicsGravitationalAccelerationModel >(
                createAccelerationModel(  bodyMap.at( "Io" ), bodyMap.at( "Jupiter" ), sphericalHarmonicGravityOnIoFromJupiterSettings,
                                          "Io", "Jupiter", bodyMap.at( "Jupiter" ), "Jupiter" ) );

    // Calculate spherical harmonic gravity of Jupiter on Io.
    sphericalHarmonicGravityOnIoFromJupiter->updateMembers( );
    Eigen::Vector3d sphericalHarmonicGravityOnIoFromJupiterAcceleration = sphericalHarmonicGravityOnIoFromJupiter->getAcceleration( );

    // Create spherical harmonic gravity of Io on Jupiter, Io-fixed (mu = Io + Jupiter)
    std::shared_ptr< AccelerationSettings > sphericalHarmonicGravityOnJupiterFromIoSettings =
            std::make_shared< SphericalHarmonicAccelerationSettings >( 2, 2 );
    std::shared_ptr< SphericalHarmonicsGravitationalAccelerationModel > sphericalHarmonicGravityOnJupiterFromIo =
            std::dynamic_pointer_cast< SphericalHarmonicsGravitationalAccelerationModel >(
                createAccelerationModel( bodyMap.at( "Jupiter" ), bodyMap.at( "Io" ), sphericalHarmonicGravityOnJupiterFromIoSettings,
                                         "Jupiter", "Io", bodyMap.at( "Io" ), "Io" ) );

    // Calculate spherical harmonic gravity of Io on Jupiter.
    sphericalHarmonicGravityOnJupiterFromIo->updateMembers( );
    Eigen::Vector3d sphericalHarmonicGravityOnJupiterFromIoAcceleration = sphericalHarmonicGravityOnJupiterFromIo->getAcceleration( );

    // Create mutual spherical harmonic gravity between Io and Jupiter on Io, Jupiter fixed (mu = Io + Jupiter)
    std::shared_ptr< AccelerationSettings > mutualDirectJupiterIoShGravitySettings =
            std::make_shared< MutualSphericalHarmonicAccelerationSettings >( 7, 7, 2, 2 );
    std::shared_ptr< MutualSphericalHarmonicsGravitationalAccelerationModel > mutualDirectJupiterIoShGravity =
            std::dynamic_pointer_cast< MutualSphericalHarmonicsGravitationalAccelerationModel >(
                createAccelerationModel( bodyMap.at( "Io" ), bodyMap.at( "Jupiter" ), mutualDirectJupiterIoShGravitySettings,
                                         "Io", "Jupiter", bodyMap.at( "Jupiter" ), "Jupiter" ) );

    // Calculate mutual spherical harmonic gravity between Io and Jupiter on Io.
    mutualDirectJupiterIoShGravity->updateMembers( );
    Eigen::Vector3d mutualDirectJupiterIoShGravityAcceleration = mutualDirectJupiterIoShGravity->getAcceleration( );

    // Calculate expected mutual spherical harmonic gravity from sub-accelerations.
    Eigen::Vector3d expectedAcceleration = -centralGravityAcceleration + sphericalHarmonicGravityOnIoFromJupiterAcceleration -
            sphericalHarmonicGravityOnJupiterFromIoAcceleration;

    // Test against directly calculated mutual spherical harmonic gravity.
    for( unsigned int i = 0; i < 3; i++ )
    {
        BOOST_CHECK_SMALL( std::fabs( expectedAcceleration( i ) - mutualDirectJupiterIoShGravityAcceleration( i ) ),
                           15.0 * std::numeric_limits< double >::epsilon( ) * expectedAcceleration.norm( ) );
    }

    // Create mutual spherical harmonic gravity between Io and Jupiter on Jupiter, Io fixed (mu = Io + Jupiter)
    std::shared_ptr< AccelerationSettings > mutualDirectJupiterIoShGravitySettings2 =
            std::make_shared< MutualSphericalHarmonicAccelerationSettings >( 2, 2, 7, 7 );
    std::shared_ptr< MutualSphericalHarmonicsGravitationalAccelerationModel > mutualDirectJupiterIoShGravity2 =
            std::dynamic_pointer_cast< MutualSphericalHarmonicsGravitationalAccelerationModel >(
                createAccelerationModel( bodyMap.at( "Jupiter" ), bodyMap.at( "Io" ), mutualDirectJupiterIoShGravitySettings2,
                                         "Jupiter", "Io", bodyMap.at( "Io" ), "Io" ) );

    // Calculate mutual spherical harmonic gravity between Io and Jupiter on Jupiter.
    mutualDirectJupiterIoShGravity2->updateMembers( );
    Eigen::Vector3d mutualDirectJupiterIoShGravityAcceleration2 = mutualDirectJupiterIoShGravity2->getAcceleration( );

    expectedAcceleration = centralGravityAcceleration - sphericalHarmonicGravityOnIoFromJupiterAcceleration + sphericalHarmonicGravityOnJupiterFromIoAcceleration;

    // Test against directly calculated mutual spherical harmonic gravity.
    for( unsigned int i = 0; i < 3; i++ )
    {
        BOOST_CHECK_SMALL( std::fabs( expectedAcceleration( i ) - mutualDirectJupiterIoShGravityAcceleration2( i ) ),
                           12.0 * std::numeric_limits< double >::epsilon( ) * expectedAcceleration.norm( ) );
    }

    // Test against directly calculated mutual spherical harmonic gravity.
    for( unsigned int i = 0; i < 3; i++ )
    {
        BOOST_CHECK_SMALL( std::fabs( mutualDirectJupiterIoShGravityAcceleration( i ) + mutualDirectJupiterIoShGravityAcceleration2( i ) ),
                           12.0 * std::numeric_limits< double >::epsilon( ) * mutualDirectJupiterIoShGravityAcceleration.norm( ) );
    }

    // Create 3rd body mutual spherical harmonics between Io and Europa on Europa, Jupiter fixed (mu = Io)
    std::shared_ptr< AccelerationSettings > mutualThirdBodyIoOnEuropaShGravitySettings =
            std::make_shared< MutualSphericalHarmonicAccelerationSettings >( 2, 2, 4, 4, 7, 7 );
    std::shared_ptr< ThirdBodyMutualSphericalHarmonicsGravitationalAccelerationModel > mutualThirdBodyIoOnEuropaShGravity =
            std::dynamic_pointer_cast< ThirdBodyMutualSphericalHarmonicsGravitationalAccelerationModel >(
                createAccelerationModel( bodyMap.at( "Europa" ), bodyMap.at( "Io" ), mutualThirdBodyIoOnEuropaShGravitySettings,
                                         "Europa", "Io", bodyMap.at( "Jupiter" ), "Jupiter" ) );

    // Calculate 3rd body mutual spherical harmonics between Io and Europa on Europa.
    mutualThirdBodyIoOnEuropaShGravity->updateMembers( );
    Eigen::Vector3d mutualThirdBodyIoOnEuropaShGravityAcceleration = mutualThirdBodyIoOnEuropaShGravity->getAcceleration( );

    // Create mutual spherical harmonics between Io and Europa on Europa, Io fixed (mu = Io + Europa)
    std::shared_ptr< MutualSphericalHarmonicsGravitationalAccelerationModel > mutualDirectIoOnEuropaShGravity =
            std::dynamic_pointer_cast< MutualSphericalHarmonicsGravitationalAccelerationModel >(
                createAccelerationModel( bodyMap.at( "Europa" ), bodyMap.at( "Io" ), mutualThirdBodyIoOnEuropaShGravitySettings,
                                         "Europa", "Io", bodyMap.at( "Io" ), "Io" ) );
    mutualDirectIoOnEuropaShGravity->updateMembers( );
    Eigen::Vector3d mutualDirectIoOnEuropaShGravityAcceleration = mutualDirectIoOnEuropaShGravity->getAcceleration( );

    // Get sub accelerations from 3rd body acceleration.
    Eigen::Vector3d directAccelerationFromThirdBodyModel = mutualThirdBodyIoOnEuropaShGravity->getAccelerationModelForBodyUndergoingAcceleration( )->getAcceleration( );
    Eigen::Vector3d centralBodyAccelerationFromThirdBodyModel = mutualThirdBodyIoOnEuropaShGravity->getAccelerationModelForCentralBody( )->
            getAcceleration( );

    for( unsigned int i = 0; i < 3; i++ )
    {
        BOOST_CHECK_SMALL(
                    ( directAccelerationFromThirdBodyModel( i ) -
                      ( ioGravityField->getGravitationalParameter( ) /
                        ( ioGravityField->getGravitationalParameter( ) + europaGravityField->getGravitationalParameter( ) ) *
                        mutualDirectIoOnEuropaShGravityAcceleration( i ) ) ),
                    ( 12.0 * std::numeric_limits< double >::epsilon( ) * directAccelerationFromThirdBodyModel.norm( ) ) );

        BOOST_CHECK_SMALL(
                    centralBodyAccelerationFromThirdBodyModel( i ) -
                    ( ioGravityField->getGravitationalParameter( ) /
                      ( ioGravityField->getGravitationalParameter( ) + jupiterGravityField->getGravitationalParameter( ) ) *
                      mutualDirectJupiterIoShGravityAcceleration2( i ) ),
                    ( 12.0 * std::numeric_limits< double >::epsilon( ) * centralBodyAccelerationFromThirdBodyModel.norm( ) ) );

        BOOST_CHECK_SMALL(
                    mutualThirdBodyIoOnEuropaShGravityAcceleration( i ) -
                    ( directAccelerationFromThirdBodyModel( i ) - centralBodyAccelerationFromThirdBodyModel( i ) ),
                    ( std::numeric_limits< double >::epsilon( ) * mutualThirdBodyIoOnEuropaShGravityAcceleration.norm( ) ) );
    }
}

BOOST_AUTO_TEST_SUITE_END( )

}

}


#pragma once

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

#include <boost/graph/graph_traits.hpp>
#include <boost/range/adaptor/transformed.hpp>
#include <boost/range/algorithm/for_each.hpp>
#include <boost/range/iterator_range.hpp>

#include "graphviz_interface.hpp"

#include "DiGraph.hpp"
#include "tran_mat_cell.hpp"
#include <fmt/format.h>
#include <nlohmann/json.hpp>

const double tuning_param = 1.0;

const size_t DEFAULT_N_SAMPLES = 1000;

enum InitialBeta { ZERO, ONE, HALF, MEAN, RANDOM };

typedef std::unordered_map<std::string, std::vector<double>>
    AdjectiveResponseMap;

typedef std::vector<std::vector<std::vector<std::vector<double>>>>
    ObservedStateSequence;

typedef std::vector<std::vector<std::vector<double>>>
    PredictedObservedStateSequence;

typedef std::pair<std::tuple<std::string, int, std::string>,
                  std::tuple<std::string, int, std::string>>
    CausalFragment;

// Access
// [ sample ][ time_step ][ vertex_name ][ indicator_name ]
typedef std::vector<std::vector<
    std::unordered_map<std::string, std::unordered_map<std::string, double>>>>
    FormattedPredictionResult;

// Access
// [ vertex_name ][ timestep ][ sample ]
typedef std::unordered_map<std::string, std::vector<std::vector<double>>>
    FormattedProjectionResult;

typedef std::tuple<std::pair<std::pair<int, int>, std::pair<int, int>>,
                   std::vector<std::string>,
                   FormattedPredictionResult>
    Prediction;

typedef boost::graph_traits<DiGraph>::edge_descriptor EdgeDescriptor;
typedef boost::graph_traits<DiGraph>::edge_iterator EdgeIterator;

typedef std::multimap<std::pair<int, int>, std::pair<int, int>>::iterator
    MMapIterator;

AdjectiveResponseMap construct_adjective_response_map(size_t n_kernels);

/**
 * The AnalysisGraph class is the main model/interface for Delphi.
 */
class AnalysisGraph {

  private:
  DiGraph graph;

  // Handle to the random number generator singleton object
  RNG* rng_instance;

  std::mt19937 rand_num_generator;

  // Uniform distribution used by the MCMC sampler
  std::uniform_real_distribution<double> uni_dist;

  // Normal distrubution used to perturb β
  std::normal_distribution<double> norm_dist;

  // Uniform discrete distribution used by the MCMC sampler
  // to perturb the initial latent state
  std::uniform_int_distribution<int> uni_disc_dist;

  /*
   ============================================================================
   Meta Data Structures
   ============================================================================
  */

  // Maps each concept name to the vertex id of the
  // vertex that concept is represented in the CAG
  // concept name --> CAG vertex id
  std::unordered_map<std::string, int> name_to_vertex = {};

  // Keeps track of indicators in CAG to ensure there are no duplicates.
  std::unordered_set<std::string> indicators_in_CAG;

  // A_beta_factors is a 2D array (std::vector of std::vectors) that keeps track
  // of the β factors involved with each cell of the transition matrix A.
  //
  // According to our current model, which uses variables and their partial
  // derivatives with respect to each other ( x --> y, βxy = ∂y/∂x ),
  // atmost half of the transition matrix cells can be affected by βs.
  // According to the way we organize the transition matrix, the cells
  // A[row][col] where row is an even index and col is an odd index
  // are such cells.
  //
  // Each cell of matrix A_beta_factors represent all the directed paths
  // starting at the vertex equal to the column index of the matrix and
  // ending at the vertex equal to the row index of the matrix.
  //
  // Each cell of matrix A_beta_factors is an object of Tran_Mat_Cell class.
  std::vector<std::vector<std::shared_ptr<Tran_Mat_Cell>>> A_beta_factors;

  // A set of (row, column) numbers of the 2D matrix A_beta_factors
  // where the cell (row, column) depends on β factors.
  std::set<std::pair<int, int>> beta_dependent_cells;

  // Maps each β to all the transition matrix cells that are dependent on it.
  std::multimap<std::pair<int, int>, std::pair<int, int>> beta2cell;

  /*
   ============================================================================
   Sampler Related Variables
   ============================================================================
  */

  // Keep track whether the model is trained.
  // Used to check whether there is a trained model before calling
  // generate_prediction()
  bool trained = false;

  int n_timesteps;
  int pred_timesteps;
  std::pair<std::pair<int, int>, std::pair<int, int>> training_range;
  std::vector<std::string> pred_range;

  double t = 0.0;
  double delta_t = 1.0;

  double log_likelihood = 0.0;
  double previous_log_likelihood = 0.0;

  // To decide whether to perturb a β or a derivative
  // If coin_flip < coin_flip_thresh perturb β else perturb derivative
  double coin_flip = 0;
  double coin_flip_thresh = 0.5;

  // Remember the old β and the edge where we perturbed the β.
  // We need this to revert the system to the previous state if the proposal
  // gets rejected.
  std::pair<EdgeDescriptor, double> previous_beta;

  // Remember the old derivative and the concept we perturbed the derivative
  int changed_derivative;
  double previous_derivative;

  // Latent state that is evolved by sampling.
  Eigen::VectorXd s0;
  Eigen::VectorXd s0_prev;

  // Transition matrix that is evolved by sampling.
  // Since variable A has been already used locally in other methods,
  // I chose to name this A_orginal. After refactoring the code, we could
  // rename this to A.
  Eigen::MatrixXd A_original;

  // Access this as
  // current_latent_state
  Eigen::VectorXd current_latent_state;

  // Access this as
  // observed_state_sequence[ time step ][ vertex ][ indicator ]
  ObservedStateSequence observed_state_sequence;

  // This is a column of the
  // this->training_latent_state_sequences
  // prediction_initial_latent_states.size() = this->res
  // TODO: If we make the code using this variable to directly fetch the values
  // from this->training_latent_state_sequences, we can get rid of this
  // std::vector<Eigen::VectorXd> prediction_initial_latent_states;

  // Access this as
  // prediction_latent_state_sequences[ sample ][ time step ]
  std::vector<std::vector<Eigen::VectorXd>> predicted_latent_state_sequences;

  // Access this as
  // predicted_observed_state_sequences
  //                            [ sample ][ time step ][ vertex ][ indicator ]
  std::vector<PredictedObservedStateSequence>
      predicted_observed_state_sequences;

  PredictedObservedStateSequence test_observed_state_sequence;

  std::vector<Eigen::MatrixXd> transition_matrix_collection;
  std::vector<Eigen::VectorXd> initial_latent_state_collection;

  std::vector<Eigen::VectorXd> synthetic_latent_state_sequence;
  // ObservedStateSequence synthetic_observed_state_sequence;
  bool synthetic_data_experiment = false;

  /*
   ============================================================================
   Private: Utilities (in graph_util.cpp)
   ============================================================================
  */

  void clear_state();

  void initialize_random_number_generator();

  void remove_node(int node_id);

  // Allocate a num_verts x num_verts 2D array (std::vector of std::vectors)
  void allocate_A_beta_factors();

  // AnalysisGraph
  // find_all_paths_for_concept(std::string concept, int depth, bool reverse);

  /**
   * Finds all the simple paths starting at the start vertex and
   * ending at the end vertex.
   * Uses find_all_paths_between_util() as a helper to recursively find the
   * paths
   */
  void find_all_paths_between(int start, int end, int cutoff);

  /**
   * Recursively finds all the simple paths starting at the start vertex and
   * ending at the end vertex. Used by find_all_paths_between()
   * Paths found are added to the Tran_Mat_Cell object that is tracking the
   * transition matrix cell (2*end, 2*start)
   *
   * @param start: Start vertex of the path
   * @param end  : End vertex of the path
   * @param path : A path starting at vettex start that is being explored
   *
   * @return void
   */
  void find_all_paths_between_util(int start,
                                   int end,
                                   std::vector<int>& path,
                                   int cutoff);

  /**
   * Utility function that converts a time range given a start date and end date
   * into an integer value.
   * At the moment returns the number of months withing the time range.
   * This should be the number of traing data time points we have
   *
   * @param start_year  : Start year of the training data sequence
   * @param start_month : Start month of the training data sequence
   * @param end_year    : End year of the training data sequence
   * @param end_month   : End month of the training data sequence
   *
   * @return            : Number of months in the training data sequence
   *                      Including both start and end months
   */
  int calculate_num_timesteps(int start_year,
                              int start_month,
                              int end_year,
                              int end_month);

  /*
   ============================================================================
   Private: Subgraphs (in subgraphs.cpp)
   ============================================================================
  */

  void get_subgraph(int vert,
                    std::unordered_set<int>& vertices_to_keep,
                    int cutoff,
                    bool inward);

  void get_subgraph_between(int start,
                            int end,
                            std::vector<int>& path,
                            std::unordered_set<int>& vertices_to_keep,
                            int cutoff);

  /*
   ============================================================================
   Private: Accessors
   ============================================================================
  */

  int num_nodes() { return boost::num_vertices(graph); }

  // int get_vertex_id(std::string concept);
  int get_vertex_id(std::string concept) {
    using namespace fmt::literals;
    try {
      return this->name_to_vertex.at(concept);
    }
    catch (const std::out_of_range& oor) {
      throw std::out_of_range("Concept \"{}\" not in CAG!"_format(concept));
    }
  }

  auto node_indices() {
    return boost::make_iterator_range(boost::vertices(this->graph));
  };

  auto nodes() {
    using boost::adaptors::transformed;
    return this->node_indices() |
           transformed([&](int v) -> Node& { return (*this)[v]; });
  };

  auto node_names() {
    using boost::adaptors::transformed;
    return this->nodes() |
           transformed([&](auto node) -> std::string { return node.name; });
  };

  int get_degree(int vertex_id) {
    return boost::in_degree(vertex_id, this->graph) +
           boost::out_degree(vertex_id, this->graph);
  };

  auto out_edges(int i) {
    return boost::make_iterator_range(boost::out_edges(i, graph));
  }

  Node& source(EdgeDescriptor e) {
    return (*this)[boost::source(e, this->graph)];
  };

  Node& target(EdgeDescriptor e) {
    return (*this)[boost::target(e, this->graph)];
  };

  auto successors(int i) {
    return boost::make_iterator_range(boost::adjacent_vertices(i, this->graph));
  }

  auto successors(std::string node_name) {
    return this->successors(this->name_to_vertex.at(node_name));
  }

  // std::vector<Node> get_successor_list(std::string node_name);
  std::vector<Node> get_successor_list(std::string node) {
    std::vector<Node> successors = {};
    for (int successor : this->successors(node)) {
      successors.push_back((*this)[successor]);
    }
    return successors;
  }

  auto predecessors(int i) {
    return boost::make_iterator_range(
        boost::inv_adjacent_vertices(i, this->graph));
  }

  auto predecessors(std::string node_name) {
    return this->predecessors(this->name_to_vertex.at(node_name));
  }

  std::vector<Node> get_predecessor_list(std::string node) {
    std::vector<Node> predecessors = {};
    for (int predecessor : this->predecessors(node)) {
      predecessors.push_back((*this)[predecessor]);
    }
    return predecessors;
  }

  double get_beta(std::string source_vertex_name,
                  std::string target_vertex_name) {

    // This is ∂target / ∂source
    return this->A_original(2 * get_vertex_id(target_vertex_name),
                            2 * get_vertex_id(source_vertex_name) + 1);
  }

  /*
   ============================================================================
   Private: Get Training Data Sequence (in train_model.cpp)
   ============================================================================
  */

  /**
   * Set the observed state sequence for a given time range from data.
   * The sequence includes both ends of the range.
   * See data.hpp::get_data_value() for missing data rules.
   * Note: units are automatically set according
   * to the parameterization of the given CAG.
   *
   * @param start_year  : Start year of the sequence of data
   * @param start_month : Start month of the sequence of data
   * @param end_year    : End year of the sequence of data
   * @param end_month   : End month of the sequenec of data
   * @param country     : Country where the data is about
   * @param state       : State where the data is about
   * @param county      : County where the data is about
   *
   */
  void
  set_observed_state_sequence_from_data(int start_year,
                                        int start_month,
                                        int end_year,
                                        int end_month,
                                        std::string country = "South Sudan",
                                        std::string state = "",
                                        std::string county = "");

  /**
   * Get the observed state (values for all the indicators)
   * for a given time point from data.
   * See data.hpp::get_data_value() for missing data rules.
   * Note: units are automatically set according
   * to the parameterization of the given CAG.
   *
   * @param year    : Year of the time point data is extracted
   * @param month   : Month of the time point data is extracted
   * @param country : Country where the data is about
   * @param state   : State where the data is about
   * @param county  : County where the data is about
   *
   * @return        : Observed state std::vector for the specified location
   *                  on the specified time point.
   *                  Access it as: [ vertex id ][ indicator id ]
   */
  std::vector<std::vector<std::vector<double>>>
  get_observed_state_from_data(int year,
                               int month,
                               std::string country,
                               std::string state = "",
                               std::string county = "");

  /*
   ============================================================================
   Private: Training by MCMC Sampling (in sampling.cpp)
   ============================================================================
  */

  // Sample elements of the stochastic transition matrix from the
  // prior distribution, based on gradable adjectives.
  void set_transition_matrix_from_betas();

  void sample_transition_matrix_collection_from_prior();

  /**
   * Utility function that sets an initial latent state from observed data.
   * This is used for the inference of the transition matrix as well as the
   * training latent state sequences.
   *
   * @param timestep: Optional setting for setting the initial state to be other
   *                  than the first time step. Not currently used.
   *                  0 <= timestep < this->n_timesteps
   */
  void set_initial_latent_state_from_observed_state_sequence();

  void set_initial_latent_from_end_of_training();

  void set_log_likelihood();

  void set_current_latent_state(int ts);

  /**
   * Run Bayesian inference - sample from the posterior distribution.
   */
  void sample_from_posterior();

  /**
   * Sample a new transition matrix from the proposal distribution,
   * given a current candidate transition matrix.
   * In practice, this amounts to:
   *    Selecting a random β.
   *    Perturbing it a bit.
   *    Updating all the transition matrix cells that are dependent on it.
   */
  // TODO: Need testng
  // TODO: Before calling sample_from_proposal() we must call
  // AnalysisGraph::find_all_paths()
  // TODO: Before calling sample_from_proposal(), we mush assign initial βs and
  // run Tran_Mat_Cell::compute_cell() to initialize the first transistion
  // matrix.
  // TODO: Update Tran_Mat_Cell::compute_cell() to calculate the proper value.
  // At the moment it just computes sum of length of all the paths realted to
  // this cell
  void sample_from_proposal();

  /**
   * Find all the transition matrix (A) cells that are dependent on the β
   * attached to the provided edge and update them.
   * Acts upon this->A_original
   *
   * @param e: The directed edge ≡ β that has been perturbed
   */
  void update_transition_matrix_cells(EdgeDescriptor e);

  double calculate_delta_log_prior();

  void revert_back_to_previous_state();

  /*
   ============================================================================
   Private: Inference (in inference.cpp)
   ============================================================================
  */

  /**
   * Sample a collection of observed state sequences from the likelihood
   * model given a collection of transition matrices.
   *
   * @param prediction_timesteps: The number of timesteps for the prediction
   * sequences.
   * @param initial_prediction_step: The initial prediction timestep relative
   *                                 to training timesteps.
   * @param total_timesteps: Total number of timesteps from the initial
   *                         training date to the end prediction date.
   */
  void sample_predicted_latent_state_sequences(int prediction_timesteps,
                                               int initial_prediction_step,
                                               int total_timesteps,
                                               bool project = false);

  /** Generate predicted observed state sequenes given predicted latent state
   * sequences using the emission model
   */
  void
  generate_predicted_observed_state_sequences_from_predicted_latent_state_sequences();

  /**
   * Format the prediction result into a format Python callers favor.
   *
   * @param pred_timestes: Number of timesteps in the predicted sequence.
   *
   * @return Re-formatted prediction result.
   *         Access it as:
   *         [ sample number ][ time point ][ vertex name ][ indicator name ]
   */
  FormattedPredictionResult format_prediction_result();

  FormattedProjectionResult format_projection_result();

  void run_model(int start_year,
                 int start_month,
                 int end_year,
                 int end_month,
                 bool project = false);

  /*
   ============================================================================
   Private: Synthetic Data Experiment (in synthetic_data.cpp)
   ============================================================================
  */

  /**
   * To help experiment with initializing βs to different values
   *
   * @param ib: Criteria to initialize β
   */
  void init_betas_to(InitialBeta ib = InitialBeta::MEAN);

  void set_random_initial_latent_state();

  void generate_synthetic_latent_state_sequence();

  void
  generate_synthetic_observed_state_sequence_from_synthetic_latent_state_sequence();

  // TODO: Need testing
  /**
   * Sample observed state std::vector.
   * This is the implementation of the emission function.
   *
   * @param latent_state: Latent state std::vector.
   *                      This has 2 * number of vertices in the CAG.
   *                      Even indices track the state of each vertex.
   *                      Odd indices track the state of the derivative.
   *
   * @return Observed state std::vector. Observed state for each indicator for
   * each vertex. Indexed by: [ vertex id ][ indicator id ]
   */
  std::vector<std::vector<double>>
  sample_observed_state(Eigen::VectorXd latent_state);

  /*
   ============================================================================
   Private: Graph Visualization (in graphviz.cpp)
   ============================================================================
  */

  std::pair<Agraph_t*, GVC_t*> to_agraph(
      bool simplified_labels =
          false, /** Whether to create simplified labels or not. */
      int label_depth =
          1, /** Depth in the ontology to which simplified labels extend */
      std::string node_to_highlight = "",
      std::string rankdir = "TB");

  public:
  AnalysisGraph() {}

  ~AnalysisGraph() {}

  std::string id;
  std::string to_json_string(int indent = 0);
  bool data_heuristic = false;

  // Sampling resolution. Default is 200
  int res = DEFAULT_N_SAMPLES;

  /*
   ============================================================================
   Constructors from INDRA-exported JSON (in constructors.cpp)
   ============================================================================
  */

  /**
   * A method to construct an AnalysisGraph object given a JSON-serialized list
   * of INDRA statements.
   *
   * @param filename: The path to the file containing the JSON-serialized INDRA
   * statements.
   */
  static AnalysisGraph
  from_indra_statements_json_dict(nlohmann::json json_data,
                                  double belief_score_cutoff = 0.9,
                                  double grounding_score_cutoff = 0.0,
                                  std::string ontology = "WM");

  static AnalysisGraph
  from_indra_statements_json_string(std::string json_string,
                                    double belief_score_cutoff = 0.9,
                                    double grounding_score_cutoff = 0.0,
                                    std::string ontology = "WM");

  static AnalysisGraph
  from_indra_statements_json_file(std::string filename,
                                  double belief_score_cutoff = 0.9,
                                  double grounding_score_cutoff = 0.0,
                                  std::string ontology = "WM");

  /**
   * A method to construct an AnalysisGraph object given from a std::vector of
   * ( subject, object ) pairs (Statements)
   *
   * @param statements: A std::vector of CausalFragment objects
   */
  static AnalysisGraph
  from_causal_fragments(std::vector<CausalFragment> causal_fragments);

  /** From internal string representation output by to_json_string */
  static AnalysisGraph from_json_string(std::string);

  /*
   ============================================================================
   Public: Integration with Uncharted's CauseMos interface
   ============================================================================
  */

  /** Construct an AnalysisGraph object from JSON exported by CauseMos. */
  static AnalysisGraph from_causemos_json_dict(nlohmann::json json_data);

  /** Construct an AnalysisGraph object from a JSON string exported by CauseMos.
   */
  static AnalysisGraph from_causemos_json_string(std::string json_string);

  /** Construct an AnalysisGraph object from a file containing JSON data from
   * CauseMos. */
  static AnalysisGraph from_causemos_json_file(std::string filename);

  /** Calculate and return a JSON string with edge weight information for
   * visualizing AnalysisGraph models in CauseMos. */
  std::string get_edge_weights_for_causemos_viz();

  /*
   ============================================================================
   Public: Accessors
   ============================================================================
  */

  /** Number of nodes in the graph */
  size_t num_vertices() { return boost::num_vertices(this->graph); }

  Node& operator[](std::string node_name) {
    return (*this)[this->get_vertex_id(node_name)];
  }

  Node& operator[](int v) { return this->graph[v]; }

  size_t num_edges() { return boost::num_edges(this->graph); }

  auto edges() { return boost::make_iterator_range(boost::edges(graph)); }

  Edge& edge(EdgeDescriptor e) { return this->graph[e]; }

  Edge& edge(int source, int target) {
    return this->graph[boost::edge(source, target, this->graph).first];
  }

  Edge& edge(int source, std::string target) {
    return this->graph
        [boost::edge(source, this->get_vertex_id(target), this->graph).first];
  }

  Edge& edge(std::string source, int target) {
    return this->graph
        [boost::edge(this->get_vertex_id(source), target, this->graph).first];
  }

  Edge& edge(std::string source, std::string target) {
    return this->graph[boost::edge(this->get_vertex_id(source),
                                   this->get_vertex_id(target),
                                   this->graph)
                           .first];
  }

  boost::range_detail::integer_iterator<unsigned long> begin() {
    return boost::vertices(this->graph).first;
  };

  boost::range_detail::integer_iterator<unsigned long> end() {
    return boost::vertices(this->graph).second;
  };

  Eigen::VectorXd& get_initial_latent_state() { return this->s0; };

  /*
   ============================================================================
   Public: Graph Building (in graph_building.cpp)
   ============================================================================
  */

  void add_node(std::string concept);

  void add_edge(CausalFragment causal_fragment);
  std::pair<EdgeDescriptor, bool> add_edge(int, int);
  std::pair<EdgeDescriptor, bool> add_edge(int, std::string);
  std::pair<EdgeDescriptor, bool> add_edge(std::string, int);
  std::pair<EdgeDescriptor, bool> add_edge(std::string, std::string);

  void remove_node(std::string concept);

  // Note:
  //      Although just calling this->remove_node(concept) within the loop
  //          for( std::string concept : concept_s )
  //      is sufficient to implement this method, it is not very efficient.
  //      It re-calculates directed simple paths for each vertex removed
  //
  //      Therefore, the code in this->remove_node() has been duplicated with
  //      slightly different flow to achive a more efficient execution.
  void remove_nodes(std::unordered_set<std::string> concepts);

  void remove_edge(std::string src, std::string tgt);

  void remove_edges(std::vector<std::pair<std::string, std::string>> edges);

  /*
   ============================================================================
   Public: Subgraph (in subgraph.cpp)
   ============================================================================
  */

  // TODO Change the name of this function to something better, like
  // restrict_to_subgraph_for_concept, update docstring

  /**
   * Returns the subgraph of the AnalysisGraph around a concept.
   *
   * @param concept: The concept to center the subgraph about.
   * @param depth  : The maximum number of hops from the concept provided
   *                 to be included in the subgraph.
   * #param inward : Sets the direction of the causal influence flow to
   *                 examine.
   *                 False - (default) A subgraph rooted at the concept
   * provided.
   *                 True  - A subgraph with all the paths ending at the concept
   * provided.
   */
  AnalysisGraph get_subgraph_for_concept(std::string concept,
                                         bool inward = false,
                                         int depth = -1);

  /**
   * Returns a new AnaysisGraph related to the source concept and the target
   * concetp provided, which is a subgraph of this graph.
   * This subgraph contains all the simple directed paths of length less than
   * or equal to the provided cutoff.
   *
   * @param source_concept: The concept where the influence starts.
   * @param target_concept: The concept where the influence ends.
   * @param cutoff        : Maximum length of a directed simple path from
   *                        the source to target to be included in the
   *                        subgraph.
   */
  AnalysisGraph get_subgraph_for_concept_pair(std::string source_concept,
                                              std::string target_concept,
                                              int cutoff = -1);

  /*
   ============================================================================
   Public: Graph Modification (in graph_modification.cpp)
   ============================================================================
  */

  void prune(int cutoff = 2);

  // Merge node n1 into node n2, with the option to specify relative polarity.
  // void
  // merge_nodes_old(std::string n1, std::string n2, bool same_polarity = true);

  /**
   * Merges the CAG nodes for the two concepts concept_1 and concept_2
   * with the option to specify relative polarity.
   */
  void merge_nodes(std::string concept_1,
                   std::string concept_2,
                   bool same_polarity = true);

  void change_polarity_of_edge(std::string source_concept,
                               int source_polarity,
                               std::string target_concept,
                               int target_polarity);

  /*
   ============================================================================
   Public: Construct Beta Pdfs (in construct_beta_pdfs.cpp)
   ============================================================================
  */

  // void construct_beta_pdfs(std::mt19937 rng);
  void construct_beta_pdfs();

  /*
   ============================================================================
   Public: Indicator Manipulation (in indicator_manipulation.cpp)
   ============================================================================
  */

  void
  set_indicator(std::string concept, std::string indicator, std::string source);

  void delete_indicator(std::string concept, std::string indicator);

  void delete_all_indicators(std::string concept);

  /**
   * Map each concept node in the AnalysisGraph instance to one or more
   * tangible quantities, known as 'indicators'.
   *
   * @param n: Int representing number of indicators to attach per node.
   * Default is 1 since our model so far is configured for only 1 indicator per
   * node.
   */
  void map_concepts_to_indicators(int n = 1, std::string country = "");

  /**
   * Parameterize the indicators of the AnalysisGraph.
   */
  void parameterize(std::string country = "South Sudan",
                    std::string state = "",
                    std::string county = "",
                    int year = -1,
                    int month = 0,
                    std::map<std::string, std::string> units = {});

  /*
   ============================================================================
   Public: Utilities (in graph_utils.cpp)
   ============================================================================
  */

  /*
   * Find all the simple paths between all the paris of nodes of the graph
   */
  void find_all_paths();

  void set_random_seed(int seed);

  void set_derivative(std::string, double);

  /*
   ============================================================================
   Public: Training the Model (in training_model.cpp)
   ============================================================================
  */

  /**
   * Train a prediction model given a CAG with indicators
   *
   * @param start_year  : Start year of the sequence of data
   * @param start_month : Start month of the sequence of data
   * @param end_year    : End year of the sequence of data
   * @param end_month   : End month of the sequenec of data
   * @param res         : Sampling resolution. The number of samples to retain.
   * @param burn        : Number of samples to throw away. Start retaining
   *                      samples after throwing away this many samples.
   * @param country     : Country where the data is about
   * @param state       : State where the data is about
   * @param county      : county where the data is about
   * @param units       : Units for each indicator. Maps
   *                      indicator name --> unit
   * @param initial_beta: Criteria to initialize β
   */
  void train_model(int start_year = 2012,
                   int start_month = 1,
                   int end_year = 2017,
                   int end_month = 12,
                   int res = 200,
                   int burn = 10000,
                   std::string country = "South Sudan",
                   std::string state = "",
                   std::string county = "",
                   std::map<std::string, std::string> units = {},
                   InitialBeta initial_beta = InitialBeta::ZERO,
                   bool use_heuristic = false);

  /*
   ============================================================================
   Public: Training by MCMC Sampling (in sampling.cpp)
   ============================================================================
  */

  void set_initial_latent_state(Eigen::VectorXd vec) { this->s0 = vec; };

  void set_default_initial_state();

  /*
   ============================================================================
   Public: Inference (in inference.cpp)
   ============================================================================
  */

  /**
   * Given a trained model, generate this->res number of
   * predicted observed state sequences.
   *
   * @param start_year  : Start year of the prediction
   *                      Should be >= the start year of training
   * @param start_month : Start month of the prediction
   *                      If training and prediction start years are equal
   *                      should be >= the start month of training
   * @param end_year    : End year of the prediction
   * @param end_month   : End month of the prediction
   *
   * @return Predicted observed state (indicator value) sequence for the
   *         prediction period including start and end time points.
   *         This is a tuple.
   *         The first element is a std::vector of std::strings with lables for
   * each time point predicted (year-month). The second element contains
   * predicted values. Access it as: [ sample number ][ time point ][ vertex
   * name ][ indicator name ]
   */
  Prediction generate_prediction(int start_year,
                                 int start_month,
                                 int end_year,
                                 int end_month);

  FormattedProjectionResult
  generate_causemos_projection(std::string json_projection);

  /**
   * this->generate_prediction() must be called before callign this method.
   * Outputs raw predictions for a given indicator that were generated by
   * generate_prediction(). Each column is a time step and the rows are the
   * samples for that time step.
   *
   * @param indicator: A std::string representing the indicator variable for
   * which we want predictions.
   * @return A (this->res, x this->pred_timesteps) dimension 2D array
   *         (std::vector of std::vectors)
   *
   */
  std::vector<std::vector<double>> prediction_to_array(std::string indicator);

  /*
   ============================================================================
   Public: Synthetic data experiment (in synthetic_data.cpp)
   ============================================================================
  */

  std::pair<PredictedObservedStateSequence, Prediction>
  test_inference_with_synthetic_data(
      int start_year = 2015,
      int start_month = 1,
      int end_year = 2015,
      int end_month = 12,
      int res = 100,
      int burn = 900,
      std::string country = "South Sudan",
      std::string state = "",
      std::string county = "",
      std::map<std::string, std::string> units = {},
      InitialBeta initial_beta = InitialBeta::HALF);

  /*
   ============================================================================
   Public: Graph Visualization (in graphviz.cpp)
   ============================================================================
  */

  std::string to_dot();

  void
  to_png(std::string filename = "CAG.png",
         bool simplified_labels =
             false, /** Whether to create simplified labels or not. */
         int label_depth =
             1, /** Depth in the ontology to which simplified labels extend */
         std::string node_to_highlight = "",
         std::string rankdir = "TB");

  /*
   ============================================================================
   Public: Printing (in printing.cpp)
   ============================================================================
  */

  void print_nodes();
  void print_edges();
  void print_name_to_vertex();
  void print_indicators();
  void print_A_beta_factors();
  void print_latent_state(const Eigen::VectorXd&);

  /*
   * Prints the simple paths found between all pairs of nodes of the graph
   * Groupd according to the starting and ending vertex.
   * find_all_paths() should be called before this to populate the paths
   */
  void print_all_paths();

  // Given an edge (source, target vertex ids - i.e. a β ≡ ∂target/∂source),
  // print all the transition matrix cells that are dependent on it.
  void print_cells_affected_by_beta(int source, int target);
};


#include <tiny.h>
#include <geometry.h>  // needed for geometry::Box and geometry::Sphere
#include <convex.h>

#define BOOST_AUTO_TEST_MAIN
#include <boost/test/auto_unit_test.hpp>
#include <boost/test/unit_test_suite.hpp>
#include <boost/test/floating_point_comparison.hpp>
#include <boost/test/test_tools.hpp>

#include <cmath>

typedef tiny::MathTypes<double>  M;
typedef M::quaternion_type       Q;
typedef M::vector3_type          V;
typedef M::real_type             T;
typedef M::coordsys_type         C;
typedef M::value_traits          VT;

BOOST_AUTO_TEST_SUITE(convex_compute_closest_points);

BOOST_AUTO_TEST_CASE(sphere_box_compile_test)
{
  {
    geometry::Sphere<V>  B;
    geometry::Box<V>     A;

    B.radius()      = 0.5;
    A.half_extent() = V::make(50, .5, 50);

    C X_A;
    C X_B;

    V p_A = V::zero();
    V p_B = V::zero();

    X_A = C::identity();

    X_B.T() = V::make(0, 1.2, 0);
    X_B.Q() = Q::identity();

    convex::compute_closest_points<M>(
                                      X_A
                                      , &A
                                      , X_B
                                      , &B
                                      , p_A
                                      , p_B
                                      );
  }
  {

    geometry::Box<V>       A;
    convex::Capsule<M>     B;

    A.half_extent() = V::make(50, 1, 50);
    B.radius()      = 0.5;

    C X_A;
    C X_B;

    V p_A = V::zero();
    V p_B = V::zero();

    X_A = C::identity();

    X_B.T() = V::make(0, 1.5, 0);
    X_B.Q() = Q::identity();

    convex::compute_closest_points<M>(
                                      X_A
                                      , &A
                                      , X_B
                                      , &B
                                      , p_A
                                      , p_B
                                      );

  }
}

BOOST_AUTO_TEST_CASE(sphere_sphere_testing)
{
  geometry::Sphere<V> const A;
  geometry::Sphere<V> const B;

  size_t const max_iterations       = 100u;
  T      const absolute_tolerance   = VT::numeric_cast(10e-6);
  T      const relative_tolerance   = VT::numeric_cast(10e-6);
  T      const stagnation_tolerance = VT::numeric_cast(10e-15);

  // Two unit-spheres placed ontop of each other
  {
    C X_A;
    C X_B;

    V p_A;
    V p_B;
    size_t iterations  = 0u;
    size_t status      = 0u;
    T      distance    = VT::infinity();

    X_A = C::identity();
    X_B = C::identity();
    X_B.T() = V::make(VT::zero(), VT::two(), VT::zero());

    convex::compute_closest_points<M>(
                                      X_A
                                      , &A
                                      , X_B
                                      , &B
                                      , p_A
                                      , p_B
                                      , distance
                                      , iterations
                                      , status
                                      , absolute_tolerance
                                      , relative_tolerance
                                      , stagnation_tolerance
                                      , max_iterations
                                      );

    T true_distance = tiny::norm( X_A.T() - X_B.T() ) - VT::two();

    BOOST_CHECK( fabs(distance - true_distance) < 10e-6);

    BOOST_CHECK( p_A(0) == VT::zero() );
    BOOST_CHECK( p_A(1) == VT::one() );
    BOOST_CHECK( p_A(2) == VT::zero() );

    BOOST_CHECK( p_B(0) == VT::zero() );
    BOOST_CHECK( p_B(1) == VT::one()  );
    BOOST_CHECK( p_B(2) == VT::zero() );
  }
  // Two unit-spheres overlapping but both placed on the x-axis
  {
    C X_A;
    C X_B;

    V p_A;
    V p_B;
    size_t iterations  = 0u;
    size_t status      = 0u;
    T      distance    = VT::infinity();

    X_A = C::identity();
    X_B = C::identity();
    X_B.T()(0) = VT::numeric_cast(2.5);

    convex::compute_closest_points<M>(
                                      X_A
                                      , &A
                                      , X_B
                                      , &B
                                      , p_A
                                      , p_B
                                      , distance
                                      , iterations
                                      , status
                                      , absolute_tolerance
                                      , relative_tolerance
                                      , stagnation_tolerance
                                      , max_iterations
                                      );

    T true_distance = tiny::norm( X_A.T() - X_B.T() ) - VT::two();

    BOOST_CHECK( fabs(distance - true_distance) < 10e-6);

    BOOST_CHECK( p_A(0) == VT::one()  );
    BOOST_CHECK( p_A(1) == VT::zero() );
    BOOST_CHECK( p_A(2) == VT::zero() );

    BOOST_CHECK( p_B(0) == VT::numeric_cast(1.5)  );
    BOOST_CHECK( p_B(1) == VT::zero() );
    BOOST_CHECK( p_B(2) == VT::zero() );
  }
  // Two unit-spheres exactly touching in one point (= one intersection point) and but both placed on the x-axis
  {
    C X_A;
    C X_B;

    V p_A;
    V p_B;
    size_t iterations  = 0u;
    size_t status      = 0u;
    T      distance    = VT::infinity();

    X_A.T().clear();
    X_A.T()(0) = -2.0;
    X_A.Q() = Q::identity();
    X_B.T().clear();
    X_B.Q() = Q::identity();

    convex::compute_closest_points<M>(
                                      X_A
                                      , &A
                                      , X_B
                                      , &B
                                      , p_A
                                      , p_B
                                      , distance
                                      , iterations
                                      , status
                                      , absolute_tolerance
                                      , relative_tolerance
                                      , stagnation_tolerance
                                      , max_iterations
                                      );

  }
  // Two unit-spheres non-overlapping but both placed on the x-axis
  {
    C X_A;
    C X_B;

    V p_A;
    V p_B;
    size_t iterations  = 0u;
    size_t status      = 0u;
    T      distance    = VT::infinity();

    X_A.T().clear();
    X_A.T()(0) = -2.5;
    X_A.Q() = Q::identity();
    X_B.T().clear();
    X_B.Q() = Q::identity();

    convex::compute_closest_points<M>(
                                      X_A
                                      , &A
                                      , X_B
                                      , &B
                                      , p_A
                                      , p_B
                                      , distance
                                      , iterations
                                      , status
                                      , absolute_tolerance
                                      , relative_tolerance
                                      , stagnation_tolerance
                                      , max_iterations
                                      );

  }
}

BOOST_AUTO_TEST_CASE(random_test)
{
  geometry::Sphere<V> const A;
  geometry::Sphere<V> const B;

  size_t const max_iterations       = 100u;
  T      const absolute_tolerance   = VT::numeric_cast(10e-6);
  T      const relative_tolerance   = VT::numeric_cast(10e-10);
  T      const stagnation_tolerance = VT::numeric_cast(0.0);

  for(size_t i=0;i<100u;++i)
  {
    C X_A;
    C X_B;

    V p_A;
    V p_B;
    size_t iterations  = 0u;
    size_t status      = 0u;
    T      distance    = VT::infinity();

    X_A.T() = V::random( -2.0, 2.0 );
    X_A.Q() = tiny::unit( Q::random() );
    X_B.T() = V::random( -2.0, 2.0 );
    X_B.Q() = tiny::unit( Q::random() );

    convex::compute_closest_points<M>(
                                      X_A
                                      , &A
                                      , X_B
                                      , &B
                                      , p_A
                                      , p_B
                                      , distance
                                      , iterations
                                      , status
                                      , absolute_tolerance
                                      , relative_tolerance
                                      , stagnation_tolerance
                                      , max_iterations
                                      );

    T const true_distance = tiny::norm( X_A.T() - X_B.T() ) - VT::two();

    if(  true_distance > absolute_tolerance )
    {
      BOOST_CHECK_CLOSE( true_distance, distance, 0.05 );

      BOOST_CHECK( status != convex::ABSOLUTE_CONVERGENCE );            // Would indicate penetration
      BOOST_CHECK( status != convex::INTERSECTION );                    // Would indicate penetration
      BOOST_CHECK( status != convex::ITERATING );                       // Would indicate internal error in CONVEX
      BOOST_CHECK( status != convex::EXCEEDED_MAX_ITERATIONS_LIMIT );   // Would indicate internal error in CONVEX
      BOOST_CHECK( status != convex::NON_DESCEND_DIRECTION );           // Would indicate internal error in CONVEX

      //BOOST_CHECK( status != convex::RELATIVE_CONVERGENCE );          // Would indicate convergence to positive distance
      //BOOST_CHECK( status != convex::SIMPLEX_EXPANSION_FAILED );      // Would indicate convergence to positive distance
      //BOOST_CHECK( status != convex::STAGNATION );                    // Would indicate convergence to positive distance
      //BOOST_CHECK( status != convex::LOWER_ERROR_BOUND_CONVERGENCE ); // Would indicate convergence to positive distance
    }
    else
    {
      BOOST_CHECK( 0.0 <= distance );
      BOOST_CHECK( distance <= absolute_tolerance );

      BOOST_CHECK( status != convex::STAGNATION );                     // Can only occur in case of positive distance
      BOOST_CHECK( status != convex::LOWER_ERROR_BOUND_CONVERGENCE );  // Can only occur in case of positive distance
      BOOST_CHECK( status != convex::RELATIVE_CONVERGENCE );           // Can only occur in case of positive distance
      BOOST_CHECK( status != convex::SIMPLEX_EXPANSION_FAILED );       // Can only occur in case of positive distance
      BOOST_CHECK( status != convex::ITERATING );                      // Would indicate internal error in CONVEX
      BOOST_CHECK( status != convex::EXCEEDED_MAX_ITERATIONS_LIMIT );  // Would indicate internal error in CONVEX
      BOOST_CHECK( status != convex::NON_DESCEND_DIRECTION );          // Would indicate internal error in CONVEX

      //BOOST_CHECK( status != convex::ABSOLUTE_CONVERGENCE );         // Indicates penetration
      //BOOST_CHECK( status != convex::INTERSECTION );                 // Indicates penetration
    }

  }
}

BOOST_AUTO_TEST_SUITE_END();


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2013 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//         Copyright 2012 - 2013 MetaScale SAS
//
//          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_ARITHMETIC_FUNCTIONS_SIMD_COMMON_FAST_REC_HPP_INCLUDED
#define BOOST_SIMD_ARITHMETIC_FUNCTIONS_SIMD_COMMON_FAST_REC_HPP_INCLUDED

#include <boost/simd/arithmetic/functions/fast_rec.hpp>
#include <boost/simd/include/functions/simd/raw_rec.hpp>
#include <boost/simd/include/functions/simd/refine_rec.hpp>
#include <boost/dispatch/attributes.hpp>

namespace boost { namespace simd { namespace ext
{
  // Single precision fast_rec is computable from the raw_rec + 1 NR step
  BOOST_DISPATCH_IMPLEMENT          (  fast_rec_, tag::cpu_
                                    , (A0)(X)
                                    , ((simd_< single_<A0>, X >))
                                    )
  {
    typedef A0 result_type;

    BOOST_FORCEINLINE result_type operator()(A0 const& a0) const
    {
      return refine_rec(a0, raw_rec( a0 ));
    }
  };
} } }

#endif


/******************************************************************************
 * Copyright (C) 2014 by Jerome Maye                                          *
 * jerome.maye@gmail.com                                                      *
 *                                                                            *
 * This program is free software; you can redistribute it and/or modify       *
 * it under the terms of the Lesser GNU General Public License as published by*
 * the Free Software Foundation; either version 3 of the License, or          *
 * (at your option) any later version.                                        *
 *                                                                            *
 * This program is distributed in the hope that it will be useful,            *
 * but WITHOUT ANY WARRANTY; without even the implied warranty of             *
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              *
 * Lesser GNU General Public License for more details.                        *
 *                                                                            *
 * You should have received a copy of the Lesser GNU General Public License   *
 * along with this program. If not, see <http://www.gnu.org/licenses/>.       *
 ******************************************************************************/

#include "aslam/calibration/error-terms/ErrorTermWheel.h"

#include <cmath>

#include <Eigen/Dense>

using namespace aslam::backend;

namespace aslam {
  namespace calibration {

/******************************************************************************/
/* Constructors and Destructor                                                */
/******************************************************************************/

    ErrorTermWheel::ErrorTermWheel(const EuclideanExpression& v_w_mw, const
        ScalarExpression& r, double measurement, const Covariance& sigma2Wheel) :
        _v_w_mw(v_w_mw),
        _r(r),
        _measurement(measurement),
        _sigma2Wheel(sigma2Wheel){
      setInvR((Eigen::Matrix<double, 1, 1>() << (1/ _sigma2Wheel)).finished());
      DesignVariable::set_t dv;
      _v_w_mw.getDesignVariables(dv);
      _r.getDesignVariables(dv);
      setDesignVariablesIterator(dv.begin(), dv.end());
    }

    ErrorTermWheel::ErrorTermWheel(const ErrorTermWheel& other) :
        ErrorTermFs<1>(other),
        _v_w_mw(other._v_w_mw),
        _r(other._r),
        _measurement(other._measurement),
        _sigma2Wheel(other._sigma2Wheel){
    }

    ErrorTermWheel& ErrorTermWheel::operator = (const ErrorTermWheel& other) {
      if (this != &other) {
        ErrorTermFs<1>::operator=(other);
       _v_w_mw= other._v_w_mw;
       _r = other._r;
       _measurement = other._measurement;
       _sigma2Wheel = other._sigma2Wheel;
      }
      return *this;
    }

    ErrorTermWheel::~ErrorTermWheel() {
    }

/******************************************************************************/
/* Accessors                                                                  */
/******************************************************************************/

    double ErrorTermWheel::getMeasurement() const {
      return _measurement;
    }

    void ErrorTermWheel::setMeasurement(double measurement) {
      _measurement = measurement;
    }

    const ErrorTermWheel::Covariance& ErrorTermWheel::getCovariance() const {
      return _sigma2Wheel;
    }

    void ErrorTermWheel::setCovariance(const Covariance& sigma2Wheel) {
      _sigma2Wheel = sigma2Wheel;
    }

/******************************************************************************/
/* Methods                                                                    */
/******************************************************************************/

    double ErrorTermWheel::evaluateErrorImplementation() {
      error_t error;
      const double v0 = _v_w_mw.toEuclidean()(0);
      const double r = _r.toScalar();
      error(0) = v0/r - _measurement;
      setError(error);
      return evaluateChiSquaredError();
    }

    void ErrorTermWheel::evaluateJacobiansImplementation(JacobianContainer&
        jacobians) {
      const double r = _r.toScalar();
      const double v0 = _v_w_mw.toValue()(0);

      Eigen::Matrix<double, 1, 3> J_w_v_mw = Eigen::Matrix<double, 1, 3>::Zero();
      J_w_v_mw(0,0) = 1.0/r;
      Eigen::Matrix<double, 1, 1> J_r;
      J_r(0,0) = -v0/(r*r);
      _v_w_mw.evaluateJacobians(jacobians, J_w_v_mw);
      _r.evaluateJacobians(jacobians, J_r);
    }

  }
}