/// Causal Dynamical Triangulations in C++ using CGAL
///
/// Copyright © 2016-2017 Adam Getchell
///
/// Data structures for operations on simplicial manifolds
///
/// \done Classify cells as (3,1), (2,2), or (1,3) based on their foliation.
/// A tuple of vectors contain cell handles to the simplices of type (3,1),
/// (2,2), and (1,3) respectively.
/// \done Classify edges as timelike or spacelike so that action can be
/// calculated.
/// \done SimplicialManifold data structure holding a std::unique_ptr to
/// the Delaunay triangulation and a std::tuple of geometry information.
/// \done Move constructor recalculates geometry.
///
/// @file  SimplicialManifold.hpp
/// @brief Data structures for simplicial manifolds
/// @author Adam Getchell
/// @todo Deprecated in favor of Manifold.hpp

#ifndef INCLUDE_SIMPLICIALMANIFOLD_HPP_
#define INCLUDE_SIMPLICIALMANIFOLD_HPP_

#include <S3Triangulation.hpp>
#include <boost/optional.hpp>
#include <map>
#include <memory>
#include <set>
#include <utility>
#include <vector>

using Facet = Delaunay3::Facet;

/// @brief A tuple of the geometric values of the Simplicial Manifold
///
/// The first element is the vector of (3,1) simplices
/// The second element is the vector of (2,2) simplices
/// The third element is the vector of (1,3) simplices
/// The fourth element is the vector of timelike edges
/// The fifth element is the vector of spacelike edges
/// The sixth element is the vector of vertices
///
/// Useful for constructing GeometryInfo, which contains this and other
/// information, and in comparing the results of moves, which will change one or
/// more elements of the Geometry_tuple.
using Geometry_tuple =
    std::tuple<std::vector<Cell_handle>, std::vector<Cell_handle>,
               std::vector<Cell_handle>, std::vector<Edge_handle>,
               std::vector<Edge_handle>, std::vector<Vertex_handle>>;

// Non-member non-friend functions

/// @brief Classifies edges
///
/// This function iterates over all edges in the triangulation
/// and classifies them as timelike or spacelike.
/// Timelike edges are stored in the **timelike_edges** vector as an Edge_handle
/// (tuple of Cell_handle, std::size_t, std::size_t) for later use by
/// ergodic moves on timelike edges. Spacelike edges are also stored as a
/// vector of Edge_handle **spacelike_edges**, for use by (4,4) moves as
/// well as the distance-finding algorithms.
/// @param[in] universe_ptr A std::unique_ptr<Delaunay> to the triangulation
/// @returns A std::pair<std::vector<Edge_handle>, std::vector<Edge_handle>> of
/// timelike edges and spacelike edges
template <typename T>
[[deprecated]] auto classify_edges(T&& universe_ptr)
{
#ifndef NDEBUG
  std::cout << "Classifying edges....\n";
#endif
  Delaunay3::Finite_edges_iterator eit;
  std::vector<Edge_handle>         timelike_edges;
  std::vector<Edge_handle>         spacelike_edges;

  // Iterate over all edges in the Delaunay triangulation
  for (eit = universe_ptr->finite_edges_begin();
       eit != universe_ptr->finite_edges_end(); ++eit)
  {
    Cell_handle ch = eit->first;
    // Get timevalues of vertices at the edge ends
    auto time1 = ch->vertex(eit->second)->info();
    auto time2 = ch->vertex(eit->third)->info();

    // Make Edge_handle
    Edge_handle thisEdge{
        ch, static_cast<std::size_t>(ch->index(ch->vertex(eit->second))),
        static_cast<std::size_t>(ch->index(ch->vertex(eit->third)))};

    if (time1 != time2)
    {  // We have a timelike edge
      timelike_edges.emplace_back(thisEdge);

#ifdef DETAILED_DEBUGGING
      std::cout << "First vertex of edge is " << std::get<1>(thisEdge)
                << " and second vertex of edge is " << std::get<2>(thisEdge)
                << std::endl;
#endif
    }
    else
    {  // We have a spacelike edge
      spacelike_edges.emplace_back(thisEdge);
    }  // endif
  }    // Finish iterating over edges

// Display results if debugging
#ifndef NDEBUG
  std::cout << "There are " << timelike_edges.size() << " timelike edges and "
            << spacelike_edges.size() << " spacelike edges.\n";
#endif
  return std::make_pair(timelike_edges, spacelike_edges);
}  // classify_edges()

/// @brief Classify simplices as (3,1), (2,2), or (1,3)
///
/// This function iterates over all cells in the triangulation
/// and classifies them as:
/// \f{eqnarray*}{
///   31 &=& (3, 1) \\
///   22 &=& (2, 2) \\
///   13 &=& (1, 3)
/// \f}
/// The vectors **three_one**, **two_two**, and **one_three** contain
/// Cell_handles to all the simplices in the triangulation of that corresponding
/// type.
///
/// @param[in] universe_ptr A std::unique_ptr<Delaunay> to the triangulation
/// @returns A std::tuple<std::vector, std::vector, std::vector> of
/// **three_one**, **two_two**, and **one_three**
template <typename T>
[[deprecated]] auto classify_simplices(T&& universe_ptr)
{
#ifndef NDEBUG
  std::cout << "Classifying simplices....\n";
#endif
  Delaunay3::Finite_cells_iterator cit;
  std::vector<Cell_handle>         three_one;
  std::vector<Cell_handle>         two_two;
  std::vector<Cell_handle>         one_three;

  // Iterate over all cells in the Delaunay triangulation
  for (cit = universe_ptr->finite_cells_begin();
       cit != universe_ptr->finite_cells_end(); ++cit)
  {
    int max_values{0};
    int min_values{0};
    // Push every time value of every vertex into a list
    int timevalues[4] = {
        cit->vertex(0)->info(),
        cit->vertex(1)->info(),
        cit->vertex(2)->info(),
        cit->vertex(3)->info(),
    };
    int max_time =
        *std::max_element(std::begin(timevalues), std::end(timevalues));
    for (auto elt : timevalues)
    {
      if (elt == max_time) { ++max_values; }
      else
      {
        ++min_values;
      }
    }

    // Classify simplex using max_values and write to cit->info()
    if (min_values == 1 && max_values == 3)
    {
      cit->info() = 13;
      one_three.emplace_back(cit);
    }
    else if (min_values == 2 && max_values == 2)
    {
      cit->info() = 22;
      two_two.emplace_back(cit);
    }
    else if (min_values == 3 && max_values == 1)
    {
      cit->info() = 31;
      three_one.emplace_back(cit);
    }
    else
    {
      throw std::runtime_error("Invalid simplex in classify_simplices()!");
    }  // endif
  }    // Finish iterating over cells

// Display results if debugging
#ifndef NDEBUG
  std::cout << "There are " << three_one.size() << " (3,1) simplices and "
            << two_two.size() << " (2,2) simplices\n";
  std::cout << "and " << one_three.size() << " (1,3) simplices.\n";
#endif
  return std::make_tuple(three_one, two_two, one_three);
}  // classify_simplices()

template <typename T>
[[deprecated]] auto classify_all_simplices(T&& universe_ptr)
{
#ifndef NDEBUG
  std::cout << "Classifying all simplices....\n";
#endif

  auto                       cells = classify_simplices(universe_ptr);
  auto                       edges = classify_edges(universe_ptr);
  std::vector<Vertex_handle> vertices;
  for (auto vit = universe_ptr->finite_vertices_begin();
       vit != universe_ptr->finite_vertices_end(); ++vit)
  { vertices.emplace_back(vit); }
  return std::make_tuple(std::get<0>(cells), std::get<1>(cells),
                         std::get<2>(cells), edges.first, edges.second,
                         vertices);
}

/// @struct
/// @brief A struct containing detailed geometry information
///
/// GeometryInfo contains information about the geometry of
/// a triangulation. In addition, it defines convenient functions to
/// retrieve commonly used values. This is to save the expense of
/// calculating manually from the triangulation. GeometryInfo() is
/// recalculated using the move assignment operator anytime a
/// SimplicialManifold() is move constructed.
/// The default constructor, destructor, move constructor, copy
/// constructor, and copy assignment operator are explicitly defaulted.
/// See http://en.cppreference.com/w/cpp/language/rule_of_three
struct [[deprecated]] GeometryInfo
{
 private:
  /// @brief (3,1) cells in the foliation
  std::vector<Cell_handle> three_one;

  /// @brief (2,2) cells in the foliation
  std::vector<Cell_handle> two_two;

  /// @brief (1,3) cells in the foliation
  std::vector<Cell_handle> one_three;

  /// @brief Edges spanning two adjacent time slices in the foliation
  std::vector<Edge_handle> timelike_edges;

  /// @brief Non-spanning edges in the foliation
  std::vector<Edge_handle> spacelike_edges;

  /// @brief Vertices of the foliation
  std::vector<Vertex_handle> vertices;

  /// @brief Spacelike facets for each timeslice
  boost::optional<std::multimap<std::size_t, Facet>> spacelike_facets;

  /// @brief Actual timevalues of simulation
  boost::optional<std::set<std::size_t>> timevalues;

 public:
  /// @brief Getter for spacelike facets
  /// @return The multimap of facets
  const boost::optional<std::multimap<std::size_t, Facet>>&
  getSpacelike_facets() const
  {
    return spacelike_facets;
  }

  /// @brief Setter for spacelike facets
  /// @param spacelike_facets The multimap of facets
  void setSpacelike_facets(
      const boost::optional<std::multimap<std::size_t, Facet>>&
          spacelike_facets)
  {
    GeometryInfo::spacelike_facets = spacelike_facets;
  }

  /// @brief Getter for timevalues
  /// @return The set of timevalues
  const boost::optional<std::set<std::size_t>>& getTimevalues() const
  {
    return timevalues;
  }

  /// @brief Setter for timevalues
  /// @param timevalues The set of timevalues
  void setTimevalues(const boost::optional<std::set<std::size_t>>& timevalues)
  {
    GeometryInfo::timevalues = timevalues;
  }

  /// @brief Default constructor
  GeometryInfo() = default;

  /// @brief Constructor from Geometry_tuple
  ///
  /// This is usually called as a result of classify_all_simplices(),
  /// which itself takes a std::unique_ptr<Delaunay>
  /// @param geometry Geometry_tuple initializing values
  /// @return A populated GeometryInfo{}
  explicit GeometryInfo(const Geometry_tuple&& geometry) noexcept
      : three_one{std::get<0>(geometry)}
      , two_two{std::get<1>(geometry)}
      , one_three{std::get<2>(geometry)}
      , timelike_edges{std::get<3>(geometry)}
      , spacelike_edges{std::get<4>(geometry)}
      , vertices{std::get<5>(geometry)}
  {}

  /// @brief Default destructor
  ~GeometryInfo() = default;

  /// @brief Default move constructor
  GeometryInfo(GeometryInfo&&) = default;

  /// @brief Default move assignment operator
  //  GeometryInfo& operator=(GeometryInfo&&) = default;
  GeometryInfo& operator=(Geometry_tuple&& other)
  {
#ifndef NDEBUG
    std::cout << "GeometryInfo move assignment operator." << std::endl;
#endif
    three_one       = std::move(std::get<0>(other));
    two_two         = std::move(std::get<1>(other));
    one_three       = std::move(std::get<2>(other));
    timelike_edges  = std::move(std::get<3>(other));
    spacelike_edges = std::move(std::get<4>(other));
    vertices        = std::move(std::get<5>(other));
    return *this;
  }

  /// @brief Default copy constructor
  GeometryInfo(const GeometryInfo&) = default;

  /// @brief Default copy assignment operator
  GeometryInfo& operator=(const GeometryInfo&) = default;

  /// @brief Timelike edges
  /// @return The number of edges spanning timeslices
  auto N1_TL() { return static_cast<std::size_t>(timelike_edges.size()); }

  /// @brief Spacelike edges
  /// @return The number of edges on same timeslice
  auto N1_SL() { return static_cast<std::size_t>(spacelike_edges.size()); }

  /// @brief (3,1) simplices
  /// @return The total number of simplices with 3 vertices on the t
  /// timeslice and 1 vertex on the t+1 timeslice
  auto N3_31() { return static_cast<std::size_t>(three_one.size()); }

  /// @brief (1,3) simplices
  /// @return The total number of simplices with 1 vertex on the t timeslice
  /// and 3 vertices on the t+1 timeslice
  auto N3_13() { return static_cast<std::size_t>(one_three.size()); }

  /// @brief (3,1) and (1,3) simplices
  /// @return The total number of simplices with 3 vertices on one
  /// timeslice and 1 vertex on the adjacent timeslice. Used to
  /// calculate the change in action.
  auto N3_31_13() { return N3_31() + N3_13(); }

  /// @brief (2,2) simplices
  /// @return The total number of simplices with 2 vertices on one
  /// timeslice and 2 vertices on the adjacent timeslice. Used to
  /// calculate the change in action.
  auto N3_22() { return static_cast<std::size_t>(two_two.size()); }

  /// @brief Number of cells
  ///
  /// This should be the equivalent of
  /// SimplicialManifold::triangulation->number_of_finite_cells(),
  /// and is used as a check to ensure that GeometryInfo{} matches.
  /// @return The number of cells in GeometryInfo{}
  auto number_of_cells() { return N3_31() + N3_22() + N3_13(); }

  /// @brief Number of edges
  ///
  /// This should be the equivalent of
  /// SimplicialManifold::triangulation->number_of_finite_edges(),
  /// and is used as a check to ensure that GeometryInfo{} matches.
  /// @return The number of edges in the triangulation
  auto number_of_edges() { return N1_TL() + N1_SL(); }

  //  auto max_timevalue() { return *timevalues.crbegin();}
  boost::optional<std::size_t> max_timevalue()
  {
    return timevalues ? *timevalues->crbegin() : 0;
  }

  boost::optional<std::size_t> min_timevalue()
  {
    return timevalues ? *timevalues->begin() : 0;
  }
  /// @brief Number of vertices
  /// @return The number of vertices in the triangulation
  auto N0() { return static_cast<std::size_t>(vertices.size()); }

  template <typename T1, typename T2>
  friend auto make_23_move(T1&& universe, T2&& attempted_moves)
      -> decltype(universe);

  template <typename T1, typename T2>
  friend auto make_32_move(T1&& universe, T2&& attempted_moves)
      -> decltype(universe);

  template <typename T1, typename T2>
  friend auto make_26_move(T1&& universe, T2&& attempted_moves)
      -> decltype(universe);

  template <typename T1, typename T2>
  friend auto make_62_move(T1&& universe, T2&& attempted_moves)
      -> decltype(universe);

  template <typename T1, typename T2>
  friend auto make_44_move(T1&& universe, T2&& attempted_moves)
      -> decltype(universe);

  template <typename T>
  friend auto VolumePerTimeslice(T&& manifold) -> decltype(manifold);
};

/// @struct
/// @brief A struct to hold triangulation and geometry information
///
/// SimplicialManifold contains information about the triangulation and
/// its geometry. In addition, it defines convenient constructors.
struct [[deprecated]] SimplicialManifold
{
  /// @brief Owning pointer to the Delaunay triangulation
  std::unique_ptr<Delaunay3> triangulation;

  /// @brief Owning pointer to GeometryInfo
  std::unique_ptr<GeometryInfo> geometry;

  /// @brief Default constructor
  SimplicialManifold()
      : triangulation{std::make_unique<Delaunay3>()}
      , geometry{std::make_unique<GeometryInfo>()}
  {
#ifndef NDEBUG
    std::cout << "SimplicialManifold default ctor.\n";
#endif
  }

  /// @brief Constructor with std::unique_ptr<Delaunay>
  ///
  /// Constructor taking a std::unique_ptr<Delaunay> which should be created
  /// using make_triangulation(). If you wish to default initialize a
  /// SimplicialManifold with no values, use the default
  /// constructor SimplicialManifold::SimplicialManifold() instead.
  /// Non-static data members are initialized in the order they are declared,
  /// (see http://open-std.org/JTC1/SC22/WG21/docs/papers/2016/n4594.pdf,
  ///  \f$\S\f$ 12.6.2.13.3), so **geometry** depending upon **triangulation**
  /// is fine.
  /// @param manifold A std::unique_ptr<Delaunay>
  /// @return A SimplicialManifold{}
  explicit SimplicialManifold(std::unique_ptr<Delaunay3>&& manifold)
      : triangulation{std::move(manifold)}
      , geometry{std::make_unique<GeometryInfo>(
            classify_all_simplices(triangulation))}
  {
#ifndef NDEBUG
    std::cout << "SimplicialManifold std::unique_ptr<Delaunay> ctor.\n";
#endif
  }

  /// @brief make_triangulation constructor
  ///
  /// Constructor that initializes **triangulation** by calling
  /// make_triangulation() and **geometry** by calling
  /// classify_all_simplices().
  /// @param simplices The number of desired simplices in the triangulation
  /// @param timeslices The number of timeslices in the triangulation
  /// @return A populated SimplicialManifold{}
  SimplicialManifold(std::size_t simplices, std::size_t timeslices)
      : triangulation{make_triangulation(simplices, timeslices)}
      , geometry{std::make_unique<GeometryInfo>(
            classify_all_simplices(triangulation))}
  {
#ifndef NDEBUG
    std::cout << "SimplicialManifold make_triangulation ctor.\n";
#endif
  }

  /// @brief Destructor
  ~SimplicialManifold()
  {
#ifndef NDEBUG
    std::cout << "SimplicialManifold dtor.\n";
#endif
    this->triangulation = nullptr;
    this->geometry      = nullptr;
  }

  /// @brief Move constructor
  /// @param other The SimplicialManifold to be move-constructed from
  /// @return A moved-to SimplicialManifold{}
  SimplicialManifold(SimplicialManifold&& other)
      : triangulation{std::move(other.triangulation)}
      , geometry{std::make_unique<GeometryInfo>(
            classify_all_simplices(triangulation))}
  //      , geometry{std::move(other.geometry)}
  {
#ifndef NDEBUG
    std::cout << "SimplicialManifold move ctor.\n";
#endif
  }
  //  SimplicialManifold(SimplicialManifold&&) = default;

  /// @brief Move assignment operator
  /// @param other The SimplicialManifold to be moved from
  /// @return A moved-assigned SimplicialManifold{}
  SimplicialManifold& operator=(SimplicialManifold&& other) noexcept
  {
#ifndef NDEBUG
    std::cout << "SimplicialManifold move assignment operator.\n";
#endif
    triangulation = std::move(other.triangulation);
    geometry      = std::make_unique<GeometryInfo>(
        classify_all_simplices(std::move(triangulation)));
    //      geometry = std::move(other.geometry);
    return *this;
  }
  //  SimplicialManifold& operator=(SimplicialManifold&&) = default;

  /// @brief SimplicialManifold copy constructor
  /// @param other The SimplicialManifold to copy
  /// @return A copied SimplicialManifold{}
  SimplicialManifold(const SimplicialManifold& other)
      : triangulation{std::make_unique<Delaunay3>(*(other.triangulation))}
      , geometry{std::make_unique<GeometryInfo>(*(other.geometry))}
  {
#ifndef NDEBUG
    std::cout << "SimplicialManifold copy ctor.\n";
#endif
  }

  /// @brief Copy assignment operator
  /// @return A copy-assigned SimplicialManifold{}
  //    SimplicialManifold& operator=(const SimplicialManifold&) = default;
  SimplicialManifold& operator=(const SimplicialManifold& other) noexcept
  {
#ifndef NDEBUG
    std::cout << "SimplicialManifold copy assignment operator.\n";
#endif
    SimplicialManifold temp(other);
    swap(triangulation, temp.triangulation);
    swap(geometry, temp.geometry);
    return *this;
  }

  /// @brief Exception-safe swap
  /// @param first  The first SimplicialManifold to be swapped
  /// @param second The second SimplicialManifold to be swapped with.
  friend void swap(SimplicialManifold& first,
                   SimplicialManifold& second) noexcept
  {
#ifndef NDEBUG
    std::cout << "SimplicialManifold swapperator.\n";
#endif
    using std::swap;
    swap(first.triangulation, second.triangulation);
    swap(first.geometry, second.geometry);
  }

  bool reconcile()
  {
    return (this->triangulation->number_of_vertices() == this->geometry->N0() &&
            this->triangulation->number_of_finite_edges() ==
                this->geometry->number_of_edges() &&
            this->triangulation->number_of_finite_cells() ==
                this->geometry->number_of_cells());
  }

  void update()
  {
    geometry = std::make_unique<GeometryInfo>(
        classify_all_simplices(std::move(triangulation)));
  }
};

#endif  // INCLUDE_SIMPLICIALMANIFOLD_HPP_


/*
 *  Copyright 2010 CNRS
 *
 *  Florent Lamiraux
 */

#include <boost/format.hpp>
#include <boost/numeric/ublas/io.hpp>

#include <dynamic-graph/factory.h>
#include <dynamic-graph/command-setter.h>
#include <dynamic-graph/command-getter.h>
#include "dynamic-graph/tutorial/inverted-pendulum.hh"
#include "command-increment.hh"
#include "constant.hh"

using namespace dynamicgraph;
using namespace dynamicgraph::tutorial;

const double Constant::gravity = 9.81;

// Register new Entity type in the factory
// Note that the second argument is the type name of the python class
// that will be created when importing the python module.
DYNAMICGRAPH_FACTORY_ENTITY_PLUGIN(InvertedPendulum, "InvertedPendulum");

InvertedPendulum::InvertedPendulum(const std::string& inName)
    : Entity(inName),
      forceSIN(NULL, "InvertedPendulum(" + inName + ")::input(double)::force"),
      stateSOUT("InvertedPendulum(" + inName + ")::output(vector)::state"),
      cartMass_(1.0),
      pendulumMass_(1.0),
      pendulumLength_(1.0),
      viscosity_(0.1) {
  // Register signals into the entity.
  signalRegistration(forceSIN);
  signalRegistration(stateSOUT);

  // Set signals as constant to size them
  Vector state(4);
  state.fill(0.);
  double input = 0.;
  stateSOUT.setConstant(state);
  forceSIN.setConstant(input);

  // Commands
  std::string docstring;

  // Incr
  docstring =
      "\n"
      "    Integrate dynamics for time step provided as input\n"
      "\n"
      "      take one floating point number as input\n"
      "\n";
  addCommand(std::string("incr"), new command::Increment(*this, docstring));

  // setCartMass
  docstring =
      "\n"
      "    Set cart mass\n"
      "\n";
  addCommand(std::string("setCartMass"), new ::dynamicgraph::command::Setter<InvertedPendulum, double>(
                                             *this, &InvertedPendulum::setCartMass, docstring));

  // getCartMass
  docstring =
      "\n"
      "    Get cart mass\n"
      "\n";
  addCommand(std::string("getCartMass"), new ::dynamicgraph::command::Getter<InvertedPendulum, double>(
                                             *this, &InvertedPendulum::getCartMass, docstring));

  // setPendulumMass
  docstring =
      "\n"
      "    Set pendulum mass\n"
      "\n";
  addCommand(std::string("setPendulumMass"), new ::dynamicgraph::command::Setter<InvertedPendulum, double>(
                                                 *this, &InvertedPendulum::setPendulumMass, docstring));

  // getPendulumMass
  docstring =
      "\n"
      "    Get pendulum mass\n"
      "\n";
  addCommand(std::string("getPendulumMass"), new ::dynamicgraph::command::Getter<InvertedPendulum, double>(
                                                 *this, &InvertedPendulum::getPendulumMass, docstring));

  // setPendulumLength
  docstring =
      "\n"
      "    Set pendulum length\n"
      "\n";
  addCommand(std::string("setPendulumLength"), new ::dynamicgraph::command::Setter<InvertedPendulum, double>(
                                                   *this, &InvertedPendulum::setPendulumLength, docstring));

  // getPendulumLength
  docstring =
      "\n"
      "    Get pendulum length\n"
      "\n";
  addCommand(std::string("getPendulumLength"), new ::dynamicgraph::command::Getter<InvertedPendulum, double>(
                                                   *this, &InvertedPendulum::getPendulumLength, docstring));
}

InvertedPendulum::~InvertedPendulum() {}

Vector InvertedPendulum::computeDynamics(const Vector& inState, const double& inControl, double inTimeStep) {
  if (inState.size() != 4)
    throw dynamicgraph::ExceptionSignal(dynamicgraph::ExceptionSignal::GENERIC, "state signal size is ",
                                        "%d, should be 4.", inState.size());

  double dt = inTimeStep;
  double dt2 = dt * dt;
  double g = Constant::gravity;
  double x = inState(0);
  double th = inState(1);
  double dx = inState(2);
  double dth = inState(3);
  double F = inControl;
  double m = pendulumMass_;
  double M = cartMass_;
  double l = pendulumLength_;
  double lambda = viscosity_;
  double l2 = l * l;
  double dth2 = dth * dth;
  double sth = sin(th);
  double cth = cos(th);
  double sth2 = sth * sth;

  double b1 = F - m * l * dth2 * sth - lambda * dx;
  double b2 = m * l * g * sth - lambda * dth;

  double det = m * l2 * (M + m * sth2);

  double ddx = (b1 * m * l2 + b2 * m * l * cth) / det;
  double ddth = ((M + m) * b2 + m * l * cth * b1) / det;

  Vector nextState(4);
  nextState(0) = x + dx * dt + .5 * ddx * dt2;
  nextState(1) = th + dth * dt + .5 * ddth * dt2;
  nextState(2) = dx + dt * ddx;
  nextState(3) = dth + dt * ddth;

  return nextState;
}

void InvertedPendulum::incr(double inTimeStep) {
  int t = stateSOUT.getTime();
  Vector nextState = computeDynamics(stateSOUT(t), forceSIN(t), inTimeStep);
  stateSOUT.setConstant(nextState);
  stateSOUT.setTime(t + 1);
  forceSIN(t + 1);
}


#include <iostream>
#include <memory>
#include <string>
#include <chrono>

#include <boost/program_options.hpp>
#include <algorithms/Grover.hpp>
#include <algorithms/QFT.hpp>
#include <algorithms/Entanglement.hpp>

#include "Simulator.hpp"
#include "QFRSimulator.hpp"
#include "GroverSimulator.hpp"
#include "ShorFastSimulator.hpp"
#include "ShorSimulator.hpp"

int main(int argc, char** argv) {
    namespace po = boost::program_options;
    unsigned long long seed;

    po::options_description description("JKQ DDSIM by https://iic.jku.at/eda/ -- Allowed options");
    description.add_options()
            ("help,h", "produce help message")
            ("seed", po::value<unsigned long long>(&seed)->default_value(0), "seed for random number generator (default zero is possibly directly used as seed!)")
            ("shots", po::value<unsigned int>()->default_value(0), "number of measurements (if the algorithm does not contain non-unitary gates, weak simulation is used)")
            ("display_vector", "display the state vector")
            ("ps", "print simulation stats (applied gates, sim. time, and maximal size of the DD)")
            ("verbose", "Causes some simulators to print additional information to STDERR")
            ("benchmark", "print simulation stats in a single CSV style line (overrides --ps and suppresses most other output, please don't rely on the format across versions)")

            ("simulate_file", po::value<std::string>(), "simulate a quantum circuit given by file (detection by the file extension)")
            ("simulate_qft", po::value<unsigned int>(), "simulate Quantum Fourier Transform for given number of qubits")
            ("simulate_ghz", po::value<unsigned int>(), "simulate state preparation of GHZ state for given number of qubits")
            ("step_fidelity", po::value<double>()->default_value(1.0), "target fidelity for each approximation run (>=1 = disable approximation)")
            ("steps", po::value<unsigned int>()->default_value(1), "number of approximation steps")
            ("initial_reorder", po::value<int>()->default_value(0), "Try to find a good initial variable order (0=None, 1=Most affected qubits to the top, 2=Most affected targets to the top)")
            ("dynamic_reorder", po::value<int>()->default_value(0), "Apply reordering strategy during simulation (0=None, 1=Sifting, 2=Move2Top)")
            ("post_reorder", po::value<int>()->default_value(0), "Apply a reordering strategy after simulation (0=None, 1=Sifting)")

            ("simulate_grover", po::value<unsigned int>(), "simulate Grover's search for given number of qubits with random oracle")
            ("simulate_grover_emulated", po::value<unsigned int>(), "simulate Grover's search for given number of qubits with random oracle and emulation")
            ("simulate_grover_oracle_emulated", po::value<std::string>(), "simulate Grover's search for given number of qubits with given oracle and emulation")

            ("simulate_shor", po::value<unsigned int>(), "simulate Shor's algorithm factoring this number")
            ("simulate_shor_coprime", po::value<unsigned int>()->default_value(0), "coprime number to use with Shor's algorithm (zero randomly generates a coprime)")
            ("simulate_shor_no_emulation", "Force Shor simulator to do modular exponentiation instead of using emulation (you'll usually want emulation)")

            ("simulate_fast_shor", po::value<unsigned int>(), "simulate Shor's algorithm factoring this number with intermediate measurements")
            ("simulate_fast_shor_coprime", po::value<unsigned int>()->default_value(0), "coprime number to use with Shor's algorithm (zero randomly generates a coprime)")
            ("noise_effects", po::value<std::string>(),
             "Noise effects (A (=amplitude damping),D (=depolarization),P (=phase flip)) in the form of a character string describing the noise effects (default=\" \")")
            ("noise_prob", po::value<double>(), "Probability for applying noise (default=0.001)")
            ("confidence", po::value<double>(),
             "Confidence in the error bound of the stochastic simulation (default= 0.05)")
            ("error_bound", po::value<double>(), "Error bound of the stochastic simulation (default=0.1)")
            ("stoch_runs", po::value<long>()->default_value(0), "Number of stochastic runs. When the value is 0 the value is calculated using the confidence, error_bound and number of tracked properties. (default = 0)")
            ("properties", po::value<std::string>(), R"(Comma separated list of tracked properties. Note that -1 is the fidelity and "-" can be used to specify a range.  (default="0-1000"))")
            ;
    po::variables_map vm;
    try {
        po::store(po::parse_command_line(argc, argv, description), vm);

        if (vm.count("help")) {
            std::cout << description;
            return 0;
        }
        po::notify(vm);
    } catch (const po::error &e) {
        std::cerr << "[ERROR] " << e.what() << "! Try option '--help' for available commandline options.\n";
        std::exit(1);
    }

    std::unique_ptr<qc::QuantumComputation> quantumComputation;
    std::unique_ptr<Simulator> ddsim{nullptr};

    if (vm.count("simulate_file")) {
        const std::string fname = vm["simulate_file"].as<std::string>();
    	quantumComputation = std::make_unique<qc::QuantumComputation>(fname);
        ddsim = std::make_unique<QFRSimulator>(quantumComputation,
                                               vm["steps"].as<unsigned int>(), vm["step_fidelity"].as<double>(),
                                               vm["initial_reorder"].as<int>(), vm["dynamic_reorder"].as<int>(), vm["post_reorder"].as<int>(),
                                               seed);
    } else {
        std::cerr << "Did not find anything to simulate. See help below.\n"
                  << description << "\n";
        return 1;
    }

    if (quantumComputation && quantumComputation->getNqubits() > dd::MAXN) {
        std::cerr << "Quantum computation contains to many qubits (limit is set to " << dd::MAXN << "). See documentation for details.\n";
        std::exit(1);
    }
        if (vm.count("noise_effects")) {
            ddsim->setNoiseEffects((char *) &vm["noise_effects"].as<std::string>()[0]);
    } else {
            ddsim->setNoiseEffects((char *) &std::string("0")[0]);
    }
    if (vm.count("noise_prob")) {
        ddsim->setAmplitudeDampingProbability(vm["noise_prob"].as<double>());
    } else {
        ddsim->setAmplitudeDampingProbability(0.001);
    }

    if (vm.count("confidence")) {
        ddsim->stoch_confidence = vm["confidence"].as<double>();
    }

    if (vm.count("properties")) {
        ddsim->setRecordedProperties((char *) &vm["properties"].as<std::string>()[0]);
    } else {
        ddsim->setRecordedProperties("0-1000");
    }

    if (vm.count("error_bound")) {
        ddsim->stoch_error_margin = vm["error_bound"].as<double>();
    } else {
        ddsim->stoch_error_margin = 0.01;
    }
    ddsim->stochastic_runs = vm["stoch_runs"].as<long>();

    auto t1 = std::chrono::high_resolution_clock::now();
    ddsim->StochSimulate();
    auto t2 = std::chrono::high_resolution_clock::now();

    std::chrono::duration<float> duration_simulation = t2-t1;

    if (vm.count("benchmark")) {
        auto more_info = ddsim->AdditionalStatistics();
        std::cout << ddsim->getName() << ", "
                  << ddsim->getNumberOfQubits() << ", "
                  //<< vm["approximate"].as<float>() << ", "
                  << std::fixed << duration_simulation.count() << std::defaultfloat << ", "
                  //<< more_info["approximation_runs"] << ","
                  //<< more_info["final_fidelity"] << ", "
                  << more_info["coprime_a"] << ", "
                  << more_info["sim_result"] << ", "
                  << more_info["polr_result"] << ", "
                  << ddsim->getSeed() << ", "
                  << ddsim->getNumberOfOps() << ", "
                  << ddsim->getMaxNodeCount()
                  << "\n";
        return 0;
    }

    std::cout << "{\n";

    if (vm.count("ps")) {
        std::cout << "  \"statistics\": {\n"
                  << "    \"simulation_time\": " << std::fixed << duration_simulation.count() << std::defaultfloat << ",\n"
                  << "    \"benchmark\": \"" << ddsim->getName() << "\",\n"
                  << "    \"shots\": " << vm["shots"].as<unsigned int>() << ",\n"
//                  << "    \"distinct_results\": " << m.size() << ",\n"
                  << "    \"n_qubits\": " << ddsim->getNumberOfQubits() << ",\n"
                  << "    \"applied_gates\": " << ddsim->getNumberOfOps() << ",\n"
                  << "    \"max_nodes\": " << ddsim->getMaxNodeCount() << ",\n"
                  ;
        for(const auto& item : ddsim->AdditionalStatistics()) {
            std::cout << "    \"" << item.first << "\": \"" << item.second << "\",\n";
        }
        std::cout << "    \"seed\": " << ddsim->getSeed() << "\n"
                  << "  },\n";
    }
    std::cout << "  \"dummy\": 0\n}\n"; // trailing element to make json printout easier
}


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

#ifndef AMA_TENSOR_DETAIL_TENSOR_OUTER_HPP
#define AMA_TENSOR_DETAIL_TENSOR_OUTER_HPP 1

#include <ama/tensor/detail/tensor_base.hpp>
#include <boost/mpl/advance.hpp>
#include <boost/mpl/assert.hpp>
#include <boost/mpl/begin_end.hpp>
#include <boost/mpl/bool.hpp>
#include <boost/mpl/fold.hpp>
#include <boost/mpl/placeholders.hpp>
#include <boost/mpl/plus.hpp>
#include <boost/mpl/push_back.hpp>
#include <boost/mpl/vector/vector0_c.hpp>

namespace ama
{
  namespace tensor_
  {

    /* forward declaration*/
    template <typename LEFT, typename RIGHT> class tensor_outer;

    /* traits definition */
    template <typename LEFT, typename RIGHT>
    struct tensor_traits< tensor_outer<LEFT, RIGHT> >
    {
      typedef typename LEFT::value_type value_type;

      typedef typename LEFT::dimension_type dimension_type;

      typedef typename ::boost::mpl::plus<
                             typename LEFT::controvariant_type
                           , typename RIGHT::controvariant_type
                           >::type controvariant_type;
      typedef typename ::boost::mpl::plus<
                             typename LEFT::covariant_type
                           , typename RIGHT::covariant_type
                           >::type covariant_type;

      typedef ::boost::mpl::false_ is_assignable;
      typedef ::boost::mpl::true_ is_temporary;
    };

    /* class declaration */
    template <typename LEFT, typename RIGHT>
    class tensor_outer:
        public tensor_base< tensor_outer<LEFT, RIGHT> >
    {
    protected:
      typedef tensor_base< tensor_outer<LEFT, RIGHT> > base_type;
      typedef tensor_outer<LEFT, RIGHT> derived_type;

    protected:
      typedef LEFT left_operand_type;
      typedef RIGHT right_operand_type;

    public:
      typedef typename base_type::value_type value_type;

    public:
      /* costructor */
      tensor_outer(left_operand_type const & left,
                   right_operand_type const & right)
          : m_left(left),
            m_right(right) { }

    public:
      /* retrieve the value */
      template <typename ILIST>
      value_type at() const
      {
        namespace mpl = ::boost::mpl;

        /* iterator */
        typedef typename mpl::begin<ILIST>::type begin1;
        typedef typename mpl::advance<begin1, typename LEFT::controvariant_type>::type begin2;
        typedef typename mpl::advance<begin2, typename RIGHT::controvariant_type>::type begin3;
        typedef typename mpl::advance<begin3, typename LEFT::covariant_type>::type begin4;
        typedef typename mpl::advance<begin4, typename RIGHT::covariant_type>::type begin5;
        typedef typename mpl::end<ILIST>::type end;

        /* check for no errors in the code */
        BOOST_MPL_ASSERT((mpl::equal_to<mpl::distance<end,begin5>,mpl::size_t<0> >));

        /* split multi-index */
        typedef typename mpl::fold<
              mpl::iterator_range<begin1, begin2>
            , mpl::vector0_c<size_t>
            , mpl::push_back<mpl::_1, mpl::_2>
            >::type controvariant_left;

        typedef typename mpl::fold<
              mpl::iterator_range<begin2, begin3>
            , mpl::vector0_c<size_t>
            , mpl::push_back<mpl::_1, mpl::_2>
            >::type controvariant_right;

        typedef typename mpl::fold<
              mpl::iterator_range<begin3, begin4>
            , mpl::vector0_c<size_t>
            , mpl::push_back<mpl::_1, mpl::_2>
            >::type covariant_left;

        typedef typename mpl::fold<
              mpl::iterator_range<begin4, begin5>
            , mpl::vector0_c<size_t>
            , mpl::push_back<mpl::_1, mpl::_2>
            >::type covariant_right;

        /* append the splittend multi-index */
        typedef typename mpl::fold<
              covariant_left
            , controvariant_left
            , mpl::push_back<mpl::_1, mpl::_2>
            >::type left_ilist;

        typedef typename mpl::fold<
              covariant_right
            , controvariant_right
            , mpl::push_back<mpl::_1, mpl::_2>
            >::type right_ilist;

        return (m_left.template at<left_ilist>()) *
            (m_right.template at<right_ilist>());
      }

    protected:
      /* if the operand is temporary we save a copy, otherwise a reference */
      typedef typename LEFT::is_temporary left_operand_is_temporary;
      typedef typename ::boost::mpl::if_<
            left_operand_is_temporary
          , left_operand_type const
          , left_operand_type const &
          >::type const_left_operand_type;

      /* if the operand is temporary we save a copy, otherwise a reference */
      typedef typename RIGHT::is_temporary right_operand_is_temporary;
      typedef typename ::boost::mpl::if_<
            right_operand_is_temporary
          , right_operand_type const
          , right_operand_type const &
          >::type const_right_operand_type;

      /* members */
      const_left_operand_type m_left;
      const_right_operand_type m_right;
    };

  }
}

#endif /* AMA_TENSOR_DETAIL_TENSOR_OUTER_HPP */


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

//-----------------------------------------------------------------------------
/*!\file    SelectivelyDampedLeastSquaresSolver.cpp
 *
 * \author  Stefan Scherzinger <scherzin@fzi.de>
 * \date    2020/03/27
 *
 */
//-----------------------------------------------------------------------------

// this package
#include <cartesian_controller_base/SelectivelyDampedLeastSquaresSolver.h>

// Pluginlib
#include <pluginlib/class_list_macros.h>

// other
#include <boost/algorithm/clamp.hpp>

/**
 * \class cartesian_controller_base::SelectivelyDampedLeastSquaresSolver
 *
 * Users may explicitly specify this solver with \a "selectively_damped_least_squares" as \a
 * ik_solver in their controllers.yaml configuration file for each controller:
 *
 * \code{.yaml}
 * <name_of_your_controller>:
 *     type: "<type_of_your_controller>"
 *     ik_solver: "selectively_damped_least_squares"
 *     ...
 * \endcode
 *
 */
PLUGINLIB_EXPORT_CLASS(cartesian_controller_base::SelectivelyDampedLeastSquaresSolver, cartesian_controller_base::IKSolver)

namespace cartesian_controller_base{

  SelectivelyDampedLeastSquaresSolver::SelectivelyDampedLeastSquaresSolver()
  {
  }

  SelectivelyDampedLeastSquaresSolver::~SelectivelyDampedLeastSquaresSolver(){}

  trajectory_msgs::JointTrajectoryPoint SelectivelyDampedLeastSquaresSolver::getJointControlCmds(
        ros::Duration period,
        const ctrl::Vector6D& net_force)
  {
    // Compute joint Jacobian
    m_jnt_jacobian_solver->JntToJac(m_current_positions,m_jnt_jacobian);

    Eigen::JacobiSVD<Eigen::Matrix<double, 6, Eigen::Dynamic> > JSVD(
      m_jnt_jacobian.data, Eigen::ComputeFullU | Eigen::ComputeFullV);
    Eigen::Matrix<double, 6, 6> U = JSVD.matrixU();
    Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> V = JSVD.matrixV();
    Eigen::Matrix<double, Eigen::Dynamic, 1> s = JSVD.singularValues();

    // Default recommendation by Buss and Kim.
    const double gamma_max = 3.141592653 / 4;

    Eigen::Matrix<double, Eigen::Dynamic, 1> sum_phi =
      Eigen::Matrix<double, Eigen::Dynamic, 1>::Zero(m_number_joints);

    // Compute each joint velocity with the SDLS method.  This implements the
    // algorithm as described in the paper (but for only one end-effector).
    // Also see Buss' own implementation:
    // https://www.math.ucsd.edu/~sbuss/ResearchWeb/ikmethods/index.html

    for (int i = 0; i < m_number_joints; ++i)
    {
      double alpha = U.col(i).transpose() * net_force;

      double N = U.col(i).head(3).norm();
      double M = 0;
      for (int j = 0; j < m_number_joints; ++j)
      {
        double rho = m_jnt_jacobian.data.col(j).head(3).norm();
        M += std::abs(V.col(i)[j]) * rho;
      }
      M *= 1.0 / s[i];

      double gamma = std::min(1.0, N / M) * gamma_max;

      Eigen::Matrix<double, Eigen::Dynamic, 1> phi = clampMaxAbs(1.0 / s[i] * alpha * V.col(i), gamma);
      sum_phi += phi;
    }

    m_current_velocities.data = clampMaxAbs(sum_phi, gamma_max);

    // Integrate once, starting with zero motion
    m_current_positions.data = m_last_positions.data + 0.5 * m_current_velocities.data * period.toSec();

    // Make sure positions stay in allowed margins
    applyJointLimits();

    // Apply results
    trajectory_msgs::JointTrajectoryPoint control_cmd;
    for (int i = 0; i < m_number_joints; ++i)
    {
      control_cmd.positions.push_back(m_current_positions(i));
      control_cmd.velocities.push_back(m_current_velocities(i));

      // Accelerations should be left empty. Those values will be interpreted
      // by most hardware joint drivers as max. tolerated values. As a
      // consequence, the robot will move very slowly.
    }
    control_cmd.time_from_start = period; // valid for this duration

    return control_cmd;
  }

  bool SelectivelyDampedLeastSquaresSolver::init(ros::NodeHandle& nh,
                                      const KDL::Chain& chain,
                                      const KDL::JntArray& upper_pos_limits,
                                      const KDL::JntArray& lower_pos_limits)
  {
    IKSolver::init(nh, chain, upper_pos_limits, lower_pos_limits);

    m_jnt_jacobian_solver.reset(new KDL::ChainJntToJacSolver(m_chain));
    m_jnt_jacobian.resize(m_number_joints);

    return true;
  }

  Eigen::Matrix<double, Eigen::Dynamic, 1> SelectivelyDampedLeastSquaresSolver::clampMaxAbs(
    const Eigen::Matrix<double, Eigen::Dynamic, 1>& w, double d)
  {
    if (w.cwiseAbs().maxCoeff() <= d)
    {
      return w;
    }
    else
    {
      return d * w / w.cwiseAbs().maxCoeff();
    }
  }

} // namespace


#include "BootstrapRunner.h"
#include <vector>
#include <boost/timer.hpp>
using namespace std;

BootstrapRunner::BootstrapRunner(BoostrapConfig* bootstrapConfig, CftConfig cftConfig)
{
    this->bootstrapConfig = bootstrapConfig;
    this->cftData = new CftData(cftConfig);
    this->correlatorSet = new CorrelatorSet(cftData, bootstrapConfig->NumberOfScalarsToBootstrap);
}

BootstrapRunner::~BootstrapRunner()
{
    if (this->cftData != 0) {
        delete cftData;
    }

    if (this->correlatorSet != 0) {
        delete this->correlatorSet;
    }
}

float_type BootstrapRunner::ConstraintCost(CftData* cftData)
{
    float_type ret = .0;
    float_type lastDim = .0, dim;
    int lastSpin = -1;

    for (uint i = 1; i <= cftData->MaxPrimaryId(); i++) {
        dim = cftData->GetPrimaryDim(i);
        if (cftData->GetPrimarySpin(i) != lastSpin) {
            lastSpin = cftData->GetPrimarySpin(i);
            if (lastSpin == 0) {
                if (dim < cftData->D / 2.0 - 1) ret += (cftData->D / 2.0 - 1 - dim);
            } else {
                if (dim < lastSpin + cftData->D - 2) ret += (lastSpin + cftData->D - 2 - dim);
            }
        } else if (dim < lastDim) {
            ret += (lastDim - dim);
        }

        lastDim = dim;
    }

    return ret * bootstrapConfig->ConstraintFactor;
}

void BootstrapRunner::ConstraintCostDerivative(CftData* cftData, vector<int>& ops, vector<float_type>& opsDerivatives)
{
    assert(ops.size() == opsDerivatives.size());

    float_type dim, lastDim, nextDim;

    for (uint i = 0; i < ops.size(); i++) {
        int l = cftData->GetPrimarySpin(ops[i]);
        dim = cftData->GetPrimaryDim(ops[i]);

        if (ops[i] > 1 && cftData->GetPrimarySpin(ops[i] - 1) == l) {
            lastDim = cftData->GetPrimaryDim(ops[i] - 1);
            if (dim < lastDim) {
                opsDerivatives[i] -= bootstrapConfig->ConstraintFactor;
            }
        }

        if (ops[i] < cftData->MaxPrimaryId() && cftData->GetPrimarySpin(ops[i] + 1) == l) {
            nextDim = cftData->GetPrimaryDim(ops[i] + 1);
            if (dim > nextDim) {
                opsDerivatives[i] += bootstrapConfig->ConstraintFactor;
            }
        }

        if (ops[i] == 1 || cftData->GetPrimarySpin(ops[i] - 1) != l) {
            if (l == 0 && dim <= cftData->D / 2.0 - 1) {
                opsDerivatives[i] -= bootstrapConfig->ConstraintFactor;
            }

            if (l > 0 && dim <= cftData->D + l - 2.0) {
                opsDerivatives[i] -= bootstrapConfig->ConstraintFactor;
            }
        }
    }
}

void BootstrapRunner::Run()
{
    vector<OpeCoefficientKey> coefs;
    vector<int> ops;

    for (uint op = 1; op <= cftData->MaxPrimaryId(); op++) {
        if (op != cftData->StressTensorId) {
            ops.push_back(op);
        }
    }

    for (uint op1 = 1; op1 <= bootstrapConfig->NumberOfScalarsToBootstrap; op1++) {
        for (uint op2 = op1; op2 <= bootstrapConfig->NumberOfScalarsToBootstrap; op2++) {
            for (uint op3 = op2; op3 <= cftData->MaxPrimaryId(); op3++) {
                if (op3 != cftData->StressTensorId) {
                    coefs.push_back(OpeCoefficientKey(op1, op2, op3));
                }
            }
        }
    }

    vector<cpx_t> zs;
    zs.reserve(bootstrapConfig->SamplesEachStep);
    for (uint i = 0; i < bootstrapConfig->SamplesEachStep; i++) {
        float_type dr = random(0.0, 0.2);
        float_type theta = random(0.0, 2 * acos(-1.0));
        zs.push_back(cpx_t(dr * cos(theta), dr * sin(theta)) + 0.5);
    }

    vector<float_type> opsDerivatives(ops.size(), .0);
    vector<float_type> opsDerivativeTot(ops.size(), .0);
    vector<float_type> coefsDerivatives(coefs.size(), .0);
    vector<float_type> coefsDerivativeTot(coefs.size(), .0);

    float_type lastCost = std::numeric_limits<float_type>::max();
    vector<int> scalarIds;
    for (uint i = 1; i <= cftData->MaxScalarId(); i++) scalarIds.push_back(i);

    boost::timer stopwatch;

    while (true) {
        std::cout << "Time elapsed: " << stopwatch.elapsed() << " seconds." << endl;
        cftData->Output(scalarIds);

        opsDerivativeTot.resize(ops.size(), .0);
        coefsDerivativeTot.resize(coefs.size(), .0);

        float_type cost = .0;

        for (uint i = 0; i < zs.size(); i++) {
            this->correlatorSet->CostDerivatives(zs[i], ops, coefs, opsDerivatives, coefsDerivatives);
            cost += this->correlatorSet->Cost(zs[i]);

            for (uint j = 0; j < opsDerivatives.size(); j++) {
                opsDerivativeTot[j] += opsDerivatives[j];
            }

            for (uint j = 0; j < coefsDerivatives.size(); j++) {
                coefsDerivativeTot[j] += coefsDerivatives[j];
            }
        }

        cost /= zs.size();
        std::cout << "pure cost = " << cost << "\t";
        cost += ConstraintCost(cftData);
        std::cout << "total cost = " << cost << std::endl;

        if (cost < bootstrapConfig->Accuracy) break;

        float_type gradNorm = .0;
        for (uint j = 0; j < coefsDerivativeTot.size(); j++) {
            coefsDerivativeTot[j] /= zs.size();
            gradNorm += coefsDerivativeTot[j] * coefsDerivativeTot[j];
        }

        std::cout << "Coefs Derivative: " << coefsDerivativeTot << std::endl;

        for (uint j = 0; j < opsDerivativeTot.size(); j++) {
            opsDerivativeTot[j] /= zs.size();
        }

        std::cout << "Ops Derivative: " << opsDerivativeTot << std::endl;

        ConstraintCostDerivative(cftData, ops, opsDerivativeTot);

        std::cout << "Ops Derivative with constraints: " << opsDerivativeTot << std::endl;

        for (uint j = 0; j < opsDerivativeTot.size(); j++) {
            gradNorm += opsDerivativeTot[j] * opsDerivativeTot[j];
        }

        float_type factor = gradNorm / ((ops.size() + coefs.size()));
        factor = sqrt(factor);

        for (uint j = 0; j < coefsDerivativeTot.size(); j++) {
            coefsDerivativeTot[j] /= factor;
        }

        for (uint j = 0; j < opsDerivativeTot.size(); j++) {
            opsDerivativeTot[j] /= factor;
        }

        std::cout << "Coefs Derivative normalized: " << coefsDerivativeTot << std::endl;
        std::cout << "Ops Derivative normalized: " << opsDerivativeTot << std::endl;

        for (uint j = 0; j < opsDerivativeTot.size(); j++) {
            float_type dim = cftData->GetPrimaryDim(ops[j]);
            cftData->SetPrimaryDim(ops[j], dim - opsDerivativeTot[j] * bootstrapConfig->Step);
        }

        for (uint j = 0; j < coefsDerivativeTot.size(); j++) {
            float_type coef = cftData->GetOpeCoefficient(coefs[j]);
            cftData->SetOpeCoefficient(coefs[j], coef - coefsDerivativeTot[j] * bootstrapConfig->Step);
        }

        this->correlatorSet->UpdateCftData(cftData);
        lastCost = cost;
    }
}


#include "../is_valid_char.hxx"
#include "../parse_vector.hxx"
#include "../parse_number.hxx"
#include "../../../../Damped_Rational.hxx"

#include <boost/algorithm/string.hpp>

#include <algorithm>
#include <iterator>
#include <string>

const char *parse_damped_rational(const char *begin, const char *end,
                                  Damped_Rational &damped_rational)
{
  const std::string damped_literal("DampedRational[");

  // Parse constant numbers
  for(auto current = begin; current != end; ++current)
    {
      if(*current == damped_literal.front())
        {
          break;
        }
      if(is_valid_char(*current))
        {
          damped_rational.base = 1;
          auto comma(std::find(current, end, ','));
          if(comma == end)
            {
              throw std::runtime_error(
                "Missing a comma when parsing a constant as a DampedRational. "
                " The parsed string is\n\t'"
                + std::string(current, end) + "'");
            }
          damped_rational.constant = Boost_Float(parse_number(current, comma));
          return comma;
        }
    }

  auto damped_start(
    std::search(begin, end, damped_literal.begin(), damped_literal.end()));
  if(damped_start == end)
    {
      throw std::runtime_error("Could not find '" + damped_literal + "'");
    }

  auto comma(std::find(damped_start, end, ','));
  if(comma == end)
    {
      throw std::runtime_error("Missing comma after DampedRational.constant");
    }
  auto constant_start(std::next(damped_start, damped_literal.size()));
  damped_rational.constant = Boost_Float(parse_number(constant_start, comma));

  auto start_poles(std::next(comma));
  auto end_poles(parse_vector(start_poles, end, damped_rational.poles));

  comma = std::find(end_poles, end, ',');
  if(comma == end)
    {
      throw std::runtime_error("Missing comma after DampedRational.poles");
    }

  auto start_base(std::next(comma));
  comma = std::find(start_base, end, ',');
  if(comma == end)
    {
      throw std::runtime_error("Missing comma after DampedRational.base");
    }
  damped_rational.base = Boost_Float(parse_number(start_base, comma));

  auto start_variable(std::next(comma));
  const auto close_bracket(std::find(start_variable, end, ']'));
  if(close_bracket == end)
    {
      throw std::runtime_error("Missing ']' at end of DampedRational");
    }

  return std::next(close_bracket);
}


#include <fenv.h>
#include <algorithm>
#include <boost/python/numpy.hpp>
#include <iomanip>
#include <iostream>
#include <list>
#include <memory>
#include <random>
#include <stdexcept>
#include <vector>

float gen_gaussian() {
    static std::random_device r;
    static std::default_random_engine generator(r());
    static std::normal_distribution<float> distribution;
    return distribution(generator);
}

template <typename T>
int sgn(T val) {
    return (T(0) < val) - (val < T(0));
}

enum class Optimizer { Sgd, Sgdm, PowerSign };

struct Dims {
    int w;
    int h;
    int c;

    int sz() const { return w * h * c; }
};

struct Settings {
    float lr;
    int batch_size;
    int epochs;
    float l2;
    float grad_max;
    float lr_decay;
    Optimizer optimizer;

    Settings()
        : lr(0.1),
          batch_size(8),
          epochs(10),
          l2(0.0001),
          grad_max(5),
          lr_decay(1),
          optimizer(Optimizer::Sgdm) {}
};

template <class T>
class Pool {
   public:
    class PoolAlloc {
        std::unique_ptr<T> x_;
        int sz_;
        int owners_;

       public:
        PoolAlloc(std::unique_ptr<T>&& x, int sz)
            : x_(std::move(x)), sz_(sz), owners_(0) {}

        bool is_free() const { return owners_ == 0; }

        void add_owner() { ++owners_; }
        void remove_owner() {
#if _GLIBCXX_DEBUG
            if (owners_ == 0)
                throw std::runtime_error(
                    "trying to deown something that was said to have zero "
                    "wner");
#endif
            --owners_;
        }

        auto* get() { return x_.get(); }
        const auto* get() const { return x_.get(); }

        int sz() const { return sz_; }
    };

    class Recyclable {
       public:
        Recyclable(PoolAlloc* alloc) : alloc_(alloc) { alloc_->add_owner(); }

        Recyclable(Recyclable&& o) {
            alloc_ = o.alloc_;
            o.alloc_ = nullptr;
        }

        Recyclable() : alloc_(nullptr) {}

        Recyclable(const Recyclable& o) {
            alloc_ = o.alloc_;
            if (alloc_) {
                alloc_->add_owner();
            }
        }

        ~Recyclable() {
            if (alloc_) {
                alloc_->remove_owner();
            }
        }

        Recyclable& operator=(Recyclable&& o) {
            if (alloc_) {
                alloc_->remove_owner();
            }
            alloc_ = o.alloc_;
            o.alloc_ = nullptr;
            return *this;
        }

        Recyclable& operator=(const Recyclable& o) {
            if (alloc_) {
                alloc_->remove_owner();
            }
            alloc_ = o.alloc_;
            if (alloc_) {
                alloc_->add_owner();
            }
            return *this;
        }

        auto& operator*() { return *alloc_->get(); }
        const auto& operator*() const { return *alloc_->get(); }

        auto& operator-> () { return alloc_->get().operator->(); }
        const auto& operator-> () const { return alloc_->get().operator->(); }

        auto& operator[](int i) { return alloc_->get()[i]; }
        const auto& operator[](int i) const { return alloc_->get()[i]; }

        int sz() const { return alloc_->sz_; }

       private:
        PoolAlloc* alloc_;
    };

    Recyclable alloc(int sz) {
        auto found = std::find_if(pool_.begin(), pool_.end(), [=](auto& r) {
            return r.sz() == sz && r.is_free();
        });

        if (found != pool_.end()) {
            PoolAlloc* p = &*found;
            return Recyclable(p);
        }

        pool_.emplace_back(std::make_unique<T>(sz), sz);
        return Recyclable(&pool_.back());
    }

   private:
    std::list<PoolAlloc> pool_;
};

static Pool<float[]> float_pool_;

class Volume {
   public:
    Volume(int w, int h, int c)
        : w_(w), h_(h), c_(c), sz_(h * w * c), res_(float_pool_.alloc(sz_)) {}
    Volume(const Dims& sz) : Volume(sz.w, sz.h, sz.c) {}

    Volume() = default;
    Volume(Volume&& o) = default;

    Volume from_shape() const { return Volume(w_, h_, c_); }
    void zero() {
        for (int i = 0; i < sz_; ++i) {
            res_[i] = 0;
        }
    }

    Volume& operator=(Volume&& o) = default;

    float& operator[](int i) {
#if _GLIBCXX_DEBUG
        if (i >= sz_) {
            throw std::runtime_error("volume bound checks fail");
        }
#endif
        return res_[i];
    }
    float operator[](int i) const {
#if _GLIBCXX_DEBUG
        if (i >= sz_) {
            throw std::runtime_error("volume bound checks fail");
        }
#endif
        return res_[i];
    }

    int w() const { return w_; }
    int h() const { return h_; }
    int c() const { return c_; }
    int sz() const { return sz_; }

    int row_idx(int i) const { return i * w_; };
    int cha_idx(int i) const { return i * w_ * h_; }

    void show() const {
        std::cout << "[";
        for (int c = 0; c < c_; ++c) {
            std::cout << "[\n";
            for (int h = 0; h < h_; ++h) {
                std::cout << "  [";
                std::cout << res_[h_ * c * w_ + h * w_];
                for (int w = 1; w < w_; ++w) {
                    std::cout << ", " << res_[w + w_ * h + w_ * h_ * c];
                }
                std::cout << "],\n";
            }
            std::cout << "]";
        }
        std::cout << "]\n";
    }

    void share_with(Volume& v) const {
        v.res_ = res_;
        v.w_ = w_;
        v.h_ = h_;
        v.c_ = c_;
        v.sz_ = sz_;
    }

    void gaussian() {
        for (int i = 0; i < sz_; ++i) {
            res_[i] = gen_gaussian();
        }
    }

    std::pair<float, float> mean_std() const {
        float m = 0;
        for (int i = 0; i < sz_; ++i) {
            m += res_[i];
        }
        m = m / sz_;

        float s = 0;
        for (int i = 0; i < sz_; ++i) {
            s += (res_[i] - m) * (res_[i] - m);
        }
        s = std::sqrt(s / (sz_ - 1));
        return std::make_pair(m, s);
    }

    float min() const {
        float m = res_[0];
        for (int i = 1; i < sz_; ++i) {
            m = m < res_[i] ? m : res_[i];
        }
        return m;
    }

    float max() const {
        float m = res_[0];
        for (int i = 1; i < sz_; ++i) {
            m = m > res_[i] ? m : res_[i];
        }
        return m;
    }

    int nonzero() const {
        int count = 0;
        for (int i = 1; i < sz_; ++i) {
            count += res_[i] == 0 ? 0 : 1;
        }
        return count;
    }

    float norm() const {
        float len = 0;
        for (int i = 0; i < sz_; ++i) {
            len += res_[i] * res_[i];
        }
        len = std::sqrt(len + 1e-6);
        return len;
    }

   private:
    int w_;
    int h_;
    int c_;
    int sz_;
    Pool<float[]>::Recyclable res_;
};

class Param {
   public:
    Param() = default;
    Param(Param&&) = default;
    Param(Dims d) : Param(d.w, d.h, d.c) {}
    Param(int w, int h, int c) : val_(w, h, c), grad_(w, h, c), mem_(w, h, c) {
        val_.gaussian();
        grad_.zero();
        mem_.zero();
    }

    float& operator[](int i) { return val_[i]; }
    float operator[](int i) const { return val_[i]; }

    void reset_grad() { grad_.zero(); }
    Volume& grad() { return grad_; }
    const Volume& grad() const { return grad_; }

    int sz() const { return val_.sz(); }
    Volume& vol() { return val_; }
    const Volume& vol() const { return val_; }

    Param& operator=(Param&& o) = default;

    void reinit(Volume& v) {
        v.share_with(val_);
        grad_ = v.from_shape();
        grad_.zero();
        mem_ = v.from_shape();
        mem_.zero();
    }

    int w() const { return val_.w(); }
    int h() const { return val_.h(); }
    int c() const { return val_.c(); }
    int cha_idx(int c) const { return val_.cha_idx(c); }

    void descend(const Settings& s) {
        switch (s.optimizer) {
            case Optimizer::Sgdm:
                sgdm(s);
                break;
            case Optimizer::Sgd:
                sgd(s);
                break;
            case Optimizer::PowerSign:
                power_sign(s);
                break;
        }
    }

   private:
    void sgd(const Settings& s) {
        float len = grad_.norm();
        for (int i = 0; i < val_.sz(); ++i) {
            val_[i] = val_[i] - s.lr * grad_[i] - s.l2 * 2 * val_[i];
        }
        grad_.zero();
    }

    void sgdm(const Settings& s) {
        float len = grad_.norm();
        for (int i = 0; i < val_.sz(); ++i) {
            grad_[i] =
                len > s.grad_max ? grad_[i] / len * s.grad_max : grad_[i];
            mem_[i] = 0.9 * mem_[i] + 0.1 * grad_[i];
            val_[i] = val_[i] - s.lr * mem_[i] - s.l2 * 2 * val_[i];
        }
        grad_.zero();
    }

    void power_sign(const Settings& s) {
        float len = grad_.norm();
        for (int i = 0; i < val_.sz(); ++i) {
            grad_[i] =
                len > s.grad_max ? grad_[i] / len * s.grad_max : grad_[i];
            mem_[i] = 0.9 * mem_[i] + 0.1 * grad_[i];
            val_[i] -= s.lr * exp(sgn(grad_[i]) * sgn(mem_[i])) * grad_[i];
        }
        grad_.zero();
    }

    Volume val_;
    Volume grad_;
    Volume mem_;
};

class Layer {
   public:
    virtual Volume& forward(const Volume&) = 0;
    virtual Volume backward(const Volume& grad) = 0;
    virtual void update(const Settings&) {}
    virtual Dims out_shape() const = 0;
    virtual void set_train_mode(bool train) {}
};

class Tanh : public Layer {
   public:
    Tanh(Dims d) : out_shape_(d) {}

    virtual Volume& forward(const Volume& input) {
        res_ = input.from_shape();
        for (int i = 0; i < res_.sz(); ++i) {
            res_[i] = std::tanh(input[i]);
        }
        return res_;
    }
    virtual Volume backward(const Volume& pgrad) {
        Volume grad = pgrad.from_shape();
        for (int i = 0; i < res_.sz(); ++i) {
            grad[i] = pgrad[i] * (1 - res_[i] * res_[i]);
        }
        return grad;
    }
    virtual Dims out_shape() const override { return out_shape_; }

   private:
    Volume res_;
    Dims out_shape_;
};

class Verbose : public Layer {
   public:
    Verbose(Dims d, std::string name)
        : out_shape_(d), name_(name), train_(false) {}

    virtual Volume& forward(const Volume& x) override {
        if (train_) {
            fwd_.emplace_back();
            x.share_with(fwd_.back());
            return fwd_.back();
        } else {
            x.share_with(test_thru_);
            return test_thru_;
        }
    }

    virtual Volume backward(const Volume& x) override {
        bwd_.emplace_back();
        x.share_with(bwd_.back());
        Volume v;
        bwd_.back().share_with(v);
        return v;
    }

    virtual void update(const Settings&) override {
        float statsfwd[] = {0, 0, 0, 0, 0, 0};
        for (auto& v : fwd_) {
            auto local = v.mean_std();
            statsfwd[0] += local.first;
            statsfwd[1] += local.second;
            statsfwd[2] += v.max();
            statsfwd[3] += v.min();
            statsfwd[4] += float(v.nonzero()) / v.sz();
            statsfwd[5] += v.norm();
        }
        float statsbwd[] = {0, 0, 0, 0, 0, 0};
        for (auto& v : bwd_) {
            auto local = v.mean_std();
            statsbwd[0] += local.first;
            statsbwd[1] += local.second;
            statsbwd[2] += v.max();
            statsbwd[3] += v.min();
            statsbwd[4] += float(v.nonzero()) / v.sz();
            statsbwd[5] += v.norm();
        }

        std::cout << "               Statistics for " << name_ << ":\n";
        std::cout << "  Batch size: " << fwd_.size() << "\n";
        std::cout << "  Forward pass:                       Backward Pass:\n";
        std::cout << "    Shape: (" << std::setw(3) << fwd_[0].w() << ", "
                  << std::setw(3) << fwd_[0].h() << ", " << std::setw(3)
                  << fwd_[0].c() << ")            "
                  << "  Shape: (" << std::setw(3) << bwd_[0].w() << ", "
                  << std::setw(3) << bwd_[0].h() << ", " << std::setw(3)
                  << bwd_[0].c() << ")\n";

        std::cout << "    Mean:     " << std::setw(8)
                  << statsfwd[0] / fwd_.size()
                  << "                  Mean:     " << statsbwd[0] / bwd_.size()
                  << "\n";
        std::cout << "    Std:      " << std::setw(8)
                  << statsfwd[1] / fwd_.size()
                  << "                  Std:      " << statsbwd[1] / bwd_.size()
                  << "\n";
        std::cout << "    Max:      " << std::setw(8)
                  << statsfwd[2] / fwd_.size()
                  << "                  Max:      " << statsbwd[2] / bwd_.size()
                  << "\n";
        std::cout << "    Min:      " << std::setw(8)
                  << statsfwd[3] / fwd_.size()
                  << "                  Min:      " << statsbwd[3] / bwd_.size()
                  << "\n";
        std::cout << "    Non-Zero: " << std::setw(8)
                  << statsfwd[4] / fwd_.size()
                  << "                  Non-Zero: " << statsbwd[4] / fwd_.size()
                  << "\n";
        std::cout << "    Norm:     " << statsfwd[5] / fwd_.size()
                  << "                  Norm: " << statsbwd[5] / bwd_.size()
                  << "\n";

        fwd_.clear();
        bwd_.clear();
    }

    virtual Dims out_shape() const override { return out_shape_; }
    virtual void set_train_mode(bool b) override { train_ = b; }

   private:
    std::vector<Volume> fwd_;
    std::vector<Volume> bwd_;
    Dims out_shape_;
    std::string name_;
    Volume test_thru_;
    bool train_;
};

class BatchNorm : public Layer {
   public:
    BatchNorm(Dims d)
        : mu_(d), sigma_(d), out_shape_(d), beta_(d), gamma_(d), train_(false) {
        beta_.vol().zero();
        mu_.zero();
        for (int i = 0; i < gamma_.sz(); ++i) {
            gamma_.vol()[i] = 1;
            sigma_[i] = 1;
        }
    }

    virtual Volume& forward(const Volume& input) {
        if (train_) {
            memory_.emplace_back();
            input.share_with(memory_.back());
        }
        scaled_ = input.from_shape();
        centered_ = input.from_shape();
        for (int i = 0; i < scaled_.sz(); ++i) {
            centered_[i] = (input[i] - mu_[i]) / sigma_[i];
            scaled_[i] = gamma_[i] * centered_[i] + beta_[i];
        }
        return scaled_;
    }

    virtual Volume backward(const Volume& pgrad) {
        Volume out = pgrad.from_shape();
        for (int i = 0; i < out.sz(); ++i) {
            gamma_.grad()[i] += pgrad[i] * centered_[i];
            beta_.grad()[i] += pgrad[i];
            out[i] = pgrad[i] * gamma_[i] / sigma_[i];
        }
        return out;
    }

    virtual void update(const Settings& s) {
        beta_.descend(s);
        gamma_.descend(s);
        Volume avg = mu_.from_shape();
        avg.zero();
        // compute average
        for (auto& m : memory_) {
            for (int i = 0; i < m.sz(); ++i) {
                avg[i] += m[i];
            }
        }
        for (int i = 0; i < mu_.sz(); ++i) {
            avg[i] /= memory_.size();
        }

        // compute std
        Volume std = mu_.from_shape();
        std.zero();
        for (auto& m : memory_) {
            for (int i = 0; i < m.sz(); ++i) {
                float x = m[i] - avg[i];
                std[i] += x * x;
            }
        }
        for (int i = 0; i < mu_.sz(); ++i) {
            std[i] /= memory_.size() - 1;
        }

        // update
        for (int i = 0; i < mu_.sz(); ++i) {
            mu_[i] = 0.9 * mu_[i] + 0.1 * avg[i];
            sigma_[i] = 0.9 * sigma_[i] + 0.1 * std[i];
        }

        memory_.clear();
    }

    virtual Dims out_shape() const { return out_shape_; }
    virtual void set_train_mode(bool train) override { train_ = train; }

   private:
    Volume mu_;
    Volume sigma_;
    Param beta_;
    Param gamma_;
    std::vector<Volume> memory_;
    Dims out_shape_;
    Volume centered_;
    Volume scaled_;
    bool train_;
};

class Input : public Layer {
   public:
    Input(Dims d) : out_shape_(d) {}
    virtual Volume& forward(const Volume& x) override {
        x.share_with(out_);
        return out_;
    }
    virtual Volume backward(const Volume& grad) override {
        Volume out;
        grad.share_with(out);
        return out;
    }
    virtual void update(const Settings&) {}
    virtual Dims out_shape() const { return out_shape_; }
    virtual void set_train_mode(bool train) override {}

   private:
    Dims out_shape_;
    Volume out_;
};

class Relu : public Layer {
   public:
    Relu() = default;
    Relu(Dims d) : out_shape_(d){};
    virtual Volume& forward(const Volume& input) override {
        auto res = input.from_shape();
        for (int i = 0; i < input.sz(); ++i) {
            float x = input[i];
            res[i] = x < 0 ? 0 : x;
        }
        res_ = std::move(res);
        return res_;
    }

    virtual Volume backward(const Volume& pgrad) override {
        Volume grad = pgrad.from_shape();
        for (int i = 0; i < grad.sz(); ++i) {
            grad[i] = res_[i] > 0 ? pgrad[i] : 0;
        }
        return std::move(grad);
    }

    virtual Dims out_shape() const { return out_shape_; }

   private:
    Volume res_;
    Dims out_shape_;
};

class FullyConn : public Layer {
   public:
    FullyConn() = default;
    FullyConn(Dims in_sz, int c) : w_(in_sz.sz(), c, 1), b_(1, c, 1) {
        b_.vol().zero();
    }

    virtual Volume& forward(const Volume& input) override {
        input.share_with(x_);
        res_ = Volume(b_.sz(), 1, 1);

        for (int c = 0, wptr = 0; c < b_.sz(); ++c, wptr += input.sz()) {
            float sum = b_[c];
            for (int i = 0; i < input.sz(); ++i) {
                sum += input[i] * w_[wptr + i];
            }
            res_[c] = sum;
        }
        return res_;
    }

    virtual Volume backward(const Volume& pgrad) override {
        Volume x_grad = x_.from_shape();
        x_grad.zero();

        for (int c = 0, wptr = 0; c < b_.sz(); ++c, wptr += x_.sz()) {
            float d = pgrad[c];
            auto& dw = w_.grad();

            b_.grad()[c] += d;
            for (int i = 0; i < x_.sz(); ++i) {
                dw[wptr + i] += d * x_[i];
            }

            for (int i = 0; i < x_.sz(); ++i) {
                x_grad[i] += d * w_[wptr + i];
            }
        }
        return x_grad;
    }

    virtual void update(const Settings& s) override {
        w_.descend(s);
        b_.descend(s);
    }

    void set_weights(std::vector<Volume>&& w, Volume b) {
        Volume v(w[0].sz(), w.size(), 1);
        for (int i = 0; i < w.size(); ++i) {
            for (int j = 0; j < w[i].sz(); ++j) {
                v[v.w() * i + j] = w[i][j];
            }
        }
        w_.reinit(v);
        b_.reinit(b);
    }

    std::vector<Volume> grads() const {
        std::vector<Volume> gs;
        for (int i = 0; i < b_.sz(); ++i) {
            gs.emplace_back(x_.w(), x_.h(), x_.c());
            for (int j = 0; j < w_.w(); ++j) {
                gs.back()[j] = w_.grad()[w_.w() * i + j];
            }
        }
        return gs;
    }
    float bias() const { return b_.vol()[0]; }
    std::vector<Volume> weights() const {
        std::vector<Volume> gs;
        for (int i = 0; i < b_.sz(); ++i) {
            gs.emplace_back(x_.w(), x_.h(), x_.c());
            for (int j = 0; j < w_.w(); ++j) {
                gs.back()[j] = w_[w_.w() * i + j];
            }
        }
        return gs;
    }

    virtual void set_train_mode(bool train) override {}
    virtual Dims out_shape() const override { return Dims{b_.h(), 1, 1}; }

   private:
    Param w_;
    Param b_;
    Volume res_;
    Volume x_;
};

class MaxPool : public Layer {
   public:
    MaxPool(Dims d) : out_shape_(Dims{d.w / 2, d.h / 2, d.c}) {}

    virtual Volume& forward(const Volume& input) override {
        input.share_with(x_);
        res_ = Volume(input.w() / 2, input.h() / 2, input.c());
        if (!cache_) {
            cache_ = std::make_unique<int[]>(res_.sz());
            std::fill(cache_.get(), cache_.get() + res_.sz(), 0);
        }

        for (int c = 0; c < input.c(); ++c) {
            int src_row_idx = input.cha_idx(c);
            int dst_row_idx = res_.cha_idx(c);
            for (int h = 0; h < res_.h(); ++h) {
                for (int w = 0; w < res_.w(); ++w) {
                    int dst_idx = dst_row_idx + w;
                    int src_idx = src_row_idx + 2 * w;
                    int max_pos = std::max(
                        {src_idx,
                         src_idx + 1,
                         src_idx + input.w(),
                         src_idx + input.w() + 1},
                        [&](int a, int b) { return input[a] < input[b]; });
                    cache_[dst_idx] = max_pos;
                    res_[dst_idx] = input[max_pos];
                }
                src_row_idx += 2 * input.w();
                dst_row_idx += res_.w();
            }
        }

        return res_;
    }

    virtual Volume backward(const Volume& pgrad) override {
        Volume grad = x_.from_shape();
        grad.zero();
        for (int i = 0; i < res_.sz(); ++i) {
            grad[cache_[i]] = pgrad[i];
        }
        return grad;
    }

    virtual void set_train_mode(bool train) override {}
    virtual void update(const Settings&) override {}

    virtual Dims out_shape() const override { return out_shape_; }

   private:
    Volume x_;
    std::unique_ptr<int[]> cache_;
    Volume res_;
    Dims out_shape_;
};

class MSE : public Layer {
   public:
    void set_target(const Volume& target) { target.share_with(target_); }

    virtual Volume& forward(const Volume& input) override {
        input.share_with(x_);
        res_ = Volume(1, 1, 1);

        float total = 0;

        for (int i = 0; i < input.sz(); ++i) {
            float err = input[i] - target_[i];
            total += err * err;
        }

        res_[0] = total / (2 * input.sz());

        return res_;
    }

    virtual Volume backward(const Volume& pgrad) override {
        Volume grad = x_.from_shape();
        for (int i = 0; i < grad.sz(); ++i) {
            grad[i] = pgrad[0] * (x_[i] - target_[i]);
        }
        return grad;
    }

    virtual void set_train_mode(bool train) override {}
    virtual void update(const Settings&) override {}
    virtual Dims out_shape() const override { return Dims{1, 1, 1}; }

   private:
    Volume target_;
    Volume x_;
    Volume res_;
};

class Conv : public Layer {
   public:
    Conv() = default;
    Conv(Dims in, Dims kerns) : in_sz_(in), kern_sz_(kerns) {
        for (int i = 0; i < kern_sz_.c; ++i) {
            filters_.emplace_back(kern_sz_.w, kern_sz_.h, in.c);
            biases_.emplace_back(1, 1, 1);
            biases_.back().vol().zero();
        }
    }

    virtual Volume& forward(const Volume& input) override {
        input.share_with(x_);
        res_ = Volume(input.w(), input.h(), filters_.size());

        for (int filter = 0; filter < int(filters_.size()); ++filter) {
            for (int i = res_.cha_idx(filter), end = res_.cha_idx(filter + 1);
                 i < end;
                 ++i) {
                res_[i] = biases_[filter][0];
            }
            for (int wptr = 0; wptr < filters_[0].sz(); ++wptr) {
                auto& f = filters_[filter].vol();

                float weight = f[wptr];

                int channel = wptr / (f.w() * f.h());
                int hoffset = (wptr - f.cha_idx(channel)) / f.w() - f.h() / 2;
                int woffset = (wptr - f.cha_idx(channel)) % f.w() - f.w() / 2;

                int dst_row_ptr = res_.cha_idx(filter) +
                                  std::max(0, hoffset) * res_.w() +
                                  std::max(0, woffset);
                int src_row_ptr = input.cha_idx(channel) +
                                  std::max(0, -hoffset) * input.w() +
                                  std::max(0, -woffset);
                for (int h = 0; h < input.h() - std::abs(hoffset); ++h) {
                    for (int w = 0; w < input.w() - std::abs(woffset); ++w) {
                        int dst_ptr = dst_row_ptr + w;
                        int src_ptr = src_row_ptr + w;

                        res_[dst_ptr] += weight * input[src_ptr];
                    }
                    dst_row_ptr += res_.w();
                    src_row_ptr += input.w();
                }
            }
        }

        return res_;
    }

    virtual Volume backward(const Volume& pgrad) override {
        Volume dx = x_.from_shape();
        dx.zero();

        for (int filter = 0; filter < int(filters_.size()); ++filter) {
            for (int i = pgrad.cha_idx(filter), end = pgrad.cha_idx(filter + 1);
                 i < end;
                 ++i) {
                biases_[filter].grad()[0] += pgrad[i];
            }
            auto& f = filters_[filter].vol();
            auto& df = filters_[filter].grad();
            for (int wptr = 0; wptr < filters_[0].sz(); ++wptr) {
                float weight = f[wptr];

                int channel = wptr / (f.w() * f.h());
                int hoffset = (wptr - f.cha_idx(channel)) / f.w() - f.h() / 2;
                int woffset = (wptr - f.cha_idx(channel)) % f.w() - f.w() / 2;

                int dst_row_ptr = res_.cha_idx(filter) +
                                  std::max(0, -hoffset) * res_.w() +
                                  std::max(0, -woffset);
                int src_row_ptr = dx.cha_idx(channel) +
                                  std::max(0, hoffset) * dx.w() +
                                  std::max(0, woffset);
                for (int h = 0; h < dx.h() - std::abs(hoffset); ++h) {
                    for (int w = 0; w < dx.w() - std::abs(woffset); ++w) {
                        int dst_ptr = dst_row_ptr + w;
                        int src_ptr = src_row_ptr + w;

                        dx[src_ptr] += pgrad[dst_ptr] * weight;
                        df[wptr] += pgrad[dst_ptr] * x_[src_ptr];
                    }
                    dst_row_ptr += dx.w();
                    src_row_ptr += dx.w();
                }
            }
        }
        return dx;
    }

    void set_filters(std::vector<Volume>&& fs, std::vector<float>&& bs) {
        if (fs.size() != bs.size()) {
            throw std::runtime_error("filters and biases of a different size");
        }

        filters_.clear();
        for (auto& v : fs) {
            filters_.emplace_back();
            filters_.back().reinit(v);
        }

        biases_.clear();
        for (auto& v : bs) {
            biases_.emplace_back(1, 1, 1);
            biases_.back()[0] = v;
        }
    }

    std::vector<Volume> filters_grad() const {
        std::vector<Volume> vs;
        for (auto& f : filters_) {
            vs.emplace_back();
            f.grad().share_with(vs.back());
        }
        return vs;
    }

    virtual void update(const Settings& s) override {
        for (int i = 0; i < int(filters_.size()); ++i) {
            filters_[i].descend(s);
            biases_[i].descend(s);
        }
    }

    virtual Dims out_shape() const override {
        return Dims{in_sz_.w, in_sz_.h, kern_sz_.c};
    }
    virtual void set_train_mode(bool train) override {}

   private:
    int nb_f_;
    std::vector<Param> filters_;
    std::vector<Param> biases_;
    Volume res_;
    Volume x_;
    Dims in_sz_;
    Dims kern_sz_;
};

class Net {
   public:
    void set_lr(float x) { settings_.lr = x; }
    void set_batch_size(int x) { settings_.batch_size = x; }
    void set_epochs(int e) { settings_.epochs = e; }
    void set_l2(float x) { settings_.l2 = x; }
    void set_lr_decay(float x) { settings_.lr_decay = x; }
    void set_grad_max(float x) { settings_.grad_max = x; }
    void set_optimizer(Optimizer o) { settings_.optimizer = o; }
    void set_train_mode(bool train) {
        for (auto& l : layers_) {
            l->set_train_mode(train);
        }
    }

    void train(std::vector<Volume>&& xs, std::vector<Volume>&& ys) {
        xs_ = std::move(xs);
        ys_ = std::move(ys);

        set_train_mode(true);
        std::vector<float> errs(xs_.size() / settings_.batch_size + 1);
        for (int e = 0; e < settings_.epochs; ++e) {
            settings_.lr *= settings_.lr_decay;
            std::cout << "Epoch " << e << "\n";
            for (int i = 0; i < int(xs_.size()); i += settings_.batch_size) {
                errs.clear();
                for (int j = i;
                     j < std::min(int(xs_.size()), i + settings_.batch_size);
                     ++j) {
                    mse_.set_target(ys_[j]);
                    errs.push_back(forward(xs_[j])[0]);
                    backward();
                }
                update();
                float total = 0;
                for (float x : errs) {
                    total += x;
                }
                total /= errs.size();
                std::cout << "loss: " << total << "\n";
            }
        }
    }

    void conv(int kw, int kh, int nb) {
        if (layers_.empty()) {
            throw std::invalid_argument("input() must be the first layer");
        }
        auto c = std::make_unique<Conv>(layers_.back()->out_shape(),
                                        Dims{kw, kh, nb});
        layers_.emplace_back(std::move(c));
    }

    void input(int w, int h, int c) {
        if (!layers_.empty()) {
            throw std::invalid_argument("input() must be the first layer");
        }
        layers_.emplace_back(std::make_unique<Input>(Dims{w, h, c}));
    }

    void relu() {
        layers_.emplace_back(
            std::make_unique<Relu>(layers_.back()->out_shape()));
    }

    void maxpool() {
        layers_.emplace_back(
            std::make_unique<MaxPool>(layers_.back()->out_shape()));
    }

    void fc(int out) {
        layers_.emplace_back(
            std::make_unique<FullyConn>(layers_.back()->out_shape(), out));
    }

    void verbose(std::string name) {
        layers_.emplace_back(
            std::make_unique<Verbose>(layers_.back()->out_shape(), name));
    }

    void batch_norm() {
        layers_.emplace_back(
            std::make_unique<BatchNorm>(layers_.back()->out_shape()));
    }

    void tanh() {
        layers_.emplace_back(
            std::make_unique<Tanh>(layers_.back()->out_shape()));
    }

    const Volume predict(const Volume& x) {
        set_train_mode(false);
        Volume it;
        x.share_with(it);
        for (int i = 0; i < int(layers_.size()); ++i) {
            layers_[i]->forward(it).share_with(it);
        }
        return std::move(it);
    }
    const Volume& forward(const Volume& x) {
        Volume it;
        x.share_with(it);
        for (int i = 0; i < int(layers_.size()); ++i) {
            layers_[i]->forward(it).share_with(it);
        }
        return mse_.forward(it);
    }

    Volume backward() {
        Volume one(1, 1, 1);
        one[0] = 1;
        Volume it = mse_.backward(one);
        for (int i = layers_.size(); i-- > 0;) {
            layers_[i]->backward(it).share_with(it);
        }
        return it;
    }

    void update() {
        for (auto& l : layers_) {
            l->update(settings_);
        }
    }

   private:
    std::vector<std::unique_ptr<Layer>> layers_;
    std::vector<Volume> xs_;
    std::vector<Volume> ys_;
    MSE mse_;
    Settings settings_;
};

namespace p = boost::python;
namespace np = boost::python::numpy;

Volume from_array(np::ndarray arr) {
    if (arr.get_nd() != 3) {
        throw std::runtime_error("volume doesn't have 3 dimensions");
    }
    auto as_f = arr.astype(np::dtype::get_builtin<float>());
    Volume v(arr.shape(1), arr.shape(0), arr.shape(2));
    const int channel_size = v.cha_idx(1);
    const int img_size = v.w() * v.h();

    for (int dst = 0, src = 0; dst < img_size; ++dst, src += arr.shape(2)) {
        int channel_dst = dst;
        for (int j = 0; j < arr.shape(2); ++j) {
            float x = reinterpret_cast<float*>(as_f.get_data())[src + j];
            v[channel_dst] = x;
            channel_dst += channel_size;
        }
    }
    return v;
}

np::ndarray to_array(const Volume& v) {
    auto res = np::empty(p::make_tuple(v.h(), v.w(), v.c()),
                         np::dtype::get_builtin<float>());
    float* dat = reinterpret_cast<float*>(res.get_data());

    int img_sz = v.w() * v.h();
    for (int src = 0, dst = 0; src < img_sz; ++src, dst += v.c()) {
        for (int c = 0, src_c = 0; c < v.c(); ++c, src_c += v.cha_idx(1)) {
            dat[dst + c] = v[src + src_c];
        }
    }

    return res;
}

BOOST_PYTHON_MODULE(miniconv) {
    using namespace boost::python;
    class_<Conv, boost::noncopyable>("Conv")
        .def("forward",
             +[](Conv* conv, np::ndarray arr) {
                 return to_array(conv->forward(from_array(arr)));
             })
        .def("backward",
             +[](Conv* conv, np::ndarray arr) {
                 return to_array(conv->backward(from_array(arr)));
             })
        .def("set_filters",
             +[](Conv* conv, p::list kerns, p::list biases) {
                 std::vector<Volume> ks;
                 std::vector<float> bs;
                 for (int i = 0; i < p::len(kerns); ++i) {
                     ks.emplace_back(
                         from_array(p::extract<np::ndarray>(kerns[i])));
                     bs.push_back(p::extract<float>(biases[i]));
                 }
                 conv->set_filters(std::move(ks), std::move(bs));
             })
        .def("filters_grad",
             +[](Conv* conv) {
                 p::list res;
                 for (auto& vol : conv->filters_grad()) {
                     res.append(to_array(vol));
                 }
                 return res;
             })
        .def("update", &Conv::update);

    class_<MSE, boost::noncopyable>("MSE")
        .def("forward",
             +[](MSE* mse, np::ndarray arr) {
                 return to_array(mse->forward(from_array(arr)));
             })
        .def("backward",
             +[](MSE* mse, np::ndarray arr) {
                 return to_array(mse->backward(from_array(arr)));
             })
        .def("set_target",
             +[](MSE* mse, np::ndarray arr) {
                 return mse->set_target(from_array(arr));
             })
        .def("update", &MSE::update);

    class_<Relu, boost::noncopyable>("Relu")
        .def("forward",
             +[](Relu* relu, np::ndarray arr) {
                 return to_array(relu->forward(from_array(arr)));
             })
        .def("backward",
             +[](Relu* relu, np::ndarray arr) {
                 return to_array(relu->backward(from_array(arr)));
             })
        .def("update", &Relu::update);

    class_<FullyConn, boost::noncopyable>("FullyConn")
        .def("forward",
             +[](FullyConn* relu, np::ndarray arr) {
                 return to_array(relu->forward(from_array(arr)));
             })
        .def("backward",
             +[](FullyConn* fc, np::ndarray arr) {
                 return to_array(fc->backward(from_array(arr)));
             })
        .def("set_weights",
             +[](FullyConn* fc, p::list weights, np::ndarray b) {
                 std::vector<Volume> ws;
                 for (int i = 0; i < p::len(weights); ++i) {
                     ws.emplace_back(
                         from_array(p::extract<np::ndarray>(weights[i])));
                 }
                 fc->set_weights(std::move(ws), from_array(b));
             })
        .def("grads",
             +[](FullyConn* fc) {
                 p::list ws;
                 auto ws2 = fc->grads();
                 for (auto& w : ws2) {
                     ws.append(to_array(w));
                 }
                 return ws;
             })
        .def("bias", +[](FullyConn* fc) { return fc->bias(); })
        .def("weights",
             +[](FullyConn* fc) {
                 p::list ws;
                 auto ws2 = fc->weights();
                 for (auto& w : ws2) {
                     ws.append(to_array(w));
                 }
                 return ws;
             })
        .def("update", &FullyConn::update);
    class_<Net, boost::noncopyable>("Net")
        .def("conv", &Net::conv)
        .def("fc", &Net::fc)
        .def("backward", +[](Net* net) { return to_array(net->backward()); })
        .def("forward",
             +[](Net* net, np::ndarray arr) {
                 return to_array(net->forward(from_array(arr)));
             })
        .def("predict",
             +[](Net* net, np::ndarray arr) {
                 return to_array(net->predict(from_array(arr)));
             })
        .def("input", &Net::input)
        .def("maxpool", &Net::maxpool)
        .def("relu", &Net::relu)
        .def("verbose", &Net::verbose)
        .def("tanh", &Net::tanh)
        .def("batch_norm", &Net::batch_norm)
        .def("set_batch_size", &Net::set_batch_size)
        .def("set_l2", &Net::set_l2)
        .def("set_epochs", &Net::set_epochs)
        .def("set_lr", &Net::set_lr)
        .def("set_lr_decay", &Net::set_lr_decay)
        .def("set_optimizer", &Net::set_optimizer)
        .def("set_grad_max", &Net::set_grad_max)
        .def("set_train_mode", &Net::set_train_mode)
        .def("train",
             +[](Net* net, p::list xs, p::list ys) {
                 std::vector<Volume> vxs;
                 std::vector<Volume> vys;
                 if (p::len(xs) != p::len(ys)) {
                     throw std::invalid_argument(
                         "ys and xs must be of equal size");
                 }
                 for (int i = 0; i < p::len(xs); ++i) {
                     vxs.emplace_back(
                         from_array(p::extract<np::ndarray>(xs[i])));
                     vys.emplace_back(
                         from_array(p::extract<np::ndarray>(ys[i])));
                 }
                 net->train(std::move(vxs), std::move(vys));
             })
        .def("update", &Net::update);

    class_<Settings>("Settings")
        .def_readwrite("lr", &Settings::lr)
        .def_readwrite("lr_decay", &Settings::lr_decay)
        .def_readwrite("l2", &Settings::l2)
        .def_readwrite("epochs", &Settings::epochs)
        .def_readwrite("grad_max", &Settings::grad_max)
        .def_readwrite("batch_size", &Settings::batch_size);

    enum_<Optimizer>("Optimizer")
        .value("sgd", Optimizer::Sgd)
        .value("sgdm", Optimizer::Sgdm)
        .value("powersign", Optimizer::PowerSign);

    np::initialize();
    //    feenableexcept(FE_ALL_EXCEPT & ~FE_INEXACT);
}