/**
 * Source file for class Ledger
 * Please see Ledger.h for documentation
 */

#include "statsdcc/ledger.h"

#include <boost/lexical_cast.hpp>

#include <algorithm>
#include <cmath>
#include <cstring>

#include "statsdcc/chrono.h"
#include "statsdcc/configs/aggregator_config.h"
#include "statsdcc/logger.h"

namespace statsdcc {

// buffers the metric
void Ledger::buffer(const std::string& metric) {
  const char *metric_csty = metric.c_str();

  char name_buffer[MAX_TOKEN_BUFFER_SIZE] = {0};
  char type_buffer[3] = {0};
  char metric_value_buffer[MAX_TOKEN_BUFFER_SIZE] = {0};
  double metric_value = -std::numeric_limits<double>::max();
  double sample_rate = 1;
  MetricType type = MetricType::counter;

  // record number of calls to buffer method
  ++this->statsd_metrics["metrics_processed"];

  sscanf(metric_csty,
         "%" MAX_TOKEN_STRING_LENGTH "[^:]:" \
         "%" MAX_TOKEN_STRING_LENGTH "[^|]|" \
         "%2[^|]|"                           \
         "@%lf",
         name_buffer, metric_value_buffer, type_buffer, &sample_rate);

  // trim leading spaces if any
  char *tmp_type = type_buffer;
  while (*tmp_type == ' ' && *tmp_type != '\0')
    ++tmp_type;
  std::string metric_name(std::move(name_buffer));
  std::string metric_type_str(tmp_type);

  // track bad lines
  // using find() intead of == so that we can ignore trailing spaces
  bool bad_line = false;
  if ("" == metric_type_str) {
    bad_line = true;
  } else if (metric_type_str.find("ms") == 0) {
    type = MetricType::timer;
    try {
      metric_value = boost::lexical_cast<double>(metric_value_buffer);
    } catch(...) {
      bad_line = true;
    }
  } else if (metric_type_str.find("c") == 0) {
    type = MetricType::counter;
    try {
      metric_value = boost::lexical_cast<double>(metric_value_buffer);
    } catch(...) {
      bad_line = true;
    }
  } else if (metric_type_str.find("g") == 0) {
    type = MetricType::gauge;
    try {
      metric_value = boost::lexical_cast<double>(metric_value_buffer);
    } catch(...) {
      bad_line = true;
    }
  } else if (metric_type_str.find("s") == 0) {
    type = MetricType::set;
    metric_value = boost::lexical_cast<double>(metric_value_buffer);
    bad_line = false;
  } else {
    try {
      metric_value = boost::lexical_cast<double>(metric_value_buffer);
    } catch(...) {
      bad_line = true;
    }
  }

  if (bad_line) {
    ::logger->info("Bad line: " + metric);
    // setup the names for the stats stored in counters
    std::string bad_lines_seen = ::config->name + ".bad_lines_seen";
    addBadLine(bad_lines_seen);
    return;
  }

  // ++this->frequency[metric_name];
  auto metric_it = metrics.find(metric_name);
  if (metric_it != metrics.end())
  {
    if (type == MetricType::gauge)
    {
      Gauge *gauge = dynamic_cast<Gauge *>(metric_it->second.get());
      // check if +/- is specified
      char char_after_colon = metric_csty[metric.find_first_of(":") + 1];
      if (('+' == char_after_colon) || ('-' == char_after_colon))
      {
        gauge->incrementalUpdate(metric_value, sample_rate);
      }
      else
      {
        gauge->update(metric_value, sample_rate);
      }
    }
    else
    {
      buffer(metric_it->second, metric_value, sample_rate);
    }
  }
  else
  {
    switch (type)
    {
      case MetricType::timer:
      {
        if (::config->percentiles.empty())
        {
          // This timer cannot be used to compute percecntiles, but has a much smaller memory footprint
          auto ret = metrics.emplace(metric_name, std::shared_ptr<Metric>(new LeanTimer));
          ret.first->second->update(metric_value, sample_rate);
        }
        else {
          auto ret = metrics.emplace(metric_name, std::shared_ptr<Metric>(new FullTimer));
          ret.first->second->update(metric_value, sample_rate);
        }
        break;
      }
      case MetricType::gauge:
      {
        auto ret = metrics.emplace(metric_name, std::shared_ptr<Metric>(new Gauge));
        ret.first->second->update(metric_value, sample_rate);
        break;
      }
      case MetricType::set:
      {
        auto ret = metrics.emplace(metric_name, std::shared_ptr<Metric>(new Set));
        ret.first->second->update(metric_value, sample_rate);
        break;
      }
      default:
        auto ret = metrics.emplace(metric_name, std::shared_ptr<Metric>(new Counter));
        ret.first->second->update(metric_value, sample_rate);
        break;
    }
  }
}

void Ledger::buffer(const std::shared_ptr<Metric> &metric, double metric_value, double sample_rate)
{
  metric->update(metric_value, sample_rate);
}

std::shared_ptr<Metric> Ledger::buffer(const std::string &metric_name, double metric_value,
                    const std::string &metric_type)
{
  MetricType type = MetricType::counter;
  double sample_rate = 1;

  // record number of calls to buffer method
  // ++this->statsd_metrics["metrics_processed"];

  // track bad lines
  bool bad_line = false;
  if (metric_type == "ms")
  {
    type = MetricType::timer;
  }
  else if (metric_type == "c")
  {
    type = MetricType::counter;
  }
  else if (metric_type == "g")
  {
    type = MetricType::gauge;
  }
  else if (metric_type == "s")
  {
    type = MetricType::set;
  }
  else
  {
    bad_line = true;
  }

  if (bad_line) {
    ::logger->info("Bad line: " + metric_name + "|" + metric_type);
    // setup the names for the stats stored in counters
    std::string bad_lines_seen = ::config->name + ".bad_lines_seen";
    addBadLine(bad_lines_seen);
    return nullptr;
  }

  // ++this->frequency[metric_name];

  auto metric_it = metrics.find(metric_name);
  if (metric_it != metrics.end())
  {
    buffer(metric_it->second, metric_value);
    return metric_it->second;
  }
  else
  {
  switch (type) {
    case MetricType::timer:
    {
      std::shared_ptr<Metric> ptr;
      if (::config->percentiles.empty())
      {
        // This timer cannot be used to compute percecntiles, but has a much smaller memory footprint
        ptr = std::make_shared<LeanTimer>();
      }
      else
      {
        ptr = std::make_shared<FullTimer>();
      }
      metrics[metric_name] = ptr;
      ptr->update(metric_value, sample_rate);
      return ptr;
    }
    case MetricType::gauge:
      {
        /// @todo +/- is specified
        auto ptr = std::make_shared<Gauge>();
        metrics[metric_name] = ptr;
        ptr->update(metric_value, sample_rate);
        return ptr;
      }
    case MetricType::set:
    {
      /// @todo sets
      auto ptr = std::make_shared<Set>();
      metrics[metric_name] = ptr;
      ptr->update(metric_value, sample_rate);
      return  ptr;
    }
    default:
      auto ptr = std::make_shared<Counter>();
      metrics[metric_name] = ptr;
      ptr->update(metric_value, sample_rate);
      return ptr;
  }
  }
}

// Aggregates buffered metrics
void Ledger::process() {
  std::uint64_t start_time = chrono::unixtime_ms();

  // process counters
  // calculate "per second" rate

  for (auto &m : this->metrics) {
    const std::string &key = m.first;
    Metric* metric = m.second.get();
    Counter* counter = dynamic_cast<Counter*>(metric);
    FullTimer* timer = dynamic_cast<FullTimer*>(metric);
    LeanTimer* lean_timer = dynamic_cast<LeanTimer*>(metric);
    if (counter)
    {
      counter->counter_rate_ = counter->counter_ / ::config->frequency;
    }
    else if (timer)
    {
    // process timers
    std::unordered_map<std::string, double> &current_timer_data = timer->timer_data_;

    if (key.length() <= 0) {
      current_timer_data["count"] = current_timer_data["count_ps"] = 0;
    } else {
      // get sorted values
      std::vector<double> values(timer->timers_);
//      std::sort(values.begin(), values.end());

      // get count, sum, mean, min, and max
      int count = values.size();
      double sum = 0;
      auto minmax = std::minmax_element(values.begin(), values.end());
      double min = *minmax.first;
      double max = *minmax.second;
      for (auto value_itr = values.cbegin();
          value_itr != values.cend();
          ++value_itr) {
        sum += *value_itr;
      }
      double mean = sum / count;

      // initialize sum, mean, and threshold boundary
      double threshold_boundary = max;

      for (auto threshold_itr = ::config->percentiles.cbegin();
          threshold_itr != ::config->percentiles.cend();
          ++threshold_itr) {
        double pct =
          (*threshold_itr < 0) ? -(*threshold_itr) : (*threshold_itr);
         // Int is potentially too short since we could have more than 65K metrics in the flush window
         unsigned long num_in_threshold = 0;

        if (count > 1) {
          num_in_threshold = round(pct / 100 * count);
          if (0 == num_in_threshold) {
            continue;
          }
          threshold_boundary = values[num_in_threshold - 1];
        }

        double pct_sum = 0;
        for (unsigned long i = 0; i < num_in_threshold; ++i) {
            pct_sum += values[i];
        }
        double pct_mean = pct_sum / num_in_threshold;

        // generate pct name
        char clean_pct[17] = {0};
        snprintf(clean_pct, sizeof(clean_pct), "%g", pct);
        for (unsigned int i = 0; i < strlen(clean_pct); ++i) {
          clean_pct[i] = ('.' == clean_pct[i]) ? '_' : clean_pct[i];
        }

        current_timer_data[std::string("upper_") + std::string(clean_pct)] = threshold_boundary;
        current_timer_data[std::string("count_") + std::string(clean_pct)] = num_in_threshold;
        current_timer_data[std::string("sum_") + std::string(clean_pct)] = pct_sum;
        current_timer_data[std::string("mean_") + std::string(clean_pct)] = pct_mean;

      }  // foreach percentile

      current_timer_data["upper"] = max;
      current_timer_data["lower"] = min;
      current_timer_data["count"] = timer->counter_;
      current_timer_data["mean"] = mean;
      timer->timer_data_ = current_timer_data;
    }
    }
    else if (lean_timer)
    {
      std::unordered_map<std::string, double> &current_timer_data = lean_timer->timer_data_;

      if (key.length() <= 0)
      {
        current_timer_data["count"] = current_timer_data["count_ps"] = 0;
      }
      else
      {
        current_timer_data["upper"] = (!std::isinf(lean_timer->max_) ?
                                           lean_timer->max_ :
                                           std::numeric_limits<double>::quiet_NaN());
        current_timer_data["lower"] = (!std::isinf(lean_timer->min_) ?
                                         lean_timer->min_ :
                                         std::numeric_limits<double>::quiet_NaN());
        current_timer_data["count"] = lean_timer->count_;
        current_timer_data["mean"] = lean_timer->sum_ / double(lean_timer->count_);
      }
      lean_timer->timer_data_ = current_timer_data;
    }
  }  // foreach metric
  this->statsd_metrics["processing_time"] =
    chrono::unixtime_ms() - start_time;
}

void Ledger::setProcTime(std::int64_t value)
{
  this->statsd_metrics["metrics_processing_time"] =
      this->statsd_metrics["metrics_processing_time"] + value;
}

void Ledger::addBadLine(const std::string &name)
{
  // Only emplaces first time
  auto metric_it = metrics.find(name);
  if (metric_it != metrics.end())
  {
    metric_it->second->update(1, 1);
  }
  else
  {
    auto ret = metrics.emplace(name, std::shared_ptr<Metric>(new Counter));
    ret.first->second->update(1, 1);
  }
}

}  // namespace statsdcc


#ifndef __STAGEWISE_METHOD_H__
#define __STAGEWISE_METHOD_H__

#include <string>
#include <vector>

#include <armadillo>

#include <besiq/method/method.hpp>
#include <besiq/stats/log_scale.hpp>
#include <besiq/stats/closed_form_models.hpp>

/**
 * This class is responsible for intializing and repeatedly
 * executing a stepwise method where p-values for a series
 * of more complex models are computed.
 */
class stagewise_method
: public method_type
{
public:
    /**
     * Constructor.
     *
     * @param data Additional data required by all methods, such as
     *             covariates.
     */
    stagewise_method(method_data_ptr data, const std::string &model);

    /**
     * @see method_type::init.
     */
    virtual std::vector<std::string> init();

    /**
     * Returns the number of usable samples.
     *
     * @param row1 The first snp.
     * @param row2 The second snp.
     * @param phenotype The phenotype.
     *
     * @return The number of usable samples.
     */
    unsigned int num_ok_samples(const snp_row &row1, const snp_row &row2, const arma::vec &phenotype);

    /**
     * Computes a likelihood ratio statistics that measures
     * whether two SNPs are in linkage disequilibrium or not. It
     * should follow a chi^2(1) distribution if there is no LD.
     *
     * @param row1 The first snp.
     * @param row2 The second snp.
     *
     * @return The estimated likelihood ratio statistic.
     */
    double compute_ld_lr(const snp_row &row1, const snp_row &row2);

    /**
     * @see method_type::run.
     */
    virtual double run(const snp_row &row1, const snp_row &row2, float *output);

private:
    /**
     * Type of model.
     */
    std::string m_model;
    /**
     * A weight > 0 associated with each sample, that allows for
     * covariate adjustment.
     */
    arma::vec m_weight;

    /**
     * The models used.
     */
    std::vector<closed_form_model *> m_models;

    /**
     * Number of usable samples in the last call to run.
     */
    unsigned int m_num_usable_samples;
};

#endif /* End of __STAGEWISE_METHOD_H__ */


//
//  occupancy_grid.cpp
//  robotcontrol2
//
//  Created by Sebastian Theophil on 24.09.15.
//  Copyright © 2015 Sebastian Theophil. All rights reserved.
//

#include "occupancy_grid.h"
#include "robot_configuration.h" // Data structures and configuration data shared with Arduino controller
#include "error_handling.h"

#include <cstdlib>
#include <cmath>
#include <algorithm>

#include <boost/range/iterator_range.hpp>
#include <boost/range/size.hpp>
#include <opencv2/imgproc.hpp>

template<typename Derived>
COccupancyGridBaseT<Derived>::COccupancyGridBaseT()
:   m_matfMapLogOdds(c_nMapExtent, c_nMapExtent, CV_32FC1, cv::Scalar(0.0f))
{}

template<typename Derived>
COccupancyGridBaseT<Derived>::COccupancyGridBaseT(COccupancyGridBaseT<Derived> const& occgrid)
:   m_matfMapLogOdds(occgrid.m_matfMapLogOdds.clone())
{}

template<typename Derived>
COccupancyGridBaseT<Derived>::COccupancyGridBaseT(COccupancyGridBaseT<Derived>&& occgrid)
:   m_matfMapLogOdds(occgrid.m_matfMapLogOdds)
{}

template<typename Derived>
COccupancyGridBaseT<Derived>& COccupancyGridBaseT<Derived>::operator=(COccupancyGridBaseT<Derived> const& occgrid) {
    m_matfMapLogOdds=occgrid.m_matfMapLogOdds.clone();
    return *this;
}

template<typename Derived>
COccupancyGridBaseT<Derived>& COccupancyGridBaseT<Derived>::operator=(COccupancyGridBaseT<Derived>&& occgrid) {
    m_matfMapLogOdds=occgrid.m_matfMapLogOdds;
    return *this;
}

template<typename Derived>
bool COccupancyGridBaseT<Derived>::occupied(rbt::point<int> const& pt) const {
	assert(is_inside(pt));
    return c_fFreeThreshold<m_matfMapLogOdds.at<float>(pt.y, pt.x);
}

template<typename Derived>
bool COccupancyGridBaseT<Derived>::is_inside(rbt::point<int> const& pt) const {
	auto const sz = m_matfMapLogOdds.size();
	return 0<=pt.x && 0<=pt.y && pt.x < sz.width && pt.y < sz.height;
}

template<typename Derived>
void COccupancyGridBaseT<Derived>::internalUpdatePerObstacle(rbt::point<double> const& ptf, rbt::point<double> const& ptfObstacle) {
    cv::LineIterator itpt(
        m_matfMapLogOdds,
        ToGridCoordinate(ptf),
        ToGridCoordinate(ptfObstacle)
    );
    for(int i = 0; i < itpt.count; i++, ++itpt) {
        auto const pt = rbt::point<int>(itpt.pos());
        assert(is_inside(pt)); // LineIterator clips to line boundaries

        auto const fDeltaValue = i<itpt.count-1
            ? c_fFreeDelta // free
            : c_fOccupiedDelta; // occupied

        auto const fOdds = m_matfMapLogOdds.at<float>(pt.y, pt.x) + fDeltaValue;
        m_matfMapLogOdds.at<float>(pt.y, pt.x) = fOdds;

        static_cast<Derived*>(this)->updateGrid(pt, fOdds);
    }
}

template<typename Derived>
void COccupancyGridBaseT<Derived>::internalUpdatePerPose(rbt::pose<double> const& pose) {
    static_cast<Derived*>(this)->updateGridPoly(
    	RenderRobotPose(m_matfMapLogOdds, pose, cv::Scalar(c_fOccupancyRover)),
    	c_fOccupiedDelta
    );
}

template<typename Derived>
void COccupancyGridBaseT<Derived>::update(rbt::pose<double> const& pose, double fRadAngle, int nDistance) {
    internalUpdatePerObstacle(pose.m_pt, Obstacle(pose, fRadAngle, nDistance));
    internalUpdatePerPose(pose);
}

template<typename Derived>
void COccupancyGridBaseT<Derived>::update(rbt::pose<double> const& pose, std::vector<rbt::point<double>> const& vecptf) {
    boost::for_each(vecptf, [&](auto const& ptf) {
        internalUpdatePerObstacle(pose.m_pt, ptf);
    });
    internalUpdatePerPose(pose);
}


#include <boost/regex.hpp>
#include <boost/lexical_cast.hpp>

#include <fstream>

#include "atomistic/fundamental.hpp"
#include "atomistic/units.hpp"
#include "io.hpp"
#include "la.hpp"
#include "types.hpp"

namespace atomistic {

typedef boost::sregex_iterator RegIt;
typedef boost::smatch Match;
typedef const boost::regex Regex;

using boost::lexical_cast;
using namespace types;

EnergyLevels::EnergyLevels(const EnergyLevels& e){
    this->levels = e.levels;
    this->fermi = e.fermi;
}

EnergyLevels::EnergyLevels(std::vector<types::Real> levels, types::Real fermi){
    this->levels = levels;
    this->fermi = fermi;
    this->sort();
}

EnergyLevels & EnergyLevels::operator*=(Real factor){
    for(std::vector<Real>::iterator it = this->levels.begin(); it != this->levels.end(); it++) {
        *it *= factor;
    }
    fermi *= factor;
    return *this;
}

void EnergyLevels::sort() {
    std::sort(this->levels.begin(), this->levels.end());
}

Uint EnergyLevels::count() const {
    return this->levels.size();
}

Uint EnergyLevels::countOccupied() const {
    Uint n = 0;
    for(std::vector<Real>::const_iterator it = this->levels.begin();
            it != this->levels.end(); it++) {
        if(*it < fermi) ++n;
    }
    return n;
}

void EnergyLevels::join(EnergyLevels e2){
        e2.shift(fermi - e2.fermi);
        levels.insert( levels.end(), e2.levels.begin(), e2.levels.end());
        this->sort();
}

bool EnergyLevels::rangeCheck(Uint i) const {
    if( i <= this->levels.size()) return true;
#ifdef ATOMISTIC_FUNDAMENTAL_STRICT
    else throw std::range_error("Level index out of bounds.");
#else
    else return false;
#endif
}

void EnergyLevels::print() const {
    std::cout << "Fermi energy: " << this->fermi << std::endl;
    for(std::vector<Real>::const_iterator it = this->levels.begin(); it != this->levels.end(); it++) {
        std::cout << *it << std::endl;
    }
}

void EnergyLevels::shift(Real deltaE) {
    for(std::vector<Real>::iterator it = this->levels.begin(); it != this->levels.end(); it++) {
        *it += deltaE;
    }
    fermi += deltaE;
}

void EnergyLevels::setFermiZero() {
    this->shift(-fermi);
}

Real Atom::distance(const  Atom& a) const {
    std::vector<Real>::const_iterator
        it  = this->coordinates.begin(),
        end = this->coordinates.end(),
        ait = a.coordinates.begin();
    Real sum = 0;
    while(it != end){
        sum += (*it - *ait) * (*it - *ait);
        ++it; ++ait;
    }
    return sqrt(sum);
}

//LDOS::LDOS(
//        const std::vector< std::vector<Real> > &levels,
//        Real eMin,
//        Ream eMax,
//        Real deltaE,
//        Real broadening){
//
//    // Setup grid
//    Real e = eMin;
//    std::vector<Real> energies;
//    for(Real e = eMin; e <= eMax; e += deltaE){
//        energies.push_back(e);
//    }
//    std>>vector<Real> densities(energies.size(), 0.0);
//
//    // Gaussian stuff
//    // \$ \frac{1}{\sigma \sqrt{2\pi} e^{-\frac{(x-\mu)^2}{2\sigma^2}} \$
//    // = a e^{c(x-b)^2}
//    types::Real a = 1.0/(broadening * std::sqrt(2 * M_PI));
//    types::Real c = 1.0/(2 * broadening * broadening);
//
//
//    std::vector< std::vector<Real > >::const_iterator levelIt = levels.begin(),
//        levelEnd = levels.end;
//    while(levelIt != levelEnd){
//        std::vector<Real>::const_iterator eIt = energies.begin(),
//            eEnd = energies.end();
//        for(Uint i = 0; i < energies.size(); ++i){
//            densities[i] += *levelIt[1] * a * std::exp(c * (energies[i]-*levelIt[0]));
//        }
//        ++levelIt;
//    }
//
//    this->densities = densities;
//    this->energies = energies;
//}

}


#include <scitbx/array_family/boost_python/flex_fwd.h>

#include <scitbx/array_family/sort.h>
#include <boost/python/def.hpp>
#include <boost/python/args.hpp>

namespace scitbx { namespace af { namespace boost_python {

namespace {

  template <typename ElementType>
  void
  sort_permutation_wrapper()
  {
    using namespace boost::python;
    def("sort_permutation",
      // cast needed for Intel C++ 9.1
      (shared<std::size_t>(*)(const_ref<ElementType> const&, bool, bool))
        af::sort_permutation, (
          arg("data"), arg("reverse")=false, arg("stable")=false));
  }

} // namespace <anonymous>

  void wrap_flex_sort()
  {
    sort_permutation_wrapper<int>();
    sort_permutation_wrapper<std::size_t>();
    sort_permutation_wrapper<double>();
  }

}}} // namespace scitbx::af::boost_python


#include <iostream>
#include <string>
#include <boost/xpressive/xpressive.hpp>
#include <boost/xpressive/regex_actions.hpp>
#include <map>
#include <utility>

#include "VECTOR.h"
#include "object.h"
#include "shape.h"
#include "globals.h"
#include "environment.h"
#include "heightfield.h"
#include "entity.h"

using namespace boost::xpressive;

#include "env_functions.h"

sregex OBJECTS_STRUCT;
sregex OBJECTS;
sregex OBJECT;
sregex BOX1_X, BOX2_X;
sregex SPHERE_X;
sregex STAIRS_X;
sregex PLANE_X;
sregex CUSTOM_X;
sregex HEIGHTFIELD_X;
sregex FRICTION;
sregex DEFINITION;
sregex DEFINITIONS;
sregex ENTITIES;
sregex ENTITY;
sregex NAME;

void Environment::setFromFile(std::string fname)
{
	std::ifstream ifs;
	ifs.open(fname.c_str());
	if (!ifs)
	{
		std::cerr << "Cannot open controller script: " << fname << std::endl;
		return;
	}
	std::string input( (std::istreambuf_iterator<char>(ifs) ), (std::istreambuf_iterator<char>()));

    std::vector<obj_info*>* obj_info_ptr;

	sregex INT, FLOAT, COMMENT, FNAME, WORD;

	COMMENT = "#" >> *~_n >> _n;
	INT = sregex::compile("-?\\d+\\.?");
	FLOAT = sregex::compile("-?\\d*\\.\\d+");
    FNAME = +_w >> ".hf";
    WORD = +_w;
    NAME = +_w;

	BOX1_X = ("Box" >> *space >> "(" >> *space >> (s1= INT|FLOAT) >> *space >> ","
 	 	     >> *space >> (s2= INT|FLOAT) >> *space >> ","
		     >> *space >> (s3= INT|FLOAT) >> *space >> ","
		     >> *space >> (s4= INT|FLOAT) >> *space >> ","
		     >> *space >> (s5= INT|FLOAT) >> *space >> ")")
		[ pushBox1(boost::xpressive::ref(obj_info_ptr), as<float>(s1), as<float>(s2), as<float>(s3), as<float>(s4), as<float>(s5) ) ];

	BOX2_X = ("Box" >> *space >> "(" >> *space >> (s1= INT|FLOAT) >> *space >> ","
 	 	     >> *space >> (s2= INT|FLOAT) >> *space >> ","
		     >> *space >> (s3= INT|FLOAT) >> *space >> ","
		     >> *space >> (s4= INT|FLOAT) >> *space >> ","
		     >> *space >> (s5= INT|FLOAT) >> *space >> ","
		     >> *space >> (s6= INT|FLOAT) >> *space >> ","
		     >> *space >> (s7= INT|FLOAT) >> *space >> ")")
		[ pushBox2(boost::xpressive::ref(obj_info_ptr), as<float>(s1), as<float>(s2), as<float>(s3), as<float>(s4), as<float>(s5), as<float>(s6), as<float>(s7) ) ];

	SPHERE_X = ("Sphere" >> *space >> "(" >> *space >> (s1= INT|FLOAT) >> *space >> ","
 	 	     >> *space >> (s2= INT|FLOAT) >> *space >> ","
		     >> *space >> (s3= INT|FLOAT) >> *space >> ","
		     >> *space >> (s4= INT|FLOAT) >> *space >> ","
		     >> *space >> (s5= INT|FLOAT) >> *space >> ")")
		[ pushSphere(boost::xpressive::ref(obj_info_ptr), as<float>(s1), as<float>(s2), as<float>(s3), as<float>(s4), as<float>(s5) ) ];

	STAIRS_X = ("Stairs" >> *space >> "(" >> *space >> (s1= INT|FLOAT) >> *space >> ","
			>> *space >> (s2= INT|FLOAT) >> *space >> ","
			>> *space >> (s3= INT|FLOAT) >> *space >> ","
			>> *space >> (s4= INT) >> *space >> ")")
		[ pushStairs(boost::xpressive::ref(obj_info_ptr), as<float>(s1), as<float>(s2), as<float>(s3), as<int>(s4)) ];

	PLANE_X = ("2DPlane" >> *space >> "(" >> *space >> (s1= INT|FLOAT) >> *space >> ","
			>> *space >> (s2= INT|FLOAT) >> *space >> ","
			>> *space >> (s3= INT|FLOAT) >> *space >> ","
			>> *space >> (s4= INT|FLOAT) >> *space >> ")")
		[ pushPlane(boost::xpressive::ref(obj_info_ptr), as<float>(s1), as<float>(s2), as<float>(s3), as<float>(s4)) ];

    HEIGHTFIELD_X = ("Heightfield" >> *space >> "(" >> *space >> (s1 = FNAME) >> *space >> ")" ) [ pushHeightfield(boost::xpressive::ref(obj_info_ptr), as<std::string>(s1)) ];

    obj_info* custom = new obj_info;

    std::map<std::string, obj_info*> customMap;

    CUSTOM_X = ((s1= WORD) >> *space >> "(" >> *space >> (s2= INT|FLOAT) >> *space >> "," >> *space >> (s3= INT|FLOAT) >> *space >> "," >> *space >> (s4= INT|FLOAT) >> *space >> ")") [ pushCustom(boost::xpressive::ref(obj_info_ptr), boost::xpressive::ref(customMap), as<std::string>(s1), as<float>(s2), as<float>(s3), as<float>(s4)) ];

	FRICTION = (".friction" >> *space >> "(" >> *space >> (s1= INT|FLOAT) >> *space >> ")") [ setFriction(boost::xpressive::ref(obj_info_ptr), as<float>(s1)) ];

	OBJECT = (BOX1_X
		| BOX2_X
		| SPHERE_X
		| STAIRS_X
        | PLANE_X
        | HEIGHTFIELD_X
        | CUSTOM_X ) >> !(FRICTION);

    std::vector<std::string> names;

    ENTITY = (s1= NAME) [ push_back(boost::xpressive::ref(names), as<std::string>(s1)) ] >> *space >> "=" >> *space >> OBJECT;

	OBJECTS = ( ((OBJECT >> *space >> ";")|COMMENT) >> *space >> by_ref(OBJECTS))
		| OBJECT >> *space >> ";"
		| COMMENT;

    ENTITIES = ( ((ENTITY >> *space >> ";")|COMMENT) >> *space >> by_ref(ENTITIES))
        | ENTITY >> *space >> ";"
        | COMMENT;

	OBJECTS_STRUCT = ("ENTITIES" >> *space >> "{" >> *space >> ENTITIES >> *space >> "}") [ fillObjects(boost::xpressive::ref(names), boost::xpressive::ref(entities), boost::xpressive::ref(obj_info_ptr)) ];

    //map of custom object names to their obj_info ptr's.
//    custom->children = new std::vector<obj_info*>;
//    obj_info_ptr = custom->children;

    // fills in custom's children thru OBJECTS call, then calls createCustom to name it
    DEFINITION = ("DEFINE" >> *space >> (s1= WORD) [ newObjInfVecPtr(boost::xpressive::ref(obj_info_ptr)) ] >> *space >> "{" >> *space >> OBJECTS >> *space >> "}") [ createCustom(as<std::string>(s1), boost::xpressive::ref(obj_info_ptr), boost::xpressive::ref(customMap)) ] ;

    DEFINITIONS = (DEFINITION >> *space >> by_ref(DEFINITION))
            | DEFINITION;

	smatch what;
	if (!regex_search(input, what, DEFINITIONS))
	{
	//	return;
	}

    std::vector<obj_info*> obj_info_vec;
    obj_info_ptr = &obj_info_vec;

	if (!regex_search(input, what, OBJECTS_STRUCT))
	{
		std::cerr << "Encountered syntax error while reading in script " << fname << std::endl;
		return;
	}

	return;

}

void Environment::addEntity(Entity* e)
{

    entities.insert(std::pair<std::string, Entity*>(e->getName(), e));

}
void Environment::deleteEntity(std::string s)
{

    delete entities.find(s)->second;

    entities.erase(s);

}


#pragma once

#include <boost/sml.hpp>
#include <bitset>
#include <iostream>
#include <vector>

static bool printDebug = false;
static const int decoderwidth = 12;

struct decoderOutput
{
  std::bitset<decoderwidth> shortPress{};
  std::bitset<decoderwidth> longPress{};
  std::bitset<decoderwidth> doubleShortPress{};
  std::bitset<decoderwidth> doubleLongPress{};
  std::bitset<decoderwidth> shortLongPress{};
  std::bitset<decoderwidth> longShortPress{};

  std::vector<unsigned long> shortPressTime = std::vector<unsigned long>(decoderwidth, 0);
};

struct touchDecoderTimingConfig
{
  unsigned long minReleaseTime{};
  unsigned long shortPressTime{};
  unsigned long longPressTime{};
  unsigned long maxIdleShortTime{};
  unsigned long maxIdleLongTime{};
};

class state_machine
{
  using Self = state_machine;

public:
  explicit state_machine(const touchDecoderTimingConfig &tc) : tc{tc} {}

  // States
  struct Idle
  {
  };

  struct IdleShort
  {
    unsigned long entry_time{};
    int keyid{};
    unsigned long press_time{};
  };

  struct IdleLong
  {
    unsigned long entry_time{};
    int keyid{};
  };

  struct Touched
  {
    unsigned long entry_time{};
  };

  struct TouchedShort
  {
    unsigned long entry_time{};
  };

  struct TouchedLong
  {
    unsigned long entry_time{};
  };

  // Events
  struct no_event
  {
    int keyid{};
    unsigned long event_time{};
  };

  struct touch_event
  {
    int keyid{};
    unsigned long event_time{};
  };
  struct release_event
  {
    int keyid{};
    unsigned long event_time{};
  };
  // Transition table
  auto operator()() const
  {
    using namespace boost::sml;

    auto guard_min_touch_time = [this](const auto &event, Touched &state) {
      bool timeout = false;
      unsigned long timeDiff = event.event_time - state.entry_time;
      if (timeDiff < tc.minReleaseTime)
      {
        timeout = true;
      }
      if (printDebug)
        std::cout << "Key " << event.keyid << " - guard_min_touch_time (" << timeDiff << ") = " << timeout << std::endl;
      return timeout;
    };
    auto guard_short_press_time = [this](const auto &event, decoderOutput &dcOutput, Touched &state, IdleShort &dstState) {
      bool result = false;
      unsigned long timeDiff = event.event_time - state.entry_time;
      if (timeDiff >= tc.minReleaseTime && timeDiff < tc.shortPressTime)
      {
        dstState.press_time = timeDiff;
        result = true;
      }
      if (printDebug)
        std::cout << "Key " << event.keyid << " - guard_short_press_time (" << timeDiff << ") = " << result << std::endl;
      return result;
    };
    auto guard_long_press_time = [this](const auto &event, Touched &state) {
      bool result = false;
      unsigned long timeDiff = event.event_time - state.entry_time;
      if (timeDiff >= tc.longPressTime)
      {
        result = true;
      }
      if (printDebug)
        std::cout << "Key " << event.keyid << " - guard_long_press_time (" << timeDiff << ") = " << result << std::endl;
      return result;
    };

    auto guard_short_min_touch_time = [this](const auto &event, TouchedShort &state) {
      bool timeout = false;
      unsigned long timeDiff = event.event_time - state.entry_time;
      if (timeDiff < tc.minReleaseTime)
      {
        timeout = true;
      }
      if (printDebug)
        std::cout << "Key " << event.keyid << " - guard_short_min_touch_time (" << timeDiff << ") = " << timeout << std::endl;
      return timeout;
    };
    auto guard_short_short_press_time = [this](const auto &event, TouchedShort &state) {
      bool result = false;
      unsigned long timeDiff = event.event_time - state.entry_time;
      if (timeDiff >= tc.minReleaseTime && timeDiff < tc.shortPressTime)
      {
        result = true;
      }
      if (printDebug)
        std::cout << "Key " << event.keyid << " - guard_short_short_press_time (" << timeDiff << ") = " << result << std::endl;
      return result;
    };
    auto guard_short_long_press_time = [this](const auto &event, TouchedShort &state) {
      bool result = false;
      unsigned long timeDiff = event.event_time - state.entry_time;
      if (timeDiff >= tc.longPressTime)
      {
        result = true;
      }
      if (printDebug)
        std::cout << "Key " << event.keyid << " - guard_short_long_press_time (" << timeDiff << ") = " << result << std::endl;
      return result;
    };

    auto guard_long_min_touch_time = [this](const auto &event, TouchedLong &state) {
      bool timeout = false;
      unsigned long timeDiff = event.event_time - state.entry_time;
      if (timeDiff < tc.minReleaseTime)
      {
        timeout = true;
      }
      if (printDebug)
        std::cout << "Key " << event.keyid << " - guard_long_min_touch_time (" << timeDiff << ") = " << timeout << std::endl;
      return timeout;
    };
    auto guard_long_short_press_time = [this](const auto &event, TouchedLong &state) {
      bool result = false;
      unsigned long timeDiff = event.event_time - state.entry_time;
      if (timeDiff >= tc.minReleaseTime && timeDiff < tc.shortPressTime)
      {
        result = true;
      }
      if (printDebug)
        std::cout << "Key " << event.keyid << " - guard_long_short_press_time (" << timeDiff << ") = " << result << std::endl;
      return result;
    };
    auto guard_long_long_press_time = [this](const auto &event, TouchedLong &state) {
      bool result = false;
      unsigned long timeDiff = event.event_time - state.entry_time;
      if (timeDiff >= tc.longPressTime)
      {
        result = true;
      }
      if (printDebug)
        std::cout << "Key " << event.keyid << " - guard_long_long_press_time (" << timeDiff << ") = " << result << std::endl;
      return result;
    };

    auto guard_max_idleshort_time = [this](const auto &event, IdleShort &state, unsigned long &currentTime) {
      bool timeout = false;
      unsigned long timeDiff = currentTime - state.entry_time;
      if (timeDiff > tc.maxIdleShortTime)
      {
        timeout = true;
      }
      if (printDebug)
        std::cout << "Key " << state.keyid << " - guard_max_idleshort_time (" << timeDiff << ") = " << timeout << std::endl;
      return timeout;
    };

    auto guard_max_idlelong_time = [this](const auto &event, IdleLong &state, unsigned long &currentTime) {
      bool timeout = false;
      unsigned long timeDiff = currentTime - state.entry_time;
      if (timeDiff > tc.maxIdleLongTime)
      {
        timeout = true;
      }
      if (printDebug)
        std::cout << "Key " << state.keyid << " - guard_max_idlelong_time (" << timeDiff << ") = " << timeout << std::endl;
      return timeout;
    };

    const auto touch_action = [](const auto &event, Touched &state) {
      if (printDebug)
        std::cout << "Key " << event.keyid << " - Transition due to touch_event" << std::endl;
      state.entry_time = event.event_time;
    };
    const auto touch_short_action = [](const auto &event, TouchedShort &state) {
      if (printDebug)
        std::cout << "Key " << event.keyid << " - Transition due to touch_short_event" << std::endl;
      state.entry_time = event.event_time;
    };
    const auto touch_long_action = [](const auto &event, TouchedLong &state) {
      if (printDebug)
        std::cout << "Key " << event.keyid << " - Transition due to touch_long_event" << std::endl;
      state.entry_time = event.event_time;
    };

    const auto release_to_idle_short = [](decoderOutput &dcOutput, IdleShort &state, const auto &event) {
      if (printDebug)
        std::cout << "Key " << event.keyid << " - Transition due to release_to_idle_short" << std::endl;
      state.entry_time = event.event_time;
      state.keyid = event.keyid;
    };

    const auto release_to_idle_long = [](decoderOutput &dcOutput, IdleLong &state, const auto &event) {
      if (printDebug)
        std::cout << "Key " << event.keyid << " - Transition due to release_to_idle_long" << std::endl;
      state.entry_time = event.event_time;
      state.keyid = event.keyid;
    };

    const auto release_short_action = [](decoderOutput &dcOutput, IdleShort &state, const auto &event) {
      if (printDebug)
        std::cout << "Key " << state.keyid << " - Transition due to release_short_action" << std::endl;
      dcOutput.shortPressTime[state.keyid] = state.press_time;
      dcOutput.shortPress.set(state.keyid);
    };

    const auto release_double_short_action = [](decoderOutput &dcOutput, TouchedShort &state, const auto &event) {
      if (printDebug)
        std::cout << "Key " << event.keyid << " - Transition due to release_double_short_action" << std::endl;
      dcOutput.doubleShortPress.set(event.keyid);
    };

    const auto release_short_long_action = [](decoderOutput &dcOutput, TouchedShort &state, const auto &event) {
      if (printDebug)
        std::cout << "Key " << event.keyid << " - Transition due to release_short_long_action" << std::endl;
      dcOutput.shortLongPress.set(event.keyid);
    };

    const auto release_long_action = [](decoderOutput &dcOutput, IdleLong &state, const auto &event) {
      if (printDebug)
        std::cout << "Key " << state.keyid << " - Transition due to release_long_action" << std::endl;
      dcOutput.longPress.set(state.keyid);
    };

    const auto release_long_short_action = [](decoderOutput &dcOutput, TouchedLong &state, const auto &event) {
      if (printDebug)
        std::cout << "Key " << event.keyid << " - Transition due to release_long_short_action" << std::endl;
      dcOutput.longShortPress.set(event.keyid);
    };

    const auto release_double_long_action = [](decoderOutput &dcOutput, TouchedLong &state, const auto &event) {
      if (printDebug)
        std::cout << "Key " << event.keyid << " - Transition due to release_double_long_action" << std::endl;
      dcOutput.doubleLongPress.set(event.keyid);
    };

    const auto min_touch_time_cancel_action = [](decoderOutput &dcOutput, const auto &event) {
      if (printDebug)
        std::cout << "Key " << event.keyid << " - Transition due to min_touch_time_cancel_action" << std::endl;
    };

    // clang-format off
    return make_transition_table(
        * state<Idle>          + event<touch_event>                                  / (touch_action)                = state<Touched>
        , state<Touched>       + event<release_event>   [guard_min_touch_time]       / (min_touch_time_cancel_action)              = state<Idle>
        , state<Touched>       + event<release_event>   [guard_short_press_time]     / (release_to_idle_short)       = state<IdleShort>
        , state<Touched>       + event<release_event>   [guard_long_press_time]      / (release_to_idle_long)        = state<IdleLong>
        , state<IdleShort>                              [guard_max_idleshort_time]   / (release_short_action)        = state<Idle>
        , state<IdleLong>                               [guard_max_idlelong_time]    / (release_long_action)         = state<Idle>

        , state<IdleShort>     + event<touch_event>                                  / (touch_short_action)               = state<TouchedShort>
        , state<IdleLong>      + event<touch_event>                                  / (touch_long_action)                = state<TouchedLong>

        , state<TouchedShort>  + event<release_event>   [guard_short_min_touch_time]       / (min_touch_time_cancel_action)              = state<Idle>
        , state<TouchedShort>  + event<release_event>   [guard_short_short_press_time]     / (release_double_short_action) = state<Idle>
        , state<TouchedShort>  + event<release_event>   [guard_short_long_press_time]      / (release_short_long_action)   = state<Idle>

        , state<TouchedLong>  + event<release_event>   [guard_long_min_touch_time]       / (min_touch_time_cancel_action)               = state<Idle>
        , state<TouchedLong>  + event<release_event>   [guard_long_short_press_time]     / (release_long_short_action)    = state<Idle>
        , state<TouchedLong>  + event<release_event>   [guard_long_long_press_time]      / (release_double_long_action)   = state<Idle>
    );
    // clang-format on
  }

  touchDecoderTimingConfig tc{};
};

class TouchDecoder
{
public:
  TouchDecoder() : statemachine(state_machine{touchDecoderTimingConfig{30, 300, 300, 500, 1000}})
  {
    for (int i = 0; i < decoderwidth; i++)
    {
      vsm.push_back(new boost::sml::sm<state_machine>{statemachine, state_machine::Touched{0}, output, currentTime});
    }
  }

  TouchDecoder(touchDecoderTimingConfig);

  ~TouchDecoder()
  {
    for (int i = 0; i < decoderwidth; i++)
    {
      delete vsm[i];
    }
  }

  state_machine statemachine;

  std::vector<boost::sml::sm<state_machine> *> vsm;

  void push(uint16_t touchstate, unsigned long pushTime)
  {
    std::bitset<decoderwidth> newState = std::bitset<decoderwidth>(touchstate);
    _released = current & ~newState;
    _touched = ~current & newState;

    // reset output event registers to zero
    output = decoderOutput{};

    if (printDebug)
      std::cout << "PRE-SM  - current: " << current
                << ", touchstate: " << touchstate
                << ", _released: " << _released
                << ", _touched: " << _touched
                << ", shortPress: " << output.shortPress
                << ", longPress: " << output.longPress
                << ", doubleShort: " << output.doubleShortPress
                << ", shortLongPress: " << output.shortLongPress
                << ", longShortPress: " << output.longShortPress
                << ", doubleLong: " << output.doubleLongPress
                << std::endl;

    current = newState;
    currentTime = pushTime;

    for (int i = 0; i < decoderwidth; i++)
    {
      if (_touched.test(i))
      {
        vsm[i]->process_event(state_machine::touch_event{i, pushTime});
      }
    }

    for (int i = 0; i < decoderwidth; i++)
    {
      if (_released.test(i))
      {
        vsm[i]->process_event(state_machine::release_event{i, pushTime});
      }
    }

    for (int i = 0; i < decoderwidth; i++)
    {
      if (~_released.test(i) && ~_touched.test(i))
      {
        vsm[i]->process_event(state_machine::no_event{i, pushTime});
      }
    }
    if (printDebug)
      std::cout << "POST-SM - current: " << current
                << ", touchstate: " << touchstate
                << ", _released: " << _released
                << ", _touched: " << _touched
                << ", shortPress: " << output.shortPress
                << ", longPress: " << output.longPress
                << ", doubleShort: " << output.doubleShortPress
                << ", shortLongPress: " << output.shortLongPress
                << ", longShortPress: " << output.longShortPress
                << ", doubleLong: " << output.doubleLongPress
                << std::endl;
    int j = 3;
  }

  std::bitset<decoderwidth> released()
  {
    return _released;
  }

  std::bitset<decoderwidth> shortPress()
  {
    return output.shortPress;
  }

  std::vector<unsigned long> shortPressTime()
  {
    return output.shortPressTime;
  }

  std::bitset<decoderwidth> longPress()
  {
    return output.longPress;
  }

  std::bitset<decoderwidth> doubleShortPress()
  {
    return output.doubleShortPress;
  }

  std::bitset<decoderwidth> doubleLongPress()
  {
    return output.doubleLongPress;
  }

  std::bitset<decoderwidth> shortLongPress()
  {
    return output.shortLongPress;
  }
  std::bitset<decoderwidth> longShortPress()
  {
    return output.longShortPress;
  }
  int width()
  {
    return decoderwidth;
  }

private:
  unsigned long currentTime;
  int empMaxTouched; // max count of touches in buffer. If exceeded will clear buffer and set empDetected = true
  int _empTouchCount = 0;
  int bufSize;

  touchDecoderTimingConfig tc{10, 20};

  decoderOutput output;
  std::bitset<decoderwidth> current = 0;
  std::bitset<decoderwidth> _released;
  std::bitset<decoderwidth> _touched;

  // actions
};

TouchDecoder::TouchDecoder(touchDecoderTimingConfig tc) : statemachine{tc}
{
  for (int i = 0; i < decoderwidth; i++)
  {
    vsm.push_back(new boost::sml::sm<state_machine>{statemachine, state_machine::Touched{0}, output, currentTime});
  }
}


#include <geometry_msgs/PoseStamped.h>
#include <kr_mav_msgs/SO3Command.h>
#include <nav_msgs/Odometry.h>
#include <nodelet/nodelet.h>
#include <ros/ros.h>
#include <rosflight_msgs/Command.h>
#include <sensor_msgs/Imu.h>
#include <std_msgs/Float64.h>
#include <tf/transform_datatypes.h>

#include <Eigen/Geometry>

class SO3CmdToRosflight : public nodelet::Nodelet
{
 public:
  void onInit(void);

 private:
  void so3_cmd_callback(const kr_mav_msgs::SO3Command::ConstPtr &msg);
  void odom_callback(const nav_msgs::Odometry::ConstPtr &odom);
  void imu_callback(const sensor_msgs::Imu::ConstPtr &pose);

  bool odom_set_, imu_set_, so3_cmd_set_;
  Eigen::Quaterniond odom_q_, imu_q_;
  int num_props_;
  double max_prop_force_;

  ros::Publisher command_pub_;

  ros::Subscriber so3_cmd_sub_;
  ros::Subscriber odom_sub_;
  ros::Subscriber imu_sub_;

  double so3_cmd_timeout_;
  ros::Time last_so3_cmd_time_;
  kr_mav_msgs::SO3Command last_so3_cmd_;
};

void SO3CmdToRosflight::odom_callback(const nav_msgs::Odometry::ConstPtr &odom)
{
  if(!odom_set_)
    odom_set_ = true;

  odom_q_ = Eigen::Quaterniond(odom->pose.pose.orientation.w, odom->pose.pose.orientation.x,
                               odom->pose.pose.orientation.y, odom->pose.pose.orientation.z);
}

void SO3CmdToRosflight::imu_callback(const sensor_msgs::Imu::ConstPtr &pose)
{
  if(!imu_set_)
    imu_set_ = true;

  imu_q_ = Eigen::Quaterniond(pose->orientation.w, pose->orientation.x, pose->orientation.y, pose->orientation.z);

  if(so3_cmd_set_ && ((ros::Time::now() - last_so3_cmd_time_).toSec() >= so3_cmd_timeout_))
  {
    ROS_INFO("so3_cmd timeout. %f seconds since last command", (ros::Time::now() - last_so3_cmd_time_).toSec());
    const auto last_so3_cmd_ptr = boost::make_shared<kr_mav_msgs::SO3Command>(last_so3_cmd_);

    so3_cmd_callback(last_so3_cmd_ptr);
  }
}

void SO3CmdToRosflight::so3_cmd_callback(const kr_mav_msgs::SO3Command::ConstPtr &msg)
{
  if(!so3_cmd_set_)
    so3_cmd_set_ = true;

  // grab desired forces and rotation from so3
  const Eigen::Vector3d f_des(msg->force.x, msg->force.y, msg->force.z);

  const Eigen::Quaterniond q_des(msg->orientation.w, msg->orientation.x, msg->orientation.y, msg->orientation.z);

  tf::Quaternion q_des_tf =
      tf::Quaternion(msg->orientation.x, msg->orientation.y, msg->orientation.z, msg->orientation.w);

  // check psi for stability
  const Eigen::Matrix3d R_des(q_des);
  const Eigen::Matrix3d R_cur(odom_q_);

  const float Psi = 0.5f * (3.0f - (R_des(0, 0) * R_cur(0, 0) + R_des(1, 0) * R_cur(1, 0) + R_des(2, 0) * R_cur(2, 0) +
                                    R_des(0, 1) * R_cur(0, 1) + R_des(1, 1) * R_cur(1, 1) + R_des(2, 1) * R_cur(2, 1) +
                                    R_des(0, 2) * R_cur(0, 2) + R_des(1, 2) * R_cur(1, 2) + R_des(2, 2) * R_cur(2, 2)));

  if(Psi > 1.0f)  // Position control stability guaranteed only when Psi < 1
  {
    ROS_WARN_THROTTLE(1, "Psi %2.2f > 1.0, orientation error is too large!", Psi);
  }

  double force = f_des(0) * R_cur(0, 2) + f_des(1) * R_cur(1, 2) + f_des(2) * R_cur(2, 2);

  double throttle = force / (num_props_ * max_prop_force_);

  // clamp from 0.0 to 1.0
  throttle = std::min(1.0, throttle);
  throttle = std::max(0.0, throttle);

  if(!msg->aux.enable_motors)
    throttle = 0;

  // publish messages
  auto setpoint_msg = boost::make_shared<rosflight_msgs::Command>();
  setpoint_msg->header = msg->header;
  setpoint_msg->mode = rosflight_msgs::Command::MODE_ROLL_PITCH_YAWRATE_THROTTLE;
  setpoint_msg->ignore = rosflight_msgs::Command::IGNORE_NONE;

  double roll_sp, pitch_sp, yaw_sp;
  tf::Matrix3x3(q_des_tf).getRPY(roll_sp, pitch_sp, yaw_sp);

  // Require conversion to FRD frame for rosflight
  setpoint_msg->x = roll_sp;
  setpoint_msg->y = -pitch_sp;
  setpoint_msg->z = -msg->angular_velocity.z;  // yaw rate

  // No field in rosflight_msgs/Command for rate setpoints at same time as attitude?
  // setpoint_msg->mode = rosflight_msgs::Command::MODE_ROLLRATE_PITCHRATE_YAWRATE_THROTTLE;
  // setpoint_msg->x = msg->angular_velocity.x;
  // setpoint_msg->y = msg->angular_velocity.y;
  // setpoint_msg->z = msg->angular_velocity.z;

  setpoint_msg->F = throttle;

  command_pub_.publish(setpoint_msg);

  // save last so3_cmd
  last_so3_cmd_ = *msg;
  last_so3_cmd_time_ = ros::Time::now();
}

void SO3CmdToRosflight::onInit(void)
{
  ros::NodeHandle priv_nh(getPrivateNodeHandle());

  // get thrust scaling parameters
  if(priv_nh.getParam("num_props", num_props_))
    ROS_INFO("Got number of props: %d", num_props_);
  else
    ROS_ERROR("Must set num_props param");

  if(priv_nh.getParam("max_prop_force", max_prop_force_))
    ROS_INFO("Using max_prop_force=%g ", max_prop_force_);
  else
    ROS_ERROR("Must set max_prop_force param for thrust scaling");

  ROS_ASSERT_MSG(max_prop_force_ > 0, "max_prop_force must be positive. max_prop_force = %g", max_prop_force_);

  // get param for so3 command timeout duration
  priv_nh.param("so3_cmd_timeout", so3_cmd_timeout_, 0.25);

  odom_set_ = false;
  imu_set_ = false;
  so3_cmd_set_ = false;

  command_pub_ = priv_nh.advertise<rosflight_msgs::Command>("command", 10);

  so3_cmd_sub_ =
      priv_nh.subscribe("so3_cmd", 10, &SO3CmdToRosflight::so3_cmd_callback, this, ros::TransportHints().tcpNoDelay());

  odom_sub_ =
      priv_nh.subscribe("odom", 10, &SO3CmdToRosflight::odom_callback, this, ros::TransportHints().tcpNoDelay());

  imu_sub_ = priv_nh.subscribe("imu", 10, &SO3CmdToRosflight::imu_callback, this, ros::TransportHints().tcpNoDelay());
}

#include <pluginlib/class_list_macros.h>
PLUGINLIB_EXPORT_CLASS(SO3CmdToRosflight, nodelet::Nodelet);


/**
 * \ file InPointerFilter.cpp
 */

#include <cmath>

#include <boost/math/constants/constants.hpp>

#include <ATK/Core/InPointerFilter.h>
#include <ATK/Core/OutPointerFilter.h>

#include <ATK/Mock/SimpleSinusGeneratorFilter.h>
#include <ATK/Mock/TriangleCheckerFilter.h>

#include <ATK/Tools/SumFilter.h>

#include <gtest/gtest.h>

constexpr gsl::index PROCESSSIZE = 100;

TEST(InPointerFloat, sin1k_test)
{
  std::array<float, PROCESSSIZE> data;
  for(gsl::index i = 0; i < PROCESSSIZE; ++i)
  {
    data[i] = std::sin(2 * boost::math::constants::pi<float>() * (i+1.)/48000 * 1000);
  }

  ATK::InPointerFilter<float> generator(data.data(), 1, PROCESSSIZE, false);
  generator.set_output_sampling_rate(48000);

  ATK::SimpleSinusGeneratorFilter<float> check;
  check.set_output_sampling_rate(48000);
  check.set_amplitude(-1);
  check.set_frequency(1000);

  ATK::SumFilter<float> sum;
  sum.set_input_sampling_rate(48000);
  sum.set_input_port(0, &generator, 0);
  sum.set_input_port(1, &check, 0);

  ATK::TriangleCheckerFilter<float> checker;
  checker.set_input_sampling_rate(48000);
  checker.set_amplitude(0);
  checker.set_frequency(1000);
  checker.set_input_port(0, &sum, 0);

  checker.process(2);
  checker.process(PROCESSSIZE - 2);
}

TEST(InPointerFloat, sin1k_test_set_pointer)
{
  std::array<float, PROCESSSIZE> data;
  for(gsl::index i = 0; i < PROCESSSIZE; ++i)
  {
    data[i] = std::sin(2 * boost::math::constants::pi<float>() * (i+1.)/48000 * 1000);
  }
  std::array<float, PROCESSSIZE> data2;

  ATK::InPointerFilter<float> generator(data2.data(), 1, PROCESSSIZE, false);
  generator.set_output_sampling_rate(48000);
  generator.set_pointer(data.data(), PROCESSSIZE);

  ATK::SimpleSinusGeneratorFilter<float> check;
  check.set_output_sampling_rate(48000);
  check.set_amplitude(-1);
  check.set_frequency(1000);

  ATK::SumFilter<float> sum;
  sum.set_input_sampling_rate(48000);
  sum.set_input_port(0, &generator, 0);
  sum.set_input_port(1, &check, 0);

  ATK::TriangleCheckerFilter<float> checker;
  checker.set_input_sampling_rate(48000);
  checker.set_amplitude(0);
  checker.set_frequency(1000);
  checker.set_input_port(0, &sum, 0);

  checker.process(2);
  checker.process(PROCESSSIZE - 2);
}

TEST(InPointerDouble, sin1k_test)
{
  std::array<double, PROCESSSIZE> data;
  for(gsl::index i = 0; i < PROCESSSIZE; ++i)
  {
    data[i] = std::sin(2 * boost::math::constants::pi<float>() * (i+1.)/48000 * 1000);
  }

  ATK::InPointerFilter<double> generator(data.data(), 1, PROCESSSIZE, false);
  generator.set_output_sampling_rate(48000);

  ATK::SimpleSinusGeneratorFilter<double> check;
  check.set_output_sampling_rate(48000);
  check.set_amplitude(-1);
  check.set_frequency(1000);

  ATK::SumFilter<double> sum;
  sum.set_input_sampling_rate(48000);
  sum.set_input_port(0, &generator, 0);
  sum.set_input_port(1, &check, 0);

  ATK::TriangleCheckerFilter<double> checker;
  checker.set_input_sampling_rate(48000);
  checker.set_amplitude(0);
  checker.set_frequency(1000);
  checker.set_input_port(0, &sum, 0);

  checker.process(2);
  checker.process(PROCESSSIZE - 2);
}

TEST(InPointerFloat, sin1k2k_interleaved_test)
{
  std::array<float, 2*PROCESSSIZE> data;
  for(gsl::index i = 0; i < PROCESSSIZE; ++i)
  {
    data[2*i] = std::sin(2 * boost::math::constants::pi<float>() * (i+1.)/48000 * 1000);
    data[2*i+1] = std::sin(2 * boost::math::constants::pi<float>() * (i+1.)/48000 * 2000);
  }

  ATK::InPointerFilter<float> generator(data.data(), PROCESSSIZE, 2, true);
  generator.set_output_sampling_rate(48000);

  ATK::SimpleSinusGeneratorFilter<float> check1;
  check1.set_output_sampling_rate(48000);
  check1.set_amplitude(-1);
  check1.set_frequency(1000);

  ATK::SimpleSinusGeneratorFilter<float> check2;
  check2.set_output_sampling_rate(48000);
  check2.set_amplitude(-1);
  check2.set_frequency(2000);

  ATK::SumFilter<float> sum1;
  sum1.set_input_sampling_rate(48000);
  sum1.set_input_port(0, &generator, 0);
  sum1.set_input_port(1, &check1, 0);
  ATK::SumFilter<float> sum2;
  sum2.set_input_sampling_rate(48000);
  sum2.set_input_port(0, &generator, 1);
  sum2.set_input_port(1, &check2, 0);
  ATK::SumFilter<float> sum;
  sum.set_input_sampling_rate(48000);
  sum.set_input_port(0, &sum1, 0);
  sum.set_input_port(1, &sum2, 0);

  ATK::TriangleCheckerFilter<float> checker;
  checker.set_input_sampling_rate(48000);
  checker.set_amplitude(0);
  checker.set_frequency(1000);
  checker.set_input_port(0, &sum, 0);

  checker.process(2);
  checker.process(PROCESSSIZE - 2);
}

TEST(InPointerFloat, sin1k2k_noninterleaved_test)
{
  std::array<float, 2*PROCESSSIZE> data;
  for(gsl::index i = 0; i < PROCESSSIZE; ++i)
  {
    data[i] = std::sin(2 * boost::math::constants::pi<float>() * (i+1.)/48000 * 1000);
  }
  for(gsl::index i = 0; i < PROCESSSIZE; ++i)
  {
    data[i+PROCESSSIZE] = std::sin(2 * boost::math::constants::pi<float>() * (i+1.)/48000 * 2000);
  }

  ATK::InPointerFilter<float> generator(data.data(), 2, PROCESSSIZE, false);
  generator.set_output_sampling_rate(48000);

  ATK::SimpleSinusGeneratorFilter<float> check1;
  check1.set_output_sampling_rate(48000);
  check1.set_amplitude(-1);
  check1.set_frequency(1000);

  ATK::SimpleSinusGeneratorFilter<float> check2;
  check2.set_output_sampling_rate(48000);
  check2.set_amplitude(-1);
  check2.set_frequency(2000);

  ATK::SumFilter<float> sum1;
  sum1.set_input_sampling_rate(48000);
  sum1.set_input_port(0, &generator, 0);
  sum1.set_input_port(1, &check1, 0);
  ATK::SumFilter<float> sum2;
  sum2.set_input_sampling_rate(48000);
  sum2.set_input_port(0, &generator, 1);
  sum2.set_input_port(1, &check2, 0);
  ATK::SumFilter<float> sum;
  sum.set_input_sampling_rate(48000);
  sum.set_input_port(0, &sum1, 0);
  sum.set_input_port(1, &sum2, 0);

  ATK::TriangleCheckerFilter<float> checker;
  checker.set_input_sampling_rate(48000);
  checker.set_amplitude(0);
  checker.set_frequency(1000);
  checker.set_input_port(0, &sum, 0);

  checker.process(2);
  checker.process(PROCESSSIZE - 2);
}

TEST(InPointerFloat, sin1k_overflow_interleaved_test)
{
  std::array<float, 2*PROCESSSIZE> data;
  std::array<float, 2*PROCESSSIZE> output;
  for(gsl::index i = 0; i < PROCESSSIZE; ++i)
  {
    data[2*i] = std::sin(2 * boost::math::constants::pi<float>() * (i+1.)/48000 * 1000);
    data[2*i+1] = std::sin(2 * boost::math::constants::pi<float>() * (i+1.)/48000 * 1000);
  }

  ATK::InPointerFilter<float> generator(data.data(), PROCESSSIZE, 2, true);
  generator.set_output_sampling_rate(48000);

  ATK::OutPointerFilter<float> getter(output.data(), 1, 2 * PROCESSSIZE, false);
  getter.set_input_sampling_rate(48000);
  getter.set_output_sampling_rate(48000);
  getter.set_input_port(0, generator, 0);

  getter.process(PROCESSSIZE*2);

  for(gsl::index i = 0; i < PROCESSSIZE; ++i)
  {
    ASSERT_EQ(data[2 * i], output[i]);
  }
  for(gsl::index i = 0; i < PROCESSSIZE; ++i)
  {
    ASSERT_EQ(0, output[i + PROCESSSIZE]);
  }
}


#include "solver.h"

#include <boost/format.hpp>
#include <cmath>
#include <cstdlib>
#include <fstream>
#include <functional>
#include <string>
#include <unordered_set>
#include <utility>
#include "../injector.h"
#include "../omp_hash_map/src/omp_hash_map.h"
#include "../omp_hash_map/src/reducer.h"
#include "../vector_stats.h"
#include "davidson_util.h"
#include "spin_det_util.h"

#define ENERGY_FORMAT "%.12f"
#define TABLE_FORMAT_ITEM_NAME "%20s"
#define TABLE_FORMAT_F "%17.10f"
#define TABLE_FORMAT_D "%'17d"
#define TABLE_FORMAT_LL "%'17llu"
#define MAX_HASH_LOAD 1.618
#define MAX_DETS_PER_FILE 20000000

class UncertainResult {
 public:
  double value = 0.0;
  double uncertainty = 0.0;
};

class SolverImpl : public Solver {
 public:
  ~SolverImpl();

  SolverImpl(
      Session* const session,
      Connections* const connections,
      AbstractSystem* const abstrct_system);

  void setup_hf() override;

  void variation(const double eps_var) override;

  void save_variation_result(const std::string& filename) override;

  bool load_variation_result(const std::string& filename) override;

  std::vector<double> apply_hamiltonian(
      const std::vector<double>& vec, const bool first_iteration) override;

  void perturbation(
      const int n_orbs_var,
      const double eps_var,
      const std::vector<int>& n_orbs_pts,
      const std::vector<double>& eps_pts) override;

 private:
  int n_up = 0;

  int n_dn = 0;

  double energy_hf = 0.0;

  double energy_var = 0.0;

  std::unordered_set<std::string> var_dets_set;

  bool verbose = false;

  void print_var_result() const;

  std::vector<data::Determinant> tmp_dets;

  std::vector<double> get_energy_pts_pre_dtm(
      const std::vector<int>& n_orbs_pts);

  std::vector<UncertainResult> get_energy_pts_dtm(
      const std::vector<int>& n_orbs_pts,
      const std::vector<double>& energy_pts_pre_dtm);

  std::vector<std::vector<UncertainResult>> get_energy_pts_stc(
      const std::vector<int>& n_orbs_pts,
      const std::vector<double>& eps_pts,
      const std::vector<UncertainResult>& energy_pts_dtm);

  double get_weight(const int i) const;
};

SolverImpl::SolverImpl(
    Session* const session,
    Connections* const connections,
    AbstractSystem* const abstract_system)
    : Solver(session, connections, abstract_system) {
  Parallel* const parallel = session->get_parallel();
  Config* const config = session->get_config();
  verbose = parallel->is_master();
  tmp_dets.resize(parallel->get_n_threads());
  n_up = config->get_int("n_up");
  n_dn = config->get_int("n_dn");
}

SolverImpl::~SolverImpl() {}

void SolverImpl::setup_hf() {
  // Create or clear wavefunction.
  abstract_system->dets.clear();
  abstract_system->coefs.clear();

  // Add a single term with coef 1.0 and no diffs.
  auto& det_hf = tmp_dets[0];
  det_hf.Clear();
  det_hf.mutable_up()->set_n_hf_elecs(n_up);
  det_hf.mutable_dn()->set_n_hf_elecs(n_dn);
  abstract_system->dets.push_back(det_hf.SerializeAsString());
  abstract_system->coefs.push_back(1.0);

  // Update HF and variational energy.
  energy_hf = energy_var = abstract_system->hamiltonian(&det_hf, &det_hf);

  // Clear connections.
  connections->clear();

  if (verbose) printf("HF det setup.\n");
};

void SolverImpl::variation(const double eps_var) {
  Timer* const timer = session->get_timer();

  // Contruct variational determinant hash set.
  const int n_start_dets = abstract_system->dets.size();
  var_dets_set.clear();
  std::vector<double> prev_coefs;
  var_dets_set.reserve(n_start_dets);
  prev_coefs.resize(n_start_dets, 0.0);
  for (int i = 0; i < n_start_dets; i++) {
    var_dets_set.insert(abstract_system->dets[i]);
  }

  // Variation iterations.
  bool converged = false;
  int iteration = 0;
  int n_new_dets = 0;

  // Add to the var wf if not exists and set its coef to zero.
  const auto& connected_det_handler =
      [&](const data::Determinant* const connected_det) {
        std::string det_code = connected_det->SerializeAsString();
        if (var_dets_set.count(det_code) == 0) {
          var_dets_set.insert(det_code);
          abstract_system->dets.push_back(det_code);
          abstract_system->coefs.push_back(0.0);
          n_new_dets++;
        }
      };

  // Apply hamiltonian to a vec (with partially cached hamiltonian).
  const auto& apply_hamiltonian_func = std::bind(
      &SolverImpl::apply_hamiltonian,
      this,
      std::placeholders::_1,
      std::placeholders::_2);

  while (!converged) {
    timer->start("loop " + std::to_string(iteration));

    // Find new dets.
    n_new_dets = 0;
    const int n_old_dets = abstract_system->dets.size();
    for (int i = 0; i < n_old_dets; i++) {
      const double coef = abstract_system->coefs[i];
      auto& det = tmp_dets[0];
      det.ParseFromString(abstract_system->dets[i]);
      if (std::abs(coef) <= std::abs(prev_coefs[i])) continue;
      abstract_system->find_connected_dets(
          &det, eps_var / std::abs(coef), connected_det_handler);
    }
    const int n_total_dets = abstract_system->dets.size();
    if (verbose) {
      printf("New / total dets: %'d / %'d\n", n_new_dets, n_total_dets);
    }
    timer->checkpoint("found new dets");

    // Update prev coefs.
    prev_coefs.resize(n_total_dets);
    for (int i = 0; i < n_total_dets; i++) {
      prev_coefs[i] = abstract_system->coefs[i];
    }

    // Diagonalize.
    connections->update();
    timer->checkpoint("updated connections");
    std::vector<double> diagonal(n_total_dets, 0.0);
#pragma omp parallel for schedule(static, 1)
    for (int i = 0; i < n_total_dets; i++) {
      const int thread_id = omp_get_thread_num();
      auto& tmp_det = tmp_dets[thread_id];
      tmp_det.ParseFromString(abstract_system->dets[i]);
      diagonal[i] = abstract_system->hamiltonian(&tmp_det, &tmp_det);
    }
    const auto& diagonalization_result = DavidsonUtil::diagonalize(
        prev_coefs, diagonal, apply_hamiltonian_func, 5, verbose);
    const double energy_var_new = diagonalization_result.first;
    const auto& new_coefs = diagonalization_result.second;
    for (int i = 0; i < n_total_dets; i++) {
      abstract_system->coefs[i] = new_coefs[i];
    }

    // Determine convergence.
    if (std::abs(energy_var_new - energy_var) < 1.0e-7) {
      converged = true;
    }

    energy_var = energy_var_new;
    iteration++;

    if (verbose) print_var_result();

    timer->end();  // iteration.
  }
};

void SolverImpl::save_variation_result(const std::string& filename) {
  // Write summary.
  std::fstream var_file(
      filename, std::ios::out | std::ios::trunc | std::ios::binary);
  data::VariationResult res;
  data::Wavefunction trunk_wf;
  res.set_energy_hf(energy_hf);
  res.set_energy_var(energy_var);
  const int n_dets = abstract_system->dets.size();
  res.set_n_dets(n_dets);
  res.SerializeToOstream(&var_file);
  var_file.close();
  if (verbose) {
    printf("Saved summary to: %s\n", filename.c_str());
  }

  // Write wavefunction trunks. (Due to the 2GB serialization limit of Protobuf)
  int trunk_n_dets = 0;
  int trunk_id = 0;
  for (int i = 0; i < n_dets; i++) {
    data::Term* const new_term = trunk_wf.add_terms();
    new_term->set_coef(abstract_system->coefs[i]);
    new_term->mutable_det()->ParseFromString(abstract_system->dets[i]);
    trunk_n_dets++;
    if (trunk_n_dets >= MAX_DETS_PER_FILE || i == n_dets - 1) {
      const auto& wf_filename = filename + ".part" + std::to_string(trunk_id);
      std::fstream wf_file(
          wf_filename, std::ios::out | std::ios::trunc | std::ios::binary);
      trunk_wf.SerializeToOstream(&wf_file);
      wf_file.close();
      if (verbose) {
        printf("Saved %'d dets to: %s\n", trunk_n_dets, wf_filename.c_str());
      }
      trunk_wf.Clear();
      trunk_n_dets = 0;
      trunk_id++;
    }
  }
}

bool SolverImpl::load_variation_result(const std::string& filename) {
  // Load summary.
  std::fstream var_file(filename, std::ios::in | std::ios::binary);
  data::VariationResult res;
  if (!res.ParseFromIstream(&var_file)) {
    return false;
  }
  var_file.close();
  energy_hf = res.energy_hf();
  energy_var = res.energy_var();
  const int n_dets_total = res.n_dets();
  if (verbose) {
    printf("Loaded summary from: %s\n", filename.c_str());
  }

  // Load wavefunctions from trunks.
  abstract_system->dets.clear();
  abstract_system->coefs.clear();
  int n_dets_read = 0;
  int trunk_id = 0;
  data::Wavefunction trunk_wf;
  while (n_dets_read < n_dets_total) {
    const auto& wf_filename = filename + ".part" + std::to_string(trunk_id);
    trunk_id++;
    std::fstream var_file(wf_filename, std::ios::in | std::ios::binary);
    if (!trunk_wf.ParseFromIstream(&var_file)) {
      throw std::runtime_error("variational results corrupted");
    }
    const int trunk_n_dets = trunk_wf.terms_size();
    for (int i = 0; i < trunk_n_dets; i++) {
      const auto& term = trunk_wf.terms(i);
      abstract_system->coefs.push_back(term.coef());
      abstract_system->dets.push_back(term.det().SerializeAsString());
      n_dets_read++;
    }
    if (verbose) {
      printf("Loaded %'d dets from: %s\n", trunk_n_dets, wf_filename.c_str());
    }
  }

  if (verbose) {
    print_var_result();
  }

  return true;
}

void SolverImpl::print_var_result() const {
  printf("Number of dets: %'zu\n", abstract_system->dets.size());
  printf("Variation energy: " ENERGY_FORMAT " Ha\n", energy_var);
  const double energy_corr = energy_var - energy_hf;
  printf("Correlation energy (variation): " ENERGY_FORMAT " Ha\n", energy_corr);
}

std::vector<double> SolverImpl::apply_hamiltonian(
    const std::vector<double>& vec, const bool first_iteration) {
  const int n_dets = vec.size();
  Parallel* const parallel = session->get_parallel();
  const int proc_id = parallel->get_proc_id();
  const int n_procs = parallel->get_n_procs();
  const int n_threads = parallel->get_n_threads();
  std::vector<std::vector<double>> res(n_threads);
  for (int i = 0; i < n_threads; i++) res[i].resize(n_dets, 0.0);
  std::vector<unsigned long long> n_nonzero_elems(n_threads, 0);

#pragma omp parallel for schedule(static, 1)
  for (int i = proc_id; i < n_dets; i += n_procs) {
    const int thread_id = omp_get_thread_num();
    const auto& conns = connections->get_connections(i);
    for (const auto conn : conns) {
      const int j = conn.first;
      const double H_ij = conn.second;
      res[thread_id][i] += H_ij * vec[j];
      if (i != j) {
        res[thread_id][j] += H_ij * vec[i];
        n_nonzero_elems[thread_id] += 2;
      } else {
        n_nonzero_elems[thread_id]++;
      }
    }
  }

  for (int i = 1; i < n_threads; i++) {
    n_nonzero_elems[0] += n_nonzero_elems[i];
    for (int j = 0; j < n_dets; j++) {
      res[0][j] += res[i][j];
    }
  }
  parallel->reduce_to_sum(res[0]);
  parallel->reduce_to_sum(n_nonzero_elems[0]);
  if (verbose && first_iteration) {
    printf("Number of non-zero elements: %'llu\n", n_nonzero_elems[0]);
  }

  session->get_timer()->checkpoint("hamiltonian applied");

  return res[0];
};

void SolverImpl::perturbation(
    const int n_orbs_var,
    const double eps_var,
    const std::vector<int>& n_orbs_pts,
    const std::vector<double>& eps_pts) {
  // Check if the results already exists.
  const auto result_filename =
      str(boost::format("pt_%d_%#.4g.csv") % n_orbs_var % eps_var);
  if (std::ifstream(result_filename)) {
    if (verbose) printf("PT results found in: %s\n", result_filename.c_str());
    return;
  }

  // Clean variation variables.
  connections->clear();

  // Construct var dets set.
  var_dets_set.clear();
  var_dets_set.reserve(abstract_system->dets.size());
  for (const auto& det_string : abstract_system->dets) {
    var_dets_set.insert(det_string);
  }

  // Get Deterministic PT correction.
  const auto& energy_pts_pre_dtm = get_energy_pts_pre_dtm(n_orbs_pts);
  const auto& energy_pts_dtm =
      get_energy_pts_dtm(n_orbs_pts, energy_pts_pre_dtm);

  // Check eps_pts validity.
  const double eps_dtm_pt = config->get_double("eps_dtm_pt");
  const int n_eps_pts = eps_pts.size();
  for (int i = 1; i < n_eps_pts; i++) {
    assert(eps_pts[i] <= eps_dtm_pt);
    assert(eps_pts[i] < eps_pts[i - 1]);
  }

  // Get Stochastic PT correction.
  const auto& energy_pts_stc =
      get_energy_pts_stc(n_orbs_pts, eps_pts, energy_pts_dtm);

  // Record results.
  std::ofstream result_file(result_filename);
  result_file << "n_orbs_var,eps_var,n_orbs_pt,eps_pt,energy_corr,uncert,"
              << "energy_hf,energy_var,energy_pt" << std::endl;

  const int n_n_orbs_pts = n_orbs_pts.size();
  for (int i = 0; i < n_eps_pts; i++) {
    for (int j = 0; j < n_n_orbs_pts; j++) {
      const double eps_pt = eps_pts[i];
      const int n_orbs_pt = n_orbs_pts[j];
      const double energy_pt_dtm = energy_pts_dtm[j].value;
      const double energy_pt_stc = energy_pts_stc[i][j].value;
      const double energy_pt = energy_pt_dtm + energy_pt_stc;
      const double uncert = std::sqrt(
          std::pow(energy_pts_stc[i][j].uncertainty, 2) +
          std::pow(energy_pts_dtm[j].uncertainty, 2));
      const double energy_corr = energy_var + energy_pt - energy_hf;
      result_file << str(boost::format("%d, %#.4g, %d, %#.4g, %#.15g, %#.15g, "
                                       "%#.15g, %#.15g, %#.15g") %
                         n_orbs_var % eps_var % n_orbs_pt % eps_pt %
                         energy_corr % uncert % energy_hf % energy_var %
                         energy_pt)
                  << std::endl;
    }
  }

  if (verbose) {
    printf("Results saved to: %s\n", result_filename.c_str());
  }

  result_file.close();
}

std::vector<double> SolverImpl::get_energy_pts_pre_dtm(
    const std::vector<int>& n_orbs_pts) {
  // Cache commonly used variables.
  const int n_n_orbs_pts = n_orbs_pts.size();
  const int n_var_dets = var_dets_set.size();
  const double eps_pre_dtm_pt = config->get_double("eps_pre_dtm_pt");
  std::hash<std::string> string_hasher;

  // Construct partial sums store and pt results store.
  omp_hash_map<std::string, double> partial_sums_pre;
  partial_sums_pre.set_max_load_factor(MAX_HASH_LOAD);
  std::vector<double> energy_pts_pre_dtm(n_n_orbs_pts, 0.0);
  std::vector<unsigned long long> n_pt_dets_pre_dtm(n_n_orbs_pts, 0);

  timer->start("pre_dtm");
  if (verbose) printf(">>> eps_pre_dtm_pt %#.4g\n", eps_pre_dtm_pt);
  timer->start("search");
  double target_progress = 0.25;
#pragma omp parallel for schedule(dynamic, 5)
  for (int i = 0; i < n_var_dets; i++) {
    auto& var_det = tmp_dets[omp_get_thread_num()];
    var_det.ParseFromString(abstract_system->dets[i]);
    const double coef = abstract_system->coefs[i];

    const auto& pt_det_handler = [&](const auto& det_a) {
      const auto& det_a_code = det_a->SerializeAsString();
      const size_t det_a_hash = string_hasher(det_a_code);
      if (var_dets_set.count(det_a_code) == 1) return;
      if (det_a_hash % n_procs != proc_id) return;
      const double H_ai = abstract_system->hamiltonian(&var_det, det_a);
      const double partial_sum_term = H_ai * coef;
      partial_sums_pre.set(
          det_a_code, [&](double& value) { value += partial_sum_term; }, 0.0);
    };

    abstract_system->find_connected_dets(
        &var_det, eps_pre_dtm_pt / std::abs(coef), pt_det_handler);

    // Report progress. (every time progress reaches 2^n %).
    const double current_progress = i * 100.0 / n_var_dets;
    if (target_progress <= current_progress) {
      if (omp_get_thread_num() == 0) {
        if (verbose) {
          printf("Master node PT dets: %'zu\n", partial_sums_pre.get_n_keys());
        }
        std::string event =
            str(boost::format("Progress: %.2f %%") % target_progress);
        timer->checkpoint(event);
        target_progress *= 2.0;
      }
    }
  }
  while (target_progress <= 100) {
    std::string event =
        str(boost::format("Progress: %.2f %%") % target_progress);
    timer->checkpoint(event);
    target_progress *= 2.0;
  }
  timer->end();  // Search.

  timer->start("accumulate");
  partial_sums_pre.apply([&](const std::string& det_code, const double value) {
    const int thread_id = omp_get_thread_num();
    auto& det = tmp_dets[thread_id];
    det.ParseFromString(det_code);
    const int n_orbs_used = SpinDetUtil::get_n_orbs_used(det);
    const double H_aa = abstract_system->hamiltonian(&det, &det);
    const double contribution = value * value / (energy_var - H_aa);
    for (int i = 0; i < n_n_orbs_pts; i++) {
      if (n_orbs_used > n_orbs_pts[i]) continue;
#pragma omp atomic
      n_pt_dets_pre_dtm[i]++;
#pragma omp atomic
      energy_pts_pre_dtm[i] += contribution;
    }
  });
  partial_sums_pre.clear();
  parallel->reduce_to_sum(n_pt_dets_pre_dtm);
  parallel->reduce_to_sum(energy_pts_pre_dtm);

  if (verbose) {
    printf(TABLE_FORMAT_ITEM_NAME, "# orbitals PT:");
    for (int i = 0; i < n_n_orbs_pts; i++) {
      printf(TABLE_FORMAT_D, n_orbs_pts[i]);
    }
    printf("\n");
    printf(TABLE_FORMAT_ITEM_NAME, "Pre DTM energy PT:");
    for (int i = 0; i < n_n_orbs_pts; i++) {
      printf(TABLE_FORMAT_F, energy_pts_pre_dtm[i]);
    }
    printf("\n");
    printf(TABLE_FORMAT_ITEM_NAME, "EST. CORR. energy:");
    for (int i = 0; i < n_n_orbs_pts; i++) {
      const double energy_corr = energy_pts_pre_dtm[i] + energy_var - energy_hf;
      printf(TABLE_FORMAT_F, energy_corr);
    }
    printf("\n");
    printf(TABLE_FORMAT_ITEM_NAME, "# DTM PT dets:");
    for (int i = 0; i < n_n_orbs_pts; i++) {
      printf(TABLE_FORMAT_LL, n_pt_dets_pre_dtm[i]);
    }
    printf("\n");
  }
  timer->end();  // Accumulate.
  timer->end();  // pre dtm.

  return energy_pts_pre_dtm;
}

std::vector<UncertainResult> SolverImpl::get_energy_pts_dtm(
    const std::vector<int>& n_orbs_pts,
    const std::vector<double>& energy_pts_pre_dtm) {
  // Cache commonly used variables.
  const int n_n_orbs_pts = n_orbs_pts.size();
  const int n_var_dets = var_dets_set.size();
  const size_t n_pt_batches_dtm = config->get_int("n_batches_dtm_pt");
  const double eps_pre_dtm_pt = config->get_double("eps_pre_dtm_pt");
  const double eps_dtm_pt = config->get_double("eps_dtm_pt");
  std::hash<std::string> string_hasher;

  // Construct partial sums store and pt results store.
  omp_hash_map<std::string, double> partial_sums_pre;
  omp_hash_map<std::string, double> partial_sums;
  partial_sums.set_max_load_factor(MAX_HASH_LOAD);
  partial_sums_pre.set_max_load_factor(MAX_HASH_LOAD);
  std::vector<UncertainResult> energy_pts_dtm(n_n_orbs_pts);
  std::vector<std::vector<double>> energy_pts_dtm_batches(n_n_orbs_pts);
  std::vector<unsigned long long> n_pt_dets_dtm(n_n_orbs_pts, 0);
  const double target_error = config->get_double("target_error");

  timer->start("dtm");
  if (verbose) printf(">>> eps_dtm_pt %#.4g\n", eps_dtm_pt);
  // Process batch by batch to achieve a larger run in a constrained mem env.
  for (size_t b = 0; b < n_pt_batches_dtm; b++) {
    timer->start(str(boost::format("%d/%d") % (b + 1) % n_pt_batches_dtm));

    double target_progress = 0.25;
    timer->start("search");
#pragma omp parallel for schedule(dynamic, 5)
    for (int i = 0; i < n_var_dets; i++) {
      auto& var_det = tmp_dets[omp_get_thread_num()];
      var_det.ParseFromString(abstract_system->dets[i]);
      const double coef = abstract_system->coefs[i];

      // Process only if:
      // 1. is not a var det, and
      // 2. belongs to the current proc_id, and
      // 3. belongs to the current batch.
      const auto& pt_det_handler = [&](const auto& det_a) {
        const auto& det_a_code = det_a->SerializeAsString();
        const size_t det_a_hash = string_hasher(det_a_code);
        if (var_dets_set.count(det_a_code) == 1) return;
        if (det_a_hash % n_procs != proc_id) return;
        if ((det_a_hash / n_procs) % n_pt_batches_dtm != b) return;
        const double H_ai = abstract_system->hamiltonian(&var_det, det_a);
        const double partial_sum_term = H_ai * coef;
        if (std::abs(partial_sum_term) >= eps_pre_dtm_pt) {
          partial_sums_pre.set(
              det_a_code,
              [&](double& value) { value += partial_sum_term; },
              0.0);
        }
        partial_sums.set(
            det_a_code, [&](double& value) { value += partial_sum_term; }, 0.0);
      };

      abstract_system->find_connected_dets(
          &var_det, eps_dtm_pt / std::abs(coef), pt_det_handler);

      // Report progress. (every time progress reaches 2^n %).
      const double current_progress = i * 100.0 / n_var_dets;
      if (target_progress <= current_progress) {
        if (omp_get_thread_num() == 0) {
          if (verbose) {
            printf("Master node PT dets: %'zu\n", partial_sums.get_n_keys());
          }
          std::string event =
              str(boost::format("Progress: %.2f %%") % target_progress);
          timer->checkpoint(event);
          target_progress *= 2.0;
        }
      }
    }
    while (target_progress <= 100) {
      std::string event =
          str(boost::format("Progress: %.2f %%") % target_progress);
      timer->checkpoint(event);
      target_progress *= 2.0;
    }
    timer->end();  // Search.

    timer->start("accumulate");

    // For reduction after each batch.
    std::vector<double> energy_pts_dtm_batch(n_n_orbs_pts, 0.0);

    // Apply to each key, value pair:
    // 1. Calculate the contribution.
    // 2. For each n_orbs_pt interested:
    //      If the n_orbs used by the key is less than n_orbs_pt:
    //        Increase n_pt_dets by one and add contribution to energy_pt.
    partial_sums.apply([&](const std::string& det_code, const double value) {
      const int thread_id = omp_get_thread_num();
      auto& det = tmp_dets[thread_id];
      det.ParseFromString(det_code);
      const int n_orbs_used = SpinDetUtil::get_n_orbs_used(det);
      const double H_aa = abstract_system->hamiltonian(&det, &det);
      const double value_pre =
          partial_sums_pre.get_copy_or_default(det_code, 0.0);
      partial_sums_pre.unset(det_code);
      const double contribution =
          (value * value - value_pre * value_pre) / (energy_var - H_aa);
      for (int i = 0; i < n_n_orbs_pts; i++) {
        if (n_orbs_used > n_orbs_pts[i]) continue;
#pragma omp atomic
        n_pt_dets_dtm[i]++;
#pragma omp atomic
        energy_pts_dtm_batch[i] += contribution;
      }
    });

    // Aggregate the results from each proc.
    parallel->reduce_to_sum(energy_pts_dtm_batch);
    double max_uncert = 0.0;
    for (int i = 0; i < n_n_orbs_pts; i++) {
      energy_pts_dtm_batches[i].push_back(energy_pts_dtm_batch[i]);
    }
    for (int i = 0; i < n_n_orbs_pts; i++) {
      energy_pts_dtm[i].value =
          get_avg(energy_pts_dtm_batches[i]) * n_pt_batches_dtm +
          energy_pts_pre_dtm[i];
      if (b == n_pt_batches_dtm - 1) {
        energy_pts_dtm[i].uncertainty = 0.0;  // All batches finished.
      } else {
        energy_pts_dtm[i].uncertainty =
            get_stdev(energy_pts_dtm_batches[i]) * n_pt_batches_dtm / b;
      }
      max_uncert = std::max(max_uncert, energy_pts_dtm[i].uncertainty);
    }

    // Print batch result and estimated total correction.
    if (verbose) {
      printf(TABLE_FORMAT_ITEM_NAME, "# orbitals PT:");
      for (int i = 0; i < n_n_orbs_pts; i++) {
        printf(TABLE_FORMAT_D, n_orbs_pts[i]);
      }
      printf("\n");
      printf(TABLE_FORMAT_ITEM_NAME, "Batch DTM energy PT:");
      for (int i = 0; i < n_n_orbs_pts; i++) {
        printf(TABLE_FORMAT_F, energy_pts_dtm_batch[i]);
      }
      printf("\n");
      printf(TABLE_FORMAT_ITEM_NAME, "EST. DTM energy PT:");
      for (int i = 0; i < n_n_orbs_pts; i++) {
        printf(TABLE_FORMAT_F, energy_pts_dtm[i].value - energy_pts_pre_dtm[i]);
      }
      printf("\n");
      printf(TABLE_FORMAT_ITEM_NAME, "(including pre dtm):");
      for (int i = 0; i < n_n_orbs_pts; i++) {
        printf(TABLE_FORMAT_F, energy_pts_dtm[i].value);
      }
      printf("\n");
      printf(TABLE_FORMAT_ITEM_NAME, "Uncertainty:");
      for (int i = 0; i < n_n_orbs_pts; i++) {
        printf(TABLE_FORMAT_F, energy_pts_dtm[i].uncertainty);
      }
      printf("\n");
      printf(TABLE_FORMAT_ITEM_NAME, "EST. CORR. energy:");
      for (int i = 0; i < n_n_orbs_pts; i++) {
        const double energy_corr =
            energy_pts_dtm[i].value + energy_var - energy_hf;
        printf(TABLE_FORMAT_F, energy_corr);
      }
      printf("\n");
    }

    timer->end();  // Accumulation.

    partial_sums.clear();
    partial_sums_pre.clear();
    timer->end();  // Batch.

    if (b > 0 && b < n_pt_batches_dtm - 2 && max_uncert < 0.2 * target_error) {
      if (verbose) {
        printf(
            "\n>>> Skip remaining batches since uncertainty is significantly "
            "smaller than target error.\n");
      }
      break;
    }
  }

  parallel->reduce_to_sum(n_pt_dets_dtm);

  // Print total PT dets.
  if (verbose) {
    printf(TABLE_FORMAT_ITEM_NAME, "# DTM PT dets:");
    for (int i = 0; i < n_n_orbs_pts; i++) {
      printf(TABLE_FORMAT_LL, n_pt_dets_dtm[i]);
    }
    printf("\n");
  }

  timer->end();

  return energy_pts_dtm;
}

std::vector<std::vector<UncertainResult>> SolverImpl::get_energy_pts_stc(
    const std::vector<int>& n_orbs_pts,
    const std::vector<double>& eps_pts,
    const std::vector<UncertainResult>& energy_pts_dtm) {
  const int n_n_orbs_pts = n_orbs_pts.size();
  const int n_eps_pts = eps_pts.size();

  // Construct results store.
  std::vector<std::vector<UncertainResult>> energy_pts_stc(n_eps_pts);
  std::vector<std::vector<std::vector<double>>> energy_pts_loops(n_eps_pts);
  std::vector<omp_hash_map<std::string, double>> partial_sums(n_eps_pts);
  omp_hash_map<std::string, double> partial_sums_dtm;
  for (int i = 0; i < n_eps_pts; i++) {
    energy_pts_stc[i].resize(n_n_orbs_pts);
    energy_pts_loops[i].resize(n_n_orbs_pts);
    partial_sums[i].set_max_load_factor(MAX_HASH_LOAD);
  }
  partial_sums_dtm.set_max_load_factor(MAX_HASH_LOAD);

  // Construct probabilities from normalized weights.
  const int n_var_dets = var_dets_set.size();
  double sum_weights = 0.0;
  for (int i = 0; i < n_var_dets; i++) {
    sum_weights += get_weight(i);
  }
  std::vector<double> probs(n_var_dets);
  std::vector<double> cum_probs(n_var_dets);  // For sampling.
  for (int i = 0; i < n_var_dets; i++) {
    probs[i] = get_weight(i) / sum_weights;
    if (i == 0) {
      cum_probs[i] = probs[i];
    } else {
      cum_probs[i] = probs[i] + cum_probs[i - 1];
    }
  }

  // Stochastic iterations.
  int iteration = 0;
  const int max_n_iterations = config->get_int("max_n_iterations_stc_pt");
  std::unordered_map<int, int> stc_pt_sample_dets;
  std::vector<int> stc_pt_sample_dets_list;
  const int n_samples = config->get_int("n_samples_stc_pt");
  srand(time(NULL));
  const double eps_dtm_pt = config->get_double("eps_dtm_pt");
  const double eps_pts_min = eps_pts.back();
  const size_t n_batches = config->get_int("n_batches_stc_pt");
  std::hash<std::string> string_hasher;
  const double target_error = config->get_double("target_error");

  timer->start("stc");
  while (iteration < max_n_iterations) {
    timer->start(str(boost::format("loop %d") % iteration));

    // Cleanup.
    for (int i = 0; i < n_eps_pts; i++) {
      partial_sums[i].clear();
    }
    partial_sums_dtm.clear();
    stc_pt_sample_dets.clear();
    stc_pt_sample_dets_list.clear();

    std::vector<std::vector<double>> energy_pts_loop(n_eps_pts);
    std::vector<std::vector<unsigned long long>> n_pt_dets_loop(n_eps_pts);
    for (int i = 0; i < n_eps_pts; i++) {
      energy_pts_loop[i].resize(n_n_orbs_pts, 0.0);
      n_pt_dets_loop[i].resize(n_n_orbs_pts, 0);
    }

    // Generate random samples.
    for (int i = 0; i < n_samples; i++) {
      const double rand_01 = ((double)rand() / (RAND_MAX));
      const int sample_det_id =
          std::lower_bound(cum_probs.begin(), cum_probs.end(), rand_01) -
          cum_probs.begin();
      if (stc_pt_sample_dets.count(sample_det_id) == 0) {
        stc_pt_sample_dets_list.push_back(sample_det_id);
      }
      stc_pt_sample_dets[sample_det_id]++;
    }

    // Use one batch to estimate contribution from all batches.
    const size_t selected_batch = rand() % n_batches;

    // Search PT dets from selected sample var dets.
    const int n_stc_pt_sample_dets = stc_pt_sample_dets_list.size();

#pragma omp parallel for schedule(dynamic, 2)
    for (int s = 0; s < n_stc_pt_sample_dets; s++) {
      const int var_det_id = stc_pt_sample_dets_list[s];
      const double prob = probs[var_det_id];
      const double cnt = static_cast<double>(stc_pt_sample_dets[var_det_id]);
      auto& var_det = tmp_dets[omp_get_thread_num()];
      var_det.ParseFromString(abstract_system->dets[var_det_id]);
      const double coef = abstract_system->coefs[var_det_id];

      const auto& pt_det_handler = [&](const auto& det_a) {
        const auto& det_a_code = det_a->SerializeAsString();
        const size_t det_a_hash = string_hasher(det_a_code);
        if (var_dets_set.count(det_a_code) == 1) return;
        if (det_a_hash % n_procs != proc_id) return;
        if ((det_a_hash / n_procs) % n_batches != selected_batch) {
          return;  // Only use one batch.
        }
        const double H_ai = abstract_system->hamiltonian(&var_det, det_a);
        const double H_aa = abstract_system->hamiltonian(det_a, det_a);
        const double factor =
            n_batches / ((energy_var - H_aa) * n_samples * (n_samples - 1));
        const double partial_sum_term = H_ai * coef;
        const double contrib_1 = cnt * partial_sum_term / prob * sqrt(-factor);

        // Add to dtm partial sums.
        if (std::abs(partial_sum_term) >= eps_dtm_pt) {
          partial_sums_dtm.set(
              det_a_code, [&](double& value) { value += contrib_1; }, 0.0);
        }

        // Add to partial sums.
        for (int i = 0; i < n_eps_pts; i++) {
          if (std::abs(partial_sum_term) >= eps_pts[i]) {
            partial_sums[i].set(
                det_a_code, [&](double& value) { value += contrib_1; }, 0.0);
            break;
          }
        }

        // Calculate 2nd order contribution if not belong to dtm pt .
        if (std::abs(partial_sum_term) >= eps_dtm_pt) return;

        const double contrib_2 =
            (cnt * (n_samples - 1) / prob - (cnt * cnt) / (prob * prob)) *
            partial_sum_term * partial_sum_term * factor;
        const int n_orbs_used = SpinDetUtil::get_n_orbs_used(*det_a);
        for (int i = 0; i < n_eps_pts; i++) {
          if (std::abs(partial_sum_term) < eps_pts[i]) continue;
          for (int j = 0; j < n_n_orbs_pts; j++) {
            if (n_orbs_pts[j] < n_orbs_used) continue;
#pragma omp atomic
            energy_pts_loop[i][j] += contrib_2;
          }
        }
      };

      abstract_system->find_connected_dets(
          &var_det, eps_pts_min / std::abs(coef), pt_det_handler);
    }

    // Accumulate and get PT correction from batch.
    for (int i = 0; i < n_eps_pts; i++) {
      partial_sums[i].apply(
          [&](const std::string& det_code, const double value) {
            const int thread_id = omp_get_thread_num();
            auto& det = tmp_dets[thread_id];
            det.ParseFromString(det_code);

            double cum_value = value;
            const double value_dtm =
                partial_sums_dtm.get_copy_or_default(det_code, 0.0);
            partial_sums_dtm.unset(det_code);

            const int n_orbs_used = SpinDetUtil::get_n_orbs_used(det);
            for (int j = i; j < n_eps_pts; j++) {
              if (j > i) {
                const double value_j =
                    partial_sums[j].get_copy_or_default(det_code, 0.0);
                partial_sums[j].unset(det_code);
                cum_value += value_j;
              }

              const double contrib_1 =
                  -cum_value * cum_value + value_dtm * value_dtm;

              for (int k = 0; k < n_n_orbs_pts; k++) {
                if (n_orbs_pts[k] < n_orbs_used) continue;
#pragma omp atomic
                energy_pts_loop[j][k] += contrib_1;
#pragma omp atomic
                n_pt_dets_loop[j][k]++;
              }
            }
          });
    }

    // Append loop results store.
    for (int i = 0; i < n_eps_pts; i++) {
      parallel->reduce_to_sum(energy_pts_loop[i]);
      parallel->reduce_to_sum(n_pt_dets_loop[i]);
      for (int j = 0; j < n_n_orbs_pts; j++) {
        energy_pts_loops[i][j].push_back(energy_pts_loop[i][j]);
      }
    }

    // Calculate statistics.
    double max_uncert = 0.0;
    for (int i = 0; i < n_eps_pts; i++) {
      for (int j = 0; j < n_n_orbs_pts; j++) {
        const double avg = get_avg(energy_pts_loops[i][j]);
        const double stdev = get_stdev(energy_pts_loops[i][j]);
        energy_pts_stc[i][j].value = avg;
        energy_pts_stc[i][j].uncertainty = stdev / sqrt(iteration);
        max_uncert = std::max(energy_pts_stc[i][j].uncertainty, max_uncert);
      }
    }

    // Report results.
    if (verbose) {
      printf(TABLE_FORMAT_ITEM_NAME, "# orbitals PT:");
      for (int j = 0; j < n_n_orbs_pts; j++) {
        printf(TABLE_FORMAT_D, n_orbs_pts[j]);
      }
      printf("\n");
      for (int i = 0; i < n_eps_pts; i++) {
        printf(">>> eps_stc_pt %#.4g\n", eps_pts[i]);
        printf(TABLE_FORMAT_ITEM_NAME, "Loop # STC PT dets:");
        for (int j = 0; j < n_n_orbs_pts; j++) {
          printf(TABLE_FORMAT_LL, n_pt_dets_loop[i][j]);
        }
        printf("\n");
        printf(TABLE_FORMAT_ITEM_NAME, "Loop STC energy PT:");
        for (int j = 0; j < n_n_orbs_pts; j++) {
          printf(TABLE_FORMAT_F, energy_pts_loops[i][j][iteration]);
        }
        printf("\n");
        printf(TABLE_FORMAT_ITEM_NAME, "EST. STC energy PT:");
        for (int j = 0; j < n_n_orbs_pts; j++) {
          printf(TABLE_FORMAT_F, energy_pts_stc[i][j].value);
        }
        printf("\n");
        printf(TABLE_FORMAT_ITEM_NAME, "Uncertainty:");
        for (int j = 0; j < n_n_orbs_pts; j++) {
          printf(TABLE_FORMAT_F, energy_pts_stc[i][j].uncertainty);
        }
        printf("\n");
        printf(TABLE_FORMAT_ITEM_NAME, "EST. energy PT:");
        for (int j = 0; j < n_n_orbs_pts; j++) {
          printf(
              TABLE_FORMAT_F,
              energy_pts_stc[i][j].value + energy_pts_dtm[j].value);
        }
        printf("\n");
        printf(TABLE_FORMAT_ITEM_NAME, "EST. CORR. energy:");
        for (int j = 0; j < n_n_orbs_pts; j++) {
          const double energy_corr = energy_pts_stc[i][j].value +
                                     energy_pts_dtm[j].value + energy_var -
                                     energy_hf;
          printf(TABLE_FORMAT_F, energy_corr);
        }
        printf("\n");
      }
    }

    timer->end();

    if (iteration >= 10 && max_uncert <= target_error) break;

    iteration++;
  }

  timer->end();

  return energy_pts_stc;
}

double SolverImpl::get_weight(const int i) const {
  const double coef = abstract_system->coefs[i];
  const double abs_coef = std::abs(coef);
  return abs_coef;
}

Solver* Injector::new_solver(
    Session* const session,
    Connections* const connections,
    AbstractSystem* const abstract_system) {
  return new SolverImpl(session, connections, abstract_system);
}