#include <bits/types/FILE.h>
#include <fstream>
#include <iostream>
#include <stdio.h>
#include <vector>
#include <string>
#include <tgmath.h>
#include "image_ppm.h"
#include <filesystem>
#include <Eigen/Dense>

using namespace std;
using namespace Eigen;

typedef Matrix<u_char, Dynamic, Dynamic> MatrixImg;
typedef Matrix<double, Dynamic, Dynamic> TempMatrixImg;
typedef Vector<u_char,Dynamic> ImgLine;
typedef Vector<double,Dynamic> TempImgLine;

auto max(auto a, auto b){
    if (a<b) return b;
    else return a;
}

auto min(auto a, auto b){
    if (a>b) return b;
    else return a;
}

vector<double> projectOnEigenSpace(vector<TempImgLine> eigenfaces, TempImgLine imToProj,int K){
    vector<double> res = vector<double>();
    for (int i=0;i<K;i++){
        res.push_back(eigenfaces[i].dot(imToProj));
    }
    return res;
}

ImgLine octToVec(OCTET* im, int nH, int nW){
    ImgLine res(nH*nW);
    for (int i=0; i<nH*nW;i++){
        res(i)=im[i];
    }
    return res;
}

double eigenProjsDistance(vector<double> proj1, vector<double> proj2){

    double res =0.0;
    for (int i=0; i<min(proj1.size(),proj2.size());i++){
        res+=(proj1[i]-proj2[i])*(proj1[i]-proj2[i]);
    }
    return res;
}

//db is fileName ; names is labels ; values : values[person][imN][doubleM]
void processDB(string db, vector<string>& names, vector<vector<vector<double>>>& values){
    ifstream file;
    file.open(db);
    values = vector<vector<vector<double>>>();
    string line ;
    while(getline(file, line)){
        size_t exclam = line.find('!');
        if (exclam != std::string::npos){
            names.push_back(line.substr(exclam+1));
            values.push_back(vector<vector<double>>());
        }
        else{
            if (line.size()>1){
                vector<double> res;
                stringstream ss(line);

                double num;
                while(ss>>num){
                    res.push_back(num);
                }
                values[values.size()-1].push_back(res);
            }
        }
    }
}

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

    if (argc<3){
        cout<<"usage : eigenfaces im.pgm"<<endl;
    }

    vector<TempImgLine> eigenfaces;
    int K =42; int nH,nW;

    for (int i=0; i<K;i++){
        OCTET* im;
        char name[100];

        sprintf(name,"/im%d.pgm",i);
        string eigenFolder = string(string(argv[1])+string(name));

        lire_nb_lignes_colonnes_image_pgm(eigenFolder.c_str(),&nH,&nW);
        allocation_tableau(im,OCTET,nH*nW);
        lire_image_pgm(eigenFolder.c_str(),im,nH*nW);
        TempImgLine eigenFace(nH*nW) ;
        double sum=0.0;
        for (int j=0;j<nH*nW;j++){
            eigenFace(j)=im[j]-127;
            sum+=(double)(im[j]-127)*(double)(im[j]-127);
        }
        for (int j=0;j<nH*nW;j++){
            eigenFace(j)=eigenFace(j)/sqrt(sum);
            }
        eigenfaces.push_back(eigenFace);
        free(im);
    }

    OCTET* img;
    lire_nb_lignes_colonnes_image_pgm(argv[2],&nH,&nW);
    TempImgLine imgLine(nH*nW);
    allocation_tableau(img,OCTET, nH*nW);
    lire_image_pgm(argv[2],img,nH*nW);
    double sumIm=0.0;
    for (int i=0;i<nH*nW;i++){sumIm+=(double)(img[i]-127)*(double)(img[i]-127);}
    for (int i=0; i<nH*nW;i++){imgLine(i)=(double)(img[i]-127)/sumIm;}

    //registre[countDirs].push_back(projectOnEigenSpace(eigenfaces,imgLine,K));
    vector<double> imProj = projectOnEigenSpace(eigenfaces,imgLine,K);

    vector<string> labels;
    vector<vector<vector<double>>> values;
    processDB("DBEigen.txt", labels, values);
    double minDist=INFINITY;
    int closestMatch = -1;
    for (int i=0;i<values.size();i++){
        for (int j=0; j<values[i].size();j++){
            double dist = eigenProjsDistance(imProj, values[i][j]);
            //cout<<labels[i]<<" "<<dist<<endl;
            if (dist<minDist){
                //cout<<dist<<" "<<i<<endl;
                minDist=dist;
                closestMatch=i;

            }
        }
    }

     if (closestMatch>=0){cout<<"\n **************** \nCLOSEST MATCH FOUND : "<<labels[closestMatch]<<"\n ****************"<<endl;}
    else {cout<<"No match found !"<<endl;}

}


/**
 * \Description:
 * this node is used to implement the Extended kalman Filter equations, its input are the prediction and the measurement, and its output is final estimation.
 * It subscribes to:
 * - /odom2D: this topic contains the 2D pose (x,y,theta) of the turtlrbot and their covariance.
 * - /beacon distances : this topic contains three distance from the robot to each beacon in order "blue, orange/red, green".
 * It publishes to:
 * - /robot_pose ekf: it publishes the 2D pose (x,y,theta) of the turtlrbot and their corresponding covariance after modifying the predicted pose using the measurement.
 *
 * ALI&SERRANO_ECN_M1_2017
 */

//Cpp
#include <sstream>
#include <stdio.h>
#include <vector>
#include <iostream>
#include <stdlib.h>
#include <math.h>
#include <Eigen/Dense>

using namespace Eigen;
using namespace std;
using Eigen::MatrixXd;

//ROS
#include "ros/ros.h"

// Include here the ".h" files corresponding to the topic type you use.
#include <nav_msgs/Odometry.h>
#include "tf/transform_datatypes.h"

#include "std_msgs/MultiArrayLayout.h"
#include "std_msgs/MultiArrayDimension.h"
#include <std_msgs/Float64MultiArray.h>

#include <geometry_msgs/Point.h>
#include <geometry_msgs/Pose2D.h>
#include <turtle_ekf/Pose2DWithCovariance.h>

//#include <geometry_msgs/PoseWithCovariance.h>
//#include <std_msgs/Float32.h>

// global variables...
turtle_ekf::Pose2DWithCovariance odom2D;
std_msgs::Float64MultiArray dist;
bool is_odom_available= false, is_dist_available=false;

// Callback functions...
void odom_Callback( turtle_ekf::Pose2DWithCovariance odom2D_msg) {
    is_odom_available= true;
    odom2D = odom2D_msg;
}

void msurmnt_Callback ( std_msgs::Float64MultiArray  dist_msg) {
    is_dist_available= true;
    dist = dist_msg;
}

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

    //ROS Initialization
    ros::init(argc, argv, "ekf_node");

    // Define your node handles
    ros::NodeHandle nh_loc("~"), nh_glob;

    // Read the node parameters if any

    // Declare your node's subscriptions and service clients
    ros::Subscriber odom2D_sub = nh_glob.subscribe<turtle_ekf::Pose2DWithCovariance>("/odom2D",1, odom_Callback) ;
    ros::Subscriber msurmnt_sub = nh_glob.subscribe<std_msgs::Float64MultiArray>("/beacon_distances",1, msurmnt_Callback) ;

    // Declare you publishers and service servers
    ros::Publisher  pose_pub = nh_glob.advertise<turtle_ekf::Pose2DWithCovariance>("/robot_pose_ekf",1) ;

    // node initilization
    MatrixXd Q_alpha =  MatrixXd::Zero(3,3);
    MatrixXd Q_gamma =  MatrixXd::Zero(3,3);
    MatrixXd P =  MatrixXd::Zero(3,3);
    MatrixXd C =  MatrixXd::Zero(3,3);
    MatrixXd D =  MatrixXd::Zero(3,3);
    MatrixXd K =  MatrixXd::Zero(3,3);
    MatrixXd I =  MatrixXd::Zero(3,3);
    I(0,0)= 1;   I(1,1)= 1;     I(2,2)= 1;

    VectorXd X(3);
    VectorXd Y(3);
    VectorXd Y_hat(3);

    geometry_msgs::Point B1, B2, B3;
    B1.x=0;  B1.y=0; B3.x=4.40; B3.y=3.20; B2.x=0;  B2.y=5.20;

    // Init
    X(0)=1;   X(1)=1;   X(2)= 0; //M_PI/2;
// init matrices
//    for (int i=0; i++; i<3){
//        for(int j=0; j++; j<3){
//            P(i,j)= 0;
//            C(i,j)= 0;
//            D(i,j)= 0;
//            K(i,j)= 0;
//        }
//    }

    ros::Rate rate(10);

    while (ros::ok()){

        ros::spinOnce();
                if( ! is_dist_available || !is_odom_available ){
                    ROS_INFO("Waiting for odom2D and/or dist") ;
                    rate.sleep() ;
                    continue ;
                }

        // odometry part:
        // Prediction step: predicted_state Xk+1/k "extract pose from odom"
        X(0)= odom2D.pose.x;
        X(1)= odom2D.pose.y;
        X(2)= odom2D.pose.theta;

        // state_trans_uncertainty_noise: Q_alpha "extract covariance matrix from odom"
        Q_alpha(0,0)= odom2D.Covariance[0];
        Q_alpha(1,1)= odom2D.Covariance[4];
        Q_alpha(2,2)= odom2D.Covariance[8];

        // jacobian_matrix: A = I

        // pred_covariance_est: Pk+1/k = A. PK/k. A^T + Q_alpha  "P= A*P*A.transpose() + Q_alpha;" but A=I
        P=P+Q_alpha;

        // measurement part:

        // Actual measurement  Y
        Y(0)= dist.data[0]/100;
        Y(1)= dist.data[1]/100;
        Y(2)= dist.data[2]/100;

        // Expected measurement Y_hat
        Y_hat(0)= sqrt( (B1.x-X(0))*(B1.x-X(0)) + (B1.y-X(1))*(B1.y-X(1)) );
        Y_hat(1)= sqrt( (B2.x-X(0))*(B2.x-X(0)) + (B2.y-X(1))*(B2.y-X(1)) );
        Y_hat(2)= sqrt( (B3.x-X(0))*(B3.x-X(0)) + (B3.y-X(1))*(B3.y-X(1)) );

cout<< "Y:  "<< Y(0) << ", " << Y(1) << ", " <<  Y(2) << endl<< endl;
cout<< "Y_hat:  " << Y_hat(0) << ", " << Y_hat(1) << ", " <<  Y_hat(2) <<  endl ;

        // Measurement jacobian C_matrix
        C(0,0)= 2*(X(0) - B1.x) / Y_hat(0);
        C(0,1)= 2*(X(1) - B1.y) / Y_hat(0);
        C(1,0)= 2*(X(0) - B1.x) / Y_hat(1);
        C(1,1)= 2*(X(1) - B1.y) / Y_hat(1);
        C(2,0)= 2*(X(0) - B1.x) / Y_hat(2);
        C(2,1)= 2*(X(1) - B1.y) / Y_hat(2);

        // measurement covariance Q_gamma
        Q_gamma(0,0)= 10;
        Q_gamma(1,1)= 10;
        Q_gamma(2,2)= 10;

        if (Y(0)> 0.3 && Y(0)< 6)
        Q_gamma(0,0)= (.0000842*Y(0))+0.0158;
        if (Y(1)> 0.3 && Y(1)< 6)
        Q_gamma(1,1)= (.0000842*Y(1))+0.0158;
        if (Y(2)> 0.3 && Y(2)< 6)
        Q_gamma(2,2)= (.0000842*Y(2))+0.0158;

if (Y(0)> 0.3 && Y(0)< 6 && Y(1)> 0.3 && Y(1)< 6 && Y(2)> 0.3 && Y(2)< 6 ){

       // Residual covariance D is Q_(Y-Y_hat)
       D= ( C*P*C.transpose() + Q_gamma );

       // Correction propotional gain  K= P*C.transpose() * (C*P*C.transpose() + Q_gamma )^-1
        K= P*C.transpose() * D.reverse();

        // Estimation Step: estimated_state Xk+1/k+1 = Xk+1/k + Kk (Yk - Yk_hat)
        X = X + K*(Y-Y_hat);

        // Estimated_covariance_est: Pk+1/k+1 = ( I - Kk*Ck ) * Pk+1/k
        P= ( I - K*C ) * P;
}
        std::cout<<  "P" << P << endl << endl;
        std::cout<<  "C" << C << endl << endl;
        std::cout<<  "D" << D << endl << endl;
        std::cout<<  "K" << K << endl << endl;

        // publish the final pose ........................................................................................................
        turtle_ekf::Pose2DWithCovariance pose_ekf;
        pose_ekf.pose.x = X(0);
        pose_ekf.pose.y = X(1);
        pose_ekf.pose.theta =  X(2);

        pose_ekf.Covariance[0] =  P(0,0);
        pose_ekf.Covariance[1] =  0;
        pose_ekf.Covariance[2] =  0;
        pose_ekf.Covariance[3] =  0;
        pose_ekf.Covariance[4] =  P(1,1);
        pose_ekf.Covariance[5] =  0;
        pose_ekf.Covariance[6] =  0;
        pose_ekf.Covariance[7] =  0;
        pose_ekf.Covariance[8] =  P(2,2);

        pose_pub.publish(pose_ekf);

        rate.sleep();
    }

}


// Copyright (c) 2018 Evan S Weinberg
// Thick restarted and deflated Lanczos.
// Finds some number of exterior largest or
// smallest values, NOT in magnitude.
// (Doesn't try solve the interior problem,
// though there is a way to trick it with the
// operator (A-sigma)^2 ...)
// Based on arXiv:1512:08135

#include <iostream>
#include <iomanip>
#include <cmath>
#include <string>
#include <sstream>
#include <complex>
#include <random>
#include <vector>
#include <functional>

// Borrow dense matrix eigenvalue routines.
#include <Eigen/Dense>

#include "blas/generic_vector.h"

#include "../square_laplace.h"

// Operator class
#include "../operator.h"
#include "../poly_operator.h"

// Lanczos
#include "../lanczos.h"
#include "../thick_deflate_lanczos.h"

using namespace std;
using namespace Eigen;

typedef Matrix<double, Dynamic, Dynamic, ColMajor> dMatrix;
typedef Matrix<std::complex<double>, Dynamic, Dynamic, ColMajor> cMatrix;

int main(int argc, char** argv)
{
  complex<double> *rhs_cplx;

  // Set output precision to be long.
  cout << setprecision(10);

  // RNG related things.
  std::mt19937 generator (1337u); // RNG, 1337u is the seed.

  ///////////////////////////////////
  // Properties of Linear operator //
  ///////////////////////////////////
  double inv_variance = 6.0; // inverse of variance for gaussian non-compact U(1) links.

  // Basic information about the lattice.
  int length = 24;
  double m_sq = 0.001;

  // Some start-up.
  int volume = length*length;

  // Create a random compact U(1) link.
  complex<double>* gauge_links = allocate_vector<complex<double>>(2*length*length);
  gaussian_real(gauge_links, 2*length*length, generator, 1.0/inv_variance);
  polar(gauge_links, 2*length*length);

  // Vectors.
  rhs_cplx = allocate_vector<complex<double>>(volume);

  // Zero out the vector.
  zero_vector(rhs_cplx, length*length);

  // Structure which gets passed to the function.
  laplace_gauged_struct lapstr_gauged;
  lapstr_gauged.length = length;
  lapstr_gauged.m_sq = m_sq;
  lapstr_gauged.gauge_links = gauge_links;

  // Uncomment this to get the free field.
  //constant_vector(gauge_links, 1.0, 2*volume);
  //std::cout << "Free case.\n\n";

  std::cout << "Interacting case.\n\n";

  //////////////////////////////////////
  // Create a linear operator object. //
  //////////////////////////////////////

  FunctionWrapper<complex<double>> lap_fcn(square_laplacian_gauged, &lapstr_gauged, volume);

  //////////////////////////////////////
  // Properties for Lanczos algorithm //
  //////////////////////////////////////

  // Fill a struct
  TRCLStruct lanc_props;
  lanc_props.n_ev = 10; // get 10 eigenvalues
  lanc_props.m = 30; // subspace size of 20
  lanc_props.tol = 1e-10; // lock at a tolerance of 1e-8
  lanc_props.max_restarts = 500; // maximum of 10 restarts
  lanc_props.preserved_space = 10; // space preserved after restart,
                                  // set to -1 to default to m/4+1
  lanc_props.deflate = false; // don't deflate locked eigenvalues
  lanc_props.generator = &generator; // passed by reference
  lanc_props.verbose = false;

  // Get the smallest eigenvalues
  ThickRestartComplexLanczos<double,std::less<double>> lanczos(&lap_fcn, lanc_props);

  //////////////
  // Compute! //
  //////////////

  lanczos.compute();

  // Get number of converged eigenvalues
  int n_converged = lanczos.num_converged();

  // Get the eigenvalues
  double* eigenvalues = new double[n_converged];
  lanczos.ritzvalues(eigenvalues);

  // Print the Ritz values
  std::cout << "The " << n_converged << " converged eigenvalues are:\n";
  for (int i = 0; i < n_converged; i++) {
    std::cout << eigenvalues[i] << "\n";
  }

  // Get converged eigenvectors
  complex<double>** eigenvectors = new complex<double>*[n_converged];
  for (int i = 0; i < n_converged; i++) {
    eigenvectors[i] = allocate_vector<complex<double>>(volume);
  }

  // Get the Ritz vectors
  lanczos.ritzvectors(eigenvectors);

  ///////////////////////////////////////////
  // Do the polynomial accelerated version //
  ///////////////////////////////////////////

  std::cout << "\nPolynomial Accleration\n";

  // Get approximate bounds from a small Lanczos
  const int m_mini = 6;
  SimpleComplexLanczos<double> simp_lanczos(&lap_fcn, m_mini, generator);
  simp_lanczos.compute();
  double approx_eigs[m_mini];
  simp_lanczos.ritzvalues((double*)approx_eigs);
  // print the approximate eigenvalues
  std::cout << "The " << m_mini << " Ritz values are:\n";
  for (int i = 0; i < m_mini; i++) {
    std::cout << approx_eigs[i] << "\n";
  }
  std::cout << "\n";
  double approx_min = approx_eigs[0]; // get the approximate min
  double approx_max = approx_eigs[m_mini-1]*1.2; // and overshoot the max
  std::cout << "The linear op window is " << approx_min << " to " << approx_max << "\n\n";

  // Make a linear interpolation of the laplace op.
  LinearMapToUnit<complex<double>,double> linear_lap(&lap_fcn, approx_min, approx_max);

  // And make the 20th order poly accelerated form
  OneOverOnePlusX<complex<double> > poly_accel(&linear_lap, 50);

  // Fill a struct
  TRCLStruct poly_lanc_props;
  poly_lanc_props.n_ev = 10; // get 10 eigenvalues
  poly_lanc_props.m = 30; // subspace size of 20
  poly_lanc_props.tol = 1e-10; // lock at a tolerance of 1e-8
  poly_lanc_props.max_restarts = 20; // maximum of 10 restarts
  poly_lanc_props.preserved_space = 10; // space preserved after restart,
                                  // set to -1 to default to m/4+1
  poly_lanc_props.deflate = true; // deflate locked eigenvalues
  poly_lanc_props.generator = &generator; // passed by reference
  poly_lanc_props.verbose = false;

  // Get the largest eigenvalues, since we're poly acceling
  ThickRestartComplexLanczos<double,std::greater<double>> poly_lanczos(&poly_accel, poly_lanc_props);
  poly_lanczos.compute();
  int n_poly_converged = poly_lanczos.num_converged();
  std::cout << n_poly_converged << "\n";

  // Get the poly eigenvalues
  double* poly_eigenvalues = new double[n_poly_converged];
  poly_lanczos.ritzvalues(poly_eigenvalues);

  // Print the Ritz values
  std::cout << "The " << n_poly_converged << " converged eigenvalues are:\n";
  for (int i = 0; i < n_poly_converged; i++) {
    std::cout << poly_eigenvalues[i] << "\n";
  }

  // Get the poly eigenvectors
  complex<double>** poly_eigenvectors = new complex<double>*[n_poly_converged];
  for (int i = 0; i < n_poly_converged; i++) {
    poly_eigenvectors[i] = allocate_vector<complex<double>>(volume);
  }

  // Get the Ritz vectors
  poly_lanczos.ritzvectors(poly_eigenvectors);

  // Get eigenvalues of original system
  complex<double>* intermediate = allocate_vector<complex<double>>(volume);
  std::cout << "\nThe " << n_poly_converged << " reconstructed eigenvalues are:\n";
  for (int i = 0; i < n_poly_converged; i++) {
    zero_vector(intermediate, volume);
    lap_fcn(intermediate, poly_eigenvectors[i]);
    double rec_eval = re_dot(poly_eigenvectors[i], intermediate, volume);
    std::cout << rec_eval << "\n";
  }

  //////////////////////////
  // Get all eigenvalues! //
  //////////////////////////

  if (volume <= 256) {

    // Let's get the eigenvalues of the full operator!

    // Allocate a sufficiently gigantic matrix.
    cMatrix mat_cplx = cMatrix::Zero(volume, volume);

    // Form matrix elements. This is where it's important that
    // dMatrix and cMatrix are column major.
    // I should probably make this safer by using a "Map".
    for (int i = 0; i < volume; i++)
    {
      // Set a point on the rhs for a matrix element.
      zero_vector(rhs_cplx, volume);
      rhs_cplx[i] = 1.0;

      // Where we put the result of the matrix element.
      complex<double>* mptr = &(mat_cplx(i*volume));

      square_laplacian_gauged(mptr, rhs_cplx, &lapstr_gauged);
    }

    // Get the eigenvalues.
    SelfAdjointEigenSolver<cMatrix> eigsolve_cplx(volume);
    eigsolve_cplx.compute(mat_cplx);

    std::cout << "The eigenvalues are:\n" << eigsolve_cplx.eigenvalues() << "\n";

    ////////////////////////////////
    // COMPARE LOWEST EIGENVECTOR //
    ////////////////////////////////

    std::cout << "Compare results:\n\n";
    std::cout << "Lanczos Exact Ratio\n";
    for (int i = 0; i < volume; i++) {
      std::cout << eigenvectors[0][i] << " " << eigsolve_cplx.eigenvectors().col(0)(i)
                << " " << eigenvectors[0][i]/eigsolve_cplx.eigenvectors().col(0)(i) << "\n";
    }
  }

  //////////////
  // CLEAN UP //
  //////////////

  delete[] eigenvalues;
  delete[] poly_eigenvalues;
  for (int i = 0; i < n_converged; i++) {
    deallocate_vector(&eigenvectors[i]);
  }
  delete[] eigenvectors;
  for (int i = 0; i < n_poly_converged; i++) {
    deallocate_vector(&poly_eigenvectors[i]);
  }
  delete[] poly_eigenvectors;
  deallocate_vector(&intermediate);

  deallocate_vector(&rhs_cplx);
  deallocate_vector(&gauge_links);
  return 0;
}


#include "AnalysisGraph.hpp"
#include <range/v3/all.hpp>
#include <unsupported/Eigen/MatrixFunctions>
#include <boost/range/adaptors.hpp>
#include <tqdm.hpp>

using namespace std;
using Eigen::VectorXd, Eigen::MatrixXd;
using fmt::print, fmt::format;
namespace rs = ranges;
using boost::adaptors::transformed;

using fmt::print;
/*
 ============================================================================
 Private: Prediction
 ============================================================================
*/

void AnalysisGraph::generate_latent_state_sequences(
    double initial_prediction_step) {

  // Allocate memory for prediction_latent_state_sequences
  this->predicted_latent_state_sequences.clear();
  this->predicted_latent_state_sequences = vector<vector<VectorXd>>(
      this->res,
      vector<VectorXd>(this->pred_timesteps, VectorXd(this->num_vertices() * 2)));

  cout << "\nPredicting for " << this->pred_timesteps << " time steps..." << endl;
  for (int samp : tq::trange(this->res)) {
      // The sampled transition matrices would be either of matrix exponential
      // (continuous) version or discretized version depending on whether the
      // matrix exponential (continuous) version or the discretized transition
      // matrix version had been used to train the model. This allows us to
      // make the prediction routine common for both the versions, except for
      // the exponentiation of the matrix exponential (continuous) transition
      // matrices.
      MatrixXd A;

      this->generate_head_node_latent_sequences(
          samp, initial_prediction_step + this->pred_timesteps);

      if (this->continuous) {
//          // Here A = Ac = this->transition_matrix_collection[samp] (continuous)
//
//          // Evolving the system till the initial_prediction_step
//          A = (this->transition_matrix_collection[samp] *
//                   initial_prediction_step).exp();
//
//          this->predicted_latent_state_sequences[samp][0] =
//                               A * this->initial_latent_state_collection[samp];
//
//          // After jumping to time step ips - 1, we take one step of length Δt
//          // at a time.
//          // So compute the transition matrix for a single step.
//          // Computing the matrix exponential for a Δt time step.
//          // By default we are using Δt = 1
//          // A = e^{Ac * Δt)
//          A = (this->transition_matrix_collection[samp] * this->delta_t).exp();

          /////////////////
          A = this->transition_matrix_collection[samp].exp();

          // Evolving the system till the initial_prediction_step
          this->current_latent_state = this->initial_latent_state_collection[samp];

          for (int ts = 0; ts < initial_prediction_step; ts++) {
            this->update_latent_state_with_generated_derivatives(ts, ts + 1);

            // Set derivatives for frozen nodes
            for (const auto & [ v, deriv_func ] : this->external_concepts) {
              const Indicator& ind = this->graph[v].indicators[0];
              this->current_latent_state[2 * v + 1] = deriv_func(ts, ind.mean);
            }

            this->current_latent_state = A * this->current_latent_state;
          }
          //this->update_latent_state_with_generated_derivatives(0, initial_prediction_step);
          //this->current_latent_state = A * this->current_latent_state;

          this->update_latent_state_with_generated_derivatives(
              initial_prediction_step, initial_prediction_step + 1);

          // Set derivatives for frozen nodes
          for (const auto & [ v, deriv_func ] : this->external_concepts) {
            const Indicator& ind = this->graph[v].indicators[0];
            this->current_latent_state[2 * v + 1] =
                deriv_func(initial_prediction_step, ind.mean);
          }

          this->predicted_latent_state_sequences[samp][0] = this->current_latent_state;
          /////////////////
      } else {
          // Here A = Ad = this->transition_matrix_collection[samp] (discrete)
          // This is the discrete transition matrix to take a single step of
          // length Δt
          A = this->transition_matrix_collection[samp];

          // Evolving the system till the initial_prediction_step
          this->current_latent_state = this->initial_latent_state_collection[samp];

          for (int ts = 0; ts < initial_prediction_step; ts++) {
            this->update_latent_state_with_generated_derivatives(ts, ts + 1);

            // Set derivatives for frozen nodes
            for (const auto & [ v, deriv_func ] : this->external_concepts) {
              const Indicator& ind = this->graph[v].indicators[0];
              this->current_latent_state[2 * v + 1] = deriv_func(ts, ind.mean);
            }

            this->current_latent_state = A * this->current_latent_state;
          }

          this->update_latent_state_with_generated_derivatives(
              initial_prediction_step, initial_prediction_step + 1);

          // Set derivatives for frozen nodes
          for (const auto & [ v, deriv_func ] : this->external_concepts) {
            const Indicator& ind = this->graph[v].indicators[0];
            this->current_latent_state[2 * v + 1] =
                                  deriv_func(initial_prediction_step, ind.mean);
          }

          this->predicted_latent_state_sequences[samp][0] = this->current_latent_state;
      }

      // Clear out perpetual constraints residual from previous sample
      this->perpetual_constraints.clear();

      if (this->clamp_at_derivative) {
          // To clamp a latent state value to x_c at prediction step 1 via
          // clamping the derivative, we have to perturb the derivative at
          // prediction step 0, before evolving it to prediction time step 1.
          // So we have to look one time step ahead whether we have to clamp
          // at 1.
          //
          if (delphi::utils::in(this->one_off_constraints, 1)) {
              this->perturb_predicted_latent_state_at(1, samp);
          }
      }

      // Since we used ts = 0 for the initial_prediction_step - 1, this loop is
      // one ahead of the prediction time steps. That is, index 1 in the
      // prediction data structures is the 0th index for the requested
      // prediction sequence. Hence, when we are at time step ts, we should
      // check whether there are any constconstraints
        // time step index is 1
      for (int ts = 1; ts < this->pred_timesteps; ts++) {
          // When continuous: The standard matrix exponential equation is,
          //                        s_{t+Δt} = e^{Ac * Δt } * s_t
          //                  Since vector indices are integral values, and in
          //                  the implementation s is the vector, to index into
          //                  the vector we uses consecutive integers. Thus in
          //                  the implementation, the matrix exponential
          //                  equation becomes,
          //                      s_{t+1} = e^{Ac * Δt } * s_t
          //                  What this equation says is that although vector
          //                  indices advance by 1, the duration between two
          //                  predictions stored in two adjacent vector cells
          //                  need not be 1. The time duration is actually Δt.
          //                  The actual line of code represents,
          //                        s_t = e^{Ac * Δt } * s_{t-1}
          // When discrete  : s_t = Ad * s_{t-1}
          this->current_latent_state = A * this->predicted_latent_state_sequences[samp][ts - 1];
          this->update_latent_state_with_generated_derivatives(
              initial_prediction_step + ts, initial_prediction_step + ts + 1);
          this->predicted_latent_state_sequences[samp][ts] = this->current_latent_state;

          // Set derivatives for frozen nodes
          for (const auto & [ v, deriv_func ] : this->external_concepts) {
            const Indicator& ind = this->graph[v].indicators[0];
            this->predicted_latent_state_sequences[samp][ts][2 * v + 1] =
                            deriv_func(initial_prediction_step + ts, ind.mean);
          }

          if (this->clamp_at_derivative ) {
              if (ts == this->rest_derivative_clamp_ts) {
                  // We have perturbed the derivative at ts - 1. If we do not
                  // take any action, that clamped derivative will be in effect
                  // until another clamping or end of prediction. This is where
                  // we take the necessary actions.

                  if (is_one_off_constraints) {
                      // We should revert the derivative to its original value
                      // at this time step so that clamping does not affect
                      // time step ts + 1.
                      for (auto constraint : this->one_off_constraints.at(ts)) {
                          int node_id = constraint.first;

                          this->predicted_latent_state_sequences[samp][ts]
                              (2 * node_id + 1) =
                              this->initial_latent_state_collection[samp]
                              (2 * node_id + 1);
                  }
                  } else {
                      for (auto [node_id, value]: this->one_off_constraints.at(ts)) {
                          this->predicted_latent_state_sequences[samp][ts]
                                                        (2 * node_id + 1) = 0;
                      }
                  }
              }

              // To clamp a latent state value to x_c at prediction step ts + 1
              // via clamping the derivative, we have to perturb the derivative
              // at prediction step ts, before evolving it to prediction time
              // step ts + 1. So we have to look one time step ahead whether we
              // have to clamp at ts + 1 and clamp the derivative now.
              //
              if (delphi::utils::in(this->one_off_constraints, ts + 1) ||
                      !this->perpetual_constraints.empty()) {
                  this->perturb_predicted_latent_state_at(ts + 1, samp);
              }
          } else {
              // Apply constraints to latent state if any
              // Logic of this condition:
              //    Initially perpetual_constraints = ∅
              //
              //    one_off_constraints.at(ts) = ∅ => Unconstrained ∀ ts
              //
              //    one_off_constraints.at(ts) ‡ ∅
              //        => ∃ some constraints (But we do not know what kind)
              //        => Perturb latent state
              //           Call perturb_predicted_latent_state_at()
              //           one_off_constraints == true
              //                => We are applying One-off constraints
              //           one_off_constraints == false
              //                => We are applying perpetual constraints
              //                => We add constraints to perpetual_constraints
              //                => perpetual_constraints ‡ ∅
              //                => The if condition is true ∀ subsequent time steps
              //                   after ts (the first time step s.t.
              //                   one_off_constraints.at(ts) ‡ ∅
              //                => Constraints are perpetual
              //
              if (delphi::utils::in(this->one_off_constraints, ts) ||
                      !this->perpetual_constraints.empty()) {
                  this->perturb_predicted_latent_state_at(ts, samp);
              }
          }
      }
  }
}

/*
 * Applying constraints (interventions) to latent states
 * Check the data structure definition to get more descriptions
 *
 * Clamping at time step ts for value v
 * Two methods work in tandem to achieve this. Let us label them as follows:
 *      glss  : generate_latent_state_sequences()
 *      ppls@ : perturb_predicted_latent_state_at()
 *                                 ┍━━━━━━━━━━━━━━━━━━━━━━━━┑
 *                                 │           How          ┃
 *                                 ┝━━━━━━━━━━┳━━━━━━━━━━━━━┫
 *                                 │ One-off  ┃  Perpetual  ┃
 * ━━━━━━┳━━━━━━━━━━━━┯━━━━━━━━━━━━┿━━━━━━━━━━┻━━━━━━━━━━━━━┫
 *       ┃            │ Clamp at   │           v - x_{ts-1} ┃
 *       ┃            │ (by ppls@) │   ts-1 to ──────────── ┃
 *       ┃            │            │               Δt       ┃
 *       ┃ Derivative ├┈┈┈┈┈┈┈┈┈┈┈┈┼┈┈┈┈┈┈┈┈┈┈┰┈┈┈┈┈┈┈┈┈┈┈┈┈┨
 * Where ┃            │ Reset at   │ ts to ẋ₀ ┃ ts to 0     ┃
 *       ┃            │ (by glss)  │ from S₀  ┃             ┃
 *       ┣━━━━━━━━━━━━┿━━━━━━━━━━━━┿━━━━━━━━━━╋━━━━━━━━━━━━━┫
 *       ┃ Value      │ Clamp at   │ ts to v  ┃ ∀ t≥ts to v ┃
 *       ┃            │ (by ppls@) │          ┃             ┃
 * ━━━━━━┻━━━━━━━━━━━━┷━━━━━━━━━━━━┷━━━━━━━━━━┻━━━━━━━━━━━━━┛
 */
void AnalysisGraph::perturb_predicted_latent_state_at(int timestep, int sample_number) {
    // Let vertices of the CAG be v = 0, 1, 2, 3, ...
    // Then,
    //    indices 2*v keeps track of the state of each variable v
    //    indices 2*v+1 keeps track of the state of ∂v/∂t

    if (this->clamp_at_derivative) {
        for (auto [node_id, value]: this->one_off_constraints.at(timestep)) {
            // To clamp the latent state value to x_c at time step t via
            // clamping the derivative, we have to clamp the derivative
            // appropriately at time step t-1.
            // Example:
            //      Say we want to clamp the latent state at t=6 to value
            //      x_c (i.e. we want x₆ = x_c. So we have to set the
            //      derivative at t=6-1=5, ẋ₅, as follows:
            //                  x_c - x₅
            //             ẋ₅ = --------- ........... (1)
            //                     Δt
            //             x₆ = x₅ + (ẋ₅ × Δt)
            //                = x_c
            //      Thus clamping ẋ₅ (at t = 5) as described in (1) gives
            //      us the desired clamping at t = 5 + 1 = 6
            double clamped_derivative = (value -
                    this->predicted_latent_state_sequences[sample_number]
                    [timestep - 1](2 * node_id)) / this->delta_t;

            this->predicted_latent_state_sequences[sample_number]
                [timestep - 1](2 * node_id + 1) = clamped_derivative;
        }

        // Clamping the derivative at t-1 changes the value at t.
        // According to our model, derivatives never chance. So if we
        // do not revert it, clamped derivative stays till another
        // clamping or the end of prediction. Since this is a one-off
        // clamping, we have to return the derivative back to its
        // original value at time step t, before we use it to evolve
        // time step t + 1. Thus we remember the time step at which we
        // have to perform this.
        this->rest_derivative_clamp_ts = timestep;

        return;
    }

    if (this->is_one_off_constraints) {

        for (auto [node_id, value]: this->one_off_constraints.at(timestep)) {
            this->predicted_latent_state_sequences[sample_number][timestep](2 * node_id) = value;
        }
    } else { // Perpetual constraints
        if (delphi::utils::in(this->one_off_constraints, timestep)) {
            // Update any previous perpetual constraints
            for (auto [node_id, value]: this->one_off_constraints.at(timestep)) {

                this->perpetual_constraints[node_id] = value;
            }
        }

        // Apply perpetual constraints
        for (auto [node_id, value]: this->perpetual_constraints) {

            this->predicted_latent_state_sequences[sample_number][timestep](2 * node_id) = value;
        }
    }
}

void AnalysisGraph::generate_observed_state_sequences() {
  using rs::to, rs::views::transform;

  // Allocate memory for observed_state_sequences
  this->predicted_observed_state_sequences.clear();
  this->predicted_observed_state_sequences =
      vector<PredictedObservedStateSequence>(
          this->res,
          PredictedObservedStateSequence(this->pred_timesteps,
                                         vector<vector<double>>()));

  for (int samp = 0; samp < this->res; samp++) {
    vector<VectorXd>& sample = this->predicted_latent_state_sequences[samp];

    this->predicted_observed_state_sequences[samp] =
        sample | transform([this](VectorXd latent_state) {
          return this->generate_observed_state(latent_state);
        }) |
        to<vector>();
  }
}

vector<vector<double>>
AnalysisGraph::generate_observed_state(VectorXd latent_state) {
  using rs::to, rs::views::transform;

  int num_verts = this->num_vertices();

  vector<vector<double>> observed_state(num_verts);

  for (int v = 0; v < num_verts; v++) {
    vector<Indicator>& indicators = (*this)[v].indicators;

    observed_state[v] = vector<double>(indicators.size());

    observed_state[v] = indicators | transform([&](Indicator ind) {
                          return ind.mean * latent_state[2 * v];
                        }) |
                        to<vector>();
  }

  return observed_state;
}

FormattedPredictionResult AnalysisGraph::format_prediction_result() {

  // NOTE: To facilitate clamping derivatives, we start prediction one time
  //       step before the requested prediction start time. We are omitting
  //       that additional time step from the results returned to the user.
  this->pred_timesteps--;

  // Access
  // [ sample ][ time_step ][ vertex_name ][ indicator_name ]
  auto result = FormattedPredictionResult(
      this->res,
      vector<unordered_map<string, unordered_map<string, double>>>(
          this->pred_timesteps));

  for (int samp = 0; samp < this->res; samp++) {
    // NOTE: To facilitate clamping derivatives, we start prediction one time
    //       step before the requested prediction start time. We are omitting
    //       that additional time step from the results returned to the user.
    for (int ts = 1; ts <= this->pred_timesteps; ts++) {
      for (auto [vert_name, vert_id] : this->name_to_vertex) {
        for (auto [ind_name, ind_id] : (*this)[vert_id].nameToIndexMap) {
          result[samp][ts - 1][vert_name][ind_name] =
              this->predicted_observed_state_sequences[samp][ts][vert_id]
                                                      [ind_id];
        }
      }
    }
  }

  return result;
}

void AnalysisGraph::run_model(int start_year,
                              int start_month,
                              int end_year,
                              int end_month) {
  if (!this->trained) {
    print("Passed untrained Causal Analysis Graph (CAG) Model. \n",
          "Try calling <CAG>.train_model(...) first!");
    throw "Model not yet trained";
  }

  // Check for sensible prediction time step ranges.
  // NOTE: To facilitate clamping derivatives, we start prediction one time
  //       step before the requested prediction start time. Therefor, the
  //       possible valid earliest prediction time step is one time step after
  //       the training start time.
  if (start_year < this->training_range.first.first ||
      (start_year == this->training_range.first.first &&
       start_month <= this->training_range.first.second)) {
    print("The initial prediction date can't be before the "
         "initial training date. Defaulting initial prediction date "
         "to initial training date + 1 month.");
    start_month = this->training_range.first.second + 1;
    if (start_month == 13) {
        start_year = this->training_range.first.first + 1;
        start_month = 1;
    } else {
        start_year = this->training_range.first.first;
    }
  }

  /*
   *              total_timesteps
   *   ____________________________________________
   *  |                                            |
   *  v                                            v
   * start training                          end prediction
   *  |--------------------------------------------|
   *  :           |--------------------------------|
   *  :         start prediction                   :
   *  ^           ^                                ^
   *  |___________|________________________________|
   *      diff              pred_timesteps
   */
  int total_timesteps =
      this->calculate_num_timesteps(this->training_range.first.first,
                                    this->training_range.first.second,
                                    end_year,
                                    end_month);

  this->pred_timesteps = this->calculate_num_timesteps(
      start_year, start_month, end_year, end_month);

  int pred_init_timestep = total_timesteps - pred_timesteps;

  int year = start_year;
  int month = start_month;

  this->pred_range.clear();
  this->pred_range = vector<string>(this->pred_timesteps);

  for (int t = 0; t < this->pred_timesteps; t++) {
    this->pred_range[t] = to_string(year) + "-" + to_string(month);

    if (month == 12) {
      year++;
      month = 1;
    }
    else {
      month++;
    }
  }

  // NOTE: To facilitate clamping derivatives, we start prediction one time
  //       step before the requested prediction start time. When we are
  //       returning results, we have to remove the predictions at the 0th
  //       index of each predicted observed state sequence.
  //       Adding that additional time step.
  // t     = Requested prediction start time step
  //       = pred_init_timestep
  // t - 1 = Prediction start time step
  //       = pred_init_timestep - 1
  pred_init_timestep--;
  this->pred_timesteps++;

  this->generate_latent_state_sequences(pred_init_timestep);
  this->generate_observed_state_sequences();
}

void AnalysisGraph::add_constraint(int step, string concept_name, string indicator_name,
                                                double indicator_clamp_value) {
    // When constraining latent state derivatives, to reach the
    // prescribed value for an observation at projection step t, we have
    // to  clamp the derivative at projection step t-1 appropriately.
    // Therefor, to constrain at projection step 0, we have to clamp the
    // derivative at projection step -1.
    // To facilitate this, we start projection one time step before the
    // requested projection start time.
    // To correct for that and make internal vector indexing easier,
    // nicer and less error prone, we shift all the constraints by one
    // step up.
    step++;

    // Check whether this concept is in the CAG
    if (!delphi::utils::in(this->name_to_vertex, concept_name)) {
        print("Concept \"{0}\" not in CAG!\n", concept_name);
        return;
    }

    int concept_id = this->name_to_vertex.at(concept_name);
    Node& n = (*this)[concept_id];

    if (n.indicators.size() <= 0) {
        // This concept does not have any indicators attached.
        print("Concept \"{0}\" does not have any indicators attached!\n", concept_name);
        print("\tCannot set constraint\n");
        return;
    }

    // If the indicator name is empty we clamp the first indicator attached to
    // this concept
    int ind_id = 0;

    if (!indicator_name.empty()) {
        if (!delphi::utils::in(this->indicators_in_CAG, indicator_name)) {
            print("Indicator \"{0}\" is not in CAG!\n", indicator_name);
            return;
        }

        if (!delphi::utils::in(n.nameToIndexMap, indicator_name)) {
            print("Indicator \"{0}\" is not attached to {1} in CAG!\n",
                    indicator_name,
                    concept_name);
            return;
        }

        ind_id = n.nameToIndexMap.at(indicator_name);
    }

    Indicator& ind = n.indicators[ind_id];

    // We have to clamp the latent state value corresponding to this
    // indicator such that the probability where the emission Gaussian
    // emitting the requested indicator value is the highest. For a
    // Gaussian emission function, the highest probable value is its
    // mean. So we have to set:
    //      μ = ind_value
    //      latent_clamp_value * scaling_factor = ind_value
    //      latent_clamp_value = ind_value / scaling_factor
    // NOTE: In our code we incorrectly call scaling_factor as
    //       indicator mean. To avoid confusion here, I am using the
    //       correct terminology.
    //       (Gosh, when we have incorrect terminology and we know it
    //       and we have not fixed it, I have to type a lot of
    //       comments)
    double latent_clamp_value = indicator_clamp_value / ind.get_mean();

    if (this->head_nodes.find(concept_id) == this->head_nodes.end()) {
      if (!delphi::utils::in(this->one_off_constraints, step)) {
        this->one_off_constraints[step] = vector<pair<int, double>>();
      }

      this->one_off_constraints[step].push_back(
          make_pair(concept_id, latent_clamp_value));
    } else {
      step += this->pred_start_timestep;
      if (!delphi::utils::in(this->head_node_one_off_constraints, step)) {
        this->head_node_one_off_constraints[step] = vector<pair<int, double>>();
      }

      this->head_node_one_off_constraints[step].push_back(
          make_pair(concept_id, latent_clamp_value));
    }
}

/*
 ============================================================================
 Public: Prediction
 ============================================================================
*/

Prediction AnalysisGraph::generate_prediction(int start_year,
                                              int start_month,
                                              int end_year,
                                              int end_month,
                                              ConstraintSchedule constraints,
                                              bool one_off,
                                              bool clamp_deri) {
  this->is_one_off_constraints = one_off;
  this->clamp_at_derivative = clamp_deri;
  this->one_off_constraints.clear();
  this->head_node_one_off_constraints.clear();

  for (auto [step, const_vec] : constraints) {
      for (auto constraint : const_vec) {
          string concept_name = get<0>(constraint);
          string indicator_name = get<1>(constraint);
          double value = get<2>(constraint);

          this->add_constraint(step, concept_name, indicator_name, value);
      }
  }

  this->run_model(start_year, start_month, end_year, end_month);

  return make_tuple(
      this->training_range, this->pred_range, this->format_prediction_result());
}

void AnalysisGraph::generate_prediction(int pred_start_timestep,
                                        int pred_timesteps,
                                        ConstraintSchedule constraints,
                                        bool one_off,
                                        bool clamp_deri) {
  this->is_one_off_constraints = one_off;
  this->clamp_at_derivative = clamp_deri;
  this->one_off_constraints.clear();
  this->head_node_one_off_constraints.clear();

  this->pred_start_timestep = pred_start_timestep - 1;
  this->pred_timesteps = pred_timesteps + 1;

  for (auto [step, const_vec] : constraints) {
    for (auto constraint : const_vec) {
      string concept_name = get<0>(constraint);
      string indicator_name = get<1>(constraint);
      double value = get<2>(constraint);

      this->add_constraint(step, concept_name, indicator_name, value);
    }
  }

  this->generate_latent_state_sequences(this->pred_start_timestep);
  this->generate_observed_state_sequences();
}

vector<vector<double>> AnalysisGraph::prediction_to_array(string indicator) {
  int vert_id = -1;
  int ind_id = -1;

  auto result =
      vector<vector<double>>(this->res, vector<double>(this->pred_timesteps));

  // Find the vertex id the indicator is attached to and
  // the indicator id of it.
  // TODO: We can make this more efficient by making indicators_in_CAG
  // a map from indicator names to vertices they are attached to.
  // This is just a quick and dirty implementation
  for (auto [v_name, v_id] : this->name_to_vertex) {
    for (auto [i_name, i_id] : (*this)[v_id].nameToIndexMap) {
      if (indicator.compare(i_name) == 0) {
        vert_id = v_id;
        ind_id = i_id;
        goto indicator_found;
      }
    }
  }
  // Program will reach here only if the indicator is not found
  throw IndicatorNotFoundException(format(
      "AnalysisGraph::prediction_to_array - indicator \"{}\" not found!\n",
      indicator));

indicator_found:

  for (int samp = 0; samp < this->res; samp++) {
    for (int ts = 0; ts < this->pred_timesteps; ts++) {
      result[samp][ts] =
          this->predicted_observed_state_sequences[samp][ts][vert_id][ind_id];
    }
  }

  return result;
}


/*
 * Random.cpp
 *
 *  Created on: Feb 22, 2009
 *      Author: Tim Babb
 */

#include <boost/static_assert.hpp>
#include <assert.h>
#include <algorithm>
#include <limits>

#include <geomc/random/Random.h>

#define INT_SIGNBIT       (1 << std::numeric_limits<int>::digits)
#define LONG_SIGNBIT      (1L << std::numeric_limits<long>::digits)
#define LONG_LONG_SIGNBIT (1LL << std::numeric_limits<long long>::digits)

using namespace std;
using namespace geom;

typedef union fpbox {
    uint32_t i;
    float f;
    uint64_t l;
    double d;
} FPBox;

/**********************************
 * Structors                      *
 **********************************/

Random::Random():_bitpool(0),
                 _bitsleft(0) {}

Random::~Random() {}

/**********************************
 * Elementary Specializations     *
 **********************************/

namespace geom {

template <> bool Random::rand<bool>() {
    if (_bitsleft == 0) {
        _bitpool = this->rand32();
        _bitsleft = 32;
    }
    int bit = _bitpool & 1;
    _bitpool = _bitpool >> 1;
    _bitsleft--;
    return bit;
}

template <> unsigned int Random::rand<unsigned int>() {
    unsigned int bits = 0;
    const int chunks = (std::numeric_limits<unsigned int>::digits + 32 - 1) / 32; //ceiling(x/y) = (x + y - 1) / y
#if UINT_MAX > (1LL << 32) - 1
    for (int i = 0; i < chunks; i++) {
        bits = (bits << 32) | rand32();
    }
#else
    bits =(unsigned int)rand32();
#endif
    return bits;
}

template <> int Random::rand<int>() {
    return rand<unsigned int>(); //MSB is now a sign bit; covers the entire range of <int> type. might we want to change this?
}

template <> unsigned long Random::rand<unsigned long>() {
    unsigned long bits = 0;
    const int chunks = (std::numeric_limits<unsigned long>::digits + 32 - 1) / 32; //ceiling(x/y) = (x + y - 1) / y
#if ULONG_MAX > (1LL << 32) - 1
    for (int i = 0; i < chunks; i++) {
        bits = (bits << 32) | rand32();
    }
#else
    bits = (unsigned long)rand32();
#endif
    return bits;
}

template <> long Random::rand<long>() {
    return rand<unsigned long>();
}

template <> unsigned long long Random::rand<unsigned long long>() {
    unsigned long long bits = 0;
    const int chunks = (std::numeric_limits<unsigned long long>::digits + 32 - 1) / 32; //ceiling(x/y) = (x + y - 1) / y
    for (int i = 0; i < chunks; i++) {
        bits = (bits << 32) | rand32();
    }
    return bits;
}

template <> long long Random::rand<long long>() {
    return rand<unsigned long long>();
}

/* From the unpublished paper by Allen B. Downey.
 *
 * The conventional method of dividing a random number
 * by a constant excludes about 93% of the representable
 * floating point values between 0.0 and 1.0. Instead
 * we pick the bits of our number explicitly, but must
 * be careful to do so such that the distribution is
 * uniform. This method is fast, and makes efficient
 * use of our pseudorandom bits.
 *
 * The gist: Pick the mantissa, then pick the exponent
 * by iterative coin-flipping.
 */

//TODO: use numeric_limits instead of hard-coded masks
//see: boost::integer's int_t<bits>

template <> float Random::rand<float>() {
    // without IEEE754, our bit twiddling will not work.
    // generally this will only fail for "exotic" systems:
    BOOST_STATIC_ASSERT(std::numeric_limits<float>::is_iec559);

    uint32_t mant;
    int exp, hi_exp, lo_exp;
    FPBox lo, hi, ans;

    lo.f = 0.0;
    hi.f = 1.0;

    lo_exp = (lo.i >> 23) & 0xFF;
    hi_exp = (hi.i >> 23) & 0xFF;

    //not >= because exp is decremented at the end of the last loop
    for (exp = hi_exp - 1; exp > lo_exp; exp--) {
        if (rand<bool>()) break;
    }

    mant = (rand32() & 0xFFFFFE00) >> 9; //use the high quality bits

    if (mant == 0 && rand<bool>()) exp++; //border values must not be skewed towards 1

    ans.i = (((uint32_t)exp) << 23) | mant; //combine exp and mantissa
    return ans.f;
}

template <> double Random::rand<double>() {
    // without IEEE754, our bit twiddling will not work.
    // generally this will only fail for "exotic" systems:
    BOOST_STATIC_ASSERT(std::numeric_limits<double>::is_iec559);

    int exp, hi_exp, lo_exp;
    uint64_t mant;
    FPBox lo, hi, ans;

    lo.d = 0.0;
    hi.d = 1.0;

    lo_exp = (lo.l >> 52) & 0x7FF;
    hi_exp = (hi.l >> 52) & 0x7FF;

    //not >= because exp is decremented at the end of the last loop
    for (exp = hi_exp - 1; exp > lo_exp; exp--) {
        if (rand<bool>()) break;
    }

    mant = ((uint64_t)rand32() << 32) | rand32();
    mant = (mant & 0xFFFFFFFFFFFFF000LL) >> 12; //use the high quality bits

    if (mant == 0 && rand<bool>()) exp++; //border values must not be skewed towards 1

    ans.l = (((uint64_t)exp) << 52) | mant; //combine exp and mantissa
    return ans.d;
}

/**********************************
 * Positive Range Specializations *
 **********************************/

template <> unsigned int Random::rand<unsigned int>(unsigned int hi) {
    unsigned int bits, val, lo;

    do {
        bits = rand<unsigned int>();
        val = bits % hi;
        lo = bits - val;
        //while this chunk wraps around to zero:
        //i.e. lo sign bit is 1, end-of-chunk sign bit is 0
    } while ( (lo & INT_SIGNBIT) & ~(lo+(hi-1)) );
    return val;
}

template <> int Random::rand<int>(int hi) {
    if (hi < 0) {
        return -rand<unsigned int>(-hi);
    } else {
        return rand<unsigned int>(hi);
    }
}

template <> unsigned long Random::rand<unsigned long>(unsigned long hi) {
    unsigned long bits, val, lo;

    do {
        bits = rand<unsigned long>();
        val = bits % hi;
        lo = bits - val;
        //while this chunk wraps around to zero:
        //i.e. lo sign bit is 1, end-of-chunk sign bit is 0
    } while ( (lo & LONG_SIGNBIT) & ~(lo+(hi-1)) );
    return val;
}

template <> long Random::rand<long>(long hi) {
    if (hi < 0) {
        return -rand<unsigned long>(-hi);
    } else {
        return rand<unsigned long>(hi);
    }
}

template <> unsigned long long Random::rand<unsigned long long>(unsigned long long hi) {
    unsigned long long bits, val, lo;

    do {
        bits = rand<unsigned long long>();
        val = bits % hi;
        lo = bits - val;
        //while this chunk wraps around to zero:
        //i.e. lo sign bit is 1, end-of-chunk sign bit is 0
    } while ( (lo & LONG_LONG_SIGNBIT) & ~(lo+(hi-1)) );
    return val;
}

template <> long long Random::rand<long long>(long long hi) {
    if (hi < 0) {
        return -rand<unsigned long long>(-hi);
    } else {
        return rand<unsigned long long>(hi);
    }
}

template <> float Random::rand<float>(float hi) {
    return hi * rand<float>();
}

template <> double Random::rand<double>(double hi) {
    return hi * rand<double>();
}

/**********************************
 * Full Range Specializations     *
 **********************************/

// todo: some of these will fail for ranges crossing zero where lo - hi is
// greater than INT_MAX.

template <> int Random::rand<int>(int lo, int hi) {
    int lo1 = min(hi,lo);
    return rand<int>(max(hi,lo) - lo1) + lo1;
}

template <> long long Random::rand<long long>(long long lo, long long hi) {
    long long lo1 = min(hi,lo);
    return rand<long long>(max(hi,lo) - lo1) + lo1;
}

template <> float Random::rand<float>(float lo, float hi) {
    return (hi-lo)*rand<float>() + lo;
}

template <> double Random::rand<double>(double lo, double hi) {
    return (hi-lo)*rand<double>() + lo;
}

/// @} //ingroup random

} //end namespace geom


/*
 * sdf2sdf_optimizer2d.cpp
 *
 *  Created on: Mar 05, 2019
 *      Author: Fei Shan
 */

//libraries
#include <boost/python.hpp>
#include <iostream>

// local
#include "../math/transformation.hpp"
#include "sdf_2_sdf_optimizer2d.hpp"
#include "sdf_gradient_wrt_transformation2d.hpp"

namespace ropt = rigid_optimization;
namespace eig = Eigen;

namespace rigid_optimization {

Sdf2SdfOptimizer2d::Sdf2SdfOptimizer2d(
		float rate,
		int maximum_iteration_count,
		tsdf::Parameters2d tsdf_generation_parameters,
		VerbosityParameters verbosity_parameters) :
		rate(rate),
				maximum_iteration_count(maximum_iteration_count),
				tsdf_generator(tsdf_generation_parameters),
				verbosity_parameters(verbosity_parameters)
{
}
;

Sdf2SdfOptimizer2d::~Sdf2SdfOptimizer2d()
{
}
;

Sdf2SdfOptimizer2d::VerbosityParameters::VerbosityParameters(
		bool print_iteration_max_warp_update,
		bool print_iteration_energy) :
		print_iteration_max_warp_update(print_iteration_max_warp_update),
				print_iteration_energy(print_iteration_energy),
				print_per_iteration_info(
						print_iteration_max_warp_update ||
								print_iteration_energy
								)
{
}
;

eig::Matrix3f Sdf2SdfOptimizer2d::optimize(int image_y_coordinate,
		const eig::MatrixXf canonical_field,
		const eig::Matrix<unsigned short, eig::Dynamic, eig::Dynamic>& live_depth_image,
		float eta,
		const eig::Matrix4f& initial_camera_pose) {

	eig::MatrixXf canonical_weight = canonical_field.replicate(1, 1);
	for (int i = 0; i < canonical_weight.rows(); ++i) { // Determine weight based on thickness
		for (int j = 0; j < canonical_weight.cols(); ++j) {
			canonical_weight(i, j) = (canonical_field(i, j) <= -eta) ? 0.0f : 1.0f;
		}
	}

	eig::Vector3f twist = eig::Vector3f(0.0f, 0.0f, 0.0f);

	for (int iteration_count = 0; iteration_count < maximum_iteration_count; ++iteration_count) {
		eig::Matrix3f matrix_A = eig::Matrix3f::Zero(); // from sdf2sdf paper,
														// used for calculating the optimal transformation.
		eig::Vector3f vector_b = eig::Vector3f::Zero(); // from sdf2sdf paper,
														// used to calculating the optimal transformation.

		eig::Matrix<float, 6, 1> twist3d;
		twist3d << twist(0), 0.0f, twist(1), 0.0f, twist(2), 0.0f;
		eig::Matrix4f twist_matrix3d = math::transformation_vector_to_matrix3d(twist3d);
		eig::MatrixXf live_field = this->tsdf_generator.generate(live_depth_image,
				twist_matrix3d,
				image_y_coordinate);
		eig::MatrixXf live_weight = live_field.replicate(1, 1);
		for (int j = 0; j < live_weight.cols(); ++j) { // Determine weight based on thickness
			for (int i = 0; i < live_weight.rows(); ++i) {
				live_weight(i, j) = (live_field(i, j) <= -eta) ? 0.0f : 1.0f;
			}
		}
		eig::Matrix<eig::Vector3f, eig::Dynamic, eig::Dynamic> live_gradient(canonical_field.rows(),
				canonical_field.cols());
		eig::Vector3i offset(tsdf_generator.parameters.array_offset.x,
				0,
				tsdf_generator.parameters.array_offset.y);
		ropt::gradient_wrt_twist(live_field,
				twist,
				offset,
				tsdf_generator.parameters.voxel_size,
				live_gradient);

		for (int j = 0; j < live_field.cols(); ++j) {
			for (int i = 0; i < live_field.rows(); ++i) {
				matrix_A += live_gradient(i, j) * live_gradient(i, j).transpose();
				vector_b += (canonical_field(i, j) - live_field(i, j) + live_gradient(i, j).transpose() * twist)
						* live_gradient(i, j);
			}
		}

		float energy = .5f * (canonical_field.cwiseProduct(canonical_weight) -
				live_field.cwiseProduct(live_weight)).array().pow(2.f).sum();
		eig::Vector3f optimal_twist = matrix_A.inverse() * vector_b;
		twist = twist + this->rate * (optimal_twist - twist);
		if (this->verbosity_parameters.print_per_iteration_info) {
			std::cout << "[ITERATION " << iteration_count << " COMPLETED]\n";
			if (this->verbosity_parameters.print_iteration_max_warp_update) {
				std::cout << " [optimize twist:" << optimal_twist.transpose() << "]\n";
				std::cout << " [twist:" << twist.transpose() << "]\n";
			}
			if (this->verbosity_parameters.print_iteration_energy) {
				std::cout << " [energy: " << energy << "]\n\n" << std::endl;
			}

		}
	}

    eig::Matrix3f twist_matrix2d = math::transformation_vector_to_matrix2d(twist);
	return twist_matrix2d;
}

}


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

#include "patch.h"
#include <stdexcept>

#undef Success
#include <Eigen/Dense>

using namespace nurbsfit;

void FitPatch::initSurface(int dims, int order0, int order1, int cps0, int cps1, Domain roi)
{
  if( cps0 < order0 )
    cps0 = order0;

  if( cps1 < order1 )
    cps1 = order1;

  m_nurbs = ON_NurbsSurface (dims, false, order0, order1, cps0, cps1);

  double ddx = roi.width  / (m_nurbs.KnotCount(0) - 2*(order0-2) - 1);
  double ddy = roi.height / (m_nurbs.KnotCount(1) - 2*(order1-2) - 1);

  m_nurbs.MakeClampedUniformKnotVector (0, ddx);
  m_nurbs.MakeClampedUniformKnotVector (1, ddy);

  for (int i = 0; i < m_nurbs.KnotCount(0); i++)
  {
    double k = m_nurbs.Knot (0, i);
    m_nurbs.SetKnot (0, i, k + roi.x);
  }

  for (int i = 0; i < m_nurbs.KnotCount(1); i++)
  {
    double k = m_nurbs.Knot (1, i);
    m_nurbs.SetKnot (1, i, k + roi.y);
  }

  m_x = Eigen::VectorXd(m_nurbs.CVCount()*dims,1);
  m_x.setZero();
  for (int i = 0; i < m_nurbs.CVCount(0); i++)
  {
    for (int j = 0; j < m_nurbs.CVCount(1); j++)
    {
      m_nurbs.SetCV (i, j, ON_3dPoint(m_nurbs.Knot(0, i),m_nurbs.Knot(1, j),0));
    }
  }
}

void FitPatch::initSolver(const Eigen::VectorXd& param0, const Eigen::VectorXd& param1)
{
  if(param0.rows()!=param1.rows())
    throw std::runtime_error("[FitPatch::initSolver] Error, param vectors must be of same length.");

  if(m_nurbs.CVCount() <= 0)
    throw std::runtime_error("[FitPatch::initSolver] Error, surface not initialized (initSurface).");

  int dim = m_nurbs.Dimension();
  m_A = SparseMatrix( param0.rows()*dim, m_nurbs.CVCount()*dim );

  typedef Eigen::Triplet<double> Tri;
  std::vector<Tri> tripletList;
  tripletList.resize( 3*param0.rows() * m_nurbs.Order(0) * m_nurbs.Order(1) );

  if(!m_quiet)
    printf("[FitPatch::initSolver] entries: %lu  rows: %lu  cols: %d\n",
           param0.rows()* m_nurbs.Order(0) * m_nurbs.Order(1),
           param0.rows(),
           m_nurbs.CVCount());

  double *N0 = new double[m_nurbs.Order (0) * m_nurbs.Order (0)];
  double *N1 = new double[m_nurbs.Order (1) * m_nurbs.Order (1)];
  int E,F;
  int ti(0);

  for(Eigen::MatrixXd::Index row=0; row<param0.rows(); row++)
  {
    E = ON_NurbsSpanIndex (m_nurbs.m_order[0], m_nurbs.m_cv_count[0], m_nurbs.m_knot[0], param0(row), 0, 0);
    F = ON_NurbsSpanIndex (m_nurbs.m_order[1], m_nurbs.m_cv_count[1], m_nurbs.m_knot[1], param1(row), 0, 0);

    ON_EvaluateNurbsBasis (m_nurbs.Order (0), m_nurbs.m_knot[0] + E, param0(row), N0);
    ON_EvaluateNurbsBasis (m_nurbs.Order (1), m_nurbs.m_knot[1] + F, param1(row), N1);

    for (int i = 0; i < m_nurbs.Order (0); i++)
    {
      for (int j = 0; j < m_nurbs.Order (1); j++)
      {
        tripletList[ti] = Tri( dim*row, dim*lrc2gl(E, F, i, j), N0[i] * N1[j] );
        if(dim>1)
          tripletList[ti+1] = Tri( dim*row+1, dim*lrc2gl(E, F, i, j)+1, N0[i] * N1[j] );
        if(dim>2)
          tripletList[ti+2] = Tri( dim*row+2, dim*lrc2gl(E, F, i, j)+2, N0[i] * N1[j] );
        ti+=dim;
      } // j
    } // i

  } // row

  delete [] N1;
  delete [] N0;

  m_A.setFromTriplets(tripletList.begin(), tripletList.end());

  if(m_solver==NULL)
    m_solver = new SPQR();
  m_solver->compute(m_A);
  if(m_solver->info()!=Eigen::Success)
    throw std::runtime_error("[FitPatch::initSolver] decomposition failed.");

  if(!m_quiet)
    printf("[FitPatch::initSolver] decomposition done\n");
}

void FitPatch::solve(const Eigen::VectorXd& values)
{
  m_x = m_solver->solve(values);

  updateSurf();
}

void FitPatch::updateSurf()
{
  int ncp = m_nurbs.CVCount ();

  if(m_x.rows()!=ncp*m_nurbs.Dimension())
    throw std::runtime_error("[FitPatch::updateSurf] Error, number of control points does not match.");

  for (int i = 0; i < ncp; i++)
  {
    ON_3dPoint cp;

    if(m_nurbs.Dimension()==1)
    {
      cp.x = m_x(i);
    }

    if(m_nurbs.Dimension()==2)
    {
      cp.x = m_x(2*i+0);
      cp.y = m_x(2*i+1);
    }

    if(m_nurbs.Dimension()==3)
    {
      cp.x = m_x(3*i+0);
      cp.y = m_x(3*i+1);
      cp.z = m_x(3*i+2);
    }

    m_nurbs.SetCV(gl2gr(i), gl2gc(i), cp);
  }
}

Eigen::VectorXd FitPatch::getError(const Eigen::VectorXd& values)
{
  // compute A*x (i.e. points on surface)
  Eigen::VectorXd Ax(values.rows(),1);
  Ax.setZero();
  for (int k=0; k<m_A.outerSize(); ++k)
  for (SparseMatrix::InnerIterator it(m_A,k); it; ++it)
    Ax(it.row()) += it.value() * m_x(it.col());

  // return (A*x-b)
  return (Ax-values);
}

Eigen::Vector2d FitPatch::reparameterize(const Eigen::VectorXd &value, const Eigen::Vector2d &hint,
                                         int& steps, double &accuracy, int maxSteps, double minAccuracy)
{
  if(m_nurbs.CVCount() <= 0)
    throw std::runtime_error("[FitPatch::reparameterize] Error, surface not initialized (initSurface).");

  int nder = 1; // number of derivatives
  int dims = m_nurbs.Dimension();
  int nvals = dims*(nder+1)*(nder+2)/2;
  double pointAndTangents[nvals];

  Eigen::Vector2d current, delta, b, c;
  Eigen::Matrix2d A, I;
  Eigen::VectorXd p(dims), tu(dims), tv(dims), r(dims);
  I = Eigen::Matrix2d::Identity();

  double minU = m_nurbs.Knot(0,0);
  double minV = m_nurbs.Knot(1,0);
  double maxU = m_nurbs.Knot(0,m_nurbs.KnotCount(0)-1);
  double maxV = m_nurbs.Knot(1,m_nurbs.KnotCount(1)-1);
  double nu = 0.5;
  double lambda = 1.0;
  double accuracy_old(DBL_MAX);

  current = hint;

  for (steps = 0; steps < maxSteps; steps++)
  {

    m_nurbs.Evaluate (current(0), current(1), nder, dims, pointAndTangents);

    if(dims==1)
    {
      p(0) = pointAndTangents[0];
      tu(0) = pointAndTangents[1];
      tv(0) = pointAndTangents[2];
    }

    if(dims==2)
    {
      p(0) = pointAndTangents[0];
      p(1) = pointAndTangents[1];
      tu(0) = pointAndTangents[2];
      tu(1) = pointAndTangents[3];
      tv(0) = pointAndTangents[4];
      tv(1) = pointAndTangents[5];
    }

    if(dims==3)
    {
      p(0) = pointAndTangents[0];
      p(1) = pointAndTangents[1];
      p(2) = pointAndTangents[2];
      tu(0) = pointAndTangents[3];
      tu(1) = pointAndTangents[4];
      tu(2) = pointAndTangents[5];
      tv(0) = pointAndTangents[6];
      tv(1) = pointAndTangents[7];
      tv(2) = pointAndTangents[8];
    }

    r = p - value;

    b(0) = -r.dot (tu);
    b(1) = -r.dot (tv);

    A(0, 0) = tu.dot (tu);
    A(0, 1) = tu.dot (tv);
    A(1, 0) = A (0, 1);
    A(1, 1) = tv.dot (tv);

    delta = A.ldlt().solve(b);

    accuracy = delta.norm();
    if (accuracy < minAccuracy)
      return current;

    // step width control (quite heuristic)
    if(accuracy>accuracy_old*nu)
      lambda *= nu;
    accuracy_old = accuracy;

    // make step
    c = current + lambda * delta;

    // clamp to domain borders
    if (c(0) < minU)
      c(0) = minU;
    else if (c (0) > maxU)
      c(0) = maxU;
    if (c(1) < minV)
      c(1) = minV;
    else if (c(1) > maxV)
      c(1) = maxV;

    // compute real step
    delta = c - current;
    current = c;

    accuracy = delta.norm();
    if (accuracy < minAccuracy)
      return current;
  }

  printf ("[FitPatch::reparameterize] Warning: Method did not converge (%e %e %d)\n",
          accuracy, minAccuracy, maxSteps);
  printf ("[FitPatch::reparameterize]   (0: %f %f) (1: %f %f) %f %f ... %f %f\n",
          minU, maxU, minV, maxV,
          hint (0), hint (1), current (0), current (1));

  return current;
}

//int nvals = nurbs.m_dim*(nder+1)*(nder+2)/2;
//double P[nvals];
//nurbs.Evaluate (u, v, nder, nurbs.m_dim, P);

//// positions
//xx = P[0];    xy = P[1];    xz = P[2];

//// 1st derivatives (for normals)
//xu(0) = P[3];    xu(1) = P[4];    xu(2) = P[5];
//xv(0) = P[6];    xv(1) = P[7];    xv(2) = P[8];

//n = xu.cross(xv);
//n.normalize();
//v.normal = TomGine::vec3(n(0),n(1),n(2));

//// 2nd derivatives (for curvature)
//xuu(0) = P[9];     xuu(1) = P[10];    xuu(2) = P[11];
//xuv(0) = P[12];    xuv(1) = P[13];    xuv(2) = P[14];
//xvv(0) = P[15];    xvv(1) = P[16];    xvv(2) = P[17];

//#include <unsupported/Eigen/NumericalDiff>
//#include <unsupported/Eigen/NonLinearOptimization>

//struct SurfaceFunctor
//{
//  const Eigen::VectorXd& m_residuals;
//  const ON_NurbsSurface& m_nurbs;

//  SurfaceFunctor(const Eigen::VectorXd& residuals, const ON_NurbsSurface& nurbs)
//    : m_residuals(residuals), m_nurbs(nurbs){}

//  int operator()(const Eigen::VectorXd& x, Eigen::VectorXd& residuals) const
//  {

//    return 0;
//  }

//  int dx(const Eigen::VectorXd& x, Eigen::MatrixXd& jacobian)
//  {

//    return 0;
//  }

//  int inputs() const { return m_nurbs.CVCount(); }
//  int values() const { return m_residuals.rows(); }
//};

//Eigen::Vector2d FitPatch::reparameterizeLM(const Eigen::VectorXd &value, const Eigen::Vector2d &hint,
//                                         int& steps, double &accuracy, int maxSteps, double minAccuracy)
//{
//  if(m_nurbs.CVCount() <= 0)
//    throw std::runtime_error("[FitPatch::reparameterize] Error, surface not initialized (initSurface).");

//  int nder = 1; // number of derivatives
//  int dims = m_nurbs.Dimension();
//  int nvals = dims*(nder+1)*(nder+2)/2;
//  double pointAndTangents[nvals];

//  Eigen::Vector2d current, delta, b, c;
//  Eigen::Matrix2d A, I;
//  Eigen::VectorXd p(dims), tu(dims), tv(dims), r(dims);
//  I = Eigen::Matrix2d::Identity();

//  double minU = m_nurbs.Knot(0,0);
//  double minV = m_nurbs.Knot(1,0);
//  double maxU = m_nurbs.Knot(0,m_nurbs.KnotCount(0)-1);
//  double maxV = m_nurbs.Knot(1,m_nurbs.KnotCount(1)-1);

//  current = hint;

//  SurfaceFunctor func;
//  Eigen::LevenbergMarquardt<Eigen::NumericalDiff<my_functor>,double> lm(numDiff);

//  for (steps = 0; steps < maxSteps; steps++)
//  {

//    m_nurbs.Evaluate (current(0), current(1), nder, dims, pointAndTangents);

//    if(dims==1)
//    {
//      p(0) = pointAndTangents[0];
//      tu(0) = pointAndTangents[1];
//      tv(0) = pointAndTangents[2];
//    }

//    if(dims==2)
//    {
//      p(0) = pointAndTangents[0];
//      p(1) = pointAndTangents[1];
//      tu(0) = pointAndTangents[2];
//      tu(1) = pointAndTangents[3];
//      tv(0) = pointAndTangents[4];
//      tv(1) = pointAndTangents[5];
//    }

//    if(dims==3)
//    {
//      p(0) = pointAndTangents[0];
//      p(1) = pointAndTangents[1];
//      p(2) = pointAndTangents[2];
//      tu(0) = pointAndTangents[3];
//      tu(1) = pointAndTangents[4];
//      tu(2) = pointAndTangents[5];
//      tv(0) = pointAndTangents[6];
//      tv(1) = pointAndTangents[7];
//      tv(2) = pointAndTangents[8];
//    }

//    r = p - value;

//    b(0) = -r.dot (tu);
//    b(1) = -r.dot (tv);

//    A(0, 0) = tu.dot (tu);
//    A(0, 1) = tu.dot (tv);
//    A(1, 0) = A (0, 1);
//    A(1, 1) = tv.dot (tv);

//    delta = A.ldlt().solve(b);

//    accuracy = delta.norm();
//    if (accuracy < minAccuracy)
//      return current;

//    // make step
//    c = current + delta;

//    // clamp to domain borders
//    if (c(0) < minU)
//      c(0) = minU;
//    else if (c (0) > maxU)
//      c(0) = maxU;
//    if (c(1) < minV)
//      c(1) = minV;
//    else if (c(1) > maxV)
//      c(1) = maxV;

//    // compute real step
//    delta = c - current;
//    current = c;

//    accuracy = delta.norm();
//    if (accuracy < minAccuracy)
//      return current;
//  }

//  printf ("[FitPatch::reparameterize] Warning: Method did not converge (%e %e %d)\n",
//          accuracy, minAccuracy, maxSteps);
//  printf ("[FitPatch::reparameterize]   (0: %f %f) (1: %f %f) %f %f ... %f %f\n",
//          minU, maxU, minV, maxV,
//          hint (0), hint (1), current (0), current (1));

//  return current;
//}


#include "nnw_hierarchy.h"

#include <google/protobuf/stubs/casts.h>

#include <Eigen/Dense>
#include <stan/math/prim/prob.hpp>

#include "hierarchy_prior.pb.h"
#include "ls_state.pb.h"
#include "marginal_state.pb.h"
#include "matrix.pb.h"
#include "src/utils/distributions.h"
#include "src/utils/eigen_utils.h"
#include "src/utils/proto_utils.h"
#include "src/utils/rng.h"

//! \param prec_ Value to set to prec
void NNWHierarchy::set_prec_and_utilities(const Eigen::MatrixXd &prec_) {
  state.prec = prec_;

  // Update prec utilities
  prec_chol = Eigen::LLT<Eigen::MatrixXd>(prec_).matrixL().transpose();
  Eigen::VectorXd diag = prec_chol.diagonal();
  prec_logdet = 2 * log(diag.array()).sum();
}

void NNWHierarchy::initialize() {
  if (prior == nullptr) {
    throw std::invalid_argument("Hierarchy prior was not provided");
  }
  state.mean = hypers->mean;
  set_prec_and_utilities(hypers->var_scaling *
                         Eigen::MatrixXd::Identity(dim, dim));
  clear_data();
}

void NNWHierarchy::clear_data() {
  data_sum = Eigen::VectorXd::Zero(dim);
  data_sum_squares = Eigen::MatrixXd::Zero(dim, dim);
  card = 0;
  cluster_data_idx = std::set<int>();
}

void NNWHierarchy::update_summary_statistics(const Eigen::VectorXd &datum,
                                             bool add) {
  if (add) {
    data_sum += datum;
    data_sum_squares += datum * datum.transpose();
  } else {
    data_sum -= datum;
    data_sum_squares -= datum * datum.transpose();
  }
}

//! \param data                    Matrix of row-vectorial data points
//! \param mu0, lambda0, tau0, nu0 Original values for hyperparameters
//! \return                        Vector of updated values for hyperparameters
NNWHierarchy::Hyperparams NNWHierarchy::normal_wishart_update() {
  // Initialize relevant objects
  Hyperparams post_params;

  // Compute updated hyperparameters
  post_params.var_scaling = hypers->var_scaling + card;
  post_params.deg_free = hypers->deg_free + 0.5 * card;

  Eigen::VectorXd mubar = data_sum.array() / card;  // sample mean
  post_params.mean = (hypers->var_scaling * hypers->mean + card * mubar) /
                     (hypers->var_scaling + card);
  // Compute tau_n
  Eigen::MatrixXd tau_temp =
      data_sum_squares - card * mubar * mubar.transpose();
  tau_temp += (card * hypers->var_scaling / (card + hypers->var_scaling)) *
              (mubar - hypers->mean) * (mubar - hypers->mean).transpose();
  tau_temp = 0.5 * tau_temp + hypers->scale_inv;
  post_params.scale = stan::math::inverse_spd(tau_temp);
  return post_params;
}

void NNWHierarchy::update_hypers(
    const std::vector<bayesmix::MarginalState::ClusterState> &states) {
  auto &rng = bayesmix::Rng::Instance().get();
  if (prior->has_fixed_values()) {
    return;
  }

  else if (prior->has_normal_mean_prior()) {
    // Get hyperparameters
    Eigen::VectorXd mu00 =
        bayesmix::to_eigen(prior->normal_mean_prior().mean_prior().mean());
    Eigen::MatrixXd sigma00 =
        bayesmix::to_eigen(prior->normal_mean_prior().mean_prior().var());
    double lambda0 = prior->normal_mean_prior().var_scaling();
    // Compute posterior hyperparameters
    Eigen::MatrixXd sigma00inv = stan::math::inverse_spd(sigma00);
    Eigen::MatrixXd prec = Eigen::MatrixXd::Zero(dim, dim);
    Eigen::VectorXd num = Eigen::MatrixXd::Zero(dim, 1);
    for (auto &st : states) {
      Eigen::MatrixXd prec_i = bayesmix::to_eigen(st.multi_ls_state().prec());
      prec += prec_i;
      num += prec_i * bayesmix::to_eigen(st.multi_ls_state().mean());
    }
    prec = hypers->var_scaling * prec + sigma00inv;
    num = hypers->var_scaling * num + sigma00inv * mu00;
    Eigen::VectorXd mu_n = prec.llt().solve(num);
    // Update hyperparameters with posterior sampling
    hypers->mean = stan::math::multi_normal_prec_rng(mu_n, prec, rng);
  }

  else if (prior->has_ngiw_prior()) {
    // Get hyperparameters:
    // for mu0
    Eigen::VectorXd mu00 =
        bayesmix::to_eigen(prior->ngiw_prior().mean_prior().mean());
    Eigen::MatrixXd sigma00 =
        bayesmix::to_eigen(prior->ngiw_prior().mean_prior().var());
    // for lambda0
    double alpha00 = prior->ngiw_prior().var_scaling_prior().shape();
    double beta00 = prior->ngiw_prior().var_scaling_prior().rate();
    // for tau0
    double nu00 = prior->ngiw_prior().scale_prior().deg_free();
    Eigen::MatrixXd tau00 =
        bayesmix::to_eigen(prior->ngiw_prior().scale_prior().scale());
    // Compute posterior hyperparameters
    Eigen::MatrixXd sigma00inv = stan::math::inverse_spd(sigma00);
    Eigen::MatrixXd tau_n = Eigen::MatrixXd::Zero(dim, dim);
    Eigen::VectorXd num = Eigen::MatrixXd::Zero(dim, 1);
    double beta_n = 0.0;
    for (auto &st : states) {
      Eigen::VectorXd mean = bayesmix::to_eigen(st.multi_ls_state().mean());
      Eigen::MatrixXd prec = bayesmix::to_eigen(st.multi_ls_state().prec());
      tau_n += prec;
      num += prec * mean;
      beta_n +=
          (hypers->mean - mean).transpose() * prec * (hypers->mean - mean);
    }
    Eigen::MatrixXd prec_n = hypers->var_scaling * tau_n + sigma00inv;
    tau_n += tau00;
    num = hypers->var_scaling * num + sigma00inv * mu00;
    beta_n = beta00 + 0.5 * beta_n;
    Eigen::MatrixXd sig_n = stan::math::inverse_spd(prec_n);
    Eigen::VectorXd mu_n = sig_n * num;
    double alpha_n = alpha00 + 0.5 * states.size();
    double nu_n = nu00 + states.size() * hypers->deg_free;
    // Update hyperparameters with posterior random Gibbs sampling
    hypers->mean = stan::math::multi_normal_rng(mu_n, sig_n, rng);
    hypers->var_scaling = stan::math::gamma_rng(alpha_n, beta_n, rng);
    hypers->scale = stan::math::inv_wishart_rng(nu_n, tau_n, rng);
    hypers->scale_inv = stan::math::inverse_spd(hypers->scale);
  }

  else {
    throw std::invalid_argument("Unrecognized hierarchy prior");
  }
}

//! \param data Matrix of row-vectorial single data point
//! \return     Log-Likehood vector evaluated in data
double NNWHierarchy::like_lpdf(const Eigen::RowVectorXd &datum) const {
  // Initialize relevant objects
  return bayesmix::multi_normal_prec_lpdf(datum, state.mean, prec_chol,
                                          prec_logdet);
}

//! \param data Matrix of row-vectorial a single data point
//! \return     Marginal distribution vector evaluated in data
double NNWHierarchy::marg_lpdf(const Eigen::RowVectorXd &datum) const {
  // Compute dof and scale of marginal distribution
  double nu_n = 2 * hypers->deg_free - dim + 1;
  Eigen::MatrixXd sigma_n = hypers->scale_inv *
                            (hypers->deg_free - 0.5 * (dim - 1)) *
                            hypers->var_scaling / (hypers->var_scaling + 1);

  // TODO: chec if this is optimized as our bayesmix::multi_normal_prec_lpdf
  return stan::math::multi_student_t_lpdf(datum, nu_n, hypers->mean, sigma_n);
}

void NNWHierarchy::draw() {
  // Generate new state values from their prior centering distribution
  auto &rng = bayesmix::Rng::Instance().get();
  Eigen::MatrixXd tau_new =
      stan::math::wishart_rng(hypers->deg_free, hypers->scale, rng);

  // Update state
  state.mean = stan::math::multi_normal_prec_rng(
      hypers->mean, tau_new * hypers->var_scaling, rng);
  set_prec_and_utilities(tau_new);
}

//! \param data Matrix of row-vectorial data points
void NNWHierarchy::sample_given_data() {
  // Update values
  Hyperparams params = normal_wishart_update();

  // Generate new state values from their prior centering distribution
  auto &rng = bayesmix::Rng::Instance().get();
  Eigen::MatrixXd tau_new =
      stan::math::wishart_rng(params.deg_free, params.scale, rng);
  state.mean = stan::math::multi_normal_prec_rng(
      params.mean, tau_new * params.var_scaling, rng);

  // Update state
  set_prec_and_utilities(tau_new);
}

void NNWHierarchy::sample_given_data(const Eigen::MatrixXd &data) {
  data_sum = Eigen::VectorXd::Zero(data.cols());
  data_sum_squares = Eigen::MatrixXd::Zero(data.cols(), data.cols());

  for (int i = 0; i < data.rows(); i++) {
    data_sum += data.row(i);
    data_sum_squares += data.row(i).transpose() * data.row(i);
  }
  card = data.rows();
  log_card = std::log(card);
  sample_given_data();
}

void NNWHierarchy::set_state_from_proto(
    const google::protobuf::Message &state_) {
  auto &statecast = google::protobuf::internal::down_cast<
      const bayesmix::MarginalState::ClusterState &>(state_);
  state.mean = to_eigen(statecast.multi_ls_state().mean());
  set_prec_and_utilities(to_eigen(statecast.multi_ls_state().prec()));
  set_card(statecast.cardinality());
}

void NNWHierarchy::set_prior(const google::protobuf::Message &prior_) {
  auto &priorcast =
      google::protobuf::internal::down_cast<const bayesmix::NNWPrior &>(
          prior_);
  prior = std::make_shared<bayesmix::NNWPrior>(priorcast);
  hypers = std::make_shared<Hyperparams>();
  if (prior->has_fixed_values()) {
    // Set values
    hypers->mean = bayesmix::to_eigen(prior->fixed_values().mean());
    dim = hypers->mean.size();
    hypers->var_scaling = prior->fixed_values().var_scaling();
    hypers->scale = bayesmix::to_eigen(prior->fixed_values().scale());
    hypers->scale_inv = stan::math::inverse_spd(hypers->scale);
    hypers->deg_free = prior->fixed_values().deg_free();
    // Check validity
    if (hypers->var_scaling <= 0) {
      throw std::invalid_argument("Variance-scaling parameter must be > 0");
    }
    if (dim != hypers->scale.rows()) {
      throw std::invalid_argument(
          "Hyperparameters dimensions are not consistent");
    }
    if (hypers->deg_free <= dim - 1) {
      throw std::invalid_argument("Degrees of freedom parameter is not valid");
    }
  }

  else if (prior->has_normal_mean_prior()) {
    // Get hyperparameters
    Eigen::VectorXd mu00 =
        bayesmix::to_eigen(prior->normal_mean_prior().mean_prior().mean());
    dim = mu00.size();
    Eigen::MatrixXd sigma00 =
        bayesmix::to_eigen(prior->normal_mean_prior().mean_prior().var());
    double lambda0 = prior->normal_mean_prior().var_scaling();
    Eigen::MatrixXd tau0 =
        bayesmix::to_eigen(prior->normal_mean_prior().scale());
    double nu0 = prior->normal_mean_prior().deg_free();
    // Check validity
    unsigned int dim = mu00.size();
    if (sigma00.rows() != dim or tau0.rows() != dim) {
      throw std::invalid_argument(
          "Hyperparameters dimensions are not consistent");
    }
    bayesmix::check_spd(sigma00);
    if (lambda0 <= 0) {
      throw std::invalid_argument("Variance-scaling parameter must be > 0");
    }
    bayesmix::check_spd(tau0);
    if (nu0 <= dim - 1) {
      throw std::invalid_argument("Degrees of freedom parameter is not valid");
    }
    // Set initial values
    hypers->mean = mu00;
    hypers->var_scaling = lambda0;
    hypers->scale = tau0;
    hypers->scale_inv = stan::math::inverse_spd(tau0);
    hypers->deg_free = nu0;
  }

  else if (prior->has_ngiw_prior()) {
    // Get hyperparameters:
    // for mu0
    Eigen::VectorXd mu00 =
        bayesmix::to_eigen(prior->ngiw_prior().mean_prior().mean());
    dim = mu00.size();
    Eigen::MatrixXd sigma00 =
        bayesmix::to_eigen(prior->ngiw_prior().mean_prior().var());
    // for lambda0
    double alpha00 = prior->ngiw_prior().var_scaling_prior().shape();
    double beta00 = prior->ngiw_prior().var_scaling_prior().rate();
    // for tau0
    double nu00 = prior->ngiw_prior().scale_prior().deg_free();
    Eigen::MatrixXd tau00 =
        bayesmix::to_eigen(prior->ngiw_prior().scale_prior().scale());
    // for nu0
    double nu0 = prior->ngiw_prior().deg_free();
    // Check validity:
    // dimensionality
    if (sigma00.rows() != dim or tau00.rows() != dim) {
      throw std::invalid_argument(
          "Hyperparameters dimensions are not consistent");
    }
    // for mu0
    bayesmix::check_spd(sigma00);
    // for lambda0
    if (alpha00 <= 0) {
      throw std::invalid_argument("Shape parameter must be > 0");
    }
    if (beta00 <= 0) {
      throw std::invalid_argument("Rate parameter must be > 0");
    }
    // for tau0
    if (nu00 <= 0) {
      throw std::invalid_argument("Degrees of freedom parameter must be > 0");
    }
    bayesmix::check_spd(tau00);
    // check nu0
    if (nu0 <= dim - 1) {
      throw std::invalid_argument("Degrees of freedom parameter is not valid");
    }
    // Set initial values
    hypers->mean = mu00;
    hypers->var_scaling = alpha00 / beta00;
    hypers->scale = tau00 / (nu00 + dim + 1);
    hypers->scale_inv = stan::math::inverse_spd(hypers->scale);
    hypers->deg_free = nu0;
  }

  else {
    throw std::invalid_argument("Unrecognized hierarchy prior");
  }
}

void NNWHierarchy::write_state_to_proto(google::protobuf::Message *out) const {
  bayesmix::MultiLSState state_;
  bayesmix::to_proto(state.mean, state_.mutable_mean());
  bayesmix::to_proto(state.prec, state_.mutable_prec());
  auto *out_cast = google::protobuf::internal::down_cast<
      bayesmix::MarginalState::ClusterState *>(out);
  out_cast->mutable_multi_ls_state()->CopyFrom(state_);
  out_cast->set_cardinality(card);
}

void NNWHierarchy::write_hypers_to_proto(
    google::protobuf::Message *out) const {
  bayesmix::NNWPrior hypers_;
  bayesmix::to_proto(hypers->mean,
                     hypers_.mutable_fixed_values()->mutable_mean());
  hypers_.mutable_fixed_values()->set_var_scaling(hypers->var_scaling);
  hypers_.mutable_fixed_values()->set_deg_free(hypers->deg_free);
  bayesmix::to_proto(hypers->scale,
                     hypers_.mutable_fixed_values()->mutable_scale());

  google::protobuf::internal::down_cast<bayesmix::NNWPrior *>(out)
      ->mutable_fixed_values()
      ->CopyFrom(hypers_.fixed_values());
}