#ifndef ZSVM_SPHERICAL_ECG_JACOBI_VARIATIONAL_OPTIMIZER_HPP_INCLUDED
#define ZSVM_SPHERICAL_ECG_JACOBI_VARIATIONAL_OPTIMIZER_HPP_INCLUDED

// C++ standard library headers
#include <chrono>
#include <cstddef> // for std::size_t
#include <random> // for std::random_device
#include <string>
#include <vector>

// OpenMP multithreading headers
#include <omp.h>

// Boost library headers
#include <boost/multiprecision/cpp_int.hpp>
#include <boost/random.hpp> // for boost::random::independent_bits_engine etc.

// Project-specific headers
#include "Restrict.hpp"
#include "AmoebaOptimizer.hpp"
#include "RealVariationalSolver.hpp"
#include "SphericalECGJacobiContext.hpp"

namespace zsvm {

    template <typename T>
    class SphericalECGJacobiVariationalOptimizer {

        typedef boost::random::mt19937 random_generator_t;

        typedef std::array<
                random_generator_t::result_type,
                random_generator_t::state_size> random_state_t;

        typedef boost::random::independent_bits_engine<
                random_generator_t,
                std::numeric_limits<T>::digits,
                boost::multiprecision::cpp_int> random_engine_t;

    private: // ================================================================

        const std::size_t num_particles;
        const std::size_t num_parameters;
        SphericalECGJacobiContext<T> context;
        RealVariationalSolver<T> solver;
        std::vector<std::vector<T>> basis;
        std::vector<std::vector<T>> basis_matrices;
        random_engine_t random_engine;
        boost::random::normal_distribution<T> correlation_distribution;

        static random_engine_t properly_seeded_random_engine() {
            random_state_t seed_data;
            std::random_device nondet_random_source;
            std::generate(seed_data.begin(), seed_data.end(),
                          std::ref(nondet_random_source));
            seed_data[0] = static_cast<random_generator_t::result_type>(
                    std::chrono::duration_cast<std::chrono::nanoseconds>(
                            std::chrono::system_clock::now().time_since_epoch()
                    ).count() % std::numeric_limits<
                            random_generator_t::result_type>::max());
            boost::random::seed_seq seed_sequence(seed_data.begin(),
                                                  seed_data.end());
            random_engine_t random_engine(seed_sequence);
            return random_engine;
        }

    public: // =================================================================

        explicit SphericalECGJacobiVariationalOptimizer(
                long long int space_dimension,
                const std::vector<zsvm::Particle<T>> &particles,
                const std::string &mass_carrier,
                const std::string &charge_carrier)
                : num_particles(particles.size()),
                  num_parameters(num_particles * (num_particles - 1) / 2),
                  context(SphericalECGJacobiContext<T>::create(
                          particles, mass_carrier, charge_carrier,
                          space_dimension)),
                  random_engine(properly_seeded_random_engine()),
                  correlation_distribution(0, 3) {}

        T get_ground_state_energy() {
            return solver.get_eigenvalue(0);
        }

    public: // =================================================================

        const std::vector<std::vector<T>> &get_basis() {
            return basis;
        }

        void set_basis(const std::vector<std::vector<T>> &input_basis) {
            basis = input_basis;
            recompute_basis_matrices();
            recompute_solver_matrices();
        }

        T augmented_ground_state_energy(
                const T *RESTRICT new_basis_matrix) {
            if (solver.empty()) {
                T overlap_element, hamiltonian_element;
                context.evaluate_matrix_elements(
                        overlap_element, hamiltonian_element,
                        new_basis_matrix, new_basis_matrix);
                return hamiltonian_element / overlap_element;
            }
            const std::size_t basis_size = basis_matrices.size();
            std::vector<T> new_overlap_column(basis_size + 1);
            std::vector<T> new_hamiltonian_column(basis_size + 1);
            for (std::size_t i = 0; i < basis_size; ++i) {
                context.evaluate_matrix_elements(
                        new_overlap_column[i], new_hamiltonian_column[i],
                        basis_matrices[i].data(), new_basis_matrix);
            }
            context.evaluate_matrix_elements(
                    new_overlap_column[basis_size],
                    new_hamiltonian_column[basis_size],
                    new_basis_matrix, new_basis_matrix);
            return solver.minimum_augmented_eigenvalue(
                    new_overlap_column.data(),
                    new_hamiltonian_column.data());
        }

        void random_basis_element(T *RESTRICT basis_element) {
            for (std::size_t i = 0; i < num_parameters; ++i) {
                basis_element[i] = correlation_distribution(random_engine);
            }
        }

        void construct_basis_matrix(
                T *RESTRICT basis_matrix,
                const T *RESTRICT basis_element) {
            context.gaussian_parameter_matrix(basis_element, basis_matrix);
        }

        void add_basis_element(const std::vector<T> &basis_element) {
            basis.push_back(basis_element);
            std::vector<T> basis_matrix(num_parameters);
            construct_basis_matrix(basis_matrix.data(), basis_element.data());
            basis_matrices.push_back(basis_matrix);
            const std::size_t basis_size = basis.size();
            solver.set_basis_size_conservative(basis_size);
            for (std::size_t i = 0; i < basis_size - 1; ++i) {
                context.evaluate_matrix_elements(
                        solver.overlap_matrix_element(i, basis_size - 1),
                        solver.hamiltonian_matrix_element(i, basis_size - 1),
                        basis_matrices[i].data(),
                        basis_matrices[basis_size - 1].data());
                solver.overlap_matrix_element(basis_size - 1, i) =
                        solver.overlap_matrix_element(i, basis_size - 1);
                solver.hamiltonian_matrix_element(basis_size - 1, i) =
                        solver.hamiltonian_matrix_element(i, basis_size - 1);
            }
            context.evaluate_matrix_elements(
                    solver.overlap_matrix_element(
                            basis_size - 1, basis_size - 1),
                    solver.hamiltonian_matrix_element(
                            basis_size - 1, basis_size - 1),
                    basis_matrices[basis_size - 1].data(),
                    basis_matrices[basis_size - 1].data());
        }

        void replace_basis_element(const std::vector<T> &basis_element) {
            const std::size_t basis_size = basis.size();
            basis[basis_size - 1] = basis_element;
            construct_basis_matrix(basis_matrices[basis_size - 1].data(),
                                   basis_element.data());
            for (std::size_t i = 0; i < basis_size - 1; ++i) {
                context.evaluate_matrix_elements(
                        solver.overlap_matrix_element(i, basis_size - 1),
                        solver.hamiltonian_matrix_element(i, basis_size - 1),
                        basis_matrices[i].data(),
                        basis_matrices[basis_size - 1].data());
                solver.overlap_matrix_element(basis_size - 1, i) =
                        solver.overlap_matrix_element(i, basis_size - 1);
                solver.hamiltonian_matrix_element(basis_size - 1, i) =
                        solver.hamiltonian_matrix_element(i, basis_size - 1);
            }
            context.evaluate_matrix_elements(
                    solver.overlap_matrix_element(
                            basis_size - 1, basis_size - 1),
                    solver.hamiltonian_matrix_element(
                            basis_size - 1, basis_size - 1),
                    basis_matrices[basis_size - 1].data(),
                    basis_matrices[basis_size - 1].data());
        }

        std::vector<T> last_basis_element() {
            return basis.back();
        }

        bool expand_amoeba(std::size_t num_trials,
                           const T &initial_step_size,
                           std::size_t max_steps) {
            std::vector<T> new_basis_element(num_parameters);
            std::vector<T> best_basis_element;
            T best_energy = solver.empty()
                            ? std::numeric_limits<T>::max()
                            : solver.get_eigenvalue(0);
            for (std::size_t trial = 0; trial < num_trials; ++trial) {
                random_basis_element(new_basis_element.data());
                const T new_energy = refine_amoeba(
                        new_basis_element.data(), initial_step_size, max_steps);
                if (new_energy < best_energy) {
                    best_basis_element = new_basis_element;
                    best_energy = new_energy;
                }
            }
            if (!best_basis_element.empty()) {
                add_basis_element(best_basis_element);
                return true;
            } else {
                return false;
            }
        }

    public: // =================================================================

        T refine_amoeba(T *RESTRICT basis_element,
                        const T &initial_step_size,
                        std::size_t max_steps) {
            std::vector<T> basis_matrix(num_parameters);
            dznl::AmoebaOptimizer<T> amoeba(
                    basis_element, num_parameters, initial_step_size,
                    [&](const T *b) {
                        construct_basis_matrix(basis_matrix.data(), b);
                        return augmented_ground_state_energy(
                                basis_matrix.data());
                    });
            for (std::size_t step = 0; step < max_steps; ++step) {
                amoeba.step();
            }
            return amoeba.current_minimum(basis_element);
        }

        void recompute_basis_matrices() {
            basis_matrices.clear();
            for (const auto &basis_element : basis) {
                std::vector<T> basis_matrix(num_parameters);
                construct_basis_matrix(basis_matrix.data(),
                                       basis_element.data());
                basis_matrices.push_back(basis_matrix);
            }
        }

        void recompute_solver_matrices() {
            const std::size_t num_basis_matrices = basis_matrices.size();
            solver.set_basis_size_destructive(num_basis_matrices);
            for (std::size_t i = 0; i < num_basis_matrices; ++i) {
                for (std::size_t j = 0; j < num_basis_matrices; ++j) {
                    context.evaluate_matrix_elements(
                            solver.overlap_matrix_element(i, j),
                            solver.hamiltonian_matrix_element(i, j),
                            basis_matrices[i].data(),
                            basis_matrices[j].data());
                }
            }
        }

    }; // class SphericalECGJacobiVariationalOptimizer

} // namespace zsvm

#endif // ZSVM_SPHERICAL_ECG_JACOBI_VARIATIONAL_OPTIMIZER_HPP_INCLUDED


#include <iostream>
#include <boost/numeric/mtl/mtl.hpp>

using namespace std;
int main(int, char**)
{
    mtl::dense2D<double>    A(3, 4), A_t(3, 4), S(3, 3), V(3, 4), D(4,4);
    A= 0;

    A[0][0]=1;  A[0][1]=1;  A[0][2]=1;  A[0][3]=4;
    A[1][0]=1;  A[1][1]=2;  A[1][2]=2;  A[1][3]=3;
    A[2][0]=9;  A[2][1]=3;  A[2][2]=2;  A[2][3]=4;
    std::cout<<"A=\n"<< A <<"\n";

    boost::tie(S, V, D)= svd(A, 1.e-10)= svd(A);  //second argument is optional (missmatch of upper R (A= Q*R))
    std::cout<<"Matrix  S=\n"<< S <<"\n";
    std::cout<<"Matrix  V=\n"<< V <<"\n";
    std::cout<<"Matrix  D=\n"<< D <<"\n";
    A_t= S*V*trans(D);
    std::cout<<"Matrix  A=S*V*D'=\n"<< A_t <<"\n";

    return 0;
}


#ifndef H_CONVEX_MPC
#define H_CONVEX_MPC

// Standard
#include <math.h>
#include <iostream>
#include <stdio.h>

// Timer
#include <chrono>

// QP Solver
#include <ExternalSource/myOptimizer/Goldfarb/QuadProg++.hh>

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

// Gait Cycle
#include <PnC/GaitCycle/GaitCycle.hpp>
#include <memory>

// Reaction Force Schedule
#include <PnC/MPC/ReactionForceSchedule.hpp>

// #define MPC_PRINT_ALL
// #define MPC_TIME_ALL

// We are following the MPC formulation from:
// Di Carlo, Jared, et al. "Dynamic locomotion in the mit cheetah 3 through convex model-predictive control."
// 2018 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). IEEE, 2018.

// We modify the formulation to smoothen the solutions at each iteration.
//  There is also an optional setting to smoothen the current solution from the previous solution

class CMPC{
public:
	CMPC();
	~CMPC();

	double robot_mass; // kg mass of the robot
	Eigen::MatrixXd I_robot; // Robot inertia

	int horizon;// mpc horizon (number of steps);
	double mpc_dt; // mpc time interval per horizon steo
	double mu; // coefficient of friction
	double fz_max; // maximum z reaction force for one force vector.
	double control_alpha_; // Regularization term on the controls
	double delta_smooth_; // Parameter to smoothen the reaction force results

	bool smooth_from_prev_result_; // Whether to smooth the solution of the MPC iteration using the previous result

	bool custom_smoothing_; // Enable user-speicified custom smoothing matrix
	Eigen::MatrixXd Dc1; // Custom smoothing for the entire horizon

	bool rotate_inertia; // whether or not the inertia needs to be rotated to the world frame.
						 // If the inertia is expressed in body frame this needs to be true.
						 // Otherwise if the inertia is already updated to be in the world frame, this can be set to false.

	Eigen::VectorXd latest_f_vec_out; // Container holding the the latest computed force output for control (not to be confused with the force at the end of the horizon)
	Eigen::VectorXd f_prev; // Previous reaction force vector for the first horizon.

	Eigen::VectorXd X_pred; // dimension: (13*horizon) state prediction over the horizon. [x_{k+1}, x_{k+2}, ..., x_{k+horizon}]
	Eigen::VectorXd F_out;  // dimension: (n_Fr*horizon) forces to exert over the horizon. [f_{k}, f_{k+1}, ..., f_{k - 1 + horizon}]

	std::shared_ptr<GaitCycle> gait_cycle_ptr; // pointer to the gait cycle object
	std::shared_ptr<ReactionForceSchedule> reaction_force_schedule_ptr; // pointer to the reaction force schedule

  	// Vector cost for the MPC: <<  th1,  th2,  th3,  px,  py,  pz,   w1,  w2,   w3,   dpx,  dpy,  dpz,  g
	// last term is gravity and should always be 0,0
	// cost_vec << 0.25, 0.25, 10.0, 2.0, 2.0, 50.0, 0.0, 0.0, 0.30, 0.20, 0.2, 0.10, 0.0;
	Eigen::VectorXd cost_vec;

	bool use_terminal_cost; // if true, will use the custom terminal cost
	Eigen::VectorXd terminal_cost_vec;

	// Returns the state vector size of the MPC per horizon.
	int getStateVecDim(){return 13;}

	// Helper function which returns a 13-vector given the state of the robot.
	// x = [Theta, p, omega, pdot, g] \in \mathbf{R}^13
	Eigen::VectorXd getx0(const double roll, const double pitch, const double yaw,
						  const double com_x, const double com_y, const double com_z,
						  const double roll_rate, const double pitch_rate, const double yaw_rate,
						  const double com_x_rate, const double com_y_rate, const double com_z_rate);

	// Gets the latest computed ground forces for control
	Eigen::VectorXd getComputedGroundForces(){return latest_f_vec_out;}
	// Get computed ground forces in matrix form where the columns are the ground reaction forces
	// in the same order as r_feet.
	Eigen::MatrixXd getMatComputedGroundForces();

	// Gets the computed forces over the horizon
	Eigen::VectorXd getForcesOverHorizon(){return F_out;}
	// Gets the predicted state evolution over the horizon
	Eigen::VectorXd getXpredOverHorizon(){return X_pred;}

	void setRobotMass(const double robot_mass_in){ robot_mass = robot_mass_in; }
	// Whether or not the inertia needs to be rotated to the world frame.
	void rotateBodyInertia(bool rotate_inertia_in){ rotate_inertia = rotate_inertia_in; }

	void setRobotInertia(const Eigen::MatrixXd & inertia_in){ I_robot = inertia_in; }
	void setDt(const double mpc_dt_in){mpc_dt = mpc_dt_in;}	// MPC dt interval per horizon
	void setHorizon(const int & horizon_in);	// MPC horizon (number of steps)

	void setMu(const double mu_in){ mu = mu_in;} // Set the coefficient of friction
	void setMaxFz(const double fz_max_in); // Set the maximum z reaction force for one force vector
	void setSmoothFromPrevResult(const bool smooth_prev_in){smooth_from_prev_result_ = smooth_prev_in;} // Sets whether to smoothen the current solution using the previous solution

	// Set preview start time for the gate cycle
	void setPreviewStartTime(const double t_preview_start_in);
	void setCustomGaitCycle(std::shared_ptr<GaitCycle> gait_cycle_ptr_in);
	void setCustomReactionForceSchedule(std::shared_ptr<ReactionForceSchedule> reaction_force_schedule_ptr_in);

  	// Vector cost for the MPC: <<  th1,  th2,  th3,  px,  py,  pz,   w1,  w2,   w3,   dpx,  dpy,  dpz,  g
	void setCostVec(const Eigen::VectorXd & cost_vec_in); // Sets the cost vector.

  	// Set Terminal vector cost for the MPC: <<  th1,  th2,  th3,  px,  py,  pz,   w1,  w2,   w3,   dpx,  dpy,  dpz,  g
	// use_terminal_cost_in : if true, will use the terminal cost specified.
	void setTerminalCostVec(const bool use_terminal_cost_in, const Eigen::VectorXd & terminal_cost_vec_in); // Sets the cost vector.

	void setControlAlpha(const double control_alpha_in){control_alpha_ = control_alpha_in;}
	void setDeltaSmooth(const double delta_smooth_in){delta_smooth_ = delta_smooth_in;}

	void enableCustomSmoothing(const bool enable_custom_smoothing){custom_smoothing_ = enable_custom_smoothing;}
	void setCustomSmoothing(const Eigen::MatrixXd & Dc1_in){ Dc1 = Dc1_in; }

	// Human readable prints out of f_vec_out.
	void print_f_vec(int & n_Fr, const Eigen::VectorXd & f_vec_out);

	// Helper function which transforms a constant x_des to a constant trajectory reference
	// Input: x_des \in \mathbf{R}^{13}
	// Ouput: X_ref \in \mathbf{R}^{13*horizon}
	void get_constant_desired_x(const Eigen::VectorXd & x_des, Eigen::VectorXd & X_ref);

	// The main solve MPC routine
	// All the values are in world frame
	// Inputs:
	// 		x0 = [Theta, p, omega, pdot, g] \in \mathbf{R}^{13}    (Starting state of the system)
	// 		X_des \in \mathbf{R}^{13*horizon} 				     (Reference or desired state evolution)
	// 		r_feet \in \mathbf{R}^{3 x number_of_point_contacts} (Matrix of point contact locations where each column is the xyz location of the contact)
	// Outputs:
	//		x_pred \in \mathbf{R}^{13} - The predicted state after a time interval of mpc_dt has dimension
	//		f_vec_out \mathbf{R}^{3 x number_of_point_contacts} The xyz forces needed to be exerted at the contact points to track X_des.

	void solve_mpc(const Eigen::VectorXd & x0, const Eigen::VectorXd & X_des, const Eigen::MatrixXd & r_feet,
	               Eigen::VectorXd & x_pred, Eigen::VectorXd & f_vec_out);

	// For testing
	void simulate_toy_mpc();

	void assemble_vec_to_matrix(const int & n, const int & m, const Eigen::VectorXd & vec, Eigen::MatrixXd & mat_out);
	void integrate_robot_dynamics(const double & dt, const Eigen::VectorXd & x_current, const Eigen::MatrixXd & f_Mat, const Eigen::MatrixXd & r_feet,
	                              Eigen::VectorXd & x_next);

private:
	Eigen::MatrixXd r_feet_; // Store the value of r_feet locally.
	double t_preview_start; // store start time of the preview.

	double gravity_acceleration;
  	Eigen::MatrixXd R_roll(const double & phi);
	Eigen::MatrixXd R_pitch(const double & theta);
	Eigen::MatrixXd R_yaw(const double & psi);
	Eigen::MatrixXd skew_sym_mat(const Eigen::VectorXd & v);

	void cont_time_state_space(const Eigen::VectorXd & x_current,
	                           const Eigen::MatrixXd & r_feet,
	                           Eigen::MatrixXd & A, Eigen::MatrixXd & B);
	void discrete_time_state_space(const Eigen::MatrixXd & A, const Eigen::MatrixXd & B, Eigen::MatrixXd & Adt, Eigen::MatrixXd & Bdt);

	void qp_matrices(const Eigen::MatrixXd & Adt, const Eigen::MatrixXd & Bdt, Eigen::MatrixXd & Aqp, Eigen::MatrixXd & Bqp);
	void get_force_constraints(const int & n_Fr, Eigen::MatrixXd & CMat, Eigen::VectorXd & cvec);
	void get_qp_constraints(const Eigen::MatrixXd & CMat, const Eigen::VectorXd & cvec, Eigen::MatrixXd & Cqp, Eigen::VectorXd & cvec_qp);
	void get_qp_costs(const int & n, const int & m, const Eigen::VectorXd & vecS_cost, const double & control_alpha, Eigen::MatrixXd & Sqp, Eigen::MatrixXd & Kqp, Eigen::MatrixXd& D0);

	// Converts location of the feet expressed in world frame to the CoM frame. We assume the CoM orientation frame is always aligned with world.
	void convert_r_feet_to_com_frame(const Eigen::VectorXd & p_com, const Eigen::MatrixXd & r_feet, Eigen::MatrixXd & r_feet_com);

	void solve_mpc_qp(const Eigen::MatrixXd & Aqp,  const Eigen::MatrixXd & Bqp, const Eigen::VectorXd & X_ref,
					  const Eigen::VectorXd & x0,   const Eigen::MatrixXd & Sqp, const Eigen::MatrixXd & Kqp,
                      const Eigen::MatrixXd& D0, const Eigen::VectorXd& f_prev_in,
					  const Eigen::MatrixXd & Cqp, const Eigen::VectorXd & cvec_qp,
					  Eigen::VectorXd & f_vec_out);

};

#endif


#include <SBGATPolyhedronGravityModel.hpp>
#include <SBGATPolyhedronGravityModelUQ.hpp>
#include <SBGATObjWriter.hpp>

#include <vtkCleanPolyData.h>
#include <vtkOBJReader.h>

#include <json.hpp>
#include <boost/progress.hpp>

#include <armadillo>

int main(){

	std::ifstream i("input_file.json");
	nlohmann::json input_data;
	i >> input_data;

	std::string PATH_SHAPE = input_data["PATH_SHAPE"];
	double CORRELATION_DISTANCE =  input_data["CORRELATION_DISTANCE"];

	double ERROR_STANDARD_DEV  = input_data["ERROR_STANDARD_DEV"];
	double DENSITY  = input_data["DENSITY"];

	bool UNIT_IN_METERS  = input_data["UNIT_IN_METERS"];
	bool HOLD_MASS_CONSTANT  = input_data["HOLD_MASS_CONSTANT"];

	int N_MONTE_CARLO = input_data["N_MONTE_CARLO"];

	std::string OUTPUT_DIR = input_data["OUTPUT_DIR"];

	std::cout << "- Path to shape: " << PATH_SHAPE << std::endl;
	std::cout << "- Standard deviation on point coordinates (m) : " << ERROR_STANDARD_DEV << std::endl;
	std::cout << "- Correlation distance (m) : " << CORRELATION_DISTANCE << std::endl;
	std::cout << "- Density (kg/m^3) : " << DENSITY << std::endl;
	std::cout << "- Monte Carlo Draws : " << N_MONTE_CARLO << std::endl;

	// Reading
	vtkSmartPointer<vtkOBJReader> reader = vtkSmartPointer<vtkOBJReader>::New();
	reader -> SetFileName(PATH_SHAPE.c_str());
	reader -> Update();

	// An instance of SBGATPolyhedronGravityModel is created to evaluate the PGM of
	// the considered polytdata
	vtkSmartPointer<SBGATPolyhedronGravityModel> pgm_filter = vtkSmartPointer<SBGATPolyhedronGravityModel>::New();
	pgm_filter -> SetInputConnection(reader -> GetOutputPort());
	pgm_filter -> SetDensity(DENSITY);

	if(UNIT_IN_METERS){
		pgm_filter -> SetScaleMeters();
	} else{
		pgm_filter -> SetScaleKiloMeters();
	}

	std::cout << "Building pgm ...\n";
	pgm_filter -> Update();

	// An instance of SBGATPolyhedronGravityModelUQ is created to perform
	// uncertainty quantification from the PGM associated to the shape
	SBGATPolyhedronGravityModelUQ pgm_uq;
	pgm_uq.SetModel(pgm_filter);
	pgm_uq.PrecomputeMassPropertiesPartials();
	std::cout << "Populating shape covariance ...\n";

	// Populate the shape vertices covariance
	pgm_uq.ComputeVerticesCovarianceGlobal(ERROR_STANDARD_DEV,CORRELATION_DISTANCE);
	// Regularizing the covariance
	int regularized_eigen_values = pgm_uq.RegularizeCovariance();

	std::cout << regularized_eigen_values << " eigenvalues were regularized\n";

	std::cout << "Saving non-zero partition of shape covariance ...\n";

	// Save the covariance
	pgm_uq.SaveNonZeroVerticesCovariance(OUTPUT_DIR + "shape_covariance.json");

	// Saving baseline slices
	pgm_uq.TakeAndSaveSlice(0,OUTPUT_DIR + "baseline_slice_x.txt",0);
	pgm_uq.TakeAndSaveSlice(1,OUTPUT_DIR + "baseline_slice_y.txt",0);
	pgm_uq.TakeAndSaveSlice(2,OUTPUT_DIR + "baseline_slice_z.txt",0);

	std::vector<arma::vec::fixed<3> > all_positions = {
		arma::vec::fixed<3>({300,0,0}),
		arma::vec::fixed<3>({400,0,0}),
		arma::vec::fixed<3>({500,0,0}),
		arma::vec::fixed<3>({-300,0,0}),
		arma::vec::fixed<3>({-400,0,0}),
		arma::vec::fixed<3>({-500,0,0}),
		arma::vec::fixed<3>({0,300,0}),
		arma::vec::fixed<3>({0,400,0}),
		arma::vec::fixed<3>({0,500,0}),
		arma::vec::fixed<3>({0,-300,0}),
		arma::vec::fixed<3>({0,-400,0}),
		arma::vec::fixed<3>({0,-500,0}),
		arma::vec::fixed<3>({0,0,300}),
		arma::vec::fixed<3>({0,0,400}),
		arma::vec::fixed<3>({0,0,500}),
		arma::vec::fixed<3>({0,0,-300}),
		arma::vec::fixed<3>({0,0,-400}),
		arma::vec::fixed<3>({0,0,-500}),
	};

	// Analytical UQ

	std::vector<arma::mat::fixed<3,3> > analytical_covariances_acc(all_positions.size());
	std::vector<double> analytical_variances_pot(all_positions.size());

	std::cout << "Computing analytical uncertainties ... ";
	auto start = std::chrono::system_clock::now();
	#pragma omp parallel for
	for (int e = 0; e < all_positions.size(); ++e){
		analytical_variances_pot[e] = pgm_uq.GetVariancePotential(all_positions[e],HOLD_MASS_CONSTANT);
		analytical_covariances_acc[e] = pgm_uq.GetCovarianceAcceleration(all_positions[e],HOLD_MASS_CONSTANT);
	}
	auto end = std::chrono::system_clock::now();

	std::chrono::duration<double> elapsed_seconds = end-start;
	std::cout << "Done computing analytical uncertainties in " << elapsed_seconds.count() << " s\n";

	// Running a Monte Carlo to compare againnst
	std::vector<arma::vec> deviations;
	std::vector<double> densities;
	std::vector<std::vector<arma::vec::fixed<3> > >  all_accelerations;
	std::vector < std::vector<double> > all_potentials;
	std::vector<vtkSmartPointer<vtkPolyData > > saved_shapes(10);

	std::cout << "Running MC ... ";

	start = std::chrono::system_clock::now();
	SBGATPolyhedronGravityModelUQ::RunMCUQPotentialAccelerationInertial(PATH_SHAPE,DENSITY,
		UNIT_IN_METERS,
		HOLD_MASS_CONSTANT,
		pgm_uq.GetCovarianceSquareRoot(),
		N_MONTE_CARLO,
		all_positions,
		OUTPUT_DIR,
		std::min(30,N_MONTE_CARLO),
		deviations,
		densities,
		all_accelerations,
		all_potentials);

	end = std::chrono::system_clock::now();

	elapsed_seconds = end-start;

	std::cout << "Done running MC in " << elapsed_seconds.count() << " s\n";

	std::vector<double> mc_variances_pot(all_positions.size());
	std::vector<arma::mat > mc_covariances_acc(all_positions.size());

	std::cout << "Computing MC dispersions...\n";

	#pragma omp parallel for
	for (int e = 0; e < mc_variances_pot.size(); ++e){

		arma::vec potentials_mc(N_MONTE_CARLO);
		arma::mat accelerations_mc(3,N_MONTE_CARLO);

		for (int sample = 0; sample < N_MONTE_CARLO; ++sample){
			potentials_mc(sample) = all_potentials[sample][e];
			accelerations_mc.col(sample) = all_accelerations[sample][e];
		}

		mc_variances_pot[e] = arma::var(potentials_mc);
		mc_covariances_acc[e] = arma::cov(accelerations_mc.t());

	}

	std::cout << "\t After " << N_MONTE_CARLO << " MC outcomes:\n";

	for (int e = 0; e < all_positions.size(); ++e){
		all_positions[e].t().print("\t At: ");
		std::cout << "\t\tMC variance in potential: " << mc_variances_pot[e] << std::endl;
		std::cout << "\t\tAnalytical variance in potential: " << analytical_variances_pot[e] << std::endl;

		std::cout << "\t\tError (%): " << (mc_variances_pot[e] - analytical_variances_pot[e])/analytical_variances_pot[e] * 100 << std::endl;

		std::cout << "\t\tMC Covariance in acceleration: \n" << mc_covariances_acc[e] << std::endl;
		std::cout << "\t\tAnalytical covariance in acceleration: \n" << analytical_covariances_acc[e] << std::endl;

		std::cout << "\t\tError (%): " << arma::norm(mc_covariances_acc[e] - analytical_covariances_acc[e])/arma::trace(mc_covariances_acc[e]) * 100 << std::endl;
	}

	return 0;
}


#pragma once

#include <iostream>
#include <complex>
#include <cmath>
#include <vector>
#include <set>
#include <assert.h>
#include <memory>
#include <iomanip>
#include <fstream>

#include <Eigen/CXX11/Tensor>
#include <Eigen/MPRealSupport>

#include "common.hpp"
#include "kernel.hpp"
#include "piecewise_polynomial.hpp"

#include "irlib/detail/basis_impl.ipp"

namespace irlib {

/**
 * Class representing kernel Ir basis
 */
    class basis {
    public:
        /**
         * Constructor
         * @param s  statistics
         * @param Lambda Lambda
         * @param sv singular values
         * @param u_basis piecewise polynomials representing u_l(x)
         * @param v_basis piecewise polynomials representing v_l(y)
         */
        basis(statistics::statistics_type s,
            double Lambda,
            const std::vector<mpfr::mpreal> &sv,
            const std::vector<piecewise_polynomial<mpfr::mpreal, mpfr::mpreal>> &u_basis,
            const std::vector<piecewise_polynomial<mpfr::mpreal, mpfr::mpreal>> &v_basis,
            const std::vector<std::vector<mpfr::mpreal>> &u_basis_coeff_l,
            const std::vector<std::vector<mpfr::mpreal>> &v_basis_coeff_l
        ) throw(std::runtime_error) {
            statistics_ = s;
            Lambda_ = Lambda;
            sv_ = sv;
            u_basis_ = u_basis;
            v_basis_ = v_basis;
            u_basis_coeff_l_ = u_basis_coeff_l;
            v_basis_coeff_l_ = v_basis_coeff_l;
        }

    private:
        statistics::statistics_type statistics_;
        double Lambda_;
        std::vector<mpfr::mpreal> sv_;
        std::vector<piecewise_polynomial<mpfr::mpreal, mpfr::mpreal>> u_basis_, v_basis_;
        std::vector<std::vector<mpfr::mpreal>> u_basis_coeff_l_;
        std::vector<std::vector<mpfr::mpreal>> v_basis_coeff_l_;

        //mutable mp_prec_t default_prec_bak;

        mp_prec_t save_default_prec() const {
            mp_prec_t default_prec_bak = mpfr::mpreal::get_default_prec();
            mpfr::mpreal::set_default_prec(get_prec());
            return default_prec_bak;
        }

        void restore_default_prec(mp_prec_t prec) const {
            mpfr::mpreal::set_default_prec(prec);
        }

    public:
        /**
         * Compute the values of the basis functions for a given x.
         * @param x    x = 2 * tau/beta - 1  (-1 <= x <= 1)
         * @param val  results
         */
        double sl(int l) const throw(std::runtime_error) {
            assert(l >= 0 && l < dim());
            python_runtime_check(l >= 0 && l < dim(), "Index l is out of range.");
            return static_cast<double>(sv_[l]);
        }

        mpfr::mpreal sl_mp(int l) const throw(std::runtime_error) {
            assert(l >= 0 && l < dim());
            python_runtime_check(l >= 0 && l < dim(), "Index l is out of range.");
            return sv_[l];
        }

        double Lambda() const {
            return Lambda_;
        }

        /**
         * @param l  order of basis function
         * @param x  x on [-1,1]
         * @return   The value of u_l(x)
         */
        double ulx(int l, double x) const throw(std::runtime_error) {
            assert(x >= -1 && x <= 1);
            assert(l >= 0 && l < dim());
            auto bak = save_default_prec();

            //auto val = ulx_mp(l, mpfr::mpreal(x));
            //auto r = static_cast<double>(val);

            auto r = static_cast<double>(ulx_mp(l, mpfr::mpreal(x)));
            restore_default_prec(bak);
            return r;
        }

        double ulx_derivative(int l, double x, int order) const throw(std::runtime_error) {
            assert(x >= -1 && x <= 1);
            assert(l >= 0 && l < dim());
            auto bak = save_default_prec();
            auto r = static_cast<double>(ulx_derivative_mp(l, mpfr::mpreal(x), order));
            restore_default_prec(bak);
            return r;
        }

        /**
         * @param l  order of basis function
         * @param y  y on [-1,1]
         * @return   The value of v_l(y)
         */
        double vly(int l, double y) const throw(std::runtime_error) {
            assert(y >= -1 && y <= 1);
            assert(l >= 0 && l < dim());
            auto bak = save_default_prec();
            auto r = static_cast<double>(vly_mp(l, mpfr::mpreal(y)));
            restore_default_prec(bak);
            return r;
        }

        double vly_derivative(int l, double y, int order) const throw(std::runtime_error) {
            assert(y >= -1 && y <= 1);
            assert(l >= 0 && l < dim());
            auto bak = save_default_prec();
            auto r = static_cast<double>(vly_derivative_mp(l, mpfr::mpreal(y), order));
            restore_default_prec(bak);
            return r;
        }

        /**
         * Direct access to coefficients of u_l(x)
         */
        double coeff_ulx(int l, int section, int p) const {
            assert(l >= 0 && l < dim());
            return static_cast<double>(u_basis_[l].coefficient(section, p));
        }

        std::vector<mpreal> coeff_ulx_leg(int l) const {
            assert(l >= 0 && l < dim());
            return u_basis_coeff_l_[l];
        }

        /**
         * Direct access to coefficients of v_l(y)
         */
        double coeff_vly(int l, int section, int p) const {
            assert(l >= 0 && l < dim());
            return static_cast<double>(v_basis_[l].coefficient(section, p));
        }

        std::vector<mpreal> coeff_vly_leg(int l) const {
          assert(l >= 0 && l < dim());
          return v_basis_coeff_l_[l];
        }

        /**
         * Access to sections
         */
        int num_sections_ulx() const {
            return u_basis_[0].num_sections();
        }

        int num_sections_vly() const {
            return v_basis_[0].num_sections();
        }

        double section_edge_ulx(int i) const {
            return static_cast<double>(u_basis_[0].section_edge(i));
        }

        double section_edge_vly(int i) const {
            return static_cast<double>(v_basis_[0].section_edge(i));
        }

        int num_local_poly_ulx() const {
            return u_basis_[0].order() + 1;
        }

        int num_local_poly_vly() const {
            return v_basis_[0].order() + 1;
        }

        /**
         * This function should not be called outside this library
         * @param l  order of basis function
         * @param x  x on [-1,1]
         * @return   The value of u_l(x)
         */
#ifndef SWIG //DO NOT EXPOSE TO PYTHON
        mpfr::mpreal ulx_mp(int l, const mpfr::mpreal &x) const throw(std::runtime_error) {
            assert(x >= -1 && x <= 1);
            assert(l >= 0 && l < dim());
            python_runtime_check(l >= 0 && l < dim(), "Index l is out of range.");
            python_runtime_check(x >= -1 && x <= 1, "x must be in [-1,1].");

            auto bak = save_default_prec();

            mpfr:mpreal r;
            if (x >= 0) {
                r = u_basis_[l].compute_value(x);
            } else {
                r = u_basis_[l].compute_value(-x) * (l % 2 == 0 ? 1 : -1);
            }

            restore_default_prec(bak);

            return r;
        }

        mpfr::mpreal ulx_derivative_mp(int l, const mpfr::mpreal &x, int order) const throw(std::runtime_error) {
            assert(x >= -1 && x <= 1);
            assert(l >= 0 && l < dim());
            python_runtime_check(l >= 0 && l < dim(), "Index l is out of range.");
            python_runtime_check(x >= -1 && x <= 1, "x must be in [-1,1].");

            auto bak = save_default_prec();

            mpfr::mpreal r;
            if (x >= 0) {
                r = u_basis_[l].derivative(x, order);
            } else {
                r = u_basis_[l].derivative(-x, order) * ((l+order) % 2 == 0 ? 1 : -1);
            }

            restore_default_prec(bak);

            return r;
        }

        /**
         * This function should not be called outside this library
         * @param l  order of basis function
         * @param y  y on [-1,1]
         * @return   The value of v_l(y)
         */
        mpfr::mpreal vly_mp(int l, const mpfr::mpreal &y) const throw(std::runtime_error) {
            assert(y >= -1 && y <= 1);
            assert(l >= 0 && l < dim());
            python_runtime_check(l >= 0 && l < dim(), "Index l is out of range.");
            python_runtime_check(y >= -1 && y <= 1, "y must be in [-1,1].");

            auto bak = save_default_prec();

            mpfr::mpreal r;
            if (y >= 0) {
                r = v_basis_[l].compute_value(y);
            } else {
                r = v_basis_[l].compute_value(-y) * (l % 2 == 0 ? 1 : -1);
            }

            restore_default_prec(bak);

            return r;
        }

        mpfr::mpreal vly_derivative_mp(int l, const mpfr::mpreal &y, int order) const throw(std::runtime_error) {
            assert(y >= -1 && y <= 1);
            assert(l >= 0 && l < dim());
            python_runtime_check(l >= 0 && l < dim(), "Index l is out of range.");
            python_runtime_check(y >= -1 && y <= 1, "y must be in [-1,1].");

            auto bak = save_default_prec();

            mpfr::mpreal r;
            if (y >= 0) {
                r = v_basis_[l].derivative(y, order);
            } else {
                r = v_basis_[l].derivative(-y, order) * ((l+order) % 2 == 0 ? 1 : -1);
            }

            restore_default_prec(bak);

            return r;
        }
#endif

        std::string ulx_str(int l, const std::string& str_x) const throw(std::runtime_error) {
            auto bak = save_default_prec();

            auto prec = u_basis_[l].section_edge(0).get_prec();
            mpfr::mpreal x(str_x, prec);
            auto ulx = ulx_mp(l, x);

            std::ostringstream out;
            out << std::setprecision(mpfr::bits2digits(ulx.get_prec())) << ulx;
            //std::cout << "debug ulx_str " << std::setprecision(20) << str_x << " " << x << " " << ulx << " " << out.str() << std::endl;

            restore_default_prec(bak);

            return out.str();
        }

        std::string vly_str(int l, const std::string& str_y) const throw(std::runtime_error) {
            auto bak = save_default_prec();

            auto prec = v_basis_[l].section_edge(0).get_prec();
            mpfr::mpreal y(str_y, prec);
            auto vly = vly_mp(l, y);

            std::ostringstream out;
            out << std::setprecision(mpfr::bits2digits(vly.get_prec())) << vly;

            restore_default_prec(bak);

            return out.str();
        }

        /**
         * Return a reference to the l-th basis function
         * @param l l-th basis function
         * @return  reference to the l-th basis function
         */
        const piecewise_polynomial<mpfr::mpreal, mpfr::mpreal> &ul(int l) const throw(std::runtime_error) {
            assert(l >= 0 && l < dim());
            python_runtime_check(l >= 0 && l < dim(), "Index l is out of range.");
            return u_basis_[l];
        }

        const piecewise_polynomial<mpfr::mpreal, mpfr::mpreal> &vl(int l) const throw(std::runtime_error) {
            assert(l >= 0 && l < dim());
            python_runtime_check(l >= 0 && l < dim(), "Index l is out of range.");
            return v_basis_[l];
        }

        /**
         * Return number of basis functions
         * @return  number of basis functions
         */
        int dim() const { return u_basis_.size(); }

        mpfr_prec_t get_prec() const {
            return ul(0).section_edge(0).get_prec();
        }

        /// Return statistics
        irlib::statistics::statistics_type get_statistics() const {
            return statistics_;
        }

        std::string get_statistics_str() const {
            return statistics_ == statistics::FERMIONIC ? "F" : "B" ;
        }

        int get_prec_int() const {
            return static_cast<int>(ul(0).section_edge(0).get_prec());
        }

#ifndef SWIG //DO NOT EXPOSE TO PYTHON

        /**
         * Compute transformation matrix to Matsubara freq.
         * The computation may take some time. You may store the result somewhere and do not call this routine frequenctly.
         * @param n_vec  This vector must contain indices of Matsubara freqencies
         * @param Tnl    Results
         */
        void compute_Tnl(
                const std::vector<long> &n_vec,
                Eigen::Tensor<std::complex<double>, 2> &Tnl
        ) const {
            auto bak = save_default_prec();

            auto trans_to_non_negative = [&](long n) {
                if (n >= 0) {
                    return n;
                } else {
                    if (statistics_ == irlib::statistics::FERMIONIC) {
                        return -n - 1;
                    } else {
                        return -n;
                    }
                }
            };

            auto nl = dim();

            std::set<long> none_negative_n;
            for (const auto &n : n_vec) {
                none_negative_n.insert(trans_to_non_negative(n));
            }

            Eigen::Tensor<std::complex<double>, 2> Tnl_tmp;
            compute_transformation_matrix_to_matsubara<mpreal>(
                    std::vector<long>(none_negative_n.begin(), none_negative_n.end()),
                    statistics_, u_basis_, Tnl_tmp
            );

            Tnl = Eigen::Tensor<std::complex<double>, 2>(n_vec.size(), nl);
            for (int i = 0; i < n_vec.size(); ++i) {
                auto index_data = std::distance(
                        none_negative_n.begin(),
                        none_negative_n.find(trans_to_non_negative(n_vec[i]))
                );
                if (n_vec[i] >= 0) {
                    for (int l = 0; l < nl; ++l) {
                        Tnl(i, l) = Tnl_tmp(index_data, l);
                    }
                } else {
                    for (int l = 0; l < nl; ++l) {
                        Tnl(i, l) = std::conj(Tnl_tmp(index_data, l));
                    }
                }
            }

            restore_default_prec(bak);
        }

#endif

        /**
         * Compute transformation matrix to Matsubara freq.
         * The computation may take some time. You may store the result somewhere and do not call this routine frequenctly.
         * @param n_vec  This vector must contain indices of Matsubara freqencies
         * @return Results
         */
        Eigen::Tensor<std::complex<double>, 2>
        compute_Tnl(const std::vector<long> &n_vec) const {
            std::cout << "Warning: compute_Tnl() is not well tested!" << std::endl;
            Eigen::Tensor<std::complex<double>, 2> Tnl;
            compute_Tnl(n_vec, Tnl);
            return Tnl;
        }

        std::complex<double> compute_Tnl_safe(long n, int l) {
            auto bak = save_default_prec();

            auto o = (statistics_ == irlib::statistics::FERMIONIC ? 2*n+1 : 2*n);
            if (o >= 0) {
                auto r = to_dcomplex(
                        compute_Tnl_impl(u_basis_[l], l%2==0, statistics_, mpfr::const_pi() * 0.5 * o,
                                        mpfr::digits2bits(get_prec()),
                                        mpfr::digits2bits(get_prec()))
                );
                restore_default_prec(bak);
                return r;
            } else {
                auto r = to_dcomplex(
                        compute_Tnl_impl(u_basis_[l], l%2==0, statistics_, -mpfr::const_pi() * 0.5 * o,
                                        mpfr::digits2bits(get_prec()),
                                        mpfr::digits2bits(get_prec()))
                );
                restore_default_prec(bak);
                return std::conj(r);
            }
        }

    };

    inline basis compute_basis(statistics::statistics_type s,
                        double Lambda,
                        int max_dim = 1000,
                        double cutoff = 1e-8,
                        const std::string& fp_mode="mp",
                        double r_tol = 1e-8,
                        long prec = 64,
                        int n_local_poly = 10,
                        int num_nodes_gauss_legendre = 24,
                        bool verbose = true
        ) throw(std::runtime_error) {
        std::vector<mpfr::mpreal> sv;
        std::vector<piecewise_polynomial<mpfr::mpreal, mpfr::mpreal>> u_basis;
        std::vector<piecewise_polynomial<mpfr::mpreal, mpfr::mpreal>> v_basis;
        std::vector<std::vector<mpfr::mpreal>> u_basis_coeff_l;
        std::vector<std::vector<mpfr::mpreal>> v_basis_coeff_l;

        // Increase default precision if needed
        auto min_prec = std::max(
                mpfr::digits2bits(std::log10(1/cutoff)+2*std::log10(1/r_tol)),
                long(64)//At least 19 digits
        );
        //min_prec = std::max(min_prec, mpfr::digits2bits(std::log10(1/r_tol))+10);
        min_prec = std::max(min_prec, prec);
        if (min_prec > mpfr::mpreal::get_default_prec()) {
            mpfr::mpreal::set_default_prec(min_prec);
        }
        if (verbose) {
            std::cout << "Using default precision = " << min_prec << " bits." << std::endl;
        }

        if (fp_mode == "mp") {
            if (s == statistics::FERMIONIC) {
                std::tie(sv, u_basis, v_basis, u_basis_coeff_l, v_basis_coeff_l) = generate_ir_basis_functions<mpfr::mpreal>(
                        fermionic_kernel<mpfr::mpreal>(Lambda), max_dim, cutoff, verbose, r_tol, n_local_poly, num_nodes_gauss_legendre);
            } else if (s == statistics::BOSONIC) {
                std::tie(sv, u_basis, v_basis, u_basis_coeff_l, v_basis_coeff_l) = generate_ir_basis_functions<mpfr::mpreal>(
                        bosonic_kernel<mpfr::mpreal>(Lambda), max_dim, cutoff, verbose, r_tol, n_local_poly, num_nodes_gauss_legendre);
            }
        } else {
            throw std::runtime_error("Unknown fp_mode " + fp_mode + ". Only 'mp' is supported.");
        }

        return basis(s, Lambda, sv, u_basis, v_basis, u_basis_coeff_l, v_basis_coeff_l);
    }

    inline void savetxt(const std::string& fname, const basis& b) throw(std::runtime_error) {
        std::ofstream ofs(fname);

        int version = 2;
        ofs << version << std::endl;
        ofs << b.get_statistics() << std::endl;
        ofs << b.Lambda() << std::endl;
        ofs << b.dim() << std::endl;

        ofs << b.sl_mp(0).get_prec() << std::endl;
        for (int l=0; l<b.dim(); ++l) {
            auto sl = b.sl_mp(l);
            ofs << std::setprecision(mpfr::bits2digits(sl.get_prec())) << sl << std::endl;
        }
        for (int l=0; l<b.dim(); ++l) {
            ofs << b.ul(l);
        }
        for (int l=0; l<b.dim(); ++l) {
            ofs << b.vl(l);
        }
        for (int l=0; l<b.dim(); ++l) {
            auto coeffs = b.coeff_ulx_leg(l);
            for (int c=0; c<coeffs.size(); ++c) {
                ofs << coeffs[c] << std::endl;
            }
        }
        for (int l=0; l<b.dim(); ++l) {
            auto coeffs = b.coeff_vly_leg(l);
            for (int c=0; c<coeffs.size(); ++c) {
                ofs << coeffs[c] << std::endl;
            }
        }
    }

    //inline void dump_coeff(std::ofstream& ofs, const std::vector<mpreal>& coeff) {
      //auto num_sections = b.ul(l).num_sections();
      //
    //}

    inline basis loadtxt(const std::string& fname) throw(std::runtime_error) {
        std::ifstream ifs(fname);

        if (!ifs.is_open()) {
            throw std::runtime_error(fname + " cannot be opened!");
        }

        statistics::statistics_type s;
        double Lambda;
        int dim;

        int version;
        ifs >> version;

        if (version == 2) {
            {
                int itmp;
                ifs >> itmp;
                s = static_cast<statistics::statistics_type>(itmp);
            }
            ifs >> Lambda;
            ifs >> dim;

            mpfr_prec_t prec;
            ifs >> prec;
            std::vector<mpfr::mpreal> sv(dim);
            for (int l=0; l<dim; ++l) {
                sv[l].set_prec(prec);
                ifs >> sv[l];
            }

            std::vector<piecewise_polynomial<mpfr::mpreal, mpfr::mpreal>> u_basis(dim), v_basis(dim);

            for (int l=0; l<dim; ++l) {
                ifs >> u_basis[l];
            }

            for (int l=0; l<dim; ++l) {
                ifs >> v_basis[l];
            }

            std::vector<std::vector<mpreal>> u_basis_coeffs_leg(dim);
            for (int l = 0; l < dim; ++l) {
                u_basis_coeffs_leg[l].resize(u_basis[l].num_sections() * (u_basis[l].order()+1));
                for (int c = 0; c < u_basis_coeffs_leg[l].size(); ++c) {
                    ifs >> u_basis_coeffs_leg[l][c];
                }
            }

            std::vector<std::vector<mpreal>> v_basis_coeffs_leg(dim);
            for (int l = 0; l < dim; ++l) {
                v_basis_coeffs_leg[l].resize(v_basis[l].num_sections() * (v_basis[l].order()+1));
                for (int c = 0; c < v_basis_coeffs_leg[l].size(); ++c) {
                    ifs >> v_basis_coeffs_leg[l][c];
                }
            }

            return basis(s, Lambda, sv, u_basis, v_basis, u_basis_coeffs_leg, v_basis_coeffs_leg);
        } else {
            throw std::runtime_error("Version " + std::to_string(version) + " is not supported!");
        }
    }
}


/*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
|  Phycas: Python software for phylogenetic analysis                          |
|  Copyright (C) 2006 Mark T. Holder, Paul O. Lewis and David L. Swofford     |
|                                                                             |
|  This program is free software; you can redistribute it and/or modify       |
|  it under the terms of the GNU General Public License as published by       |
|  the Free Software Foundation; either version 2 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              |
|  GNU General Public License for more details.                               |
|                                                                             |
|  You should have received a copy of the GNU General Public License along    |
|  with this program; if not, write to the Free Software Foundation, Inc.,    |
|  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.                |
\~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~*/

#if ! defined(LARGET_SIMON_MOVE_HPP)
#define LARGET_SIMON_MOVE_HPP

#include <vector>							// for std::vector
#include <boost/shared_ptr.hpp>				// for boost::shared_ptr
#include <boost/weak_ptr.hpp>				// for boost::weak_ptr
#include "mcmc_updater.hpp"		// for base class MCMCUpdater

namespace phycas
{

class MCMCChainManager;
typedef boost::weak_ptr<MCMCChainManager>	ChainManagerWkPtr;

/*----------------------------------------------------------------------------------------------------------------------
|	Encapsulates the Larget-Simon local move. The LS move is the default move, but an exception must be made in the case
|	of a star tree, which does not have three contiguous edges required for the LS move. For star trees, the proposal
|	changes the length of just one randomly-chosen edge in the tree. An edge chosen at random is set to the value
|	Y = m*exp(`lambda'*(u - 0.5)), where m is the original length and u is a Uniform(0,1) random deviate. Under this
|	proposal scheme, the random variable Y has the following properties:
|>
|	density         f(Y) = 1/(lambda*Y)
|	cdf             F(Y) = 0.5 + (1/lambda) log(Y/m)
|	minimum         m*exp(-lambda/2)
|	maximum         m*exp(lambda/2)
|	mean            (m/lambda)[exp(lambda/2) - exp(-lambda/2)]
|	variance        (m/lambda)^2 [(lambda/2)(exp(lambda) - exp(-lambda)) - (exp(lambda/2) - exp(-lambda/2))^2]
|>
|	With a starting edge length of 1.0, the proposed edge lengths have increasing mean and variance with increasing
|	lambda, but values of lambda in the range 0.1 to 2.0 appear to be reasonable.
|>
|	lambda       mean        s.d.
|	-----------------------------
|	 0.1      1.00042     0.02888
|	 0.5      1.01045     0.14554
|	 1.0      1.04219     0.29840
|	 2.0      1.17520     0.65752
|	10.0     14.84064    29.68297
|>
*/
class LargetSimonMove : public MCMCUpdater
	{
	public:
						LargetSimonMove();
						virtual ~LargetSimonMove();

		unsigned		getWhichCase() const;
		void			setTuningParameter(double x);
		double			getTuningParameter() const;
		bool			topologyChanged() const;
		void			defaultProposeNewState();
		void			starTreeProposeNewState();
        TreeNode *      randomInternalAboveSubroot();
        TreeNode *      randomChild(TreeNode * nd);
        TreeNode *      chooseZ(TreeNode * middle);

		virtual double  sampleWorkingPrior() const;

		virtual bool    computesTopologyPrior() const {return true;}
		// These are virtual functions in the MCMCUpdater base class
		//
		double			recalcWorkingPrior() const;
		bool 			isPriorSteward() const;
		virtual bool	update();
		virtual double	getLnHastingsRatio() const;
		virtual double	getLnJacobian() const;
		virtual void	proposeNewState();
		virtual void	revert();
		virtual void	accept();

	private:

		double			    lambda;						/**< The tuning parameter for this move (the factor used in modifying backbone length) */

		TreeNode *		    ndX;						/**< Node at one end of segment involved in move; used by Revert to undo a move */
		TreeNode *		    ndY;						/**< One of two nodes in the middle of segment involved in move; used by Revert to undo a move */
		TreeNode *		    ndZ;						/**< Node at other end (from ndX) of segment involved in move; used by Revert to undo a move */
		double			    x;						    /**< Original length of ndX's branch; used by Revert to undo a move */
		double			    y;						    /**< Original length of ndX's branch; used by Revert to undo a move */
		double			    z;						    /**< Original length of ndX's branch; used by Revert to undo a move */
		TreeNode *		    swap1;						/**< First of the two nodes involved in an NNI swap; NULL if no swap was performed; used by Revert to undo a move */
		TreeNode *		    swap2;						/**< Second of the two nodes involved in an NNI swap; NULL if no swap was performed; used by Revert to undo a move */

		double			    m;							/**< Original 3-segment length; needed for computing Hastings ratio */
		double			    mstar;						/**< Modified 3-segment length; needed for computing Hastings ratio */

        double              expand_contract_factor;     /**< The factor by which the selected 3-edge segment is expanded or contracted */

		bool			    topol_changed;				/**< If true, last proposal changed topology */
		unsigned		    which_case;					/**< Which of the eight possible cases was tried last */

		std::vector<double> three_edgelens;			/**< workspace declared here to avoid unnecessary allocs/deallocs */

		void			    reset();					/**< Returns variables involved with reversing a proposed move to the state needed for the start of another proposal */

		bool			    star_tree_proposal;			/**< True if last proposed move was on a star tree (only one randomly-chosen edge changed); False if last proposed move was not on a star tree */

		// These are needed for the star tree exception
		double				orig_edge_len;	/**< Length of modified edge saved (in case revert is necessary) */
		TreeNode *			orig_node;		/**< Node owning the modified edge (in case revert is necessary) */
		std::vector<double>	one_edgelen;	/**< workspace declared here to avoid unnecessary allocs/deallocs */
	};

} // namespace phycas

#endif


//  Copyright (c) 2006 Xiaogang Zhang
//  Use, modification and distribution are subject to the
//  Boost Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_MATH_BESSEL_JY_HPP
#define BOOST_MATH_BESSEL_JY_HPP

#include <boost/math/tools/config.hpp>
#include <boost/math/special_functions/gamma.hpp>
#include <boost/math/special_functions/sign.hpp>
#include <boost/math/special_functions/hypot.hpp>
#include <boost/math/special_functions/sin_pi.hpp>
#include <boost/math/special_functions/cos_pi.hpp>
#include <boost/math/special_functions/detail/simple_complex.hpp>
#include <boost/math/special_functions/detail/bessel_jy_asym.hpp>
#include <boost/math/constants/constants.hpp>
#include <boost/math/policies/error_handling.hpp>
#include <boost/mpl/if.hpp>
#include <boost/type_traits/is_floating_point.hpp>
#include <complex>

// Bessel functions of the first and second kind of fractional order

namespace boost { namespace math {

namespace detail {

// Calculate Y(v, x) and Y(v+1, x) by Temme's method, see
// Temme, Journal of Computational Physics, vol 21, 343 (1976)
template <typename T, typename Policy>
int temme_jy(T v, T x, T* Y, T* Y1, const Policy& pol)
{
    T g, h, p, q, f, coef, sum, sum1, tolerance;
    T a, d, e, sigma;
    unsigned long k;

    BOOST_MATH_STD_USING
    using namespace boost::math::tools;
    using namespace boost::math::constants;

    BOOST_ASSERT(fabs(v) <= 0.5f);  // precondition for using this routine

    T gp = boost::math::tgamma1pm1(v, pol);
    T gm = boost::math::tgamma1pm1(-v, pol);
    T spv = boost::math::sin_pi(v, pol);
    T spv2 = boost::math::sin_pi(v/2, pol);
    T xp = pow(x/2, v);

    a = log(x / 2);
    sigma = -a * v;
    d = abs(sigma) < tools::epsilon<T>() ?
        T(1) : sinh(sigma) / sigma;
    e = abs(v) < tools::epsilon<T>() ? v*pi<T>()*pi<T>() / 2
        : 2 * spv2 * spv2 / v;

    T g1 = (v == 0) ? -euler<T>() : (gp - gm) / ((1 + gp) * (1 + gm) * 2 * v);
    T g2 = (2 + gp + gm) / ((1 + gp) * (1 + gm) * 2);
    T vspv = (fabs(v) < tools::epsilon<T>()) ? 1/constants::pi<T>() : v / spv;
    f = (g1 * cosh(sigma) - g2 * a * d) * 2 * vspv;

    p = vspv / (xp * (1 + gm));
    q = vspv * xp / (1 + gp);

    g = f + e * q;
    h = p;
    coef = 1;
    sum = coef * g;
    sum1 = coef * h;

    T v2 = v * v;
    T coef_mult = -x * x / 4;

    // series summation
    tolerance = tools::epsilon<T>();
    for (k = 1; k < policies::get_max_series_iterations<Policy>(); k++)
    {
        f = (k * f + p + q) / (k*k - v2);
        p /= k - v;
        q /= k + v;
        g = f + e * q;
        h = p - k * g;
        coef *= coef_mult / k;
        sum += coef * g;
        sum1 += coef * h;
        if (abs(coef * g) < abs(sum) * tolerance)
        {
           break;
        }
    }
    policies::check_series_iterations("boost::math::bessel_jy<%1%>(%1%,%1%) in temme_jy", k, pol);
    *Y = -sum;
    *Y1 = -2 * sum1 / x;

    return 0;
}

// Evaluate continued fraction fv = J_(v+1) / J_v, see
// Abramowitz and Stegun, Handbook of Mathematical Functions, 1972, 9.1.73
template <typename T, typename Policy>
int CF1_jy(T v, T x, T* fv, int* sign, const Policy& pol)
{
    T C, D, f, a, b, delta, tiny, tolerance;
    unsigned long k;
    int s = 1;

    BOOST_MATH_STD_USING

    // |x| <= |v|, CF1_jy converges rapidly
    // |x| > |v|, CF1_jy needs O(|x|) iterations to converge

    // modified Lentz's method, see
    // Lentz, Applied Optics, vol 15, 668 (1976)
    tolerance = 2 * tools::epsilon<T>();
    tiny = sqrt(tools::min_value<T>());
    C = f = tiny;                           // b0 = 0, replace with tiny
    D = 0.0L;
    for (k = 1; k < policies::get_max_series_iterations<Policy>() * 100; k++)
    {
        a = -1;
        b = 2 * (v + k) / x;
        C = b + a / C;
        D = b + a * D;
        if (C == 0) { C = tiny; }
        if (D == 0) { D = tiny; }
        D = 1 / D;
        delta = C * D;
        f *= delta;
        if (D < 0) { s = -s; }
        if (abs(delta - 1.0L) < tolerance)
        { break; }
    }
    policies::check_series_iterations("boost::math::bessel_jy<%1%>(%1%,%1%) in CF1_jy", k / 100, pol);
    *fv = -f;
    *sign = s;                              // sign of denominator

    return 0;
}

template <class T>
struct complex_trait
{
   typedef typename mpl::if_<is_floating_point<T>,
      std::complex<T>, sc::simple_complex<T> >::type type;
};

// Evaluate continued fraction p + iq = (J' + iY') / (J + iY), see
// Press et al, Numerical Recipes in C, 2nd edition, 1992
template <typename T, typename Policy>
int CF2_jy(T v, T x, T* p, T* q, const Policy& pol)
{
    BOOST_MATH_STD_USING

    typedef typename complex_trait<T>::type complex_type;

    complex_type C, D, f, a, b, delta, one(1);
    T tiny, zero(0.0L);
    unsigned long k;

    // |x| >= |v|, CF2_jy converges rapidly
    // |x| -> 0, CF2_jy fails to converge
    BOOST_ASSERT(fabs(x) > 1);

    // modified Lentz's method, complex numbers involved, see
    // Lentz, Applied Optics, vol 15, 668 (1976)
    T tolerance = 2 * tools::epsilon<T>();
    tiny = sqrt(tools::min_value<T>());
    C = f = complex_type(-0.5f/x, 1.0L);
    D = 0;
    for (k = 1; k < policies::get_max_series_iterations<Policy>(); k++)
    {
        a = (k - 0.5f)*(k - 0.5f) - v*v;
        if (k == 1)
        {
            a *= complex_type(T(0), 1/x);
        }
        b = complex_type(2*x, T(2*k));
        C = b + a / C;
        D = b + a * D;
        if (C == zero) { C = tiny; }
        if (D == zero) { D = tiny; }
        D = one / D;
        delta = C * D;
        f *= delta;
        if (abs(delta - one) < tolerance) { break; }
    }
    policies::check_series_iterations("boost::math::bessel_jy<%1%>(%1%,%1%) in CF2_jy", k, pol);
    *p = real(f);
    *q = imag(f);

    return 0;
}

enum
{
   need_j = 1, need_y = 2
};

// Compute J(v, x) and Y(v, x) simultaneously by Steed's method, see
// Barnett et al, Computer Physics Communications, vol 8, 377 (1974)
template <typename T, typename Policy>
int bessel_jy(T v, T x, T* J, T* Y, int kind, const Policy& pol)
{
    BOOST_ASSERT(x >= 0);

    T u, Jv, Ju, Yv, Yv1, Yu, Yu1(0), fv, fu;
    T W, p, q, gamma, current, prev, next;
    bool reflect = false;
    unsigned n, k;
    int s;

    static const char* function = "boost::math::bessel_jy<%1%>(%1%,%1%)";

    BOOST_MATH_STD_USING
    using namespace boost::math::tools;
    using namespace boost::math::constants;

    if (v < 0)
    {
        reflect = true;
        v = -v;                             // v is non-negative from here
        kind = need_j|need_y;               // need both for reflection formula
    }
    n = real_cast<unsigned>(v + 0.5L);
    u = v - n;                              // -1/2 <= u < 1/2

    if (x == 0)
    {
       *J = *Y = policies::raise_overflow_error<T>(
          function, 0, pol);
       return 1;
    }

    // x is positive until reflection
    W = T(2) / (x * pi<T>());               // Wronskian
    if (x <= 2)                           // x in (0, 2]
    {
        if(temme_jy(u, x, &Yu, &Yu1, pol))             // Temme series
        {
           // domain error:
           *J = *Y = Yu;
           return 1;
        }
        prev = Yu;
        current = Yu1;
        for (k = 1; k <= n; k++)            // forward recurrence for Y
        {
            next = 2 * (u + k) * current / x - prev;
            prev = current;
            current = next;
        }
        Yv = prev;
        Yv1 = current;
        if(kind&need_j)
        {
          CF1_jy(v, x, &fv, &s, pol);                 // continued fraction CF1_jy
          Jv = W / (Yv * fv - Yv1);           // Wronskian relation
        }
        else
           Jv = std::numeric_limits<T>::quiet_NaN(); // any value will do, we're not using it.
    }
    else                                    // x in (2, \infty)
    {
        // Get Y(u, x):
        // define tag type that will dispatch to right limits:
        typedef typename bessel_asymptotic_tag<T, Policy>::type tag_type;

        T lim;
        switch(kind)
        {
        case need_j:
           lim = asymptotic_bessel_j_limit<T>(v, tag_type());
           break;
        case need_y:
           lim = asymptotic_bessel_y_limit<T>(tag_type());
           break;
        default:
           lim = (std::max)(
              asymptotic_bessel_j_limit<T>(v, tag_type()),
              asymptotic_bessel_y_limit<T>(tag_type()));
           break;
        }
        if(x > lim)
        {
           if(kind&need_y)
           {
              Yu = asymptotic_bessel_y_large_x_2(u, x);
              Yu1 = asymptotic_bessel_y_large_x_2(u + 1, x);
           }
           else
              Yu = std::numeric_limits<T>::quiet_NaN(); // any value will do, we're not using it.
           if(kind&need_j)
           {
              Jv = asymptotic_bessel_j_large_x_2(v, x);
           }
           else
              Jv = std::numeric_limits<T>::quiet_NaN(); // any value will do, we're not using it.
        }
        else
        {
           CF1_jy(v, x, &fv, &s, pol);
           // tiny initial value to prevent overflow
           T init = sqrt(tools::min_value<T>());
           prev = fv * s * init;
           current = s * init;
           for (k = n; k > 0; k--)             // backward recurrence for J
           {
               next = 2 * (u + k) * current / x - prev;
               prev = current;
               current = next;
           }
           T ratio = (s * init) / current;     // scaling ratio
           // can also call CF1_jy() to get fu, not much difference in precision
           fu = prev / current;
           CF2_jy(u, x, &p, &q, pol);                  // continued fraction CF2_jy
           T t = u / x - fu;                   // t = J'/J
           gamma = (p - t) / q;
           Ju = sign(current) * sqrt(W / (q + gamma * (p - t)));

           Jv = Ju * ratio;                    // normalization

           Yu = gamma * Ju;
           Yu1 = Yu * (u/x - p - q/gamma);
        }
        if(kind&need_y)
        {
           // compute Y:
           prev = Yu;
           current = Yu1;
           for (k = 1; k <= n; k++)            // forward recurrence for Y
           {
               next = 2 * (u + k) * current / x - prev;
               prev = current;
               current = next;
           }
           Yv = prev;
        }
        else
           Yv = std::numeric_limits<T>::quiet_NaN(); // any value will do, we're not using it.
    }

    if (reflect)
    {
        T z = (u + n % 2);
        *J = boost::math::cos_pi(z, pol) * Jv - boost::math::sin_pi(z, pol) * Yv;     // reflection formula
        *Y = boost::math::sin_pi(z, pol) * Jv + boost::math::cos_pi(z, pol) * Yv;
    }
    else
    {
        *J = Jv;
        *Y = Yv;
    }

    return 0;
}

} // namespace detail

}} // namespaces

#endif // BOOST_MATH_BESSEL_JY_HPP


/*
 Copyright (C) 2017 Quaternion Risk Management Ltd
 All rights reserved.

 This file is part of ORE, a free-software/open-source library
 for transparent pricing and risk analysis - http://opensourcerisk.org

 ORE is free software: you can redistribute it and/or modify it
 under the terms of the Modified BSD License.  You should have received a
 copy of the license along with this program.
 The license is also available online at <http://opensourcerisk.org>

 This program is distributed on the basis that it will form a useful
 contribution to risk analytics and model standardisation, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 FITNESS FOR A PARTICULAR PURPOSE. See the license for more details.
*/

/*! \file normalsabrinterpolation.hpp
    \brief normal SABR interpolation interpolation between discrete points
*/

#pragma once

#include <qle/models/normalsabr.hpp>

#include <ql/math/interpolations/xabrinterpolation.hpp>
#include <ql/termstructures/volatility/sabr.hpp>

#include <boost/assign/list_of.hpp>
#include <boost/make_shared.hpp>

namespace QuantExt {
using namespace QuantLib;

namespace detail {

class NormalSABRWrapper {
public:
    NormalSABRWrapper(const Time t, const Real& forward, const std::vector<Real>& params,
                      const std::vector<Real>& addParams)
        : t_(t), forward_(forward), params_(params) {
        // validateSabrParameters(params[0], 0.0, params[1], params[2]);
    }
    const std::vector<Real>& params() const { return params_; }
    Real volatility(const Real x) { return normalSabrVolatility(x, forward_, t_, params_[0], params_[1], params_[2]); }

private:
    const Real t_, &forward_;
    const std::vector<Real> params_;
};

struct NormalSABRSpecs {
    Size dimension() { return 3; }
    void defaultValues(std::vector<Real>& params, std::vector<bool>&, const Real& forward, const Real expiryTime,
                       const std::vector<Real>& addParams) {
        if (params[0] == Null<Real>())
            params[0] = 0.0040;
        if (params[1] == Null<Real>())
            params[1] = std::sqrt(0.4);
        if (params[2] == Null<Real>())
            params[2] = 0.0;
    }
    void guess(Array& values, const std::vector<bool>& paramIsFixed, const Real& forward, const Real expiryTime,
               const std::vector<Real>& r, const std::vector<Real>& addParams) {
        Size j = 0;
        if (!paramIsFixed[0]) {
            values[0] = (0.01 - 2E-6) * r[j++] + 1E-6; // normal vol guess
        }
        if (!paramIsFixed[1])
            values[1] = 1.5 * r[j++] + 1E-6;
        if (!paramIsFixed[2])
            values[2] = (2.0 * r[j++] - 1.0) * (1.0 - 1E-6);
    }
    Real eps1() { return .0000001; }
    Real eps2() { return .9999; }
    Real dilationFactor() { return 0.001; }
    Array inverse(const Array& y, const std::vector<bool>&, const std::vector<Real>&, const Real) {
        Array x(3);
        x[0] = std::tan(y[0] * M_PI / 0.02 - M_PI_2);
        x[1] = std::tan(y[1] * M_PI / 5.00 - M_PI_2);
        x[2] = std::tan((y[2] + 1.0) * M_PI / 2.0 - M_PI_2);
        return x;
    }
    Array direct(const Array& x, const std::vector<bool>&, const std::vector<Real>&, const Real) {
        Array y(3);
        y[0] = 0.02 * (std::atan(x[0]) + M_PI_2) / M_PI;
        y[1] = 5.00 * (std::atan(x[1]) + M_PI_2) / M_PI;
        y[2] = 2.0 * (std::atan(x[2]) + M_PI_2) / M_PI - 1.0;
        return y;
    }
    Real weight(const Real strike, const Real forward, const Real stdDev, const std::vector<Real>& addParams) {
        return bachelierBlackFormulaStdDevDerivative(strike, forward, stdDev, 1.0);
    }
    typedef NormalSABRWrapper type;
    boost::shared_ptr<type> instance(const Time t, const Real& forward, const std::vector<Real>& params,
                                     const std::vector<Real>& addParams) {
        std::vector<Real> updatedParams(params);
        if (!addParams.empty()) {
            updatedParams[0] = normalSabrAlphaFromAtmVol(forward, t, addParams.front(), params[1], params[2]);
        }
        return boost::make_shared<type>(t, forward, updatedParams, addParams);
    }
};
} // namespace detail

//! %SABR smile interpolation between discrete volatility points.
/*! \ingroup interpolations */
class NormalSABRInterpolation : public Interpolation {
public:
    template <class I1, class I2>
    NormalSABRInterpolation(
        const I1& xBegin, // x = strikes
        const I1& xEnd,
        const I2& yBegin, // y = volatilities
        Time t,           // option expiry
        const Real& forward, Real alpha, Real nu, Real rho, bool alphaIsFixed, bool nuIsFixed, bool rhoIsFixed,
        bool vegaWeighted = true, const Size atmStrikeIndex = Null<Size>(), const bool implyAlphaFromAtmVol = false,
        const boost::shared_ptr<EndCriteria>& endCriteria = boost::shared_ptr<EndCriteria>(),
        const boost::shared_ptr<OptimizationMethod>& optMethod = boost::shared_ptr<OptimizationMethod>(),
        const Real errorAccept = 0.0002, const bool useMaxError = false, const Size maxGuesses = 50) {

        QL_REQUIRE(
            !implyAlphaFromAtmVol || atmStrikeIndex != Null<Real>(),
            "NormalSABRInterpolation: imply alpha from atm vol implies a) that the atm strike index must be given");

        std::vector<Real> addParams;
        if (implyAlphaFromAtmVol) {
            addParams.push_back(*std::next(yBegin, atmStrikeIndex));
        }

        impl_ = boost::shared_ptr<Interpolation::Impl>(
            new QuantLib::detail::XABRInterpolationImpl<I1, I2, detail::NormalSABRSpecs>(
                xBegin, xEnd, yBegin, t, forward, boost::assign::list_of(alpha)(nu)(rho),
                boost::assign::list_of(alphaIsFixed)(nuIsFixed)(rhoIsFixed), vegaWeighted, endCriteria, optMethod,
                errorAccept, useMaxError, maxGuesses, addParams));
        coeffs_ = boost::dynamic_pointer_cast<QuantLib::detail::XABRCoeffHolder<detail::NormalSABRSpecs>>(impl_);
    }
    Real expiry() const { return coeffs_->t_; }
    Real forward() const { return coeffs_->forward_; }
    Real alpha() const { return coeffs_->modelInstance_->params()[0]; }
    Real nu() const { return coeffs_->modelInstance_->params()[1]; }
    Real rho() const { return coeffs_->modelInstance_->params()[2]; }
    Real rmsError() const { return coeffs_->error_; }
    Real maxError() const { return coeffs_->maxError_; }
    const std::vector<Real>& interpolationWeights() const { return coeffs_->weights_; }
    EndCriteria::Type endCriteria() { return coeffs_->XABREndCriteria_; }

private:
    boost::shared_ptr<QuantLib::detail::XABRCoeffHolder<detail::NormalSABRSpecs>> coeffs_;
};

//! %SABR interpolation factory and traits
/*! \ingroup interpolations */
class NormalSABR {
public:
    NormalSABR(Time t, Real forward, Real alpha, Real nu, Real rho, bool alphaIsFixed, bool nuIsFixed, bool rhoIsFixed,
               bool vegaWeighted = false,
               const boost::shared_ptr<EndCriteria> endCriteria = boost::shared_ptr<EndCriteria>(),
               const boost::shared_ptr<OptimizationMethod> optMethod = boost::shared_ptr<OptimizationMethod>(),
               const Real errorAccept = 0.0002, const bool useMaxError = false, const Size maxGuesses = 50)
        : t_(t), forward_(forward), alpha_(alpha), nu_(nu), rho_(rho), alphaIsFixed_(alphaIsFixed),
          nuIsFixed_(nuIsFixed), rhoIsFixed_(rhoIsFixed), vegaWeighted_(vegaWeighted), endCriteria_(endCriteria),
          optMethod_(optMethod), errorAccept_(errorAccept), useMaxError_(useMaxError), maxGuesses_(maxGuesses) {}
    template <class I1, class I2> Interpolation interpolate(const I1& xBegin, const I1& xEnd, const I2& yBegin) const {
        return SABRInterpolation(xBegin, xEnd, yBegin, t_, forward_, alpha_, nu_, rho_, alphaIsFixed_, nuIsFixed_,
                                 rhoIsFixed_, vegaWeighted_, endCriteria_, optMethod_, errorAccept_, useMaxError_,
                                 maxGuesses_);
    }
    static const bool global = true;

private:
    Time t_;
    Real forward_;
    Real alpha_, nu_, rho_;
    bool alphaIsFixed_, nuIsFixed_, rhoIsFixed_;
    bool vegaWeighted_;
    const boost::shared_ptr<EndCriteria> endCriteria_;
    const boost::shared_ptr<OptimizationMethod> optMethod_;
    const Real errorAccept_;
    const bool useMaxError_;
    const Size maxGuesses_;
};
} // namespace QuantExt


#include <iostream>
#include <vector>
#include <string>
#include <sstream>
#include <fstream>
#include <numeric>

#include <boost/filesystem.hpp>
#include <glog/logging.h>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <opencv2/opencv.hpp>
#include <opencv2/core/eigen.hpp>

#include "YamlFileIO.h"

Eigen::Matrix4d loadCameraPose(const std::string &file)
{
    Eigen::Matrix4d T = Eigen::Matrix4d::Identity();
    std::ifstream fs(file);
    if (!fs.is_open())
    {
        std::cerr << "Fail to open camera pose file " << file << std::endl;
        return T;
    }

    std::string line;
    int i = 0;
    while (!fs.eof())
    {
        line.clear();
        getline(fs, line);
        if (line.empty())
            continue;

        std::stringstream ss(line);
        std::string dummy_str;
        if (i < 4)
        {
            ss >> T(i, 0) >> dummy_str >> T(i, 1) >> dummy_str >> T(i, 2) >> dummy_str >> T(i, 3);
            i++;
        }
        else
        {
            std::cerr << "Invalid camera pose format!\n";
            break;
        }
    }
    fs.close();

    return T;
}

void evaluateBackpackExt(const std::vector<Eigen::Matrix4d> &v_extrinsics, const std::string &eval_pattern){
    Eigen::Matrix4d T_c0_c1_gt = Eigen::Matrix4d::Identity();
    Eigen::AngleAxisd c0_c1_vec = Eigen::AngleAxisd(-90*M_PI/180.0, Eigen::Vector3d(0,-1,0));
    Eigen::Vector3d t_c0_c1(-0.03543, 0, -0.03543);
    T_c0_c1_gt.block<3, 3>(0, 0) = c0_c1_vec.matrix();
    T_c0_c1_gt.block<3, 1>(0, 3) = t_c0_c1;
    Eigen::Matrix4d T_c0_c2_gt = Eigen::Matrix4d::Identity();
    Eigen::AngleAxisd c0_c2_vec = Eigen::AngleAxisd(-180*M_PI/180.0, Eigen::Vector3d(0,-1,0));
    Eigen::Vector3d t_c0_c2(0, 0, -0.07086);
    T_c0_c2_gt.block<3, 3>(0, 0) = c0_c2_vec.matrix();
    T_c0_c2_gt.block<3, 1>(0, 3) = t_c0_c2;
    Eigen::Matrix4d T_c0_c3_gt = Eigen::Matrix4d::Identity();
    Eigen::AngleAxisd c0_c3_vec = Eigen::AngleAxisd(90*M_PI/180.0, Eigen::Vector3d(0,-1,0));
    Eigen::Vector3d t_c0_c3(0.03543, 0, -0.03543);
    T_c0_c3_gt.block<3, 3>(0, 0) = c0_c3_vec.matrix();
    T_c0_c3_gt.block<3, 1>(0, 3) = t_c0_c3;

    Eigen::Matrix4d T_l0_c0_gt = Eigen::Matrix4d::Identity();
    Eigen::Matrix4d T_l1_c0_gt = Eigen::Matrix4d::Identity();
    Eigen::AngleAxisd l1_c0_vec1 = Eigen::AngleAxisd(-15*M_PI/180.0, Eigen::Vector3d(0,1,0));
    Eigen::AngleAxisd l1_c0_vec2 = Eigen::AngleAxisd(90*M_PI/180.0, Eigen::Vector3d(0,0,1));
    Eigen::Matrix3d l1_c0_vec =  l1_c0_vec1.matrix()*l1_c0_vec2.matrix();
    Eigen::Vector3d t_l1_c0(0.58192, 0.0, -0.1954);
    T_l1_c0_gt.block<3,3>(0, 0) = l1_c0_vec;
    T_l1_c0_gt.block<3, 1>(0, 3) = l1_c0_vec* t_l1_c0;
    std::cout << "T_l1_c0_gt:\n" << T_l1_c0_gt << "\n";

    Eigen::Matrix4d T_l0_l1_gt = Eigen::Matrix4d::Identity();
    Eigen::AngleAxisd l0_l1_vec1 = Eigen::AngleAxisd(-30*M_PI/180.0, Eigen::Vector3d(0,0,1));
    Eigen::AngleAxisd l0_l1_vec2 = Eigen::AngleAxisd(-75*M_PI/180.0, Eigen::Vector3d(0,1,0));
    Eigen::Matrix3d l0_l1_vec =  l0_l1_vec1.matrix()*l0_l1_vec2.matrix();
    Eigen::Vector3d t_l0_l1(-0.2604, 0, -0.34922);
    T_l0_l1_gt.block<3,3>(0, 0) = l0_l1_vec;
    T_l0_l1_gt.block<3, 1>(0, 3) = l0_l1_vec1.matrix() * t_l0_l1;

    Eigen::Matrix4d T_gt = Eigen::Matrix4d::Identity();
    if (eval_pattern == "camera0_to_camera1")
        T_gt = T_c0_c1_gt;
    if (eval_pattern == "camera0_to_camera2")
        T_gt = T_c0_c2_gt;
    if (eval_pattern == "camera0_to_camera3")
        T_gt = T_c0_c3_gt;
    if (eval_pattern == "lidar0_to_camera0")
        T_gt = T_l0_c0_gt;
    if (eval_pattern == "lidar1_to_camera0")
        T_gt = T_l1_c0_gt;
    if (eval_pattern == "lidar0_to_lidar1")
        T_gt = T_l0_l1_gt;

    double avg_roll = 0.0, avg_pitch = 0.0, avg_yaw = 0.0;
    double avg_tx = 0, avg_ty = 0, avg_tz = 0.0;
    for(auto &T_ext : v_extrinsics){
        Eigen::Matrix4d T_error = T_gt.inverse() * T_ext;

        std::cout << "T_ext:\n" << T_ext << "\n";
        Eigen::Matrix3d R = T_ext.block<3, 3>(0, 0);
        Eigen::Vector3d v_angles = R.eulerAngles(2,1,0) * 180.0/M_PI;
        avg_yaw += (std::abs(v_angles[0]) > 90) ? (180 - std::abs(v_angles[0])) : std::abs(v_angles[0]);
        avg_pitch += (std::abs(v_angles[1]) > 90) ? (180 - std::abs(v_angles[1])) : std::abs(v_angles[1]);
        avg_roll += (std::abs(v_angles[2]) > 90) ? (180 - std::abs(v_angles[2])) : std::abs(v_angles[2]);
        LOG(INFO) << eval_pattern << " euler angles(ZYX): " << v_angles.transpose() << "\n";

        Eigen::Vector3d t = T_ext.block<3,1>(0, 3);
        avg_tx += t[0]; avg_ty += t[1]; avg_tz += t[2];
        LOG(INFO) << eval_pattern << " trans vector: " << t.transpose() << "\n";

        Eigen::Matrix3d err_R = T_error.block<3, 3>(0, 0);
        Eigen::Vector3d err_t = T_error.block<3, 1>(0, 3);
        Eigen::AngleAxisd delta_R_vec(err_R);
        LOG(INFO) << " Rotation difference angle "
                    << delta_R_vec.angle() * 180.0 / M_PI << "\n";
        LOG(INFO) << " Translation difference norm is "
                    << err_t.norm() << "\n";
    }
    int size = v_extrinsics.size();
    LOG(INFO) << "Evaluate " << size  << " extrinsic parameters.";
    LOG(INFO) << "Avg euler angles: (" << avg_roll/size << ", " << avg_pitch/size << ", " << avg_yaw/size << ")";
    LOG(INFO) << "Avg trans: (" << avg_tx/size << ", " << avg_ty/size << ", " << avg_tz/size << ")";
}

void evaluateAntmanExt(const std::vector<Eigen::Matrix4d> &v_extrinsics, const std::string &eval_pattern){
    Eigen::Matrix4d T_c0_c1_gt = Eigen::Matrix4d::Identity();
    Eigen::AngleAxisd c0_c1_vec = Eigen::AngleAxisd(-90*M_PI/180.0, Eigen::Vector3d(0,-1,0));
    Eigen::Vector3d t_c0_c1(0.05647, 0, -0.05647);
    T_c0_c1_gt.block<3, 3>(0, 0) = c0_c1_vec.matrix();
    T_c0_c1_gt.block<3, 1>(0, 3) = t_c0_c1;
    Eigen::Matrix4d T_c0_c2_gt = Eigen::Matrix4d::Identity();
    Eigen::AngleAxisd c0_c2_vec = Eigen::AngleAxisd(180*M_PI/180.0, Eigen::Vector3d(0,1,0));
    Eigen::Vector3d t_c0_c2(0, 0, -0.11294);
    T_c0_c2_gt.block<3, 3>(0, 0) = c0_c2_vec.matrix();
    T_c0_c2_gt.block<3, 1>(0, 3) = t_c0_c2;
    Eigen::Matrix4d T_c0_c3_gt = Eigen::Matrix4d::Identity();
    Eigen::AngleAxisd c0_c3_vec = Eigen::AngleAxisd(-90*M_PI/180.0, Eigen::Vector3d(0,1,0));
    Eigen::Vector3d t_c0_c3(-0.05647, 0, -0.05647);
    T_c0_c3_gt.block<3, 3>(0, 0) = c0_c3_vec.matrix();
    T_c0_c3_gt.block<3, 1>(0, 3) = t_c0_c3;

    Eigen::Matrix4d T_c3_l0_gt = Eigen::Matrix4d::Identity();
    Eigen::AngleAxisd c3_l0_vec1 = Eigen::AngleAxisd(-M_PI/4.0, -Eigen::Vector3d::UnitY());
    Eigen::AngleAxisd c3_l0_vec2 = Eigen::AngleAxisd(M_PI/2.0, Eigen::Vector3d::UnitX());
    Eigen::Matrix3d c3_l0_vec = c3_l0_vec2.matrix() * c3_l0_vec1.matrix();
    Eigen::Vector3d t_c3_l0(-0.03887, -0.03887, 0.0593);
    T_c3_l0_gt.block<3,3>(0, 0) = c3_l0_vec;
    T_c3_l0_gt.block<3,1>(0, 3) = t_c3_l0;

    Eigen::Matrix4d T_gt = Eigen::Matrix4d::Identity();
    if (eval_pattern == "camera0_to_camera1")
        T_gt = T_c0_c1_gt;
    if (eval_pattern == "camera0_to_camera2")
        T_gt = T_c0_c2_gt;
    if (eval_pattern == "camera0_to_camera3")
        T_gt = T_c0_c3_gt;
    // it is lidar0-camera0 actually
    if (eval_pattern == "camera0_to_lidar0")
        T_gt = T_c3_l0_gt;

    Eigen::Matrix3d R_gt = T_gt.block<3, 3>(0, 0);
    Eigen::Vector3d v_gt_angles = R_gt.eulerAngles(2,1,0) * 180.0/M_PI;
    LOG(INFO) << "GT Euler angles(ZYX): " << v_gt_angles.transpose() << "\n";

    double avg_roll = 0.0, avg_pitch = 0.0, avg_yaw = 0.0;
    double avg_tx = 0, avg_ty = 0, avg_tz = 0.0;
    for(size_t i = 0; i < v_extrinsics.size(); i++){
        Eigen::Matrix4d T_ext = v_extrinsics[i];

        Eigen::Matrix4d T_error = T_gt.inverse() * T_ext;

        Eigen::Matrix3d R = T_ext.block<3, 3>(0, 0);
        Eigen::Vector3d v_angles = R.eulerAngles(2,1,0) * 180.0/M_PI;
        avg_yaw += (std::abs(v_angles[0]) > 90) ? (180 - std::abs(v_angles[0])) : std::abs(v_angles[0]);
        avg_pitch += (std::abs(v_angles[1]) > 90) ? (180 - std::abs(v_angles[1])) : std::abs(v_angles[1]);
        avg_roll += (std::abs(v_angles[2]) > 90) ? (180 - std::abs(v_angles[2])) : std::abs(v_angles[2]);
        LOG(INFO) << eval_pattern << " euler angles(ZYX): " << v_angles.transpose() << "\n";

        Eigen::Vector3d t = T_ext.block<3,1>(0, 3);
        avg_tx += t[0]; avg_ty += t[1]; avg_tz += t[2];
        LOG(INFO) << eval_pattern << " trans vector: " << t.transpose() << "\n";

        Eigen::Matrix3d err_R = T_error.block<3, 3>(0, 0);
        Eigen::Vector3d err_t = T_error.block<3, 1>(0, 3);
        Eigen::AngleAxisd delta_R_vec(err_R);
        LOG(INFO) << " Rotation difference angle "
                    << delta_R_vec.angle() * 180.0 / M_PI << "\n";
        LOG(INFO) << " Translation difference norm is "
                    << err_t.norm() << "\n";
    }

    int size = v_extrinsics.size();
    LOG(INFO) << "Evaluate " << size  << " extrinsic parameters.";
    LOG(INFO) << "Avg euler angles: (" << avg_roll/size << ", " << avg_pitch/size << ", " << avg_yaw/size << ")";
    LOG(INFO) << "Avg trans: (" << avg_tx/size << ", " << avg_ty/size << ", " << avg_tz/size << ")";
}

void evaluateKaleidoExt(const std::vector<Eigen::Matrix4d> &v_extrinsics, const std::string &eval_pattern){
    Eigen::Matrix4d T_c0_c1_gt = Eigen::Matrix4d::Identity();
    Eigen::AngleAxisd c0_c1_vec = Eigen::AngleAxisd(28*M_PI/180.0, Eigen::Vector3d(1,0,0));
    Eigen::Vector3d t_c0_c1(0, 0.028, 0);
    T_c0_c1_gt.block<3, 3>(0, 0) = c0_c1_vec.matrix();
    T_c0_c1_gt.block<3, 1>(0, 3) = t_c0_c1;

    Eigen::Matrix4d T_c0_l0_gt = Eigen::Matrix4d::Identity();
    Eigen::AngleAxisd c0_l0_vec = Eigen::AngleAxisd(-14*M_PI/180.0, Eigen::Vector3d(1,0,0));
    Eigen::Vector3d t_c0_l0(0, -0.014, 0.020);
    T_c0_l0_gt.block<3, 3>(0, 0) = c0_l0_vec.matrix();
    T_c0_l0_gt.block<3, 1>(0, 3) = t_c0_l0;

    Eigen::Matrix4d T_gt = Eigen::Matrix4d::Identity();
    if (eval_pattern == "camera0_to_camera1")
        T_gt = T_c0_c1_gt;
    if (eval_pattern == "camera0_to_lidar0")
        T_gt = T_c0_l0_gt;

    double avg_roll = 0.0, avg_pitch = 0.0, avg_yaw = 0.0;
    double avg_tx = 0, avg_ty = 0, avg_tz = 0.0;
    for(auto &T_ext : v_extrinsics){
        Eigen::Matrix4d T_error = T_gt.inverse() * T_ext;

        Eigen::Matrix3d R = T_ext.block<3, 3>(0, 0);
        Eigen::Vector3d v_angles = R.eulerAngles(2,1,0) * 180.0/M_PI;
        avg_yaw += (std::abs(v_angles[0]) > 90) ? (180 - std::abs(v_angles[0])) : std::abs(v_angles[0]);
        avg_pitch += (std::abs(v_angles[1]) > 90) ? (180 - std::abs(v_angles[1])) : std::abs(v_angles[1]);
        avg_roll += (std::abs(v_angles[2]) > 90) ? (180 - std::abs(v_angles[2])) : std::abs(v_angles[2]);
        LOG(INFO) << eval_pattern << " euler angles: " << v_angles.transpose() << "\n";

        Eigen::Vector3d t = T_ext.block<3,1>(0, 3);
        avg_tx += t[0]; avg_ty += t[1]; avg_tz += t[2];
        LOG(INFO) << eval_pattern << " trans vector: " << t.transpose() << "\n";
        Eigen::Matrix3d err_R = T_error.block<3, 3>(0, 0);
        Eigen::Vector3d err_t = T_error.block<3, 1>(0, 3);
        Eigen::AngleAxisd delta_R_vec(err_R);
        LOG(INFO) << " Rotation difference angle "
                    << delta_R_vec.angle() * 180.0 / M_PI << "\n";
        LOG(INFO) << " Translation difference norm is "
                    << err_t.norm() << "\n";
    }
    int size = v_extrinsics.size();
    LOG(INFO) << "Evaluate " << size  << " extrinsic parameters.";
    LOG(INFO) << "Avg euler angles: (" << avg_roll/size << ", " << avg_pitch/size << ", " << avg_yaw/size << ")";
    LOG(INFO) << "Avg trans: (" << avg_tx/size << ", " << avg_ty/size << ", " << avg_tz/size << ")";
}

// compare several camera poses or extrinsic params, to evaluate params consistency
int main(int argc, char** argv){

    // This exe is used to analyze the exytrinsic calibration repeatability
    if(argc < 2){
        std::cerr <<"Usage: ./test_eval_extrinsics [input_folder] [file_prefix]\n";
        std::cout << "[input_folder]: contains many extrinsic calibration result.\n";
        std::cout << "[file_prefix]: extrinsic file prefix, different sensor configuration results in different extrinsic file.\n";
        return -1;
    }

    std::string dataset_folder(argv[1]);
    std::string file_prefix(argv[2]);

    std::vector<Eigen::Matrix4d> v_extrinsics;

    bool b_eval_backpack = false, b_eval_antman = false, b_eval_kaleido = false;
    if (dataset_folder.find("backpack") != std::string::npos)
        b_eval_backpack = true;
    if (dataset_folder.find("antman") != std::string::npos)
        b_eval_antman = true;
    if (dataset_folder.find("kaleido") != std::string::npos)
        b_eval_kaleido = true;

    for (const auto & entry : boost::filesystem::directory_iterator(dataset_folder)){
        if(!boost::filesystem::is_directory(entry))
            continue;

        std::string extrin_filepath = entry.path().string() + "/" +file_prefix + ".yml";
        Eigen::Matrix4d T_ext = Eigen::Matrix4d::Identity();
        if (common::loadExtFileOpencv(extrin_filepath, T_ext)){
            LOG(INFO) << "Read extrinsic parameters from " << extrin_filepath;
        }else{
            LOG(ERROR) << "Fail to load " << extrin_filepath;
            continue;
        }
        v_extrinsics.emplace_back(T_ext);
    }

    if (b_eval_backpack){
    evaluateBackpackExt(v_extrinsics, file_prefix);
    }else{
        if (b_eval_antman){
            evaluateAntmanExt(v_extrinsics, file_prefix);
        }else{
            evaluateKaleidoExt(v_extrinsics, file_prefix);
        }
    }
    return 0;
}