#include <fstream>
#include <iostream>
#include <cmath>
#include <algorithm>
#include <vector>

#include <vtkMath.h>
#include <vtkSmartPointer.h>
#include <vtkNew.h>
#include <vtkObjectFactory.h>
#include <vtkTransform.h>
#include <vtkMath.h>

#include <Eigen/SVD>
#include <Eigen/Eigenvalues>

#include "vtkConversions.h"
#include "vtkCustomTransformInterpolator.h"
#include "vtkTimeCalibration.h"
#include "statistics.h"
#include "eigenFFTCorrelation.h"
#include "interpolator1D.h"

std::string ToString(CorrelationStrategy correlationStrategy)
{
  switch (correlationStrategy)
  {
    case CorrelationStrategy::DPOS:
      return "dpos";
    case CorrelationStrategy::SPEED_WINDOW:
      return "speed_window";
    case CorrelationStrategy::ACC_WINDOW:
      return "acc_window";
    case CorrelationStrategy::JERK_WINDOW:
      return "jerk_window";
    case CorrelationStrategy::LENGTH:
      return "length";
    case CorrelationStrategy::DERIVATED_LENGTH:
      return "derivated_length";
    case CorrelationStrategy::DROT:
      return "drot";
    case CorrelationStrategy::TRAJECTORY_ANGLE:
      return "trajectory_angle";
    case CorrelationStrategy::ORIENTATION_ANGLE:
      return "orientation_angle";
    case CorrelationStrategy::DERIVATED_ORIENTATION_ARC:
      return "derivated_orientation_arc";
    default:
      return "unkown";
  }
}

Interpolator1D<double> compute_speed_window(
    const vtkSmartPointer<vtkCustomTransformInterpolator>& transform,
    double window_width)
{
  std::vector<double> times = std::vector<double>();
  std::vector<double> speeds = std::vector<double>();
  vtkSmartPointer<vtkTransform> prev = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> next = vtkSmartPointer<vtkTransform>::New();

  double minMidWindowTime = transform->GetMinimumT() + 0.5 * window_width;
  double maxMidWindowTime = transform->GetMaximumT() - 0.5 * window_width;
  double period = transform->GetPeriod();
  double time = minMidWindowTime;
  while (time < maxMidWindowTime)
  {
    transform->InterpolateTransform(time - 0.5 * window_width, prev);
    transform->InterpolateTransform(time + 0.5 * window_width, next);
    speeds.push_back((PositionVectorFromTransform(next)
			    - PositionVectorFromTransform(prev)).norm() / window_width);
    times.push_back(time);
    time = time + period;
  }

  return Interpolator1D<double>(times, speeds);
}

Interpolator1D<double> compute_acc_window(
    const vtkSmartPointer<vtkCustomTransformInterpolator>& transform,
    double window_width)
{
  std::vector<double> times = std::vector<double>();
  std::vector<double> accs = std::vector<double>();
  vtkSmartPointer<vtkTransform> prev = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> next = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> curr = vtkSmartPointer<vtkTransform>::New();

  double minMidWindowTime = transform->GetMinimumT() + 0.5 * window_width;
  double maxMidWindowTime = transform->GetMaximumT() - 0.5 * window_width;
  double period = transform->GetPeriod();
  double time = minMidWindowTime;
  while (time < maxMidWindowTime)
  {
    transform->InterpolateTransform(time - 0.5 * window_width, prev);
    transform->InterpolateTransform(time, curr);
    transform->InterpolateTransform(time + 0.5 * window_width, next);
    Eigen::Vector3d a = (PositionVectorFromTransform(next)
		    + PositionVectorFromTransform(prev)
		    - 2 * PositionVectorFromTransform(curr)) / (window_width * window_width);
    accs.push_back(a.norm());
    times.push_back(time);
    time = time + period;
  }

  return Interpolator1D<double>(times, accs);
}

Interpolator1D<double> compute_jerk_window(
    const vtkSmartPointer<vtkCustomTransformInterpolator>& transform,
    double window_width)
{
  std::vector<double> times = std::vector<double>();
  std::vector<double> jerks = std::vector<double>();
  vtkSmartPointer<vtkTransform> t1 = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> t2 = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> t3 = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> t4 = vtkSmartPointer<vtkTransform>::New();

  double minMidWindowTime = transform->GetMinimumT() + 0.5 * window_width;
  double maxMidWindowTime = transform->GetMaximumT() - 0.5 * window_width;
  double period = transform->GetPeriod();
  double time = minMidWindowTime;
  while (time < maxMidWindowTime)
  {
    transform->InterpolateTransform(time - 0.5 * window_width, t1);
    transform->InterpolateTransform(time + (- 0.5 + 1.0/3.0) * window_width, t2);
    transform->InterpolateTransform(time + (- 0.5 + 2.0/3.0) * window_width, t3);
    transform->InterpolateTransform(time + (- 0.5 + 3.0/3.0) * window_width, t4);
    Eigen::Vector3d j = (PositionVectorFromTransform(t4)
		   - 3 * PositionVectorFromTransform(t3)
		   + 3 * PositionVectorFromTransform(t2)
		   - PositionVectorFromTransform(t1))
	    / std::pow(window_width, 3.0);
    jerks.push_back(j.norm());
    times.push_back(time);
    time = time + period;
  }

  return Interpolator1D<double>(times, jerks);
}

Interpolator1D<double> compute_dPos(
    const vtkSmartPointer<vtkCustomTransformInterpolator>& transform)
{
  std::vector<std::vector<double>> transforms = transform->GetTransformList();
  std::vector<double> t = std::vector<double>(transforms.size() - 1);
  std::vector<double> x = std::vector<double>(transforms.size() - 1);
  vtkSmartPointer<vtkTransform> prev = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> curr = vtkSmartPointer<vtkTransform>::New();
  for (unsigned int i = 0; i < transforms.size() - 1; i++)
  {
    double t0 = transforms[i][0];
    double t1 = transforms[i+1][0];
    t[i] = 0.5 * (t0 + t1);
    transform->InterpolateTransform(t0, prev);
    transform->InterpolateTransform(t1, curr);
    if (std::abs(t1 - t0) < 0.0001) {
      x[i] = 0.0;
    } else {
      x[i] = (PositionVectorFromTransform(curr)
		      - PositionVectorFromTransform(prev)).norm()
	      / (t1 - t0);
    }
  }
  return Interpolator1D<double>(t, x);
}

Interpolator1D<double> compute_length(
    const vtkSmartPointer<vtkCustomTransformInterpolator>& transform)
{
  std::vector<std::vector<double>> transforms = transform->GetTransformList();
  std::vector<double> t = std::vector<double>(transforms.size());
  std::vector<double> x = std::vector<double>(transforms.size());
  vtkSmartPointer<vtkTransform> prev = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> curr = vtkSmartPointer<vtkTransform>::New();
  t[0] = transforms[0][0];
  x[0] = 0.0;
  for (unsigned int i = 1; i < transforms.size(); i++)
  {
    t[i] = transforms[i][0];
    transform->InterpolateTransform(t[i-1], prev);
    transform->InterpolateTransform(t[i], curr);
    x[i] = x[i - 1] + (PositionVectorFromTransform(curr)
		    - PositionVectorFromTransform(prev)).norm();
  }

  return Interpolator1D<double>(t, x);
}

Interpolator1D<double> compute_derivated_length(
    const vtkSmartPointer<vtkCustomTransformInterpolator>& transform,
    double window_width)
{
  Interpolator1D<double> length = compute_length(transform);
  double tMin = transform->GetMinimumT() + 0.5 * window_width;
  double tMax = transform->GetMaximumT() - 0.5 * window_width;
  int steps = (tMax - tMin) / transform->GetPeriod() + 1;
  std::vector<double> derivated_length = std::vector<double>(steps);
  std::vector<double> times = std::vector<double>(steps);
  for (int i = 0; i < steps; i++)
  {
    double time = tMin + i * transform->GetPeriod();
    times[i] = time;
    // length is an interpolator so no need to check that the sample instants
    // are not the same (they are not, even if the interpolation mode of
    // this->Reference/Aligned is "NEAREST")
    derivated_length[i] = (length.Get(time + 0.5 * window_width)
                           - length.Get(time - 0.5 * window_width))
                    / window_width;
  }

  return Interpolator1D<double>(times, derivated_length);
}

Interpolator1D<double> compute_dRot(
    const vtkSmartPointer<vtkCustomTransformInterpolator>& transform)
{
  std::vector<std::vector<double>> transforms = transform->GetTransformList();
  std::vector<double> t = std::vector<double>(transforms.size() - 1);
  std::vector<double> x = std::vector<double>(transforms.size() - 1);
  vtkSmartPointer<vtkTransform> prev = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> curr = vtkSmartPointer<vtkTransform>::New();
  for (int i = 0; i < static_cast<int>(transforms.size()) - 1; i++)
  {
    double t0 = transforms[i][0];
    double t1 = transforms[i+1][0];
    transform->InterpolateTransform(t0, prev);
    transform->InterpolateTransform(t1, curr);
    Eigen::AngleAxisd aa = Eigen::AngleAxisd(RotationMatrixFromTransform(curr)
		    * RotationMatrixFromTransform(prev).transpose());
    t[i] = 0.5 * (t0 + t1);
    if (std::abs(t1 - t0) < 0.0001) {
      x[i] = 0.0;
    } else {
      x[i] = std::abs(aa.angle()) / (t1 - t0);
    }
  }
  return Interpolator1D<double>(t, x);
}

// TODO: handle case where real_sample_time(next) == real_sample_time(prev)
Interpolator1D<double> compute_trajectory_angle(
    const vtkSmartPointer<vtkCustomTransformInterpolator>& transform,
    double window_width)
{
  vtkSmartPointer<vtkTransform> prev = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> curr = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> next = vtkSmartPointer<vtkTransform>::New();
  double tMin = transform->GetMinimumT() + 0.5 * window_width;
  double tMax = transform->GetMaximumT() - 0.5 * window_width;
  int steps = (tMax - tMin) / transform->GetPeriod() + 1;
  std::vector<double> trajectory_angle = std::vector<double>(steps);
  std::vector<double> times = std::vector<double>(steps);
  for (int i = 0; i < steps; i++)
  {
    double time = tMin + i * transform->GetPeriod();
    transform->InterpolateTransform(time - 0.5 * window_width, prev);
    transform->InterpolateTransform(time, curr);
    transform->InterpolateTransform(time + 0.5 * window_width, next);
    times[i] = time;
    trajectory_angle[i] = SignedAngle(PositionVectorFromTransform(curr)
		    - PositionVectorFromTransform(prev),
		    PositionVectorFromTransform(next)
		    - PositionVectorFromTransform(curr));
  }

  return Interpolator1D<double>(times, trajectory_angle);
}

Interpolator1D<double> compute_orientation_arc(
    const vtkSmartPointer<vtkCustomTransformInterpolator>& transform)
{
  std::vector<std::vector<double>> transforms = transform->GetTransformList();
  std::vector<double> t = std::vector<double>(transforms.size());
  std::vector<double> x = std::vector<double>(transforms.size());
  vtkSmartPointer<vtkTransform> prev = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> curr = vtkSmartPointer<vtkTransform>::New();
  t[0] = transforms[0][0];
  x[0] = 0.0;
  for (unsigned int i = 1; i < transforms.size(); i++)
  {
    t[i] = transforms[i][0];
    transform->InterpolateTransform(t[i-1], prev);
    transform->InterpolateTransform(t[i], curr);
    Eigen::AngleAxisd aa = Eigen::AngleAxisd(RotationMatrixFromTransform(curr)
		    * RotationMatrixFromTransform(prev).transpose());
    x[i] = x[i - 1] + std::abs(aa.angle());
  }

  return Interpolator1D<double>(t, x);
}

Interpolator1D<double> compute_derivated_orientation_arc(
    const vtkSmartPointer<vtkCustomTransformInterpolator>& transform,
    double window_width)
{
  Interpolator1D<double> orientation_arc = compute_orientation_arc(transform);
  double tMin = transform->GetMinimumT() + 0.5 * window_width;
  double tMax = transform->GetMaximumT() - 0.5 * window_width;
  int steps = (tMax - tMin) / transform->GetPeriod() + 1;
  std::vector<double> derivated_orientation_arc = std::vector<double>(steps);
  std::vector<double> times = std::vector<double>(steps);
  for (int i = 0; i < steps; i++)
  {
    double time = tMin + i * transform->GetPeriod();
    times[i] = time;
    // length is an interpolator so no need to check that the sample instants
    // are not the same (they are not, even if the interpolation mode of
    // this->Reference/Aligned is "NEAREST")
    derivated_orientation_arc[i] =
                    (orientation_arc.Get(time + 0.5 * window_width)
                     - orientation_arc.Get(time - 0.5 * window_width))
                    / window_width;
  }

  return Interpolator1D<double>(times, derivated_orientation_arc);
}

// TODO: handle case where real_sample_time(next) == real_sample_time(prev)
Interpolator1D<double> compute_orientation_angle(
    const vtkSmartPointer<vtkCustomTransformInterpolator>& transform,
    double window_width)
{
  vtkSmartPointer<vtkTransform> prev = vtkSmartPointer<vtkTransform>::New();
  vtkSmartPointer<vtkTransform> next = vtkSmartPointer<vtkTransform>::New();
  double tMin = transform->GetMinimumT() + 0.5 * window_width;
  double tMax = transform->GetMaximumT() - 0.5 * window_width;
  int steps = (tMax - tMin) / transform->GetPeriod() + 1;
  std::vector<double> orientation_angle = std::vector<double>(steps);
  std::vector<double> times = std::vector<double>(steps);
  for (int i = 0; i < steps; i++)
  {
    double time = tMin + i * transform->GetPeriod();
    transform->InterpolateTransform(time - 0.5 * window_width, prev);
    transform->InterpolateTransform(time + 0.5 * window_width, next);
    times[i] = time;
    Eigen::AngleAxisd angleAxis(RotationMatrixFromTransform(next)
		    * RotationMatrixFromTransform(prev).transpose());
    orientation_angle[i] = angleAxis.angle();
  }

  return Interpolator1D<double>(times, orientation_angle);
}

double ComputeTimeShift(vtkSmartPointer<vtkTemporalTransforms> reference,
                      vtkSmartPointer<vtkTemporalTransforms> aligned,
                      CorrelationStrategy correlationStrategy,
                      double time_window_width,
                      bool substract_mean)
{
  vtkSmartPointer<vtkCustomTransformInterpolator> referenceInterpolator
      = reference->CreateInterpolator();
  referenceInterpolator->SetInterpolationTypeToLinear();
  vtkSmartPointer<vtkCustomTransformInterpolator> alignedInterpolator
      = aligned->CreateInterpolator();
  alignedInterpolator->SetInterpolationTypeToLinear();
  Interpolator1D<double> sig_reference;
  Interpolator1D<double> sig_aligned;
  // first, compute the signals using the chosen method
  switch (correlationStrategy)
  {
    case CorrelationStrategy::DPOS:
      sig_reference = compute_dPos(referenceInterpolator);
      sig_aligned = compute_dPos(alignedInterpolator);
      break;
    case CorrelationStrategy::SPEED_WINDOW:
      sig_reference = compute_speed_window(referenceInterpolator,
                                               time_window_width);
      sig_aligned = compute_speed_window(alignedInterpolator,
                                             time_window_width);
      break;
    case CorrelationStrategy::ACC_WINDOW:
      sig_reference = compute_acc_window(referenceInterpolator,
                                               time_window_width);
      sig_aligned = compute_acc_window(alignedInterpolator,
                                             time_window_width);
      break;
    case CorrelationStrategy::JERK_WINDOW:
      sig_reference = compute_jerk_window(referenceInterpolator,
                                               time_window_width);
      sig_aligned = compute_jerk_window(alignedInterpolator,
                                             time_window_width);
      break;
    case CorrelationStrategy::LENGTH:
      sig_reference = compute_length(referenceInterpolator);
      sig_aligned = compute_length(alignedInterpolator);
      break;
    case CorrelationStrategy::DERIVATED_LENGTH:
      sig_reference = compute_derivated_length(referenceInterpolator,
                                               time_window_width);
      sig_aligned = compute_derivated_length(alignedInterpolator,
                                             time_window_width);
      break;
    case CorrelationStrategy::DROT:
      sig_reference = compute_dRot(referenceInterpolator);
      sig_aligned = compute_dRot(alignedInterpolator);
      break;
    case CorrelationStrategy::TRAJECTORY_ANGLE:
      sig_reference = compute_trajectory_angle(referenceInterpolator,
                                               time_window_width);
      sig_aligned = compute_trajectory_angle(alignedInterpolator,
                                             time_window_width);
      break;
    case CorrelationStrategy::ORIENTATION_ANGLE:
      sig_reference = compute_orientation_angle(referenceInterpolator,
                                                time_window_width);
      sig_aligned = compute_orientation_angle(alignedInterpolator,
                                              time_window_width);
      break;
    case CorrelationStrategy::DERIVATED_ORIENTATION_ARC:
      sig_reference = compute_derivated_orientation_arc(
                              referenceInterpolator,
                              time_window_width);
      sig_aligned = compute_derivated_orientation_arc(
                              alignedInterpolator,
                              time_window_width);
      break;
    default:
      std::cerr << "unknown correlation strategy" << std::endl;
      return 0.0;
  }

  if (substract_mean)
  {
    sig_reference.ApplyValueShift(- sig_reference.Mean());
    sig_aligned.ApplyValueShift(- sig_aligned.Mean());
  }

  // We prefere the two signals to start at t = 0, so we time shift them,
  // but before that we save the information that we would lose otherwise.
  double pre_resample = sig_aligned.GetMinimumT() - sig_reference.GetMinimumT();
  sig_aligned.ApplyTimeShift(- sig_aligned.GetMinimumT());
  sig_reference.ApplyTimeShift(- sig_reference.GetMinimumT());

  // by construction we now have tMin == 0.0;
  double tMax = std::max(sig_reference.GetMaximumT(), sig_aligned.GetMaximumT());
  double period = std::min(sig_reference.GetAveragePeriod(),
		  sig_aligned.GetAveragePeriod());
  int steps = std::floor(tMax / period);
  std::vector<double> reference_resampled = std::vector<double>(steps);
  std::vector<double> aligned_resampled = std::vector<double>(steps);
  for (int i = 0; i < steps; i++)
  {
    double time = 0.0 + i * period;
    reference_resampled[i] = sig_reference.Get(time);
    aligned_resampled[i] = sig_aligned.Get(time);
  }

  std::vector<double> test = fftcorrelate(reference_resampled, aligned_resampled);

  int correlation = max_fftcorrelation(reference_resampled, aligned_resampled);
  double correction = correlation * period;
  double delta_t = pre_resample - correction;
  return delta_t;
}

void ShowTrajectoryInfo(vtkSmartPointer<vtkTemporalTransforms> reference, vtkSmartPointer<vtkTemporalTransforms> aligned)
{
  vtkSmartPointer<vtkCustomTransformInterpolator> referenceI = reference->CreateInterpolator();
  vtkSmartPointer<vtkCustomTransformInterpolator> alignedI = aligned->CreateInterpolator();
  std::cout << std::fixed;
  std::cout << std::setprecision(4);
  std::cout << "reference:                 " << referenceI->GetMaximumT() - referenceI->GetMinimumT()
            << "s from " << referenceI->GetMinimumT() << " to " << referenceI->GetMaximumT()
            << ", period is " << referenceI->GetPeriod() << std::endl;
  std::cout << "aligned:                   " << alignedI->GetMaximumT() - alignedI->GetMinimumT()
            << "s from " << alignedI->GetMinimumT() << " to " << alignedI->GetMaximumT()
            << ", period is " << alignedI->GetPeriod() << std::endl;
}

void DemoAllTimesyncMethods(vtkSmartPointer<vtkTemporalTransforms> reference, vtkSmartPointer<vtkTemporalTransforms> aligned) {
  vtkSmartPointer<vtkCustomTransformInterpolator> referenceI = reference->CreateInterpolator();
  vtkSmartPointer<vtkCustomTransformInterpolator> alignedI = aligned->CreateInterpolator();

  std::cout << std::fixed;
  std::cout << std::setprecision(4);
  ShowTrajectoryInfo(reference, aligned);
  std::cout << std::endl;
  std::cout << "dPos:                      " << ComputeTimeShift(reference, aligned, CorrelationStrategy::DPOS, 1.0) << std::endl;
  std::cout << "speed window:              " << ComputeTimeShift(reference, aligned, CorrelationStrategy::SPEED_WINDOW, 1.0) << std::endl;
  std::cout << "acceleration window:       " << ComputeTimeShift(reference, aligned, CorrelationStrategy::ACC_WINDOW, 3) << std::endl;
  std::cout << "jerk window:               " << ComputeTimeShift(reference, aligned, CorrelationStrategy::JERK_WINDOW, 6) << std::endl;
  std::cout << "derivated length:          " << ComputeTimeShift(reference, aligned, CorrelationStrategy::DERIVATED_LENGTH, 1.0) << std::endl;
  std::cout << "dRot:                      " << ComputeTimeShift(reference, aligned, CorrelationStrategy::DROT, 1.0) << std::endl;
  std::cout << "trajectory angle:          " << ComputeTimeShift(reference, aligned, CorrelationStrategy::TRAJECTORY_ANGLE, 10.0) << std::endl;
  std::cout << "orientation angle:         " << ComputeTimeShift(reference, aligned, CorrelationStrategy::ORIENTATION_ANGLE, 1.0) << std::endl;
  std::cout << "derivated orientation arc: " << ComputeTimeShift(reference, aligned, CorrelationStrategy::DERIVATED_ORIENTATION_ARC, 1.0) << std::endl;
}

double ComputeScale(vtkSmartPointer<vtkTemporalTransforms> reference,
                  vtkSmartPointer<vtkTemporalTransforms> aligned,
                  CorrelationStrategy correlationStrategy,
                  double time_window_width)
{
  const double div_epsilon = 1e-4;
  vtkSmartPointer<vtkCustomTransformInterpolator> referenceInterpolator = reference->CreateInterpolator();
  referenceInterpolator->SetInterpolationTypeToLinear();
  vtkSmartPointer<vtkCustomTransformInterpolator> alignedInterpolator = aligned->CreateInterpolator();
  alignedInterpolator->SetInterpolationTypeToLinear();
  Interpolator1D<double> sig_reference;
  Interpolator1D<double> sig_aligned;
  switch (correlationStrategy)
  {
    case CorrelationStrategy::DPOS:
      sig_reference = compute_dPos(referenceInterpolator);
      sig_aligned = compute_dPos(alignedInterpolator);
      break;
    case CorrelationStrategy::SPEED_WINDOW:
      sig_reference = compute_speed_window(referenceInterpolator,
                                               time_window_width);
      sig_aligned = compute_speed_window(alignedInterpolator,
                                             time_window_width);
      break;
    case CorrelationStrategy::ACC_WINDOW:
      sig_reference = compute_acc_window(referenceInterpolator,
                                               time_window_width);
      sig_aligned = compute_acc_window(alignedInterpolator,
                                             time_window_width);
      break;
    case CorrelationStrategy::JERK_WINDOW:
      sig_reference = compute_jerk_window(referenceInterpolator,
                                               time_window_width);
      sig_aligned = compute_jerk_window(alignedInterpolator,
                                             time_window_width);
      break;
    case CorrelationStrategy::LENGTH:
      sig_reference = compute_length(referenceInterpolator);
      sig_aligned = compute_length(alignedInterpolator);
      break;
    case CorrelationStrategy::DERIVATED_LENGTH:
      sig_reference = compute_derivated_length(referenceInterpolator,
                                               time_window_width);
      sig_aligned = compute_derivated_length(alignedInterpolator,
                                             time_window_width);
      break;
    default:
      std::cerr << "unsuported correlation strategy" << std::endl;
      return 0.0;
  }

  double tMin = std::min(sig_reference.GetMinimumT(), sig_aligned.GetMaximumT());
  double tMax = std::max(sig_reference.GetMaximumT(), sig_aligned.GetMaximumT());
  double period = std::min(sig_reference.GetAveragePeriod(),
		  sig_aligned.GetAveragePeriod());
  int steps = std::floor((tMax - tMin) / period);
  std::vector<double> ratios;
  for (int i = 0; i < steps; i++)
  {
    double time = tMin + i * period;
    double reference_resampled = sig_reference.Get(time);
    double aligned_resampled = sig_aligned.Get(time);
    if (reference_resampled < div_epsilon)
    {
      continue;
    }
    else
    {
      ratios.push_back(aligned_resampled / reference_resampled);
    }
  }

  return ComputeMedian(ratios);
}


static char help[] = "Advance_ small CMEs to benchmark intranode performance.\n\n";

#include <fstream>
#include <iomanip>
#include <petscmat.h>
#include <petscvec.h>
#include <petscviewer.h>
#include <armadillo>
#include <cmath>
#include <sys/stat.h>
#include "pacmensl_all.h"

namespace hog1p_cme {
// stoichiometric matrix of the toggle switch model
arma::Mat<PetscInt> SM{{1,-1,-1,0,0,0,0,0,0},
                       {0,0,0,1,0,-1,0,0,0},
                       {0,0,0,0,1,0,-1,0,0},
                       {0,0,0,0,0,1,0,-1,0},
                       {0,0,0,0,0,0,1,0,-1},};

// reaction parameters
const PetscReal k12{1.29},k23{0.0067},k34{0.133},k32{0.027},k43{0.0381},k21{1.0e0},kr21{0.005},kr31{
    0.45},      kr41{0.025},kr22{0.0116},kr32{0.987},kr42{0.0538},trans{0.01},gamma1{0.001},gamma2{0.0049},
// parameters for the time-dependent factors
                r1{6.9e-5},r2{7.1e-3},eta{3.1},Ahog{9.3e09},Mhog{6.4e-4};

// propensity function
inline PetscReal hog_propensity(int *X,int k)
{
  switch (k)
  {
    case 0:return k12 * double(X[0] == 0) + k23 * double(X[0] == 1) + k34 * double(X[0] == 2);
    case 1:return k32 * double(X[0] == 2) + k43 * double(X[0] == 3);
    case 2:return k21 * double(X[0] == 1);
    case 3:return kr21 * double(X[0] == 1) + kr31 * double(X[0] == 2) + kr41 * double(X[0] == 3);
    case 4:return kr22 * double(X[0] == 1) + kr32 * double(X[0] == 2) + kr42 * double(X[0] == 3);
    case 5:return trans * double(X[1]);
    case 6:return trans * double(X[2]);
    case 7:return gamma1 * double(X[3]);
    case 8:return gamma2 * double(X[4]);
    default:return 0.0;
  }
}

int propensity(const int reaction,
               const int num_species,
               const int num_states,
               const int *states,
               PetscReal *outputs,
               void *args)
{
  int (*X)[5] = ( int (*)[5] ) states;
  for (int i = 0; i < num_states; ++i)
  {
    outputs[i] = hog_propensity(X[i],reaction);
  }
  return 0;
}

// function to compute the time-dependent coefficients of the propensity functions
int t_fun(double t,int num_coefs,PetscReal *outputs,void *args)
{
  if (num_coefs != 9) return -1;

  arma::Row<double> u(outputs,9,false,true);
  u.fill(1.0);

  double h1 = (1.0 - exp(-r1 * t)) * exp(-r2 * t);

  double hog1p = pow(h1 / (1.0 + h1 / Mhog),eta) * Ahog;

  u(2) = std::max(0.0,3200.0 - 7710.0 * (hog1p));

  return 0;
}

// Function to constraint the shape of the Fsp
int lhs_constr(PetscInt num_species,PetscInt num_constrs,PetscInt num_states,PetscInt *states,int *vals,
               void *args)
{
  if (num_species != 5)
  {
    return -1;
  }
  if (num_constrs != 7)
  {
    return -1;
  }
  for (int j{0}; j < num_states; ++j)
  {
    for (int i{0}; i < 5; ++i)
    {
      vals[num_constrs * j + i] = (states[num_species * j + i]);
    }
    vals[num_constrs * j + 5] = ((states[num_species * j + 1]) + (states[num_species * j + 3]));
    vals[num_constrs * j + 6] = ((states[num_species * j + 2]) + (states[num_species * j + 4]));
  }
  return 0;
}

arma::Row<int>    rhs_constr_hyperrec{3,10,10,10,10};
arma::Row<double> expansion_factors_hyperrec{0.0,0.25,0.25,0.25,0.25};
arma::Row<int>    rhs_constr{3,10,10,10,10,10,10};
arma::Row<double> expansion_factors{0.0,0.2,0.2,0.2,0.2,0.2,0.2};
}

using arma::dvec;
using arma::Col;
using arma::Row;

using std::cout;
using std::endl;

using namespace hog1p_cme;
using namespace pacmensl;

void output_marginals(MPI_Comm comm,std::string model_name,PartitioningType fsp_par_type,
                      PartitioningApproach fsp_repart_approach,std::string constraint_type,
                      DiscreteDistribution &solution,arma::Row<int> constraints);

void output_performance(MPI_Comm comm,std::string &model_name,PartitioningType fsp_par_type,
                        PartitioningApproach fsp_repart_approach,std::string constraint_type,
                        ODESolverType ode_type,
                        FspSolverMultiSinks &fsp_solver);

int ParseOptions(MPI_Comm comm,
                 PartitioningType &fsp_par_type,
                 PartitioningApproach &fsp_repart_approach,
                 PetscBool &output_marginal,
                 PetscBool &fsp_log_events,
                 ODESolverType &ode_solver);

int main(int argc,char *argv[])
{
  Environment my_env(&argc,&argv,help);

  PetscMPIInt    ierr,myRank,num_procs;
  PetscErrorCode petsc_err;
  MPI_Comm       comm;
  MPI_Comm_dup(PETSC_COMM_WORLD,&comm);
  MPI_Comm_size(comm,&num_procs);
  PetscPrintf(comm,"Solving with %d processors.\n",num_procs);

  // Register PETSc stages
  PetscLogStage stages[2];
  petsc_err = PetscLogStageRegister("Solve with adaptive default state set shape",&stages[0]);
  CHKERRQ(petsc_err);
  petsc_err = PetscLogStageRegister("Solve with fixed default state set shape",&stages[1]);
  CHKERRQ(petsc_err);

  // Set up CME
  std::string         model_name = "hog1p";
  Model               hog1p_model(hog1p_cme::SM,hog1p_cme::t_fun,hog1p_cme::propensity,
                                  nullptr,nullptr,{2});
  PetscReal           t_final    = 60 * 3.0;
  PetscReal           fsp_tol    = 1.0e-4;
  arma::Mat<PetscInt> X0         = {0,0,0,0,0};
  X0 = X0.t();
  arma::Col<PetscReal> p0 = {1.0};

  // Default options
  PartitioningType     fsp_par_type        = PartitioningType::GRAPH;
  PartitioningApproach fsp_repart_approach = PartitioningApproach::REPARTITION;
  ODESolverType        fsp_odes_type       = CVODE;
  PetscBool            output_marginal     = PETSC_FALSE;
  PetscBool            fsp_log_events      = PETSC_FALSE;

  ierr = ParseOptions(comm,fsp_par_type,fsp_repart_approach,output_marginal,fsp_log_events,fsp_odes_type);
  CHKERRQ(ierr);

  FspSolverMultiSinks fsp_solver(comm,fsp_par_type,fsp_odes_type);
  fsp_solver.SetModel(hog1p_model);
  fsp_solver.SetInitialDistribution(X0,p0);
  fsp_solver.SetFromOptions();
  DiscreteDistribution solution;

  petsc_err = PetscLogStagePush(stages[0]);
  CHKERRQ(petsc_err);

  // Solve using adaptive default constraints
  fsp_solver.SetInitialBounds(rhs_constr_hyperrec);
  fsp_solver.SetExpansionFactors(expansion_factors_hyperrec);
  fsp_solver.SetFromOptions();
  solution = fsp_solver.Solve(t_final,fsp_tol,0);
  std::shared_ptr<const StateSetConstrained>
                 fss                   = std::static_pointer_cast<const StateSetConstrained>(fsp_solver.GetStateSet());
  arma::Row<int> final_hyperrec_constr = fss->GetShapeBounds();
  if (fsp_log_events)
  {
    output_performance(PETSC_COMM_WORLD,model_name,fsp_par_type,fsp_repart_approach,
                       std::string("adaptive_default"),fsp_odes_type,fsp_solver);
  }
  if (output_marginal)
  {
    output_marginals(PETSC_COMM_WORLD,model_name,fsp_par_type,fsp_repart_approach,
                     std::string("adaptive_default"),solution,final_hyperrec_constr);
  }
  fsp_solver.ClearState();
  PetscPrintf(comm,"\n ================ \n");

  petsc_err = PetscLogStagePop();
  CHKERRQ(petsc_err);
  petsc_err = PetscLogStagePush(stages[1]);
  CHKERRQ(petsc_err);
  // Solve using fixed default constraints
  fsp_solver.SetInitialBounds(final_hyperrec_constr);
  solution = fsp_solver.Solve(t_final,fsp_tol,0);
  if (fsp_log_events)
  {
    output_performance(PETSC_COMM_WORLD,model_name,fsp_par_type,fsp_repart_approach,
                       std::string("fixed_default"),fsp_odes_type,fsp_solver);
  }
  if (output_marginal)
  {
    output_marginals(PETSC_COMM_WORLD,model_name,fsp_par_type,fsp_repart_approach,
                     std::string("fixed_default"),solution,final_hyperrec_constr);
  }
  fsp_solver.ClearState();
  PetscPrintf(comm,"\n ================ \n");
  return ierr;
}

int ParseOptions(MPI_Comm comm,
                 PartitioningType &fsp_par_type,
                 PartitioningApproach &fsp_repart_approach,
                 PetscBool &output_marginal,
                 PetscBool &fsp_log_events,
                 ODESolverType &ode_solver)
{
  std::string part_type;
  std::string part_approach;
  part_type     = part2str(fsp_par_type);
  part_approach = partapproach2str(fsp_repart_approach);

  // Read options for fsp
  char      opt[100];
  PetscBool opt_set;
  int       ierr;
  ierr = PetscOptionsGetString(NULL,PETSC_NULL,"-fsp_partitioning_type",opt,100,&opt_set);
  CHKERRQ(ierr);
  if (opt_set)
  {
    fsp_par_type = str2part(std::string(opt));
  }

  ierr = PetscOptionsGetString(NULL,PETSC_NULL,"-fsp_repart_approach",opt,100,&opt_set);
  CHKERRQ(ierr);
  if (opt_set)
  {
    fsp_repart_approach = str2partapproach(std::string(opt));
  }

  ierr = PetscOptionsGetString(NULL,PETSC_NULL,"-fsp_output_marginal",opt,100,&opt_set);
  CHKERRQ(ierr);
  if (opt_set)
  {
    if (strcmp(opt,"1") == 0 || strcmp(opt,"true") == 0)
    {
      output_marginal = PETSC_TRUE;
    }
  }

  ierr = PetscOptionsGetString(NULL,PETSC_NULL,"-fsp_log_events",opt,100,&opt_set);
  CHKERRQ(ierr);
  if (opt_set)
  {
    if (strcmp(opt,"1") == 0 || strcmp(opt,"true") == 0)
    {
      fsp_log_events = PETSC_TRUE;
    }
  }

  ierr = PetscOptionsGetString(NULL,PETSC_NULL,"-fsp_use_solver",opt,100,&opt_set);
  CHKERRQ(ierr);
  if (opt_set)
  {
    if (strcmp(opt,"krylov") == 0)
    {
      ode_solver = KRYLOV;
    } else
    {
      ode_solver = CVODE;
    }
  }

  PetscPrintf(comm,"Partitiniong option %s \n",part2str(fsp_par_type).c_str());
  PetscPrintf(comm,"Repartitoning option %s \n",partapproach2str(fsp_repart_approach).c_str());
  return 0;
}

void output_performance(MPI_Comm comm,std::string &model_name,PartitioningType fsp_par_type,
                        PartitioningApproach fsp_repart_approach,std::string constraint_type,
                        ODESolverType ode_type,
                        FspSolverMultiSinks &fsp_solver)
{
  int myRank,num_procs;
  MPI_Comm_rank(comm,&myRank);
  MPI_Comm_size(comm,&num_procs);

  std::string ode;
  if (ode_type == KRYLOV)
  {
    ode = "krylov";
  } else
  {
    ode = "cvode";
  }

  std::string part_type;
  std::string part_approach;
  part_type     = part2str(fsp_par_type);
  part_approach = partapproach2str(fsp_repart_approach);

  // Output time breakdowns
  FspSolverComponentTiming    sum_times, min_times, max_times;
  sum_times = fsp_solver.ReduceComponentTiming("sum");
  min_times = fsp_solver.ReduceComponentTiming("min");
  max_times = fsp_solver.ReduceComponentTiming("max");

  if (myRank == 0)
  {
    struct stat buffer;
    int fstat;

    std::string   filename =
                      model_name + "_time_breakdown.dat";

    fstat = stat (filename.c_str(), &buffer);

    std::ofstream file;
    file.open(filename,std::ios_base::app);

    if (fstat != 0){
      file << "ncpu,partitioner,fsp_shape,ode_solver,min_cput,max_cput,avg_cput,mat_gen_time,ode_time,state_expand_time,min_flops,max_flops,avg_flops \n";
    }

    file << num_procs << ","
         << part_type << ","
         << constraint_type << ","
         << ode << ","
         << min_times.TotalTime << ","
         << max_times.TotalTime << ","
         << sum_times.TotalTime/num_procs << ","
         << sum_times.MatrixGenerationTime/num_procs << ","
         << sum_times.ODESolveTime/num_procs << ","
         << sum_times.StatePartitioningTime/num_procs  << ","
         << min_times.TotalFlops << ","
         << max_times.TotalFlops << ","
         << sum_times.TotalFlops/num_procs << "\n"
        ;
    file.close();
  }

  FiniteProblemSolverPerfInfo perf_info   = fsp_solver.GetSolverPerfInfo();

  if (myRank == 0){
    std::string filename =
                    model_name + "_perf_info_" + std::to_string(num_procs) + "_" + part_type + "_" + part_approach + "_" +
                        constraint_type + ".dat";
    std::ofstream file;
    file.open(filename);
    file << "Model time, ODEs size, Average processor time (sec) \n";
    for (auto i{0}; i < perf_info.n_step; ++i)
    {
      file << perf_info.model_time[i] << "," << perf_info.n_eqs[i] << "," << perf_info.cpu_time[i] << "\n";
    }
    file.close();
  }
}

void output_marginals(MPI_Comm comm,std::string model_name,PartitioningType fsp_par_type,
                      PartitioningApproach fsp_repart_approach,std::string constraint_type,
                      DiscreteDistribution &solution,arma::Row<int> constraints)
{
  int myRank,num_procs;
  MPI_Comm_rank(comm,&myRank);
  MPI_Comm_size(comm,&num_procs);

  std::string part_type;
  std::string part_approach;
  part_type     = part2str(fsp_par_type);
  part_approach = partapproach2str(fsp_repart_approach);

  /* Compute the marginal distributions */
  std::vector<arma::Col<PetscReal>> marginals(solution.states_.n_rows);
  for (PetscInt                     i{0}; i < marginals.size(); ++i)
  {
    marginals[i] = Compute1DMarginal(solution,i);
  }

  MPI_Comm_rank(PETSC_COMM_WORLD,&myRank);
  if (myRank == 0)
  {
    for (PetscInt i{0}; i < marginals.size(); ++i)
    {
      std::string filename =
                      model_name + "_marginal_" + std::to_string(i) + "_" + std::to_string(num_procs) + "_" +
                          part_type + "_" + part_approach + "_" + constraint_type + ".dat";
      marginals[i].save(filename,arma::raw_ascii);
    }
    std::string   filename =
                      model_name + "_constraint_bounds_" + std::to_string(num_procs) + "_" + part_type + "_"
                          + part_approach +
                          "_" + constraint_type + ".dat";
    constraints.save(filename,arma::raw_ascii);
  }
}


//
// Created by huangkun on 2020/9/18.
//
#include <Eigen/Eigen>
#include <opencv2/core/eigen.hpp>
#include <algorithm>
#include <limits>

#include <cv_calib.hpp>
#include <dbscan.h>
#include <opengv2/utility/utility.hpp>
#include <opengv2/event_camera_calib/CirclesEventFrame.hpp>
#include <opengv2/event_camera_calib/CalibCircle.hpp>
#include <opengv2/sensor/PinholeCamera.hpp>

opengv2::CirclesEventFrame::CirclesEventFrame(EventContainer::Ptr container, const std::pair<double, double> &duration,
                                              CirclePatternParameters::Ptr pattern, Params params) :
        EventFrame(container, duration), pattern_(pattern), params_(params) {
    auto camera = dynamic_cast<CameraBase *>(sensor_.get());
    if (pattern_->isAsymmetric) {
        circleRadiusThreshold_ = std::min(
                std::max(camera->size()[0], camera->size()[1]) /
                std::max(pattern_->rows, 2 * pattern_->cols),
                std::min(camera->size()[0], camera->size()[1]) /
                std::min(pattern_->rows, 2 * pattern_->cols)) / pattern->squareSize * pattern->circleRadius * 1.5;
    } else {
        circleRadiusThreshold_ = std::min(
                std::max(camera->size()[0], camera->size()[1]) /
                std::max(pattern_->rows, pattern_->cols),
                std::min(camera->size()[0], camera->size()[1]) /
                std::min(pattern_->rows, pattern_->cols)) / pattern->squareSize * pattern->circleRadius * 1.5;
    }
}

opengv2::CirclesEventFrame::Params::Params() {
    dbscan_eps = 4;
    dbscan_startMinSample = 2;
    clusterMinSample = 5;
    knn_num = 3;
}

opengv2::CirclesEventFrame::Params::Params(const cv::FileStorage &node) {
    node["dbscan_eps"] >> dbscan_eps;
    node["dbscan_startMinSample"] >> dbscan_startMinSample;
    node["clusterMinSample"] >> clusterMinSample;
    node["knn_num"] >> knn_num;
    node["fitCircle"] >> fitCircle;
}

struct CircleErr {
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW

    CircleErr(double fitErr, double radius, const Eigen::Ref<const Eigen::Vector2d> &center) :
            fitErr(fitErr), radius(radius), center(center) {}

    double fitErr;
    double radius;
    Eigen::Vector2d center;
};

bool opengv2::CirclesEventFrame::extractFeatures() {
    if (positiveEvents_.empty() || negativeEvents_.empty()) {
        return false;
    }

    auto dbscan = DBSCAN<Eigen::Vector2d, double>();
    dbscan.Run(&positiveEvents_, 2, params_.dbscan_eps, params_.dbscan_startMinSample);
    auto p_noise = std::move(dbscan.Noise);
    auto p_clusters = std::move(dbscan.Clusters);
    dbscan.Run(&negativeEvents_, 2, params_.dbscan_eps, params_.dbscan_startMinSample);
    auto n_noise = std::move(dbscan.Noise);
    auto n_clusters = std::move(dbscan.Clusters);

    /** Draw the result and filter data **/
    auto camera = dynamic_cast<CameraBase *>(sensor_.get());
    image_ = cv::Mat(camera->size()[1], camera->size()[0], CV_8UC3, cv::Vec3b(0, 0, 0));

    //debug
    eventImage = image_.clone();
    for (const Eigen::Vector2d &p: positiveEvents_) {
        cv::Point loc(p[0], p[1]);
        eventImage.at<cv::Vec3b>(loc) = cv::Vec3b(0, 0, 255);
    }
    for (const Eigen::Vector2d &p: negativeEvents_) {
        cv::Point loc(p[0], p[1]);
        eventImage.at<cv::Vec3b>(loc) = cv::Vec3b(0, 255, 0);
    }

    for (auto itr = p_clusters.begin(); itr != p_clusters.end();) {
        // remove cluster with too few samples
        if (itr->size() < params_.clusterMinSample) {
            itr = p_clusters.erase(itr);
        } else {
            unsigned int color = 20 * (itr - p_clusters.begin());
            for (unsigned int idx : *itr) {
                cv::Point loc(positiveEvents_[idx][0], positiveEvents_[idx][1]);
                image_.at<cv::Vec3b>(loc) = cv::Vec3b(color / 256, color % 256, 200);
            }

            itr++;
        }
    }

    for (auto itr = n_clusters.begin(); itr != n_clusters.end();) {
        // remove cluster with too few samples
        if (itr->size() < params_.clusterMinSample) {
            itr = n_clusters.erase(itr);
        } else {
            unsigned int color = 20 * (itr - n_clusters.begin());
            for (unsigned int idx : *itr) {
                cv::Point loc(negativeEvents_[idx][0], negativeEvents_[idx][1]);
                image_.at<cv::Vec3b>(loc) = cv::Vec3b(color / 256, color % 256, 100);
            }

            itr++;
        }
    }

    // store
    nClusters_ = n_clusters;
    pClusters_ = p_clusters;

    // debug
    clusterImage = image_.clone();

    // too few events
    if (p_clusters.size() < pattern_->rows * pattern_->cols || n_clusters.size() < pattern_->rows * pattern_->cols) {
        return false;
    }

    /** find candidate circles **/
    std::vector<std::pair<size_t, size_t>> candidates;
    vectorofEigenMatrix<Eigen::Vector2d> candidateCenters;
    std::vector<double> candidatesRadius;

    // calculate each cluster center (median)
    std::vector<uint> p_centers, n_centers;
    auto p_compare_fun = [&](uint lhs, uint rhs) { return positiveEvents_[lhs].norm() < positiveEvents_[rhs].norm(); };
    auto n_compare_fun = [&](uint lhs, uint rhs) { return negativeEvents_[lhs].norm() < negativeEvents_[rhs].norm(); };
    for (auto &pCluster: p_clusters) {
        std::nth_element(pCluster.begin(), pCluster.begin() + pCluster.size() / 2, pCluster.end(), p_compare_fun);
        p_centers.push_back(pCluster[pCluster.size() / 2]);
    }
    for (auto &nCluster: n_clusters) {
        std::nth_element(nCluster.begin(), nCluster.begin() + nCluster.size() / 2, nCluster.end(), n_compare_fun);
        n_centers.push_back(nCluster[nCluster.size() / 2]);
    }

    /*// draw cluster centers
    for (auto center: p_centers) {
        cv::Point loc(positiveEvents_[center][0], positiveEvents_[center][1]);
        cv::circle(image_, loc, 2, cv::Vec3b(255, 255, 255));
    }
    for (auto center: n_centers) {
        cv::Point loc(negativeEvents_[center][0], negativeEvents_[center][1]);
        cv::circle(image_, loc, 2, cv::Vec3b(255, 255, 255));
    }*/

    // building k-d tree
    vectorofEigenMatrix<Eigen::Vector2d> pCenters(p_centers.size()), nCenters(n_centers.size());
    for (int i = 0; i < p_centers.size(); ++i) {
        pCenters[i] = positiveEvents_[p_centers[i]];
    }
    for (int i = 0; i < n_centers.size(); ++i) {
        nCenters[i] = negativeEvents_[n_centers[i]];
    }
    KDTreeVectorOfVectorsAdaptor<vectorofEigenMatrix<Eigen::Vector2d>, double, 2, nanoflann::metric_L2_Simple>
            p_kdTree(2, pCenters, 10), n_kdTree(2, nCenters, 10);

    // Store fitError for Reusing
    std::unordered_map<std::pair<size_t, size_t>, CircleErr, pair_hash> fitErrandRadiusMap;

    const size_t num_results = params_.fitCircle ? params_.knn_num : 1;
    std::vector<size_t> n_idx(num_results);
    std::vector<size_t> p_idx(num_results);
    std::vector<double> out_dists_sqr(num_results);
    vectorofEigenMatrix<Eigen::Vector2d> centers(num_results);
    std::vector<double> radius(num_results);
    std::vector<double> fitErrs;
    if (params_.fitCircle) {
        for (size_t pi = 0; pi < p_centers.size(); ++pi) {
            size_t realNum_results = num_results;
            // search knn in opposite polarity cluster
            n_kdTree.query(positiveEvents_[p_centers[pi]].data(), num_results, n_idx.data(), out_dists_sqr.data());

            // remove cluster too far away
            for (int oi = 0; oi < out_dists_sqr.size(); ++oi) {
                if (out_dists_sqr[oi] > out_dists_sqr[0] * 4 ||
                    out_dists_sqr[oi] > 4 * circleRadiusThreshold_ * circleRadiusThreshold_) {
                    realNum_results = oi;
                    break;
                }
            }
            if (realNum_results == 0)
                continue;

            // connect two centers as diameter, verify circle fitting error
            fitErrs.assign(realNum_results, 0);
            for (int j = 0; j < realNum_results; ++j) {
                auto itr = fitErrandRadiusMap.find(std::pair<size_t, size_t>(pi, n_idx[j]));
                if (itr == fitErrandRadiusMap.end()) {
                    fitCircle(p_clusters[pi], n_clusters[n_idx[j]], centers[j], radius[j]);

                    double approxRadius =
                            (positiveEvents_[p_centers[pi]] - negativeEvents_[n_centers[n_idx[j]]]).norm() / 2;
                    if (radius[j] > circleRadiusThreshold_ || radius[j] > 2 * approxRadius) {
                        fitErrs[j] = std::numeric_limits<double>::max();
                    } else {
                        for (auto pEvent: p_clusters[pi]) {
                            fitErrs[j] += std::abs((positiveEvents_[pEvent] - centers[j]).norm() - radius[j]);
                        }
                        for (auto nEvent: n_clusters[n_idx[j]]) {
                            fitErrs[j] += std::abs((negativeEvents_[nEvent] - centers[j]).norm() - radius[j]);
                        }

                        fitErrs[j] /= (p_clusters[pi].size() + n_clusters[n_idx[j]].size()) * radius[j];
                    }
                    fitErrandRadiusMap.emplace(std::pair<size_t, size_t>(pi, n_idx[j]),
                                               CircleErr(fitErrs[j], radius[j], centers[j]));
                } else {
                    fitErrs[j] = itr->second.fitErr;
                    radius[j] = itr->second.radius;
                    centers[j] = itr->second.center;
                }
            }
            int n_minIdx = std::min_element(fitErrs.begin(), fitErrs.end()) - fitErrs.begin();

            // Do circle check for the minimum, if the fitting error greater than Expectation(assume gaussian noise)
            const double gaussianNoise = 2 / radius[n_minIdx]; // TODO: considering ellipse
            if (fitErrs[n_minIdx] < gaussianNoise) {
                /* Double Direction Check */
                realNum_results = num_results;
                p_kdTree.query(negativeEvents_[n_centers[n_idx[n_minIdx]]].data(), num_results, p_idx.data(),
                               out_dists_sqr.data());
                for (int oi = 0; oi < out_dists_sqr.size(); ++oi) {
                    if (out_dists_sqr[oi] > out_dists_sqr[0] * 4 ||
                        out_dists_sqr[oi] > 4 * circleRadiusThreshold_ * circleRadiusThreshold_) {
                        realNum_results = oi;
                        break;
                    }
                }
                if (realNum_results == 0)
                    continue;
                fitErrs.assign(realNum_results, 0);
                for (int i = 0; i < realNum_results; ++i) {
                    auto itr = fitErrandRadiusMap.find(std::pair<size_t, size_t>(p_idx[i], n_idx[n_minIdx]));
                    if (itr == fitErrandRadiusMap.end()) {
                        fitCircle(p_clusters[p_idx[i]], n_clusters[n_idx[n_minIdx]], centers[i], radius[i]);

                        double approxRadius = (positiveEvents_[p_centers[p_idx[i]]] -
                                               negativeEvents_[n_centers[n_idx[n_minIdx]]]).norm() / 2;
                        if (radius[i] > circleRadiusThreshold_ || radius[i] > 2 * approxRadius) {
                            fitErrs[i] = std::numeric_limits<double>::max();
                        } else {
                            for (auto pEvent: p_clusters[p_idx[i]]) {
                                fitErrs[i] += std::abs((positiveEvents_[pEvent] - centers[i]).norm() - radius[i]);
                            }
                            for (auto nEvent: n_clusters[n_idx[n_minIdx]]) {
                                fitErrs[i] += std::abs((negativeEvents_[nEvent] - centers[i]).norm() - radius[i]);
                            }
                            fitErrs[i] /=
                                    (p_clusters[p_idx[i]].size() + n_clusters[n_idx[n_minIdx]].size()) * radius[i];
                        }
                        fitErrandRadiusMap.emplace(std::pair<size_t, size_t>(p_idx[i], n_idx[n_minIdx]),
                                                   CircleErr(fitErrs[i], radius[i], centers[i]));
                    } else {
                        fitErrs[i] = itr->second.fitErr;
                        radius[i] = itr->second.radius;
                        centers[i] = itr->second.center;
                    }
                }
                int p_minIdx = std::min_element(fitErrs.begin(), fitErrs.end()) - fitErrs.begin();

                // if the same, add to candidate
                if (p_idx[p_minIdx] == pi) {
                    candidates.emplace_back(pi, n_idx[n_minIdx]);
                    candidateCenters.push_back(centers[p_minIdx]);
                    candidatesRadius.push_back(radius[p_minIdx]);
                }
            }
        }
    } else {
        for (size_t pi = 0; pi < p_centers.size(); ++pi) {
            n_kdTree.query(positiveEvents_[p_centers[pi]].data(), num_results, n_idx.data(), out_dists_sqr.data());
            // remove cluster too far away
            if (out_dists_sqr[0] > 4 * circleRadiusThreshold_ * circleRadiusThreshold_) {
                continue;
            }
            p_kdTree.query(negativeEvents_[n_centers[n_idx[0]]].data(), num_results, p_idx.data(),
                           out_dists_sqr.data());
            if (p_idx[0] == pi) {
                // check circle
                Eigen::Vector2d center =
                        (positiveEvents_[p_centers[p_idx[0]]] + negativeEvents_[n_centers[n_idx[0]]]) / 2;
                double r = (positiveEvents_[p_centers[p_idx[0]]] - negativeEvents_[n_centers[n_idx[0]]]).norm() / 2;
                double fitErr = 0;
                for (auto pEvent: p_clusters[p_idx[0]]) {
                    fitErr += std::abs((positiveEvents_[pEvent] - center).norm() - r);
                }
                for (auto nEvent: n_clusters[n_idx[0]]) {
                    fitErr += std::abs((negativeEvents_[nEvent] - center).norm() - r);
                }
                fitErr /= (p_clusters[p_idx[0]].size() + n_clusters[n_idx[0]].size()) * r;

                if (fitErr < 10 / r) {
                    candidates.emplace_back(pi, n_idx[0]);
                    candidateCenters.push_back(center);
                    candidatesRadius.push_back(r);
                }
            }
        }
    }

    // draw candidate circle
    for (int i = 0; i < candidates.size(); ++i) {
        cv::Point loc(candidateCenters[i][0], candidateCenters[i][1]);
        cv::circle(image_, loc, candidatesRadius[i], cv::Vec3b(0, 255, 0));
    }

    /** use prior pattern to match the candidate **/
    std::vector<cv::Point2f> points, outCenters;
    for (const Eigen::Vector2d &center: candidateCenters) {
        points.emplace_back(center[0], center[1]);
    }
    // TODO: 1. improve the sensitivity to additional nearby outlier
    //  2. we should assume there are absence,
    //  (even wrong pairs inside the pattern (seems too hard, maybe solved previously by more samples)),
    //  also the centers are not accurate.
    //  (solution) maybe optimization liangzu;
    //  or maybe solved by mul-spline segment(just ignore the error segment) and ignore with go on;
    //  user interaction (note: find sample in raw data)
    bool isFound = cv::findCirclesGrid(points, cv::Size(pattern_->cols, pattern_->rows), outCenters,
                                       cv::CALIB_CB_ASYMMETRIC_GRID);
    if (!isFound)
        isFound = cv::findCirclesGrid(points, cv::Size(pattern_->cols, pattern_->rows), outCenters,
                                      cv::CALIB_CB_ASYMMETRIC_GRID | cv::CALIB_CB_CLUSTERING);
    /*if (isFound)
        circleExtractionImage = image_.clone();*/
    drawChessboardCorners(image_, cv::Size(pattern_->cols, pattern_->rows), cv::Mat(outCenters), isFound);
    if (isFound) {
        KDTreeVectorOfVectorsAdaptor<vectorofEigenMatrix<Eigen::Vector2d>, double, 2, nanoflann::metric_L2_Simple>
                kdTree(2, candidateCenters, 10);
        std::vector<size_t> orderIdxs(pattern_->rows * pattern_->cols);
        std::vector<double> out_dist_sqr(1);
        for (int i = 0; i < outCenters.size(); ++i) {
            Eigen::Vector2d c(outCenters[i].x, outCenters[i].y);
            kdTree.query(c.data(), 1, orderIdxs.data() + i, out_dist_sqr.data());
        }

        // features_
        for (auto idx: orderIdxs) {
            features_.push_back(std::make_shared<CalibCircle>(candidateCenters[idx], candidatesRadius[idx]));
        }

        //detectionImage = image_.clone();
    }

    return isFound;
}

void opengv2::CirclesEventFrame::fitCircle(const std::vector<uint> &pSet,
                                           const std::vector<uint> &nSet,
                                           Eigen::Ref<Eigen::Vector2d> center, double &radius) {
    double sum_x = 0, sum_y = 0;
    double sum_xx = 0, sum_yy = 0, sum_xy = 0;
    double sum_xxx = 0, sum_yyy = 0, sum_xyy = 0, sum_xxy = 0;
    for (auto p_idx: pSet) {
        const Eigen::Vector2d &sample = positiveEvents_[p_idx];

        sum_x += sample[0];
        sum_y += sample[1];

        double xx = sample[0] * sample[0];
        double yy = sample[1] * sample[1];
        double xy = sample[0] * sample[1];

        sum_xx += xx;
        sum_yy += yy;
        sum_xy += xy;

        sum_xxx += xx * sample[0];
        sum_yyy += yy * sample[1];
        sum_xyy += xy * sample[1];
        sum_xxy += sample[0] * xy;
    }
    for (auto n_idx: nSet) {
        const Eigen::Vector2d &sample = negativeEvents_[n_idx];

        sum_x += sample[0];
        sum_y += sample[1];

        double xx = sample[0] * sample[0];
        double yy = sample[1] * sample[1];
        double xy = sample[0] * sample[1];

        sum_xx += xx;
        sum_yy += yy;
        sum_xy += xy;

        sum_xxx += xx * sample[0];
        sum_yyy += yy * sample[1];
        sum_xyy += xy * sample[1];
        sum_xxy += sample[0] * xy;
    }

    Eigen::Matrix3d A;
    A << 2 * sum_x, 2 * sum_y, pSet.size() + nSet.size(),
            2 * sum_xx, 2 * sum_xy, sum_x,
            2 * sum_xy, 2 * sum_yy, sum_y;
    Eigen::Vector3d b(sum_xx + sum_yy, sum_xxx + sum_xyy, sum_xxy + sum_yyy);

    Eigen::Vector3d x = A.lu().solve(b);
    center = x.block<2, 1>(0, 0);
    radius = std::sqrt(x[0] * x[0] + x[1] * x[1] + x[2]);
}

bool
opengv2::CirclesEventFrame::rectifyFeatures(const std::unordered_set<int> &outlierIdxs,
                                            const Eigen::Ref<const Eigen::Matrix3d> &Rcw,
                                            const Eigen::Ref<const Eigen::Vector3d> &tcw) {
    KDTreeVectorOfVectorsAdaptor<vectorofEigenMatrix<Eigen::Vector2d>, double, 2, nanoflann::metric_L2_Simple>
            p_kdTree(2, positiveEvents_, 10), n_kdTree(2, negativeEvents_, 10);

    Eigen::Matrix3d Rsw = Rcw;
    Eigen::Vector3d tsw = tcw;
    auto camera = dynamic_cast<PinholeCamera *>(sensor_.get());
    cv::Mat tvec, rvec, Rvec, distCoeffs, cameraMatrix;
    cv::eigen2cv(Rsw, Rvec);
    cv::Rodrigues(Rvec, rvec);
    cv::eigen2cv(tsw, tvec);
    cv::eigen2cv(camera->distCoeffs(), distCoeffs);
    cv::eigen2cv(camera->K(), cameraMatrix);
    for (int kIdx = 0; kIdx < features_.size(); kIdx++) {
        auto f = dynamic_cast<CalibCircle *>(features_[kIdx].get());
        auto lm = f->landmark();
        Eigen::Vector3d center = lm->position();

        std::vector<cv::Point3f> objectPoints;
        objectPoints.reserve(5);
        objectPoints.emplace_back(center[0], center[1], center[2]);
        double skewR = pattern_->circleRadius / std::sqrt(2);
        // Four quadrant, considering circle distortion
        objectPoints.emplace_back(center[0] + skewR, center[1] + skewR, center[2]);
        objectPoints.emplace_back(center[0] + skewR, center[1] - skewR, center[2]);
        objectPoints.emplace_back(center[0] - skewR, center[1] - skewR, center[2]);
        objectPoints.emplace_back(center[0] - skewR, center[1] + skewR, center[2]);

        std::vector<cv::Point2f> imagePoints;
        cv::projectPoints(objectPoints, rvec, tvec, cameraMatrix, distCoeffs, imagePoints);
        vectorofEigenMatrix<Eigen::Vector2d> eigenImagePoints;
        eigenImagePoints.reserve(5);
        for (const auto &p: imagePoints) {
            eigenImagePoints.emplace_back(p.x, p.y);
        }

        // if center exceed image bound, delete feature
        if (imagePoints[0].x >= camera->size()[0] || imagePoints[0].y >= camera->size()[1] ||
            imagePoints[0].x < 0 || imagePoints[0].y < 0) {
            features_[kIdx] = nullptr;
            continue;
        }

        std::vector<double> radius;
        double maxRadius = 0;
        radius.reserve(4);
        for (int i = 1; i < 5; ++i) {
            radius.push_back((eigenImagePoints[i] - eigenImagePoints[0]).norm());
            if (radius.back() > maxRadius) {
                maxRadius = radius.back();
            }
        }

        double inlierThreshold = 3; // pixel unit
        std::vector<std::pair<size_t, double>> p_IndicesDists, n_IndicesDists;
        p_kdTree.index->radiusSearch(eigenImagePoints[0].data(), std::pow(maxRadius + inlierThreshold, 2),
                                     p_IndicesDists,
                                     nanoflann::SearchParams(32, 0, false));
        n_kdTree.index->radiusSearch(eigenImagePoints[0].data(), std::pow(maxRadius + inlierThreshold, 2),
                                     n_IndicesDists,
                                     nanoflann::SearchParams(32, 0, false));

        // collect inliers
        std::pair<std::vector<uint>, std::vector<uint>> patternCircle;
        for (const auto &pair: p_IndicesDists) {
            Eigen::Vector2d direction = positiveEvents_[pair.first] - eigenImagePoints[0];
            double distance = std::sqrt(pair.second);
            int idx = 0;
            if (direction[0] >= 0 && direction[1] >= 0) {
                idx = 0;
            } else if (direction[0] >= 0 && direction[1] <= 0) {
                idx = 1;
            } else if (direction[0] <= 0 && direction[1] <= 0) {
                idx = 2;
            } else if (direction[0] <= 0 && direction[1] >= 0) {
                idx = 3;
            }

            if (std::abs(distance - radius[idx]) <= inlierThreshold) {
                patternCircle.first.push_back(pair.first);
            }
        }
        for (const auto &pair: n_IndicesDists) {
            Eigen::Vector2d direction = negativeEvents_[pair.first] - eigenImagePoints[0];
            double distance = std::sqrt(pair.second);
            int idx = 0;
            if (direction[0] >= 0 && direction[1] >= 0) {
                idx = 0;
            } else if (direction[0] >= 0 && direction[1] <= 0) {
                idx = 1;
            } else if (direction[0] <= 0 && direction[1] <= 0) {
                idx = 2;
            } else if (direction[0] <= 0 && direction[1] >= 0) {
                idx = 3;
            }

            if (std::abs(distance - radius[idx]) <= inlierThreshold) {
                patternCircle.second.push_back(pair.first);
            }
        }

        // expand sample set by clustering
        std::set<uint> pClusterSet, nClusterSet;
        std::unordered_map<uint, uint> pSample2Set, nSample2Set;
        for (int i = 0; i < pClusters_.size(); ++i) {
            for (int j = 0; j < pClusters_[i].size(); ++j) {
                pSample2Set.emplace(pClusters_[i][j], i);
            }
        }
        for (int i = 0; i < nClusters_.size(); ++i) {
            for (int j = 0; j < nClusters_[i].size(); ++j) {
                nSample2Set.emplace(nClusters_[i][j], i);
            }
        }
        for (const auto &p : patternCircle.first) {
            if (pSample2Set.find(p) != pSample2Set.end())
                pClusterSet.insert(pSample2Set[p]);
        }
        for (const auto &n : patternCircle.second) {
            if (nSample2Set.find(n) != nSample2Set.end())
                nClusterSet.insert(nSample2Set[n]);
        }
        patternCircle.first.clear();
        for (const auto &pCluster: pClusterSet) {
            for (const auto &p: pClusters_[pCluster]) {
                patternCircle.first.push_back(p);
            }
        }
        patternCircle.second.clear();
        for (const auto &nCluster: nClusterSet) {
            for (const auto &n: nClusters_[nCluster]) {
                patternCircle.second.push_back(n);
            }
        }

        // remove feature with too few measurements
        if (patternCircle.first.size() < 5 || patternCircle.second.size() < 5) {
            features_[kIdx] = nullptr;
            continue;
        }

        Eigen::Vector2d rectifiedCenter;
        double rectifiedRadius;
        fitCircle(patternCircle.first, patternCircle.second, rectifiedCenter,
                  rectifiedRadius);

        std::nth_element(radius.begin(), radius.begin() + radius.size() / 2, radius.end());
        // if fit result not good, delete feature
        if ((rectifiedCenter - eigenImagePoints[0]).norm() > 2 * inlierThreshold ||
            std::abs(rectifiedRadius - radius[radius.size() / 2]) > 1.5 * inlierThreshold) {
            features_[kIdx] = nullptr;
            continue;
        }

        f->setLocation(rectifiedCenter);
        f->radius = rectifiedRadius;
    }

    // release no longer used member
    positiveEvents_.clear();
    negativeEvents_.clear();

    // outlier features in the edges are more important
    std::vector<std::unordered_set<int>> edgePattern(4);
    std::vector<int> score(4, 0);
    for (int i = 0; i < pattern_->cols; ++i) // 1st row
        edgePattern[0].insert(i);
    for (int i = (pattern_->rows - 1) * pattern_->cols; i < (pattern_->rows * pattern_->cols); ++i) // last row
        edgePattern[1].insert(i);
    for (int i = 0;
         i < (pattern_->rows * pattern_->cols); i += (pattern_->isAsymmetric ? 2 : 1) * pattern_->cols) // 1st col
        edgePattern[2].insert(i);
    for (int i = pattern_->isAsymmetric ? (2 * pattern_->cols - 1) : (pattern_->cols - 1);
         i < (pattern_->rows * pattern_->cols); i += (pattern_->isAsymmetric ? 2 : 1) * pattern_->cols) // last col
        edgePattern[3].insert(i);

    int counter = 0, index = 0;
    for (auto itr = features_.begin(); itr != features_.end(); index++) {
        if (*itr == nullptr) {
            itr = features_.erase(itr);
            for (int i = 0; i < 4; ++i) {
                if (edgePattern[i].find(index) != edgePattern[i].end())
                    score[i]++;
            }
            counter++;
        } else {
            itr++;
        }
    }

    // check features on edge
    if (!params_.fitCircle) {
        for (int i = 0; i < 4; ++i) {
            if (score[i] >= (int) (edgePattern[i].size() - 1))
                return false;
        }
    }

    // remove frame with too many features gone
    if (counter >= 0.2 * (pattern_->cols * pattern_->rows)) {
        return false;
    }

    for (auto fb: features_) {
        auto f = dynamic_cast<CalibCircle *>(fb.get());
        cv::Point loc(f->location()[0], f->location()[1]);
        cv::circle(image_, loc, f->radius, cv::Vec3b(255, 255, 255));
    }

    // Building KD-Tree for neighbor searching on features_, used for establish corresponds for given event.
    circles_.reserve(features_.size());
    for (const auto &f: features_) {
        circles_.push_back(f->location());
    }
    circleKdTree_ = std::make_shared<KDTreeVectorOfVectorsAdaptor<vectorofEigenMatrix<Eigen::Vector2d>, double, 2>>
            (2, circles_, 10);

    return true;
}


#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <iterator>
#include <iomanip>
#include <Eigen/Dense>
#include <cmath>
#include <random>
#include "constants.h"
#include "simulationbox.h"
#include "fileinput.h"
#include "fileoutput.h"
#include "thermostats.h"

using namespace Eigen;

class LennardJones
{
    private:
         Vector3d cutoffadjust;
        double sigma;
        double epsilon;
        double cutoff;
    public:
        LennardJones(double, double, double);
        Vector3d getForce(Vector3d);
};

LennardJones::LennardJones(double insigma, double inepsilon, double incutoff=10000) {
    // Constructs the Lennard Jones class, converting to atomic units as we go
    // We expect the sigma to be in angstrom,
    // The epsilon to be in Hartrees
    // and the cutoff to also be in angstrom

    sigma = insigma * constants::bohrAng;
    epsilon = inepsilon * constants::boltzHar;
    cutoff = incutoff * constants::bohrAng;
    double cutoff_pair_virial = 24 * epsilon / pow(cutoff, 2);
    double cutoff_sigma_6 = pow(sigma, 6) / pow(cutoff, 6);
    double cutoff_sigma_12 = pow(cutoff_sigma_6, 2);
    double cutoffadjust_value = cutoff_pair_virial * (2 * cutoff_sigma_12 - cutoff_sigma_6);
    cutoffadjust(0) = cutoffadjust_value;
    cutoffadjust(1) = cutoffadjust_value;
    cutoffadjust(2) = cutoffadjust_value;
}

Vector3d LennardJones::getForce(Vector3d distance)
{
    /* Class function for the Lennard Jones model that takes in
     * a radial distance and outputs the force vector,
     * and the energy in relation to that */
    double radiusSquared=distance.squaredNorm();
    if (radiusSquared > pow(cutoff,2)){
        // We're outside the range we wanted to look at
        // so just return no force and carry on
        Vector3d force(0.0, 0.0, 0.0);
        return force;
    }
    double sigma6 = pow(sigma, 6) / pow(radiusSquared, 3);
    double sigma12 = pow(sigma6, 2);
    //Allen and Tildesey eq 2.59-2.63, then eq 5.3
    double pair_virial = 24.0 * epsilon / radiusSquared;
    Vector3d force = (distance * pair_virial * (2 * sigma12 - sigma6)) + cutoffadjust;
    return force;
}

void implement_boundary_conditions(ArrayXXd &positions, const Box &inputBox) {
    // Mutates the input array to fix any atoms that are outside the
    // boundary conditions
    // Fails when an atom is more than twice the boundary away,
    // but something else has gone wrong then

    for (int atom=0; atom < positions.rows(); ++atom) {
        for (int dim=0; dim < positions.cols(); ++dim) {
            if (positions(atom, dim) < 0.0) {
                positions(atom, dim) += inputBox.sides(dim);
            }
            else if (positions(atom, dim) > inputBox.sides(dim)) {
                positions(atom, dim) -= inputBox.sides(dim);
            }
        }
    }
}

ArrayXXd populate_velocities(const Box &input_box, const ArrayXd &masses, double temperature){
    // Populates the velocity array with velocities sampled
    // randomly from a Gaussian distribution with mean 0 and
    // variance 1.

    // This is then rescaled to match the temperature before
    // it is used

    std::random_device device;
    std::mt19937 mt_rand(device());

    // This is the normal distribution generator
    double mean = 0.0;
    double var = 1.0;
    std::normal_distribution<double> gaussian(mean, var);

    ArrayXXd velocities = ArrayXXd::Zero(masses.rows(), input_box.dimensions);
    // Use the number of rows in the masses as a proxy for
    // number of atoms
    for (int atom=0; atom < masses.rows(); ++atom) {
        for (int dim=0; dim < input_box.dimensions; ++dim) {
            velocities(atom, dim) = gaussian(mt_rand);
        }
    }

    auto rescaler = VelocityRescale();
    rescaler.apply(velocities, masses, temperature);
    return velocities;
}

ArrayXXd populate_positions(const Box &inputBox, int num_x, int num_y, int num_z) {
    /* Fills the position array with atoms evenly spaced in the box
     * starting from (0, 0) */
    ArrayXXd positions(num_x * num_y * num_z, inputBox.dimensions);
    double x_increment = inputBox.sides(0) / static_cast<double>(num_x);
    double y_increment = inputBox.sides(1) / static_cast<double>(num_y);
    double z_increment = inputBox.sides(2) / static_cast<double>(num_z);
    int counter = 0;
    for (int row=0; row < num_x; ++row) {
        for (int col=0; col < num_y; ++col) {
            for (int plane=0; plane < num_z; ++plane) {
                positions(counter, 0) = row * x_increment;
                positions(counter, 1) = col * y_increment;
                positions(counter, 2) = plane * z_increment;
                ++counter;
            }
        }
    }
    return positions;
}

Vector3d implement_minimum_image(const Vector3d distance, const Box &box) {
    auto new_distance = distance;
    for (int dim=0; dim < box.dimensions; ++dim) {
        if (distance(dim) <= -0.5 * box.sides(dim)) {
            new_distance(dim) += box.sides(dim);
        } else if (distance(dim) > 0.5 * box.sides(dim)) {
            new_distance(dim) -= box.sides(dim);
        }
    }
    return new_distance;
}

ArrayXXd calculate_forces(ArrayXXd &atoms_array, const Box &box, LennardJones &potential_model) {
    /* Calculates the force according to a Lennard-Jones potential
     * from one atom to its nearest image, and returns
     * an array of these forces */
    ArrayXXd forces = ArrayXXd::Zero(atoms_array.rows(), atoms_array.cols());
    for (int i=0; i < atoms_array.rows(); ++i) {
        for (int j=0; j<i; ++j) {
            Vector3d distance = atoms_array.row(i) - atoms_array.row(j);
            distance = implement_minimum_image(distance, box);
            Vector3d instant_force(potential_model.getForce(distance));

            for (int dim=0; dim < box.dimensions; ++dim) {
                forces(i, dim) += instant_force(dim);
                forces(j, dim) -= instant_force(dim);
            }
        }
    }
    return forces;
}

void velocityVerlet(ArrayXXd &positions, ArrayXXd &velocities, double timestep, ArrayXd &masses,
                    const Box &box, LennardJones &potentialModel) {
    ArrayXXd accelerations = calculate_forces(positions, box, potentialModel).colwise() / masses;
    ArrayXXd positions_next = positions
                              + (velocities * timestep)
                              + (0.5 * accelerations * pow(timestep, 2));
    ArrayXXd accelerations_next = calculate_forces(positions_next, box, potentialModel).colwise() / masses;
    ArrayXXd velocities_next = velocities
                               + 0.5 * ( accelerations + accelerations_next) * timestep;
    velocities = velocities_next;
    implement_boundary_conditions(positions_next, box);
    positions = positions_next;
}

int main() {
    auto parameter_map = read_in_from_file("input.inpt");
    LennardJones potentialModel(std::stod(parameter_map["LJ_sigma"]),
                                std::stod(parameter_map["LJ_epsilon"]),
                                std::stod(parameter_map["LJ_cutoff"]));
    Box simulation_box(std::stod(parameter_map["length_x"]),
                       std::stod(parameter_map["length_y"]),
                       std::stod(parameter_map["length_z"]));
    int num_x = stoi(parameter_map["num_x"]);
    int num_y = stoi(parameter_map["num_y"]);
    int num_z = stoi(parameter_map["num_z"]);
    double temperature = std::stod(parameter_map["temperature"]);
    double timestep = std::stod(parameter_map["timestep"]);
    int rescaleStep=std::stoi(parameter_map["num_rescale"]);
    int outputStep=std::stoi(parameter_map["num_output"]);
    int numSteps=std::stoi(parameter_map["num_steps"]);

    ArrayXXd atoms = populate_positions(simulation_box, num_x, num_y, num_z);
    ArrayXd masses = ArrayXd::Zero(atoms.rows(), 1);
    masses += std::stod(parameter_map["mass"]);

    ArrayXXd velocities = populate_velocities(simulation_box, masses, temperature);

    clearOutputFiles(parameter_map["position_file"]);

    for (int step=0; step < numSteps; ++step) {
        velocityVerlet(atoms, velocities, timestep, masses, simulation_box, potentialModel);
        if (step % rescaleStep == 0) {
            rescale_velocities(velocities, masses, temperature);
        }

        if (step % outputStep == 0) {
            dumpToFile(atoms, parameter_map["position_file"], step);
            std::cout << "Step " << step << std::endl;
        }
    }
}


#pragma once
#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/IterativeLinearSolvers>

#include <tbb/tbb.h>

#include "Utilities/Utility.hpp"
#include "Observables/Energy.hpp"
#include "./utils.hpp"

namespace yannq
{
template<typename Machine, typename Hamiltonian>
class SRMat;
} //namespace yannq

namespace Eigen { //namespace Eigen
namespace internal {
	template<typename Machine, typename Hamiltonian>
	struct traits<yannq::SRMat<Machine, Hamiltonian> >
		:  public Eigen::internal::traits<Eigen::SparseMatrix<typename Machine::Scalar> > {};
}
}// namespace Eigen;

namespace yannq
{
//! \addtogroup GroundState

//! \ingroup GroundState
//! This class that generates the quantum Fisher matrix for the stochastic reconfiguration (SR) method.
template<typename Machine, typename Hamiltonian>
class SRMat
	: public Eigen::EigenBase<SRMat<Machine, Hamiltonian> >
{
public:
	using Scalar = typename Machine::Scalar;
	using RealScalar = typename remove_complex<Scalar>::type;

	using Matrix = typename Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>;
	using Vector = typename Eigen::Matrix<Scalar, Eigen::Dynamic, 1>;

	using VectorConstRef = typename Eigen::Ref<const Vector>;
	using RealMatrix = typename Eigen::Matrix<RealScalar, Eigen::Dynamic, Eigen::Dynamic>;
	using RealVector = typename Eigen::Matrix<RealScalar, Eigen::Dynamic, 1>;

	using StorageIndex = uint32_t;

	enum {
		ColsAtCompileTime = Eigen::Dynamic,
		MaxColsAtCompileTime = Eigen::Dynamic,
		IsRowMajor = false
	};

private:
	uint32_t n_;

	const Machine& qs_;
	const Hamiltonian& ham_;

	Energy<Scalar, Hamiltonian> energy_;

	RealScalar shift_;

	Matrix deltas_;
	Vector deltaMean_;

	Vector energyGrad_;
	Vector weights_;

public:
	//! \param qs Machine that describes quantum states
	//! \param ham Hamiltonian for SR
	SRMat(const Machine& qs, const Hamiltonian& ham)
	  : n_{qs.getN()}, qs_(qs), ham_(ham), energy_(ham)
	{
	}

	Eigen::Index rows() const { return qs_.getDim(); }
	Eigen::Index cols() const { return qs_.getDim(); }

	template<typename Rhs>
	Eigen::Product<SRMat<Machine, Hamiltonian>, Rhs, Eigen::AliasFreeProduct>
			operator*(const Eigen::MatrixBase<Rhs>& x) const {
	  return Eigen::Product<SRMat<Machine, Hamiltonian>, Rhs, Eigen::AliasFreeProduct>(*this, x.derived());
	}

	//! \param rs Sampling results obtained from samplers.
	template<class SamplingResult>
	void constructFromSamples(SamplingResult&& sr)
	{
		weights_.resize(0);
		int nsmp = sr.size();

		constructDelta(qs_, sr, deltas_);
		constructObs(qs_, sr, energy_);

		deltaMean_ = deltas_.colwise().mean();
		deltas_ = deltas_.rowwise() - deltaMean_.transpose();

		energyGrad_ =  deltas_.adjoint() * energy_.elocs();
		energyGrad_ /= nsmp;
	}

	template<class SamplingResult>
	void constructFromWeightSamples(const Eigen::Ref<const RealVector>& weights, SamplingResult&& sr)
	{
		weights_ = weights;

		constructDelta(qs_, sr, deltas_);
		constructObsWeights(qs_, std::forward<SamplingResult>(sr), weights, energy_);

		deltaMean_ = weights.transpose()*deltas_;
		deltas_ = deltas_.rowwise() - deltaMean_.transpose();

		energyGrad_ =  deltas_.adjoint() * weights.asDiagonal()
			* energy_.elocs();
	}

	void setShift(RealScalar shift)
	{
		shift_ = shift;
	}

	RealScalar getShift() const
	{
		return shift_;
	}

	//! return \f$\langle \nabla_{\theta} \psi_\theta(\sigma) \rangle\f$
	const Vector& oloc() const&
	{
		return deltaMean_;
	}

	Vector oloc() &&
	{
		return deltaMean_;
	}

	Matrix corrMat()
	{
		int nsmp = deltas_.rows();
		if(weights_.size() == 0)
			return (deltas_.adjoint() * deltas_)/nsmp;
		else
			return (deltas_.adjoint() * weights_.asDiagonal() * deltas_);
	}

	RealScalar eloc() const
	{
		return std::real(energy_.eloc());
	}

	Scalar elocVar() const
	{
		return std::real(energy_.elocVar());
	}

	const Vector& energyGrad() const&
	{
		return energyGrad_;
	}
	Vector energyGrad() &&
	{
		return energyGrad_;
	}

	Vector apply(const Vector& rhs) const
	{
		assert(rhs.size() == qs_.getDim());
		Vector r = deltas_*rhs;

		Vector res;

		if(weights_.size() == 0)
		{
			r /= r.rows();
			res = deltas_.adjoint()*r;
		}
		else
		{
			r.array() *= weights_.array();
			res = deltas_.adjoint() * r;
		}

		return res + Scalar(shift_)*rhs;
	}

	/*! \brief Use conjugate gradient solover to solve the optimizing vector.
	 *
	 * We solve \f$ (S + \epsilon \mathbb{1})v = f \f$ where \f$ f \f$ is the gradient of the energy expectation values.
	 * \param shift \f$ \epsilon \f$ that controls regularization
	 * \param tol tolerance for CG solver
	 */
	Vector solveCG(RealScalar shift, RealScalar tol = 1e-4)
	{
		setShift(shift);
		Eigen::ConjugateGradient<
			SRMat<Machine, Hamiltonian>,
			Eigen::Lower|Eigen::Upper,
			Eigen::IdentityPreconditioner> cg;
		cg.compute(*this);
		cg.setTolerance(tol);
		return cg.solve(energyGrad_);

	}

	/**
	 * Solve S^{-1}vec using conjugate gradient method
	 */
	Vector solveCG(const VectorConstRef& vec, double shift, double tol = 1e-4)
	{
		setShift(shift);
		Eigen::ConjugateGradient<
			SRMat<Machine, Hamiltonian>,
			Eigen::Lower|Eigen::Upper,
			Eigen::IdentityPreconditioner> cg;
		cg.compute(*this);
		cg.setTolerance(tol);
		return cg.solve(vec);
	}

	/*! \brief Solve the optimizing vector by solving the linear equation exactly..
	 *
	 * We solve \f$ (S + \epsilon \mathbb{1})v = f \f$ where \f$ f \f$ is the gradient of the energy expectation values.
	 * \param shift \f$ \epsilon \f$ that controls regularization
	 */
	Vector solveExact(RealScalar shift)
	{
		Matrix mat = corrMat();
		mat += shift*Matrix::Identity(mat.rows(),mat.cols());
		Eigen::LLT<Matrix> llt{mat};
		return llt.solve(energyGrad_);
	}

};
} //namespace yannq

namespace Eigen {
namespace internal {
	template<typename Rhs, typename Machine, typename Hamiltonian>
	struct generic_product_impl<yannq::SRMat<Machine, Hamiltonian>, Rhs, SparseShape, DenseShape, GemvProduct> // GEMV stands for matrix-vector
	: generic_product_impl_base<yannq::SRMat<Machine, Hamiltonian>, Rhs, generic_product_impl<yannq::SRMat<Machine, Hamiltonian>, Rhs> >
	{
		typedef typename Product<yannq::SRMat<Machine, Hamiltonian>, Rhs>::Scalar Scalar;
		template<typename Dest>
		static void scaleAndAddTo(Dest& dst, const yannq::SRMat<Machine, Hamiltonian>& lhs, const Rhs& rhs, const Scalar& alpha)
		{
			// This method should implement "dst += alpha * lhs * rhs" inplace,
			// however, for iterative solvers, alpha is always equal to 1, so let's not bother about it.
			assert(alpha==Scalar(1) && "scaling is not implemented");
			EIGEN_ONLY_USED_FOR_DEBUG(alpha);

			dst += lhs.apply(rhs);
		}
	};
} //namespace internal
} //namespace Eigen