#include "potentials.hpp"
#include "metropolis.hpp"
#include "distance.hpp"
#include <cmath>
#include <cassert>
#include <stdexcept>
#include <boost/algorithm/string.hpp>
#include "yaml-cpp/yaml.h"

namespace pauth {

using namespace arma;

double abstract_LJ_full_potential::_U(const metropolis &sim) const {

  const auto &molecular_ids = sim.molecular_ids();
  const auto &positions = sim.positions();
  const auto &edge_lengths = sim.edge_lengths();
  const auto N = sim.N();
  double potential = 0;

  #pragma omp parallel for reduction(+:potential)
  for (auto i = size_t{0}; i < N - 1; ++i) {
    for (auto j = i + 1; j < N; ++j) {

      const double rij2 = sim.m(positions.col(i), positions.col(j),
                                edge_lengths);
      const double rzero = get_rzero(molecular_ids[i], molecular_ids[j]);
      const double rat2 = (rzero * rzero) / rij2;
      const double rat6 = rat2 * rat2 * rat2;

      potential += 4.0 * get_well_depth(molecular_ids[i], molecular_ids[j]) *
                   (rat6 * rat6 - rat6);
    }
  }

  return potential;
}

double abstract_LJ_full_potential::_delta_U(const metropolis &sim, const size_t j,
                                       arma::vec &rn_j) const {

  const auto &molecular_ids = sim.molecular_ids();
  const auto &positions = sim.positions();
  const auto &ro_j = positions.col(j);
  const auto &edge_lengths = sim.edge_lengths();
  const auto N = sim.N();
  double dU = 0;

  #pragma omp parallel for reduction(+:dU)
  for (auto i = size_t{0}; i < N; ++i) {
    if (i == j) continue;

    const auto well_depth = get_well_depth(molecular_ids[i], molecular_ids[j]);
    double rij2 = sim.m(positions.col(i), ro_j, edge_lengths);
    double rzero = get_rzero(molecular_ids[i], molecular_ids[j]);
    double rat2 = (rzero * rzero) / rij2;
    double rat6 = rat2 * rat2 * rat2;

    const double Uo = 4.0 * well_depth * (rat6 * rat6 - rat6);

    rij2 = sim.m(positions.col(i), rn_j, edge_lengths);
    rzero = get_rzero(molecular_ids[i], molecular_ids[j]);
    rat2 = (rzero * rzero) / rij2;
    rat6 = rat2 * rat2 * rat2;

    const double Un = 4.0 * well_depth * (rat6 * rat6 - rat6);

    dU += Un - Uo;
  }

  return dU;
}

arma::vec abstract_LJ_full_potential::_forceij(const metropolis &sim, const size_t i,
                                          const size_t j) const {
  const auto &molecular_ids = sim.molecular_ids();
  const auto &positions = sim.positions();
  const auto &edge_lengths = sim.edge_lengths();
  const auto rij = sim.rij(positions.col(i), positions.col(j), edge_lengths);
  const double _rij2 = dot(rij, rij);
  const double rzero = get_rzero(molecular_ids[i], molecular_ids[j]);
  const double rat2 = (rzero * rzero) / _rij2;
  const double rat6 = rat2 * rat2 * rat2;
  const double rat8 = rat6 * rat2;
  const double well_depth = get_well_depth(molecular_ids[i], molecular_ids[j]);

  return rij / rzero * well_depth * (48.0 * rat8 * rat6 - 24.0 * rat8);
}

double abstract_LJ_cutoff_potential::_U(const metropolis &sim) const {

  const auto &molecular_ids = sim.molecular_ids();
  const auto &positions = sim.positions();
  const auto &edge_lengths = sim.edge_lengths();
  const auto N = sim.N();
  double potential = 0;

  #pragma omp parallel for reduction(+:potential)
  for (auto i = size_t{0}; i < N - 1; ++i) {
    for (auto j = i + 1; j < N; ++j) {

      const double rij2 = sim.m(positions.col(i), positions.col(j),
                                edge_lengths);

      if (rij2 <= _rc2) {

        const double rzero = get_rzero(molecular_ids[i], molecular_ids[j]);
        const double rz2 = rzero * rzero;
        const double rz4 = rz2 * rz2;
        const double rz8 = rz4 * rz4;

        const double dudr_rc = 6.0 * (rz4 * rz2 * rzero) / (_rc7)-12.0 *
                               (rz8 * rz4 * rzero) / (_rc13);
        const double u_rc = ((rz8 * rz4) / (_rc12) - (rz4 * rz2) / (_rc6));

        const double rat2 = rz2 / rij2;
        const double rat6 = rat2 * rat2 * rat2;

        potential += 4.0 * get_well_depth(molecular_ids[i], molecular_ids[j]) *
                     ((rat6 * rat6 - rat6) - u_rc -
                      (std::sqrt(rij2) - _cutoff) * dudr_rc);
      }
    }
  }

  return potential;
}

double abstract_LJ_cutoff_potential::_delta_U(const metropolis &sim,
                                              const size_t j,
                                              arma::vec &rn_j) const {

  const auto &molecular_ids = sim.molecular_ids();
  const auto &positions = sim.positions();
  const auto &ro_j = positions.col(j);
  const auto &edge_lengths = sim.edge_lengths();
  const auto N = sim.N();
  double dU = 0;

  #pragma omp parallel for reduction(+:dU)
  for (auto i = size_t{0}; i < N; ++i) {
    if (i == j) continue;

    const auto rzero = get_rzero(molecular_ids[i], molecular_ids[j]);
    const auto well_depth = get_well_depth(molecular_ids[i], molecular_ids[j]);
    const double rz2 = rzero * rzero;
    const double rz4 = rz2 * rz2;
    const double rz8 = rz4 * rz4;

    const double dudr_rc = 6.0 * (rz4 * rz2 * rzero) / (_rc7)-12.0 *
                           (rz8 * rz4 * rzero) / (_rc13);
    const double u_rc = ((rz8 * rz4) / (_rc12) - (rz4 * rz2) / (_rc6));

    double Uo = 0.0;
    double rij2 = sim.m(positions.col(i), ro_j, edge_lengths);

    if (rij2 <= _rc2) {
      const double rat2 = rz2 / rij2;
      const double rat6 = rat2 * rat2 * rat2;

      Uo = 4.0 * well_depth *
           ((rat6 * rat6 - rat6) - u_rc -
            (std::sqrt(rij2) - _cutoff) * dudr_rc);
    }

    double Un = 0.0;
    rij2 = sim.m(positions.col(i), rn_j, edge_lengths);

    if (rij2 <= _rc2) {
      const double rat2 = rz2 / rij2;
      const double rat6 = rat2 * rat2 * rat2;

      Un = 4.0 * well_depth *
           ((rat6 * rat6 - rat6) - u_rc -
            (std::sqrt(rij2) - _cutoff) * dudr_rc);
    }

    dU += Un - Uo;
  }

  return dU;
}

arma::vec abstract_LJ_cutoff_potential::_forceij(const metropolis &sim,
                                                 const size_t i,
                                                 const size_t j) const {

  const auto &molecular_ids = sim.molecular_ids();
  const auto &positions = sim.positions();
  const auto &edge_lengths = sim.edge_lengths();
  const auto _rij = sim.rij(positions.col(i), positions.col(j), edge_lengths);
  const double _rij2 = dot(_rij, _rij);

  if (_rij2 <= _rc2) {

    const double rzero = get_rzero(molecular_ids[i], molecular_ids[j]);
    const double rz2 = rzero * rzero;
    const double rz4 = rz2 * rz2;
    const double rz8 = rz4 * rz4;

    const double dudr_rc = 24.0 * (rz4 * rz2 * rzero) / (_rc7)-48.0 *
                           (rz8 * rz4 * rzero) / (_rc13);

    const double rat2 = rz2 / _rij2;
    const double rat6 = rat2 * rat2 * rat2;
    const double rat8 = rat6 * rat2;
    const double well_depth =
        get_well_depth(molecular_ids[i], molecular_ids[j]);

    return _rij / rzero * well_depth *
           (48.0 * rat8 * rat6 - 24.0 * rat8 + dudr_rc / std::sqrt(_rij2));
  }

  return arma::zeros(sim.D());
}

double abstract_spring_potential::_U(const metropolis &sim) const {

  const auto N = sim.N();
  const auto &molecular_ids = sim.molecular_ids();
  const auto& positions = sim.positions();
  double potential = 0;

  #pragma omp parallel for reduction(+:potential)
  for (auto i = size_t{0}; i < N; ++i) {
    const double r2 = dot(positions.col(i), positions.col(i));
    potential += 0.5 * get_k(molecular_ids[i]) * r2;
  }

  return potential;
}

double abstract_spring_potential::_delta_U(const metropolis &sim,
                                           const size_t j,
                                           arma::vec &rn_j) const {
  const arma::vec &ro_j = sim.positions().col(j);
  return 0.5 * get_k(sim.molecular_ids()[j]) *
         (dot(rn_j, rn_j) - dot(ro_j, ro_j));
}

const_poly_spring_potential::const_poly_spring_potential(
    const std::initializer_list<double> &coeffs)
    : pcoeffs(coeffs) {

  for (size_t i = 1; i < pcoeffs.size(); ++i)
    fcoeffs.push_back(pcoeffs[i] * 2 * i);
}

double const_poly_spring_potential::_U(const metropolis &sim) const {

  const auto N = sim.N();
  const auto& positions = sim.positions();
  double potential = 0;

  #pragma omp parallel for reduction(+:potential)
  for (auto i = size_t{0}; i < N; ++i) {
    double Ui = 0.0;
    const double r2 = dot(positions.col(i), positions.col(i));
    for (auto p = size_t{0}; p < pcoeffs.size(); ++p) {
      double x = 1;
      for (auto k = size_t{0}; k < p; ++k)
        x *= r2;
      Ui += pcoeffs[p] * x;
    }
    potential += Ui;
  }

  return potential;
}

double const_poly_spring_potential::_delta_U(const metropolis &sim,
                                             const size_t j,
                                             arma::vec &rn_j) const {
  const arma::vec &ro_j = sim.positions().col(j);

  double Uo = 0.0;
  double r2 = dot(ro_j, ro_j);
  for (auto p = size_t{0}; p < pcoeffs.size(); ++p) {
    double x = 1;
    for (auto k = size_t{0}; k < p; ++k)
      x *= r2;
    Uo += pcoeffs[p] * x;
  }

  double Un = 0.0;
  r2 = dot(rn_j, rn_j);
  for (auto p = size_t{0}; p < pcoeffs.size(); ++p) {
    double x = 1;
    for (auto k = size_t{0}; k < p; ++k)
      x *= r2;
    Un += pcoeffs[p] * x;
  }

  return Un - Uo;
}

arma::vec const_poly_spring_potential::_forceij(const metropolis &sim, const size_t,
                                                const size_t) const {
  return arma::zeros(sim.D());
}

double const_quad_spring_potential::_U(const metropolis &sim) const {

  const auto N = sim.N();
  const auto &positions = sim.positions();
  double potential = 0;

  #pragma omp parallel for reduction(+:potential)
  for (auto i = size_t{0}; i < N; ++i) {
    const double r2 = dot(positions.col(i), positions.col(i));
    potential += a * r2 * r2 + b * r2 + c;
  }

  return potential;
}

double const_quad_spring_potential::_delta_U(const metropolis &sim,
                                             const size_t j,
                                             arma::vec &rn_j) const {
  const arma::vec &ro_j = sim.positions().col(j);
  const double r2o = dot(ro_j, ro_j);
  const double r2n = dot(rn_j, rn_j);

  return a * (r2n*r2n - r2o*r2o) + b * (r2n - r2o);
}

arma::vec abstract_spring_potential::_forceij(const metropolis &sim,
                                                      const size_t,
                                                      const size_t) const {
  return arma::zeros(sim.D());
}

arma::vec const_quad_spring_potential::_forceij(const metropolis &sim, const size_t,
                                                const size_t) const {
  return arma::zeros(sim.D());
}

double twostate_int_potential::_U(const metropolis &sim) const {

  const auto N = sim.N();
  const auto &positions = sim.positions();

  double potential = 0.0;

  #pragma omp parallel for reduction(+:potential)
  for (auto i = size_t{0}; i < N; ++i) {
    assert(_check_x(positions.col(i)));
    potential += _Ui(positions.col(i));
  }

  #pragma omp parallel for reduction(+:potential)
  for (auto i = size_t{0}; i < N; ++i) {
    for (auto j = size_t{i+1}; j < N; ++j) {
      potential += _Ui(positions.col(i)) * _Ui(positions.col(j));
    }
  }

  return potential;
}

double twostate_int_potential::_delta_U(const metropolis &sim, const size_t j,
                                        arma::vec &rn_j) const {

  const auto N = sim.N();
  const auto &positions = sim.positions();

  double dUi = _Ui(rn_j) - _Ui(positions.col(j));

  double U2 = 0.0;
  #pragma omp parallel for reduction(+:U2)
  for (auto i = size_t{0}; i < N; ++i) {
    if (i == j) continue;
    U2 += _Ui(positions.col(i));
  }

  return dUi + dUi * U2;

}

arma::vec twostate_int_potential::_forceij(const metropolis &sim, const size_t,
                                           const size_t) const {
  return arma::zeros(sim.D());
}

abstract_LJ_lookup_potential::abstract_LJ_lookup_potential(const char *fname) {
  vector<string> molecular_strs;
  const auto well_params = YAML::LoadFile(fname);

  for (const auto &well_param : well_params) {
    auto molecular_pair_str = well_param.first.as<string>();
    boost::split(molecular_strs, molecular_pair_str, boost::is_any_of(","));
    for (auto &molecular_str : molecular_strs) boost::trim(molecular_str);
    if (molecular_names_to_ids.count(molecular_strs[0]) &&
        molecular_names_to_ids.count(molecular_strs[1])) {
      unordered_pair<molecular_id> key {
        molecular_names_to_ids.at(molecular_strs[0]),
        molecular_names_to_ids.at(molecular_strs[1])
      };
      _well_depth_map[key] = well_param.second["epsilon"].as<double>();
      _rzero_map[key] = well_param.second["sigma"].as<double>();
    }
    else throw runtime_error(string("Molecular pair ") +
                             well_param.first.as<string>() +
                             string(" not found."));
  }
}

double abstract_LJ_lookup_potential::_get_well_depth(const molecular_id id1,
                                                     const molecular_id id2) const {
  unordered_pair<molecular_id> key { id1, id2 };
  return _well_depth_map.at(key);
}

double abstract_LJ_lookup_potential::_get_rzero(const molecular_id id1,
                                                const molecular_id id2) const {
  unordered_pair<molecular_id> key { id1, id2 };
  return _rzero_map.at(key);
}

double dipole_electric_potential::_U(const metropolis &sim) const {

  const auto N = sim.N();
  const auto &positions = sim.positions();
  double sum = 0.0;

  #pragma omp parallel for reduction(+:sum)
  for (auto i = size_t{0}; i < N; ++i) {
    auto& positionsi = positions.col(i);
    for (auto& kv : _efield)
      sum += -(positionsi(kv.first) * kv.second(positionsi));
  }

  return sum;

}

double dipole_electric_potential::_delta_U(const metropolis &sim,
  const size_t j, arma::vec &dx) const {

  auto& positionsj = sim.positions().col(j);

  double sum = 0.0;
  for (auto& kv : _efield)
    sum += -(dx(kv.first) * kv.second(positionsj));

  return sum;

}

arma::vec dipole_electric_potential::_forceij(const metropolis &sim,
    const size_t, const size_t) const { return arma::zeros(sim.D()); }

double abstract_dipole_strain_potential::_U(const metropolis &sim) const {

  const auto N = sim.N();
  const auto &positions = sim.positions();
  const auto &ids = sim.molecular_ids();
  double sum = 0.0;

  #pragma omp parallel for reduction(+:sum)
  for (auto i = size_t{0}; i < N; ++i) {

    auto& xsi = positions.col(i);
    auto pi = p(xsi);

    // Ui = 1/2 * p . X p
    sum += arma::dot(pi, inv_chi(xsi, ids[i]) * pi);

  }

  return 0.5 * sum;

}

double abstract_dipole_strain_potential::_delta_U(const metropolis &sim,
    const size_t j, arma::vec &dx) const {

  auto& xs = sim.positions().col(j);
  const auto id = sim.molecular_ids()[j];
  auto xs_new = xs + dx;

  auto pi = p(xs);
  auto pnew = p(xs_new);

  return 0.5 * (arma::dot(pi, inv_chi(xs, id) * pi) -
                arma::dot(pnew, inv_chi(xs_new, id) * pnew));

}

arma::vec abstract_dipole_strain_potential::_forceij
  (const metropolis &sim, const size_t, const size_t) const {
  return arma::zeros(sim.D());
}

// definitions for pure virtual destructors
abstract_potential::~abstract_potential() {}
abstract_LJ_potential::~abstract_LJ_potential() {}
abstract_LJ_full_potential::~abstract_LJ_full_potential() {}
abstract_LJ_lookup_potential::~abstract_LJ_lookup_potential() {}
abstract_LJ_cutoff_potential::~abstract_LJ_cutoff_potential() {}
abstract_spring_potential::~abstract_spring_potential() {}
abstract_dipole_strain_potential::~abstract_dipole_strain_potential() {}
abstract_dipole_strain_2d_potential::~abstract_dipole_strain_2d_potential() {}
abstract_dipole_strain_3d_potential::~abstract_dipole_strain_3d_potential() {}
abstract_dipole_strain_linear_potential::~abstract_dipole_strain_linear_potential() {}

} // namespace pauth


// Copyright [2020] Yoonyoung (Jamie) Cho

#include <fmt/format.h>
#include <Eigen/Core>

#include <boost/functional/hash.hpp>
#include <boost/graph/random.hpp>
#include <boost/random.hpp>

#include <unordered_map>

#include "bglpy/graph.hpp"
#include "bglpy/rrt.hpp"

static const float SquaredDistance(const cho::graph::Node2D& src,
                                   const cho::graph::Node2D& dst) {
  const float dx = (src.x - dst.x);
  const float dy = (src.y - dst.y);
  return dx * dx + dy * dy;
}

float WorkspaceDistance(const cho::graph::Node2D& node) {
  using Point = cho::graph::Node2D;
  static constexpr const float kHipJointOffset = 0.0110;
  static constexpr const float kKneeLinkLength = 0.0285;
  static constexpr const float kHipLinkLength = 0.0175;

  static constexpr const float kSmallRadius = kKneeLinkLength - kHipLinkLength;
  static constexpr const float kLargeRadius = kKneeLinkLength + kHipLinkLength;

  const Point center_a{kHipJointOffset, 0.0f};
  const Point center_b{-kHipJointOffset, 0.0f};
  const float radius_a = std::sqrt(SquaredDistance(node, center_a));
  const float radius_b = std::sqrt(SquaredDistance(node, center_b));

  const std::array<float, 4> distances{
      {radius_a - kSmallRadius, radius_b - kSmallRadius,
       kLargeRadius - radius_a, kLargeRadius - radius_b}};
  const float dist = *std::min_element(distances.begin(), distances.end());
  return dist;
}

int main(const int argc, const char* const argv[]) {
  const cho::graph::Workspace ws{{-0.035, -0.044}, {0.035, 0.044}};
  const cho::graph::RrtSettings settings{ws, 512, 1e-3, 1e-3};
  cho::graph::Rrt rrt{settings, cho::graph::SdfFun(WorkspaceDistance)};
  const cho::graph::Node2D& source{0.0032096, 0.03201624};
  const cho::graph::Node2D& target{0.01301248, -0.01502493};
  const std::vector<cho::graph::Node2D>& path =
      rrt.GetTrajectory(source, target);
  fmt::print("Source : {} {}\n", source.x, source.y);
  fmt::print("Target : {} {}\n", target.x, target.y);
  for (const auto& point : path) {
    fmt::print("{} {}\n", point.x, point.y);
  }

  // Eigen::Matrix<std::uint8_t, Eigen::Dynamic, Eigen::Dynamic> grid =
  // decltype(grid)::Zero(3, 5); const cho::VecVector2i& path =
  // cho::graph::FindPath(grid, { 0, 0 }, { 2, 4 }); for (const auto& v : path)
  // {
  //    fmt::print("P {} {}\n", v.x(), v.y());
  //}
  return 0;
}


/**
 * @file 1dwaveabsorbingbc_main.cc
 * @brief NPDE homework "1DWaveAbsorbingBC" code
 * @author Oliver Rietmann
 * @date 08.04.2019
 * @copyright Developed at ETH Zurich
 */

#include <Eigen/Core>
#include <cstdlib>
#include <fstream>
#include <iostream>

#include "1dwaveabsorbingbc.h"

using namespace WaveAbsorbingBC1D;

const static Eigen::IOFormat CSVFormat(Eigen::FullPrecision,
                                       Eigen::DontAlignCols, ", ", "\n");

int main() {
  /* SAM_LISTING_BEGIN_1 */
  double c = 1.0;
  double T = 7.0;
  unsigned int N = 100;
  unsigned int m = 2000;
  Eigen::MatrixXd R = waveLeapfrogABC(c, T, N, m);
  Eigen::VectorXd t = Eigen::VectorXd::LinSpaced(m + 1, 0.0, T);

  // print the data, e.g. to a .csv file, in a suitable way
  std::ofstream solution_file;

  solution_file.open("solution.csv");
  Eigen::MatrixXd tR(m + 1, N + 2);
  tR.col(0) = t;
  tR.block(0, 1, m + 1, N + 1) = R;
  solution_file << tR.format(CSVFormat) << std::endl;
  solution_file.close();
  std::cout << "Generated " CURRENT_BINARY_DIR "/solution.csv" << std::endl;
  std::system("python3 " CURRENT_SOURCE_DIR "/viswave.py " CURRENT_BINARY_DIR
              "/solution.csv " CURRENT_BINARY_DIR "/solution.eps");
  /* SAM_LISTING_END_1 */

  std::pair<Eigen::VectorXd, Eigen::VectorXd> energies =
      computeEnergies(R, c, T / m);
  Eigen::VectorXd E_pot = energies.first;
  Eigen::VectorXd E_kin = energies.second;

  /* SAM_LISTING_BEGIN_2 */
  std::ofstream energies_file;
  energies_file.open("energies.csv");
  energies_file << t.transpose().format(CSVFormat) << std::endl
                << E_pot.transpose().format(CSVFormat) << std::endl
                << E_kin.transpose().format(CSVFormat) << std::endl;
  energies_file.close();
  std::cout << "Generated " CURRENT_BINARY_DIR "/energies.csv" << std::endl;
  std::system("python3 " CURRENT_SOURCE_DIR
              "/visenergies.py " CURRENT_BINARY_DIR
              "/energies.csv " CURRENT_BINARY_DIR "/energies.eps");
/* SAM_LISTING_END_2 */

  return 0;
}


/*(utf8)
Simple moving average filter
*/

#include "filter_base.hpp"
#include <boost/circular_buffer.hpp>

namespace value_filters
{

template <typename T = double>
class floating_average : public filter_base<T>
{

public:
  floating_average()
  {
    set_amount(SCALED_AMOUNT);
  }

  T operator()(T x)
  {
    T sum {0};

    this->buffer.push_back(x);

    for (T v : this->buffer)
      sum += v;

    return sum / this->buffer.size();
  }

  void set_amount(double amt)
  {
    if (amt <= 1)
      amt = 1;
    this->buffer.set_capacity(amt);
  };

private:
  boost::circular_buffer<T> buffer {0};
};

}


#ifndef BODY_HPP
#define BODY_HPP

#include <armadillo>
#include "Math.hpp"

class Body
{
public:
    Body();
    ~Body() {};

    arma::mat get_TBI();
    arma::vec get_POSITION();
    arma::vec get_VELOCITY();
    arma::vec get_ACCELERATION();
    arma::vec get_ANG_VEL();
    arma::vec get_ANGLE();
    arma::vec get_ANGLE_ACC();

protected:
    arma::vec POSITION;
    arma::vec VELOCITY;
    arma::vec ACCELERATION;
    arma::vec ANGLE;
    arma::vec ANGLE_VEL;
    arma::vec ANGLE_ACC;
    arma::mat M;
    arma::mat I;
    arma::vec FORCE;
    arma::vec TORQUE;
    arma::mat TBI;
};

class Ground : public Body
{
public:
    Ground();
    ~Ground() {};
};

class Mobilized_body : public Body
{
public:
    Mobilized_body(arma::vec PosIn, arma::vec VelIn, arma::vec AccIn, arma::vec AttIn
        , arma::vec ANG_VEL_In, arma::vec ANG_ACC_In, double MIn, arma::vec IIn
        , arma::vec F_In, arma::vec T_In);
    ~Mobilized_body();
};

#endif  //BODY_HPP


/*
 * lorenz_gmpxx.cpp
 *
 * This example demonstrates how odeint can be used with arbitrary precision types.
 *
 * Copyright 2011-2012 Karsten Ahnert
 * Copyright 2011-2012 Mario Mulansky
 *
 * Distributed under the Boost Software License, Version 1.0.
 * (See accompanying file LICENSE_1_0.txt or
 * copy at http://www.boost.org/LICENSE_1_0.txt)
 */

#include <iostream>
#include <boost/array.hpp>

#include <gmpxx.h>

#include <boost/numeric/odeint.hpp>

using namespace std;
using namespace boost::numeric::odeint;

//[ gmpxx_lorenz
typedef mpf_class value_type;
typedef boost::array< value_type , 3 > state_type;

struct lorenz
{
    void operator()( const state_type &x , state_type &dxdt , value_type t ) const
    {
        const value_type sigma( 10.0 );
        const value_type R( 28.0 );
        const value_type b( value_type( 8.0 ) / value_type( 3.0 ) );

        dxdt[0] = sigma * ( x[1] - x[0] );
        dxdt[1] = R * x[0] - x[1] - x[0] * x[2];
        dxdt[2] = -b * x[2] + x[0] * x[1];
    }
};
//]

struct streaming_observer
{
    std::ostream& m_out;

    streaming_observer( std::ostream &out ) : m_out( out ) { }

    template< class State , class Time >
    void operator()( const State &x , Time t ) const
    {
        m_out << t;
        for( size_t i=0 ; i<x.size() ; ++i ) m_out << "\t" << x[i] ;
        m_out << "\n";
    }
};

int main( int argc , char **argv )
{
    //[ gmpxx_integration
    const int precision = 1024;
    mpf_set_default_prec( precision );

    state_type x = {{ value_type( 10.0 ) , value_type( 10.0 ) , value_type( 10.0 ) }};

    cout.precision( 1000 );
    integrate_const( runge_kutta4< state_type , value_type >() ,
            lorenz() , x , value_type( 0.0 ) , value_type( 10.0 ) , value_type( value_type( 1.0 ) / value_type( 10.0 ) ) ,
            streaming_observer( cout ) );
    //]

    return 0;
}


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

#include <iostream>
#include <boost/math/tools/ulps_plot.hpp>
#include <boost/core/demangle.hpp>
#include <boost/math/distributions/kolmogorov_smirnov.hpp>

using boost::math::tools::ulps_plot;

int main() {
    using PreciseReal = long double;
    using CoarseReal = float;

    boost::math::kolmogorov_smirnov_distribution<CoarseReal> dist_coarse(10);
    auto pdf_coarse = [&, dist_coarse](CoarseReal x) {
        return boost::math::pdf(dist_coarse, x);
    };
    boost::math::kolmogorov_smirnov_distribution<PreciseReal> dist_precise(10);
    auto pdf_precise = [&, dist_precise](PreciseReal x) {
        return boost::math::pdf(dist_precise, x);
    };

    int samples = 2500;
    int width = 800;
    PreciseReal clip = 100;

    std::string filename1 = "kolmogorov_smirnov_pdf_" + boost::core::demangle(typeid(CoarseReal).name()) + ".svg";
    auto plot1 = ulps_plot<decltype(pdf_precise), PreciseReal, CoarseReal>(pdf_precise, 0.0, 1.0, samples);
    plot1.clip(clip).width(width);
    std::string title1 = "Kolmogorov-Smirnov PDF (N=10) ULP plot at " + boost::core::demangle(typeid(CoarseReal).name()) + " precision";
    plot1.title(title1);
    plot1.vertical_lines(10);
    plot1.add_fn(pdf_coarse);
    plot1.write(filename1);
}


#include <iomanip>
#include <iostream>

#include <mxpfit/approx_sph_bessel.hpp>

// #include <mxpfit/math/sph_bessel.hpp>

#include <boost/math/special_functions/bessel.hpp>

//==============================================================================
// Main
//==============================================================================

using Index              = Eigen::Index;
using Real               = double;
using Complex            = std::complex<Real>;
using RealArray          = Eigen::Array<Real, Eigen::Dynamic, 1>;
using ComplexArray       = Eigen::Array<Complex, Eigen::Dynamic, 1>;
using ExponentialSumType = mxpfit::ExponentialSum<Complex, Complex>;

void sph_bessel_kernel_error(int l, const RealArray& x,
                             const ExponentialSumType& ret,
                             bool verbose_print = false)
{
    RealArray exact(x.size());
    RealArray approx(x.size());

    for (Index i = 0; i < x.size(); ++i)
    {
        // exact(i)  = math::sph_bessel_j(l, x(i));
        exact(i)  = boost::math::sph_bessel(l, x(i));
        approx(i) = std::real(ret(x(i)));
    }

    RealArray abserr(Eigen::abs(exact - approx));

    if (verbose_print)
    {
        for (Index i = 0; i < x.size(); ++i)
        {
            std::cout << std::setw(24) << x(i) << ' '      // point
                      << std::setw(24) << exact(i) << ' '  // exact value
                      << std::setw(24) << approx(i) << ' ' // approximation
                      << std::setw(24) << abserr(i) << '\n';
        }
    }

    Index imax;
    abserr.maxCoeff(&imax);

    std::cout << "\n  abs. error in interval [" << x(0) << ","
              << x(x.size() - 1) << "]\n"
              << "    maximum : " << abserr(imax) << '\n'
              << "    averaged: " << abserr.sum() / x.size() << std::endl;
}

int main()
{
    std::cout.precision(15);
    std::cout.setf(std::ios::scientific);

    const Real threshold = 1.0e-12;
    const Real eps       = Eigen::NumTraits<Real>::epsilon();
    const Index lmax     = 20;
    const Index N        = 1000000; // # of sampling points

    std::cout
        << "# Approximation of spherical Bessel function by exponential sum\n";

    RealArray x = Eigen::pow(10.0, RealArray::LinSpaced(N, -5.0, 7.0));
    ExponentialSumType ret;
    mxpfit::ApproxSphBesselFunction<Real> sph_bessel_approx;
    for (Index l = 0; l <= lmax; ++l)
    {
        std::cout << "\n# --- order " << l;
        ret = sph_bessel_approx.compute(l, threshold);
        const auto thresh_weight =
            std::max(eps, threshold) / std::sqrt(Real(ret.size()));
        ret = mxpfit::removeIf(
            ret, [=](const Complex& /*exponent*/, const Complex& wi) {
                return std::abs(std::real(wi)) < thresh_weight &&
                       std::abs(std::imag(wi)) < thresh_weight;
            });

        const bool verbose = false;
        sph_bessel_kernel_error(l, x, ret, verbose);
        std::cout << " (" << ret.size() << " terms approximation)\n";
        std::cout << "# real(exponent), imag(exponent), real(weight), "
                     "imag(weight)\n";
        for (Index i = 0; i < ret.size(); ++i)
        {
            std::cout << std::setw(24) << std::real(ret.exponent(i)) << '\t'
                      << std::setw(24) << std::imag(ret.exponent(i)) << '\t'
                      << std::setw(24) << std::real(ret.weight(i)) << '\t'
                      << std::setw(24) << std::imag(ret.weight(i)) << '\n';
        }
        std::cout << '\n' << std::endl;

        // std::cout << "# no. of terms and (exponents, weights)\n" << ret <<
        // '\n'; std::cout << "# sum of weights: " << ret.weights().sum() <<
        // '\n';
    }

    return 0;
}


/* Copyright 2021-2022 Massachusetts Institute of Technology

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

        https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
        limitations under the License.
==============================================================================*/

#define CATCH_CONFIG_MAIN

#include <catch2/catch.hpp>

#include <iostream>
#include <Eigen/Dense>

#include <gentl/types.h>
#include <gentl/util/randutils.h>
#include <gentl/inference/particle_filter.h>

using gentl::GenerateOptions;
using gentl::UpdateOptions;

using Eigen::VectorXd;
using Eigen::MatrixXd;
using Eigen::indexing::all;
using std::valarray;

// *******************************************************************
// ** forward algorithm to compute ground truth marginal likelihood **
// *******************************************************************

double hmm_forward_alg(const VectorXd& prior,
                       const MatrixXd& emission_dists,
                       const MatrixXd& transition_dists,
                       const std::vector<size_t>& emissions) {
    assert(prior.rows() == emission_dists.cols());
    assert(prior.rows() == transition_dists.cols());
    assert(transition_dists.rows() == transition_dists.rows());
    double marginal_likelihood = 1.0;
    VectorXd alpha = prior; // copy
    for (auto emission : emissions) {
        auto likelihoods = emission_dists(emission, all).transpose();
        VectorXd prev_posterior = (alpha.array() * likelihoods.array()).matrix();
        double denom = prev_posterior.sum();
        prev_posterior /= denom;
        alpha = transition_dists * prev_posterior;
        marginal_likelihood *= denom;
    }
    return marginal_likelihood;
}

TEST_CASE("hmm forward algorithm", "[particle filtering]") {

    // test the hmm_forward_alg on a hand-calculated example
    VectorXd prior {{0.4, 0.6}};
    MatrixXd emission_dists {
            {0.1, 0.9}, // dist for state 0
            {0.7, 0.3} // dist for state 1
    };
    emission_dists.transposeInPlace();
    MatrixXd transition_dists {
            {0.5, 0.5}, // dist. for state 0
            {0.2, 0.8} // dist. for state 1
    };
    transition_dists.transposeInPlace();
    std::vector<size_t> obs {1, 0};
    double expected_marginal_likelihood = 0.0;
    // z = [0, 0]
    expected_marginal_likelihood += prior(0) * transition_dists(0, 0) * emission_dists(obs[0], 0) * emission_dists(obs[1], 0);
    // z = [0, 1]
    expected_marginal_likelihood += prior(0) * transition_dists(1, 0) * emission_dists(obs[0], 0) * emission_dists(obs[1], 1);
    // z = [1, 0]
    expected_marginal_likelihood += prior(1) * transition_dists(0, 1) * emission_dists(obs[0], 1) * emission_dists(obs[1], 0);
    // z = [1, 1]
    expected_marginal_likelihood += prior(1) * transition_dists(1, 1)* emission_dists(obs[0], 1) * emission_dists(obs[1], 1);
    auto actual_marginal_likelihood = hmm_forward_alg(prior, emission_dists, transition_dists, obs);
    REQUIRE(std::abs(actual_marginal_likelihood - expected_marginal_likelihood) < 1e-16);
}

// ************************************************************************
// ** minimum HMM implementation for testing particle filter correctness **
// ************************************************************************

class HMMParams {
    friend class HMM;
    friend class HMMTrace;
private:
    const size_t num_states_;
    const MatrixXd emission_matrix_;
    std::discrete_distribution<size_t> prior_dist_;
    std::vector<std::discrete_distribution<size_t>> transition_dists_;
public:
    HMMParams(const VectorXd& prior, const MatrixXd& emission_matrix, const MatrixXd& transition_matrix) :
            emission_matrix_(emission_matrix), num_states_(emission_matrix.cols()) {
        if (transition_matrix.rows() != num_states_ || transition_matrix.cols() != num_states_)
            throw std::logic_error("dimension mismatch");
        prior_dist_ = std::discrete_distribution<size_t>(prior.cbegin(), prior.cend());
        for (auto i = 0; i < num_states_; i++) {
            auto column_vector = transition_matrix(all, i);
            assert(std::abs(column_vector.sum() - 1.0) < 1e-16);
            transition_dists_.emplace_back(column_vector.begin(), column_vector.end());
        }
    }
};

class ParameterStore {};

class Extend {};

struct NewObservation {
    size_t value;
    NewObservation(size_t value_) : value(value_) {}
};

class HMMTrace;

class HMM {
    friend class HMMTrace;
    size_t num_time_steps_;
    HMMParams& params_;
private:
public:
    explicit HMM(HMMParams& params) : num_time_steps_(0), params_(params) {}
    std::pair<std::unique_ptr<HMMTrace>,double> generate(
            std::mt19937& rng, ParameterStore&, const NewObservation& observation,
            const GenerateOptions& options) const;

};

class HMMTrace {
    friend class HMM;
    HMM model_;
    std::vector<size_t> emissions_;
    std::vector<size_t> latents_;
private:
    HMMTrace(const HMMTrace& other) = default;
    HMMTrace(HMM model, size_t emission, size_t latent) :
            emissions_(std::initializer_list<size_t>({emission})),
            latents_(std::initializer_list<size_t>({latent})),
            model_(model) {}
public:
    double update(
            std::mt19937& rng, const Extend&, const NewObservation& observation,
            const UpdateOptions& options) {
        if (options.save() || options.precompute_gradient())
            throw std::logic_error("not implemented");
        size_t latent = model_.params_.transition_dists_[latents_.back()](rng);
        double log_weight = std::log(model_.params_.emission_matrix_(observation.value, latent));
        latents_.emplace_back(latent);
        emissions_.emplace_back(observation.value);
        model_.num_time_steps_ += 1;
        return log_weight;
    }
    std::unique_ptr<HMMTrace> fork() {
        // NOTE: this trace implementation is not efficient, it copies entire histories unnecessarily
        return std::unique_ptr<HMMTrace>(new HMMTrace(*this));
    }
};

std::pair<std::unique_ptr<HMMTrace>,double> HMM::generate(
        std::mt19937& rng, ParameterStore&, const NewObservation& observation,
        const GenerateOptions& options) const {
    if (options.precompute_gradient())
        throw std::logic_error("not implemented");
    size_t latent = params_.prior_dist_(rng);
    double log_weight = std::log(params_.emission_matrix_(observation.value, latent));
    auto trace = std::unique_ptr<HMMTrace>(new HMMTrace(*this, observation.value, latent));
    return {std::move(trace), log_weight};
}

TEST_CASE("hmm particle filter", "[particle filtering]") {

    gentl::randutils::seed_seq_fe128 seed_seq {0};
    std::mt19937 rng(seed_seq);

    VectorXd prior {{0.2, 0.3, 0.5}};
    MatrixXd emission_matrix {
            {0.1, 0.2, 0.7},
            {0.2, 0.7, 0.1},
            {0.7, 0.2, 0.1}
    };
    emission_matrix.transposeInPlace();
    MatrixXd transition_matrix {
            {0.4, 0.4, 0.2},
            {0.2, 0.3, 0.5},
            {0.9, 0.05, 0.05}
    };
    transition_matrix.transposeInPlace();
    HMMParams params {
        prior, emission_matrix, transition_matrix
    };

    HMM model {params};

    // test particle filter
    std::vector<size_t> data = {0, 0, 1, 2};
    double expected = std::log(hmm_forward_alg(prior, emission_matrix, transition_matrix, data));

    std::vector<NewObservation> observations;
    for (auto datum : data)
        observations.emplace_back(datum);

    size_t num_particles = 10000;
    gentl::smc::ParticleSystem<HMMTrace,std::mt19937> filter{num_particles, rng};
    auto observations_it = observations.cbegin();
    ParameterStore store{};
    filter.init_step(model, store, *observations_it++);
    using std::cerr, std::endl;
    while (observations_it != observations.cend()) {
        filter.step(Extend{}, *observations_it++);
        double ess = filter.effective_sample_size();
        cerr << "effective sample size: " << ess << endl;
        double log_weight = filter.resample();
        cerr << "log weight from resample: " << log_weight << endl;
    }
    double actual = filter.log_marginal_likelihood_estimate();
    cerr << "actual: " << actual << ", expected: " << expected << endl;
    REQUIRE(std::abs(actual - expected) < 0.02);
}


<%
cfg['compiler_args'] = ['-std=c++11','-DUSE_DOUBLE']
cfg['include_dirs'] = ['eigen', 'include']
cfg['sources'] = ['./HODLR_Matrix.cpp',
    './HODLR_Node.cpp',
    './HODLR_Tree.cpp',
    './HODLR_Tree_NonSPD.cpp',
    './HODLR_Tree_SPD.cpp',
    './KDTree.cpp',
    './sample_funs.cpp',
    './predict.cpp']
setup_pybind11(cfg)
%>

#include <squaredeMat.hpp> //squared exponential kernel
#include <squaredeP1Mat.hpp>
#include <random>
#include <chrono>
#include <iostream>
#include <fstream>
#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <Eigen/Dense>
#include "HODLR_Tree.hpp"
#include "sample_funs.h"

using std::gamma_distribution;
using std::uniform_real_distribution;
using std::normal_distribution;

namespace py = pybind11;

void predict_module(py::module &);

Eigen::VectorXd samplefstar_HODLR(Eigen::MatrixXd X, Eigen::MatrixXd Y, Eigen::MatrixXd Xtest, double sig, double rho, double tau, HODLR_Tree* T, double multiplier, std::function<float()> &r_std_normal, int s) {

    /*
       Sample a draw of the GP function f*|f, sig, rho, tau, x*.
       Assume a squared exponential Gaussian Process based on
       observed function f at new test points x*.
    */

    int Ntest =  Xtest.rows();
    int N =  X.rows();
    int D =  X.cols();

    double sigsq = pow(sig, 2.0);
    double tmpSSR;

    Eigen::VectorXd fstarsamp(Ntest);
    Eigen::VectorXd KobsNew(N);

    for (int i = 0; i < Ntest; i++) {
      Eigen::RowVectorXd Xtest_i = Xtest.row(i);

      // Covariance between X and Xtest
      for (int j = 0; j < N; j++) {
        Eigen::RowVectorXd tmp = X.row(j) - Xtest_i;
        tmpSSR = 0.0;
        for (int d = 0; d < D; d++) {
          tmpSSR = tmpSSR + pow(tmp(d), 2.0);
        }
        KobsNew(j) = sigsq * exp(- tmpSSR * rho);
      }

      // Variance at Xtest
      double kNewNew = sigsq + 1e-8;

      // Posterior mean and variance of f* at xtest(i)
      double sdstar = pow(kNewNew - (multiplier * KobsNew.transpose() * T->solve(tau * KobsNew))(0, 0), 0.5);
      double mustar = (multiplier * KobsNew.transpose() * T->solve(tau * Y))(0, 0);

      auto normal_samp = r_std_normal();

      fstarsamp(i) =  sdstar * normal_samp + mustar;
    }

    return fstarsamp;
}

Eigen::VectorXd get_transition_probs_byind(Eigen::VectorXi rhovec_inds, double krnl, int ind_current) {
  /*
    Creates transition probability matrix for Metropolis Hastings update of lengthscale
  */
  int nrhos = rhovec_inds.rows();
  Eigen::VectorXd prop_prob(nrhos);
  for (int i = 0; i < nrhos; i++) {
    prop_prob(i) = exp(-krnl * pow((double)rhovec_inds(i) - (double)ind_current, 2.0));
  }
  prop_prob(ind_current) = 0.0;
  double sumInv = 1.0 / prop_prob.sum();
  prop_prob = sumInv * prop_prob;
  return prop_prob;
}

py::list sampleGP_HODLR(Eigen::MatrixXd X, Eigen::VectorXd Y,
                          double sig, Eigen::VectorXd rho_choices, double tau,
                          bool regression=true,
                          bool Gibbs_ls=false, bool default_MHkernel=true,
                          double numeric_MHkernel=-1.0,
                          double a_f=1, double b_f=1,
                          double a_tau=1, double b_tau=1,
                          double tol=1e-12, int M=20,
                          bool save_fsamps=true, unsigned int seed=169,
                          int burnin=1000, int nsamps=100, int thin=10, bool verbose=false) {
  /*
    C++ sample hyperparameter and function draws
    from a squared exponential Gaussian Process
    based on observed data y having precision tau.
    X   -matrix of locations
    Y   -matrix of observed data
    tau -data precision
    sig -function std deviation
    rho -length-scale
    a_f and b_f - prior sig^(-2) ~ Ga(a_f/2, b_f/2)
    a_tau and b_tau - prior tau ~ Ga(a_tau/2, b_tau/2)
    tol -specified tolerance for accuracy of calculations
    Gibbs_ls: bool, whether to use Gibbs or MH to sample length scale
    M   -Max submatrix size
  */

  // For sampling random numbers
  std::mt19937 generator(seed);

  // For sampling standard uniforms (for MH acceptance evaluation)
  uniform_real_distribution<float> unif(0, 1);
  std::function<float()> runif = bind(unif, generator);

  // For sampling standard normals (e.g. for sampling f)
  normal_distribution<float> normal(0, 1);
  std::function<float()> r_std_normal = bind(normal, generator);

  // Data sizes and initializing fsamp and params / transformed params
  int N = X.rows();

  Eigen::VectorXd KobsNew(N);
  Eigen::VectorXd fsamp(N);

  double sqrt_tau;
  if (!regression) { // If you observe non-noisy data
    fsamp = Y;
  }
  else {
    sqrt_tau = pow(tau, 0.5);
  }

  double sigsq = pow(sig, 2.0);
  double prec_f = 1 / sigsq;

  // Set up HODLR details
  int n_levels = log(N / M) / log(2);
  bool is_sym = true;
  bool is_pd  = true;

  // Create rho and matrices for each value option
  int nrhos = rho_choices.rows();
  Eigen::MatrixXd transition_probs_mat(nrhos,nrhos);
  bool fix_ls;
  int rho_ind;

  Eigen::VectorXd logdetK_all(nrhos);
  SQRExponential_Kernel* K_tmp;
  HODLR_Tree* S_tmp;
  std::vector<HODLR_Tree*> Svec;
  Svec.reserve(nrhos);

  if (nrhos != 1) {
    fix_ls = false;
    rho_ind = round(nrhos / 2);
    if (!Gibbs_ls) {
      // Make transition probability matrix from each "view" of the data
      // (i.e., from each index at which you could make a proposal)
      bool cond1 = (numeric_MHkernel > 0);
      double krnl;
      if (cond1 | default_MHkernel) {
        if (cond1) {
          krnl = numeric_MHkernel;
        }

        if (default_MHkernel) {
          krnl = 1.0 / pow((double)nrhos, 1.0);
        }

        Eigen::VectorXi rhovec_inds(nrhos);
        for (int i = 0; i < nrhos; i++) {
          rhovec_inds(i) = i;
        }
        for (int i = 0; i < nrhos; i++) { // Each col is "view" from that index.
          transition_probs_mat.col(i) = get_transition_probs_byind(rhovec_inds, krnl, i);
        }
      } else{
        printf("default_MHkernel must be TRUE or numeric_MHkernel needs to be set > 0(see function description for details).");
      }
    }
  }
  else {
    fix_ls = true;
    rho_ind = 0;
    if (verbose) {
      printf("Length scale fixed to provided input: %.2f.\n", rho_choices(rho_ind));
      printf("To sample length scale from discrete options provide vector rho_choices.\n");
    }
  }

  // Svec here is K in the paper
  for (int i = 0; i < nrhos; i++) {
    // Set up squared exponential kernel K with sigf fixed to 1.
    K_tmp  = new SQRExponential_Kernel(X, N, 1, rho_choices(i));
    S_tmp  = new HODLR_Tree(n_levels, tol, K_tmp); // Without noise (i.e., Sigma)
    Svec.push_back(S_tmp);
    Svec[i]->assembleTree(is_sym, is_pd);
    Svec[i]->factorize();
    logdetK_all(i) = Svec[i]->logDeterminant();
  }

  // Initialize first sample of rho
  double rho = rho_choices(rho_ind);

  // Set up storage
  int samp_count = 0;
  int total_draws = burnin + nsamps * thin;
  Eigen::VectorXd tau_save(nsamps);
  Eigen::VectorXd sigf_save(nsamps);
  Eigen::VectorXd rho_save(nsamps);
  int nsamps_f = nsamps;
  if (!save_fsamps) {
    nsamps_f = 0;
  }
  Eigen::MatrixXd f_save(N, nsamps_f);

  for (int s = 0; s < total_draws; s++) {
    bool save_samps = (s >= burnin) & (((s + 1) % thin) == 0);

    // Non-noisy data
    if (regression) {

      // Assemble the kernel with noise term tau
      SQRExponentialP1_Kernel* L  = new SQRExponentialP1_Kernel(X, N, sig, rho, tau);
      HODLR_Tree* T = new HODLR_Tree(n_levels, tol, L); // With noise (i.e. Sigma + I/tau)
      T->assembleTree(is_sym, is_pd);
      T->factorize();

      // Sample f
      fsamp = samplef_HODLR(X, Y, sig, tau, T, Svec[rho_ind], r_std_normal);

      delete T;
      delete L;
    }

    // Sample function variance
    prec_f = sample_prec_f(Svec[rho_ind], fsamp, a_f, b_f, generator);
    sig = pow(prec_f, -0.5);
    sigsq = pow(sig, 2.0);

    // Sample length scale
    if (!fix_ls) {
      if (Gibbs_ls) {
        rho_ind = sample_rho_gibbs(prec_f, fsamp, Svec, logdetK_all, generator);
      }
      else {
        rho_ind = sample_rho_mh(rho_ind, transition_probs_mat, prec_f, fsamp, Svec, logdetK_all, generator, runif);
      }
      rho = rho_choices(rho_ind);
    }

    // Sample noise precision
    if (regression) {
      tau = sample_tau(Y, fsamp, a_tau, b_tau, generator);
      sqrt_tau = pow(tau, 0.5);
    }

    if ((s == burnin) and verbose) {
      printf("Finished with burnin, beginning sampling. \n");
    }
    if ((10 * (s + 1) % total_draws == 0) and verbose) {
      printf("%d%% done \n", 100 * (s + 1) / total_draws);
    }

    // Save samples if you're past burnin
    if (save_samps) {
      if (regression) {
        tau_save(samp_count) = tau;
      }
      sigf_save(samp_count) = sig;
      rho_save(samp_count) = rho;

      if (save_fsamps) {
        f_save.col(samp_count) = fsamp;
      }

      samp_count = samp_count + 1;
    }

  } // END for (int s=0; s < total_draws; s++)

  delete S_tmp;
  delete K_tmp;
  std::vector<HODLR_Tree*>().swap(Svec);

  py::list out;
  out.append(tau_save);
  out.append(sigf_save);
  out.append(rho_save);
  out.append(f_save);
  return out;
}

PYBIND11_MODULE(sample, m) {
    m.doc() = "C++ Implementation of FIFA-GP.";
    m.def("samplegp", &sampleGP_HODLR, "1-d GP.");
    predict_module(m);
}


#include "CGAL_point_transform_function.hpp"
#include <CGAL/Exact_predicates_inexact_constructions_kernel.h>
#include <CGAL/Regular_triangulation_3.h>
#include <CGAL/Regular_triangulation_euclidean_traits_3.h>
#include <CGAL/Fixed_alpha_shape_3.h>
#include <CGAL/Triangulation_vertex_base_with_info_3.h>
#include <CGAL/Fixed_alpha_shape_vertex_base_3.h>
#include <CGAL/Fixed_alpha_shape_cell_base_3.h>
#include <boost/iterator/counting_iterator.hpp>
#include <boost/iterator/transform_iterator.hpp>
#include <utility>
#include <iterator>
#include <stdexcept>

#include <iostream>

typedef CGAL::Exact_predicates_inexact_constructions_kernel K;
typedef CGAL::Regular_triangulation_euclidean_traits_3<K> Gt;
typedef CGAL::Triangulation_vertex_base_with_info_3<unsigned int, Gt> Tb;
typedef CGAL::Fixed_alpha_shape_vertex_base_3<Gt, Tb> Vb;
typedef CGAL::Fixed_alpha_shape_cell_base_3<Gt> Fb;
typedef CGAL::Triangulation_data_structure_3<Vb,Fb> Tds;
typedef CGAL::Regular_triangulation_3<Gt,Tds> Triangulation_3;
typedef CGAL::Fixed_alpha_shape_3<Triangulation_3> Fixed_alpha_shape_3;
typedef Fixed_alpha_shape_3::Cell_handle Cell_handle;
typedef Fixed_alpha_shape_3::Vertex_handle Vertex_handle;
typedef Fixed_alpha_shape_3::Facet Facet;
typedef Fixed_alpha_shape_3::Edge Edge;
typedef Gt::Weighted_point Weighted_point;
typedef Gt::Bare_point Bare_point;

template <typename T>
T* init_mem(unsigned int size, const T& init_val){
	T* mem = (T*)malloc(sizeof(T)*size);
	if(!mem){
		throw std::runtime_error("Insufficient memory");
	}
	for(unsigned int i = 0; i < size; ++i){
		mem[i] = init_val;
	}
	return mem;
}

void alpha_shapes(
	unsigned int num_points, double alpha,
	const double* points, const double* weights,
	unsigned int* num_edges, unsigned int** edges,
	unsigned int* num_triangles, unsigned int** triangles,
	unsigned int* num_tetrahedra, unsigned int** tetrahedra
){
	busv::CGALPointTransformFunction<Gt> tfunc(points, weights);
	unsigned int i;

	std::cout << "alpha " << alpha << std::endl;

	//build one alpha_shape with alpha=0
	Fixed_alpha_shape_3 as(
		boost::make_transform_iterator(boost::counting_iterator<unsigned int>(0), tfunc),
		boost::make_transform_iterator(boost::counting_iterator<unsigned int>(num_points), tfunc),
		alpha
	);

	i = 0;
	*num_edges = std::distance(as.finite_edges_begin(), as.finite_edges_end());
	*edges = init_mem<unsigned int>((*num_edges)*2, 0);
	for(Fixed_alpha_shape_3::Finite_edges_iterator it = as.finite_edges_begin(), eit = as.finite_edges_end(); it!= eit; ++it){
		Fixed_alpha_shape_3::Edge e = *it;
		//if(as.classify(e) == Fixed_alpha_shape_3::SINGULAR){
		if(as.classify(e) != Fixed_alpha_shape_3::EXTERIOR){
			(*edges)[i] = e.first->vertex(e.second)->info();
			(*edges)[i+1] = e.first->vertex(e.third)->info();
			i += 2;
		}
//		(*edges)[i] = it->vertex(0)->info();
//		(*edges)[i+1] = it->vertex(1)->info();
	}
	*num_edges = i/2;

	i = 0;
	*num_triangles = std::distance(as.finite_facets_begin(), as.finite_facets_end());
	*triangles = init_mem<unsigned int>((*num_triangles)*3, 0);
	for(Fixed_alpha_shape_3::Finite_facets_iterator it = as.finite_facets_begin(), eit = as.finite_facets_end(); it!= eit; ++it){
		Fixed_alpha_shape_3::Facet f = *it;
		//if(as.classify(f) == Fixed_alpha_shape_3::SINGULAR){
		if(as.classify(f) != Fixed_alpha_shape_3::EXTERIOR){
			(*triangles)[i] = f.first->vertex((f.second+1)%4)->info();
			(*triangles)[i+1] = f.first->vertex((f.second+2)%4)->info();
			(*triangles)[i+2] = f.first->vertex((f.second+3)%4)->info();
			i += 3;
		}
	}
	*num_triangles = i/3;

	i = 0;
	*num_tetrahedra = std::distance(as.finite_cells_begin(), as.finite_cells_end());
	*tetrahedra = init_mem<unsigned int>((*num_tetrahedra)*4, 0);
	for(Fixed_alpha_shape_3::Finite_cells_iterator it = as.finite_cells_begin(), eit = as.finite_cells_end(); it!= eit; ++it){
		Fixed_alpha_shape_3::Cell c = *it;
		//if(as.classify(it) == Fixed_alpha_shape_3::INTERIOR){
//		if(as.classify(it) != Fixed_alpha_shape_3::EXTERIOR){
			(*tetrahedra)[i] = c.vertex(0)->info();
			(*tetrahedra)[i+1] = c.vertex(1)->info();
			(*tetrahedra)[i+2] = c.vertex(2)->info();
			(*tetrahedra)[i+3] = c.vertex(3)->info();
			i += 4;
//		}
	}
	*num_tetrahedra = i/4;
}


#ifndef TASK_VQEPARAMETERGENERATOR_HPP_
#define TASK_VQEPARAMETERGENERATOR_HPP_

#include "XACC.hpp"
#include <Eigen/Dense>
#include <boost/math/constants/constants.hpp>

namespace xacc {
namespace vqe {

class VQEParameterGenerator {

public:

	static Eigen::VectorXd generateParameters(const int nParameters, std::shared_ptr<Communicator> comm) {

		if (xacc::optionExists("vqe-parameters")) {
			if (xacc::getOption("vqe-task") == "sweep-1d") {
				auto paramStr = xacc::getOption("vqe-parameters");

				// HERE WE COULD HAVE SOMETHING LIKE EITHER
				// 50:-3.14,3.14 or
				// -3.14,3.14
				std::vector<std::string> split, colon;
				boost::split(split, paramStr, boost::is_any_of(","));
				int nSteps = 50;
				std::string minVal;
				if (boost::contains(paramStr, ":")) {
					boost::split(colon, split[0], boost::is_any_of(":"));
					nSteps = std::stoi(colon[0]);
					minVal = colon[1];
				} else {
					minVal = split[0];
				}

				return Eigen::VectorXd::LinSpaced(nSteps, std::stod(minVal),
						std::stod(split[1]));
			} else {
				auto paramStr = xacc::getOption("vqe-parameters");
				std::vector<std::string> split;
				boost::split(split, paramStr, boost::is_any_of(","));
				Eigen::VectorXd params(nParameters);
				for (int i = 0; i < split.size(); i++) {
					params(i) = std::stod(split[i]);
				}

				return params;
			}
		} else {
			std::srand(time(0));
			auto pi = boost::math::constants::pi<double>();
			Eigen::VectorXd rand;

			std::vector<double> data;

			// Random parameters between -pi and pi
			if (comm->rank() == 0) {
				rand = -1.0 * pi * Eigen::VectorXd::Ones(nParameters)
						+ (Eigen::VectorXd::Random(nParameters) * 0.5
								+ Eigen::VectorXd::Ones(nParameters) * 0.5)
								* (pi - (-1 * pi));
				data.resize(rand.size());
				Eigen::VectorXd::Map(&data[0], rand.size()) = rand;
			}

			comm->broadcast(data, 0);

			if (comm->rank() != 0) {
				rand = Eigen::Map<Eigen::VectorXd>(data.data(), data.size());
			}
			return rand;
		}
	}

};
}
}

#endif /* VQETASK_STATEPREPARATIONEVALUATOR_HPP_ */


#include <iostream>
#include <g2o/core/base_vertex.h>
#include <g2o/core/base_unary_edge.h>
#include <g2o/core/block_solver.h>
#include <g2o/core/optimization_algorithm_levenberg.h>
#include <g2o/core/optimization_algorithm_gauss_newton.h>
#include <g2o/core/optimization_algorithm_dogleg.h>
#include <g2o/solvers/dense/linear_solver_dense.h>
#include <Eigen/Core>
#include <opencv2/core/core.hpp>
#include <cmath>
#include <chrono>

using namespace std;

/* Curve fitting for function "y = e^(a*x*x + b*x + c)"
Only need a single vertex with three parameter a,b,c in it.
All edges are unary edges connected to the only vertex
*/

// vertices as optimization target
// inherent from g2o::BaseVertex, predefined Vertex
class CurveFittingVertex: public g2o::BaseVertex<3, Eigen::Vector3d>
{
    public:
        EIGEN_MAKE_ALIGNED_OPERATOR_NEW
        virtual void setToOriginImpl(){
            _estimate << 0, 0, 0;
        }

        virtual void oplusImpl( const double* update ) // 更新
        {
            _estimate += Eigen::Vector3d(update);
        }

        // read and write leave empty
        virtual bool read( istream& in ) {}
        virtual bool write( ostream& out ) const {}
};

// edges to compute error
class CurveFittingEdge: public g2o::BaseUnaryEdge<1, double, CurveFittingVertex>
{
    public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    CurveFittingEdge( double x ): BaseUnaryEdge(), _x(x) {}

    void computeError(){
        const CurveFittingVertex* v = static_cast<const CurveFittingVertex*>(_vertices[0]);
        const Eigen::Vector3d abc = v->estimate();
        _error(0, 0) = _measurement - std::exp(abc(0, 0)*_x*_x + abc(1, 0)*_x + abc(2, 0));
    }
    virtual bool read( istream& in ) {}
    virtual bool write( ostream& out ) const {}
    public:
        double _x;
};

int main(){
    cv::RNG rng;
    double a=1.0, b=2.0, c=1.0;   // Real parameters
    int N = 200;
    double w_sigma=1.0;
<<<<<<< HEAD
=======

>>>>>>> d3aa7eab7a66633ee6b6716534884a999f4bfb7a
    vector<double> x_data, y_data;      // generated data

    cout<<"generating data: "<<endl;
    for ( int i=0; i<N; i++ )
    {
        double x = i/100.0;
        x_data.push_back ( x );
        y_data.push_back (
            exp ( a*x*x + b*x + c ) + rng.gaussian ( w_sigma )
        );
        cout<<x_data[i]<<" "<<y_data[i]<<endl;
    }

    double parameters[3] = {0,0,0};  // Parameters need to be solved

    typedef g2o::BlockSolver<g2o::BlockSolverTraits<3, 1>> Block;   // input 3 dims, output 1 dims
    Block::LinearSolverType* linearSolver = new g2o::LinearSolverDense<Block::PoseMatrixType>();

    Block *solver_ptr = new Block(linearSolver);

    g2o::OptimizationAlgorithmLevenberg *solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);

    g2o::SparseOptimizer optimizer;
    optimizer.setAlgorithm (solver);
    optimizer.setVerbose(true);

    // add Vertex as optimization target

    CurveFittingVertex *v = new CurveFittingVertex();
    v->setEstimate(Eigen::Vector3d(0, 0, 0));
    v->setId(0);
    optimizer.addVertex(v);

    for(int i=0;i<N;i++){
        CurveFittingEdge *edge = new CurveFittingEdge(x_data[i]);
        edge->setId(i);
        edge->setVertex(0, v);
        edge->setMeasurement(y_data[i]);
        edge->setInformation(Eigen::Matrix<double, 1, 1>::Identity()*1/(w_sigma*w_sigma));
        optimizer.addEdge(edge);
    }

    // Execute optimization.
    chrono::steady_clock::time_point t1 = chrono::steady_clock::now();
    optimizer.initializeOptimization();
    optimizer.optimize(100);
    chrono::steady_clock::time_point t2 = chrono::steady_clock::now();

    chrono::duration<double> time_used = chrono::duration_cast<chrono::duration<double>>( t2-t1 );
    cout<<"solve time cost = "<<time_used.count()<<" seconds. "<<endl;

    Eigen::Vector3d parameter_estimate = v->estimate();
    cout << "Estimate model" << parameter_estimate.transpose() << endl;

}