/*
// Copyright (c) 2000-2009, Texas Engineering Experiment Station (TEES), a
// component of the Texas A&M University System.

// All rights reserved.

// The information and source code contained herein is the exclusive
// property of TEES and may not be disclosed, examined or reproduced
// in whole or in part without explicit written authorization from TEES.
*/

#ifndef STAPL_BENCHMARK_LONESTAR_AGGLOMERATIVECLUSTERING_HPP
#define STAPL_BENCHMARK_LONESTAR_AGGLOMERATIVECLUSTERING_HPP

#include <iostream>
#include <sstream>
#include <map>

#include <stapl/utility/do_once.hpp>
#include <stapl/containers/array/array.hpp>
#include <stapl/views/array_view.hpp>

#include <stapl/containers/graph/hierarchical_graph.hpp>
#include <stapl/containers/graph/views/hgraph_view.hpp>
#include <stapl/containers/graph/dynamic_graph.hpp>
#include <stapl/containers/graph/graph.hpp>
#include <stapl/containers/graph/views/graph_view.hpp>
#include <stapl/containers/graph/algorithms/hierarchical_view.hpp>
#include <stapl/views/repeated_view.hpp>
#include <stapl/views/native_view.hpp>
#include <stapl/containers/graph/algorithms/create_level.hpp>
#include <stapl/containers/graph/algorithms/graph_io.hpp>
#include <stapl/containers/graph/views/property_maps.hpp>

#include <stapl/algorithms/algorithm.hpp>
#include <stapl/skeletons/utility/tags.hpp>
#include <boost/random.hpp>

using namespace stapl;

//////////////////////////////////////////////////////////////////////
/// @brief Holds the actual points to be clustered.
//////////////////////////////////////////////////////////////////////
struct my_point
{
  double m_x, m_y, m_z;

  my_point(double val=0.0)
    : m_x(val), m_y(val), m_z(val)
  { }

  void set(double x, double y, double z)
  {
    m_x = x;
    m_y = y;
    m_z = z;
  }

  void set(my_point const& other)
  {
    m_x = other.m_x;
    m_y = other.m_y;
    m_z = other.m_z;
  }

  void scale(double factor)
  {
    m_x *= factor;
    m_y *= factor;
    m_z *= factor;
  }

  void sub(my_point const& other)
  {
    m_x -= other.m_x;
    m_y -= other.m_y;
    m_z -= other.m_z;
  }

  void add(my_point const& other)
  {
    m_x += other.m_x;
    m_y += other.m_y;
    m_z += other.m_z;
  }

  double get_Magnitude() const
  {
    return m_x*m_x + m_y*m_y + m_z*m_z;
  }

  double get_SquaredEuclideanDistance(my_point const& other) const
  {
    return pow(other.m_x - m_x, 2)
         + pow(other.m_y - m_y, 2)
         + pow(other.m_z - m_z, 2);
  }

  double get_EuclideanDistance(my_point const& other) const
  {
    return sqrt(get_SquaredEuclideanDistance(other));
  }

  double get_ManhattanDistance(my_point const& other) const
  {
    return std::abs(other.m_x - m_x)
         + std::abs(other.m_y - m_y)
         + std::abs(other.m_z - m_z);
  }

  double get_DotProduct(my_point const& other) const
  {
    return other.m_x * m_x
         + other.m_y * m_y
         + other.m_z * m_z;
  }

  double get_CosineSimilarity(my_point const& other) const
  {
    return get_DotProduct(other) / (get_Magnitude() * other.get_Magnitude());
  }

  double get_Distance(my_point const& other) const
  {
    return get_SquaredEuclideanDistance(other);
  }

  void define_type(stapl::typer& t)
  {
    t.member(m_x);
    t.member(m_y);
    t.member(m_z);
  }
};

//////////////////////////////////////////////////////////////////////
/// @brief Stores the points in a full cluster, weight, and nearest neighbor.
//////////////////////////////////////////////////////////////////////
class my_vertex_property
{
private:
  /// @todo Currently only a single point average of all
  /// the true points is stored. This is naive. It will later be changed.
  my_point m_centroid;

  /// The number of points in the cluster.
  int m_count;

  /// The ID of the closest cluster.
  size_t m_my_nearest_neighbor;

public:
  typedef int property_type;

  my_vertex_property()
    : m_centroid(0.0), m_count(0), m_my_nearest_neighbor(0)
  { }

  my_vertex_property(my_point const& other)
    : m_centroid(other), m_count(1), m_my_nearest_neighbor(0)
  { }

  my_point get_Centroid() const
  {
    return m_centroid;
  }

  void set_Centroid(double new_x, double new_y, double new_z)
  {
    m_centroid.set(new_x, new_y, new_z);
  }

  void set_Centroid(my_vertex_property const& other)
  {
    m_centroid = other.m_centroid;
    m_count = other.m_count;
  }

  size_t get_nearest_neighbor() const
  {
    return m_my_nearest_neighbor;
  }

  void set_nearest_neighbor(size_t nearest_neighbor_id)
  {
    m_my_nearest_neighbor = nearest_neighbor_id;
  }

  double get_Distance(const my_vertex_property& other) const
  {
    return m_centroid.get_Distance(other.m_centroid);
  }

  int get_Count() const
  {
    return m_count;
  }

  //////////////////////////////////////////////////////////////////////
  /// @todo Merges this vertex with another and returns a
  /// new my_vertex_property. This is currently naive, using centroids.
  /// @param other The other cluster to be merged.
  /// @return my_vertex_property The new, merged, cluster.
  //////////////////////////////////////////////////////////////////////
  my_vertex_property mergeVertex(const my_vertex_property& other) const
  {
    //calculate centroid
    my_point p1(m_centroid);
    my_point p2(other.m_centroid);

    p1.scale(m_count);
    p2.scale(other.m_count);

    p1.add(p2);
    p1.scale(1.0 / (other.m_count + m_count));

    my_vertex_property vp_out(p1);
    vp_out.m_count = other.m_count + m_count;

    return vp_out;
  }

  my_vertex_property operator=(my_vertex_property const& vertex)
  {
    m_centroid = vertex.m_centroid;
    m_count = vertex.m_count;
    m_my_nearest_neighbor = vertex.m_my_nearest_neighbor;
    return my_vertex_property(vertex);
  }

  void define_type(stapl::typer& t)
  {
    t.member(m_centroid);
    t.member(m_count);
    t.member(m_my_nearest_neighbor);
  }
};

namespace stapl
{
  STAPL_PROXY_HEADER(my_vertex_property)
  {
    STAPL_PROXY_DEFINES(my_vertex_property)
    STAPL_PROXY_METHOD_RETURN(get_Centroid, my_point)
    STAPL_PROXY_METHOD(set_Centroid, double, double, double)
    STAPL_PROXY_METHOD(set_Centroid, my_vertex_property)
    STAPL_PROXY_METHOD_RETURN(get_nearest_neighbor, size_t)
    STAPL_PROXY_METHOD(set_nearest_neighbor, size_t)
    STAPL_PROXY_METHOD_RETURN(get_Distance, my_vertex_property&, double)
    STAPL_PROXY_METHOD_RETURN(get_Count, int)
    STAPL_PROXY_METHOD_RETURN(mergeVertex, my_vertex_property,
      my_vertex_property)
    STAPL_PROXY_METHOD_RETURN(operator=, my_vertex_property,
      my_vertex_property)
    STAPL_PROXY_METHOD(define_type, stapl::typer)
  };
}

//////////////////////////////////////////////////////////////////////
/// @brief Work function to merge vertices.
//////////////////////////////////////////////////////////////////////
struct functor_merge_vertex
{
  typedef my_vertex_property result_type;

  result_type operator()(my_vertex_property& vertex1,
                         my_vertex_property const& vertex2)
  {
    return vertex1.mergeVertex(vertex2);
  }
};

//////////////////////////////////////////////////////////////////////
/// @brief Creates pairs of the vertices to be merged in the next level.
//////////////////////////////////////////////////////////////////////
struct make_node_pairs
{
  typedef std::pair<size_t, double> result_type;

  my_vertex_property m_vertex_prop;
  size_t m_vertex_descriptor;

  make_node_pairs(my_vertex_property clust_prop, size_t vertex_id)
    : m_vertex_prop(clust_prop), m_vertex_descriptor(vertex_id)
  { }

  template<typename GraphNode>
  result_type operator()(GraphNode element) const
  {
    if (element.descriptor() != m_vertex_descriptor) {
      double distance =
         m_vertex_prop.get_Distance(element.property().property);
      return std::make_pair(element.descriptor(), distance);
    }
    return std::make_pair(0, 0.0);
  }

  void define_type(typer &t)
  {
    t.member(m_vertex_prop);
    t.member(m_vertex_descriptor);
  }
};

//////////////////////////////////////////////////////////////////////
/// @brief Reduces two single element pairs to the smaller value.
/// Multi-element pairs are preferred.
//////////////////////////////////////////////////////////////////////
struct pick_smaller_node
{
  typedef std::pair<size_t, double> result_type;

  result_type operator()(result_type el1, result_type el2) const
  {
    const result_type null_value = result_type(0, 0.0);

    if (el1 != null_value && el2 != null_value) {
      if (el1.second < el2.second) {
        return el1;
      } else {
        return el2;
      }
    } else if (el1 == null_value) {
      return el2;
    } else {
      return el1;
    }
  }
};

//////////////////////////////////////////////////////////////////////
/// @brief Work function to determine the nearest neighbors
/// for all of the clusters.
//////////////////////////////////////////////////////////////////////
struct find_nearest_node
{
  typedef void result_type;

  find_nearest_node()
  { }

  template<typename Node, typename GraphView>
  result_type operator()(Node node, const GraphView g_view)
  {
    make_node_pairs pairs_wf(
            node.property().property, node.descriptor());
    size_t mynearestnode =
        (nc_map_reduce(pairs_wf, pick_smaller_node(), g_view)).first;
    node.property().property.set_nearest_neighbor(mynearestnode);
  }
};

//////////////////////////////////////////////////////////////////////
/// @brief Gives an empty property. This is needed for create_level.
//////////////////////////////////////////////////////////////////////
struct empty_functor
{
  typedef int value_type;

  template<class Graph>
  properties::no_property operator()(Graph& g, size_t lvl) const
  {
    return properties::no_property();
  }

  template<class EdgePropertyType>
  properties::no_property operator()(EdgePropertyType& p,
                                     properties::no_property const& ref) const
  {
    return properties::no_property();
  }
};

//////////////////////////////////////////////////////////////////////
/// @brief Creates a property map of the graph, indicating which vertices are
/// to be clustered with which other vertices.
//////////////////////////////////////////////////////////////////////
struct create_property_map
{
  typedef void result_type;

  create_property_map()
  { }

  template<typename Vertex, typename PropMap, typename GraphView>
  result_type operator()
    (Vertex element, PropMap& property_map, GraphView whole_graph)
  {
    size_t element_descriptor = element.descriptor();
    size_t element_nearest_node =
      element.property().property.get_nearest_neighbor();
    if (whole_graph[element_nearest_node].property()
         .property.get_nearest_neighbor() == element_descriptor) {
      // Nodes are to be clustered
      if (element_descriptor < element_nearest_node) {
        property_map[element_descriptor] = element_descriptor;
      } else {
        property_map[element_descriptor] = element_nearest_node;
      }
    } else {
      // Node is not to be clustered
      property_map[element_descriptor] = element_descriptor;
    }
  }
};

//////////////////////////////////////////////////////////////////////
/// @brief Used to fill the graph used in @ref agglomerative_clustering
/// with starting points.
//////////////////////////////////////////////////////////////////////
struct fill_graph
{
  typedef void result_type;

  template<typename Vertex, typename Point>
  void operator()(Vertex element, Point point) {
    element.property() = my_vertex_property(point);
  }
};

//////////////////////////////////////////////////////////////////////
/// @brief Wrapper class for agglomerative clustering algorithm
/// Class needed to store so as to contain typedef for return_type
//////////////////////////////////////////////////////////////////////
class agglomerative_clustering
{
private:
  typedef stapl::static_array<size_t> array_type;
  typedef array_view<array_type> array_vw;
  typedef graph<stapl::DIRECTED, stapl::MULTIEDGES,
    super_vertex_property<my_vertex_property>,
    super_edge_property<properties::no_property> > graph_type;
  typedef graph_view<graph_type> graph_vw;
  typedef graph_external_property_map<graph_vw, size_t, array_vw>
    graph_ext_prop_map;

public:
  typedef std::vector<graph_vw> return_type;

  //////////////////////////////////////////////////////////////////////
  /// @brief Agglomeratively clusters an array of points.
  /// @param inputView The view of the graph to be clustered.
  //////////////////////////////////////////////////////////////////////
  template<typename View>
  return_type cluster(View inputView) {
    graph_type* g = new graph_type(inputView.size());
    graph_vw g_vw(g);

    map_func(fill_graph(), g_vw, inputView);

    std::vector<graph_vw> g_level_vw;
    g_level_vw.push_back(g_vw);

    while (g_vw.size() > 1)
    {
      //Generates the array of nearest neighbors
      nc_map_func(find_nearest_node(), g_vw, make_repeat_view(g_vw));

      array_type group_array(g_vw.size());
      array_vw group_array_vw(group_array);
      nc_map_func(create_property_map(), g_vw, make_repeat_view(group_array_vw),
        make_repeat_view(g_vw));

      graph_ext_prop_map group_id_prop_map(g_vw, group_array_vw);

      g_level_vw.push_back(create_level(g_vw, group_id_prop_map,
        functor_merge_vertex(), empty_functor()));

      g_vw = g_level_vw.back();
    }
    return g_level_vw;
  }

  #ifdef SHOW_RESULTS
  //////////////////////////////////////////////////////////////////////
  /// @brief Converts a point to its RGB hex representation.
  //////////////////////////////////////////////////////////////////////
  std::string ConvertPointToRGBColor(my_point const& point) {
    char b[10];
    sprintf(&b[0], "%02X", (unsigned char)(point.m_x * 255));
    sprintf(&b[2], "%02X", (unsigned char)(point.m_y * 255));
    sprintf(&b[4], "%02X", (unsigned char)(point.m_z * 255));
    return std::string(b);
  }

  //////////////////////////////////////////////////////////////////////
  /// @brief Writes the graph to a dot file (colors points using RGB)
  /// @todo Make more efficient version that does not redundantly build strings
  /// during the creation of dot file.
  //////////////////////////////////////////////////////////////////////
  template<typename Clusters>
  void printgraph2dotfile(Clusters clusters, char const* file="cluster.dot")
  {
    std::map<size_t, std::map<size_t, std::string> > mymap;

    do_once([&clusters, &mymap, file, this](void) {
      for (size_t lvl=0; lvl<clusters.size(); ++lvl) {
        for (size_t n=0; n<clusters[lvl].size(); ++n) {
          auto v = clusters[lvl][n];
          int v_id = v.descriptor();
          if (v.property().children.size() == 1) {
            mymap[lvl][v_id] =  mymap[lvl-1][v.property().children[0]];
            continue;
          }
          std::stringstream ss;
          ss << "subgraph cluster_" << lvl << "_" << v_id << " {\n";

          for (size_t k=0; k<v.property().children.size(); ++k) {
            ss << mymap[lvl-1][v.property().children[k]] << "\n";
          }
          if (lvl != 0) {
            ss << "label = \"";
          } else {
            ss << "\"[" << v_id<< "]\\n";
          }
          my_point p = v.property().property.get_Centroid();
          ss << p.m_x << "\\n" << p.m_y << "\\n" << p.m_z << "\";\n";
          ss << "style=\"filled\";fillcolor=\"#";
          ss << ConvertPointToRGBColor(p) << "\";\n}\n";

          mymap[lvl][v_id] = ss.str();
        }
      }
      std::ofstream ofile;
      ofile.open(file, std::fstream::out);
      ofile << "digraph pGraph {\n";
      ofile << mymap[clusters.size()-1][0];
      ofile << "}\n";
      ofile.close();
    });
    rmi_fence();
  }
  #endif
};

#endif /* STAPL_BENCHMARK_LONESTAR_AGGLOMERATIVECLUSTERING_HPP */


#pragma once
#pragma ide diagnostic ignored "modernize-use-nodiscard"
#pragma ide diagnostic ignored "NotImplementedFunctions"
#pragma ide diagnostic ignored "OCUnusedGlobalDeclarationInspection"
#pragma ide diagnostic ignored "OCUnusedTypeAliasInspection"

#include <array>
#include <map>
#include <vector>

#include <Eigen/Geometry>
#include "observation.hpp"
#include "transform3_with_covariance.hpp"

namespace fvlam
{
  class Logger;

// ==============================================================================
// Marker class
// ==============================================================================

  class Marker
  {
  public:
    static constexpr size_t ArraySize = 4;
    using Element2 = Translate2;
    using Element3 = Translate3;
    using Array2 = std::array<Element2, ArraySize>;
    using Array3 = std::array<Element3, ArraySize>;

  private:
    // The id of the marker
    std::uint64_t id_;

    // The pose of the marker in some world frame. Which frame is world depends on the context.
    Transform3WithCovariance t_map_marker_;

    // Prevent modification if true
    bool is_fixed_{false};

    inline static Array2 unit_corners2_f_marker()
    {
      return Array2{Translate2{-1, 1},
                    Translate2{1, 1},
                    Translate2{1, -1},
                    Translate2{-1, -1}};
    }

    inline static Array3 unit_corners3_f_marker()
    {
      return Array3{Translate3{-1, 1, 0},
                    Translate3{1, 1, 0},
                    Translate3{1, -1, 0},
                    Translate3{-1, -1, 0}};
    }

    inline static Array2 calc_corners2_f_marker(double marker_length)
    {
      auto half_marker_length{marker_length / 2.0};
      auto unit_corners2{unit_corners2_f_marker()};
      return Array2{unit_corners2[0] * half_marker_length,
                    unit_corners2[1] * half_marker_length,
                    unit_corners2[2] * half_marker_length,
                    unit_corners2[3] * half_marker_length};
    }

    inline static Array3 calc_corners3_f_marker(double marker_length)
    {
      auto half_marker_length{marker_length / 2.0};
      auto unit_corners3{unit_corners3_f_marker()};
      return Array3{unit_corners3[0] * half_marker_length,
                    unit_corners3[1] * half_marker_length,
                    unit_corners3[2] * half_marker_length,
                    unit_corners3[3] * half_marker_length};
    }

  public:
    inline Array3 calc_corners3_f_world(double marker_length) const
    {
      auto corners_f_marker = calc_corners3_f_marker(marker_length);
      return Array3{t_map_marker_.tf() * corners_f_marker[0],
                    t_map_marker_.tf() * corners_f_marker[1],
                    t_map_marker_.tf() * corners_f_marker[2],
                    t_map_marker_.tf() * corners_f_marker[3]};
    }

  public:
   Marker() :
      id_(0), t_map_marker_(), is_fixed_(false)
    {}

    Marker(std::uint64_t id, Transform3WithCovariance t_map_marker, bool is_fixed = false) :
      id_(id), t_map_marker_(std::move(t_map_marker)), is_fixed_(is_fixed)
    {}

    auto is_valid() const
    { return t_map_marker_.is_valid(); }

    auto id() const
    { return id_; }

    auto is_fixed() const
    { return is_fixed_; }

    const auto &t_map_marker() const
    { return t_map_marker_; }

    template<class T>
    static Marker from(T &other);

    template<class T>
    T to() const;

    template<class T>
    void to(T &other) const;

    // Return the marker corners in the marker frame (z=0)
    template<class T>
    static T corners_f_marker(double marker_length);

    // Return the marker corners in the world frame given the t_world_marker transform.
    template<class T>
    T corners_f_world(double marker_length) const; //
    template<class T>
    void corners_f_world(double marker_length, T &other) const;

    std::string to_string(bool also_cov = false) const; //
    std::string to_id_string() const; //
    std::string to_corners_f_world_string(double marker_length) const; //

    bool equals(const Marker &other, double tol = 1.0e-9, bool check_relative_also = true) const;

    using ProjectFunction = std::function<Observation(const Marker &marker)>;

    template<class TCameraCalibration>
    static ProjectFunction project_t_world_marker(const TCameraCalibration &camera_calibration,
                                                  const Transform3 &t_world_camera,
                                                  double marker_length);

    template<class TCameraCalibration>
    static ProjectFunction project_t_camera_marker(const TCameraCalibration &camera_calibration,
                                                   double marker_length)
    {
      return project_t_world_marker<TCameraCalibration>(camera_calibration, Transform3{}, marker_length);
    }
  };

// ==============================================================================
// MapEnvironment class
// ==============================================================================

  class MapEnvironment
  {
    std::string description_;
    int marker_dictionary_id_;
    double marker_length_;

  public:
    MapEnvironment() :
      description_{}, marker_dictionary_id_{0}, marker_length_{0.0}
    {}

    MapEnvironment(std::string description,
                   int marker_dictionary_id,
                   double marker_length) :
      description_{std::move(description)},
      marker_dictionary_id_{marker_dictionary_id},
      marker_length_{marker_length}
    {}

    const auto &description() const
    { return description_; }

    const auto &marker_dictionary_id() const
    { return marker_dictionary_id_; }

    const auto &marker_length() const
    { return marker_length_; }

    template<class T>
    static MapEnvironment from(T &other);

    template<class T>
    T to() const;

    template<class T>
    void to(T &other) const;

    std::string to_string() const;

    bool equals(const MapEnvironment &other, double tol = 1.0e-9, bool check_relative_also = true) const;
  };

// ==============================================================================
// MarkerMap class
// ==============================================================================

  class MarkerMap
  {
    MapEnvironment map_environment_;
    std::map<std::uint64_t, Marker> m_{};

  public:
    explicit MarkerMap() :
      map_environment_{}
    {}

    explicit MarkerMap(MapEnvironment map_environment) :
      map_environment_{std::move(map_environment)}
    {}

    const auto &map_environment() const
    { return map_environment_; }

    auto marker_length() const
    { return map_environment_.marker_length(); }

    auto &m_mutable()
    { return m_; }

    auto &m() const
    { return m_; }

    auto size() const
    { return m_.size(); }

    auto empty() const
    { return m_.empty(); }

    template<class T>
    static MarkerMap from(T &other);

    template<class T>
    T to() const;

    template<class T>
    void to(T &other) const;

    std::string to_string(bool also_cov = false) const;

    bool equals(const MarkerMap &other, double tol = 1.0e-9, bool check_relative_also = true) const;

    Marker *find_marker(int id)
    {
      auto marker_pair = m_.find(id);
      return marker_pair == m_.end() ? nullptr : &marker_pair->second;
    }

    const Marker *find_marker_const(std::uint64_t id) const
    {
      auto marker_pair = m_.find(id);
      return marker_pair == m_.end() ? nullptr : &marker_pair->second;
    }

    void add_marker(Marker marker)
    {
      m_.emplace(marker.id(), std::move(marker));
    }

    void save(const std::string &filename, Logger &logger) const; //
    static MarkerMap load(const std::string &filename, Logger &logger); //
  };
}


#ifndef KSA_H
#define KSA_H
#include <boost/multiprecision/cpp_int.hpp>
namespace mp = boost::multiprecision;

class Kahan
{
   public:
      double sum;
      double c;
      mp::cpp_int counter;
      Kahan();
      Kahan(const double& val);
      Kahan& operator+=(const double& val);
      double get_av()const;
};

Kahan::Kahan()
{
   sum = 0.0;
   c   = 0.0;
   counter=0;
}

Kahan::Kahan(const double& val)
{
   sum = val;
   c   = 0.0;
   ++counter;
}

double Kahan::get_av()const
{
   const auto res = sum/(counter.convert_to<double>());
   return res;
}

#pragma clang optimize off
inline Kahan& Kahan::operator+=(const double& val)
{
   const auto y = val - c;
   const auto t = (this->sum) + y;
   this->c      = (t-(this->sum))-y;
   this->sum    = t;
   ++counter;
   return *this;
}
#pragma clang optimize on

#endif


//---------------------------------------------------------------------------//
//!
//! \file   tstUniformDistribution.cpp
//! \author Alex Robinson
//! \brief  Uniform distribution unit tests.
//!
//---------------------------------------------------------------------------//

// Std Lib Includes
#include <iostream>

// Boost Includes
#include <boost/units/systems/si.hpp>
#include <boost/units/systems/cgs.hpp>
#include <boost/units/io.hpp>

// Trilinos Includes
#include <Teuchos_UnitTestHarness.hpp>
#include <Teuchos_RCP.hpp>
#include <Teuchos_Array.hpp>
#include <Teuchos_ParameterList.hpp>
#include <Teuchos_XMLParameterListCoreHelpers.hpp>
#include <Teuchos_VerboseObject.hpp>

// FRENSIE Includes
#include "Utility_UnitTestHarnessExtensions.hpp"
#include "Utility_OneDDistribution.hpp"
#include "Utility_UniformDistribution.hpp"
#include "Utility_RandomNumberGenerator.hpp"
#include "Utility_PhysicalConstants.hpp"
#include "Utility_UnitTraits.hpp"
#include "Utility_QuantityTraits.hpp"
#include "Utility_ElectronVoltUnit.hpp"

using boost::units::quantity;
using namespace Utility::Units;
namespace si = boost::units::si;
namespace cgs = boost::units::cgs;

//---------------------------------------------------------------------------//
// Testing Variables
//---------------------------------------------------------------------------//

Teuchos::RCP<Teuchos::ParameterList> test_dists_list;

Teuchos::RCP<Utility::TabularOneDDistribution> tab_distribution(
			  new Utility::UniformDistribution( -1.0, 1.0, 2.0 ) );

Teuchos::RCP<Utility::OneDDistribution> distribution = tab_distribution;

Teuchos::RCP<Utility::UnitAwareTabularOneDDistribution<si::energy,si::amount> >
  unit_aware_tab_distribution;

Teuchos::RCP<Utility::UnitAwareOneDDistribution<si::energy,si::amount> >
  unit_aware_distribution;

//---------------------------------------------------------------------------//
// Tests.
//---------------------------------------------------------------------------//
// Check that the distribution can be evaluated
TEUCHOS_UNIT_TEST( UniformDistribution, evaluate )
{
  TEST_EQUALITY_CONST( distribution->evaluate( -2.0 ), 0.0 );
  TEST_EQUALITY_CONST( distribution->evaluate( -1.0 ), 2.0 );
  TEST_EQUALITY_CONST( distribution->evaluate( 0.0 ), 2.0 );
  TEST_EQUALITY_CONST( distribution->evaluate( 1.0 ), 2.0 );
  TEST_EQUALITY_CONST( distribution->evaluate( -2.0 ), 0.0 );
}

//---------------------------------------------------------------------------//
// Check that the unit aware distribution can be evaluated
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, evaluate )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluate( -1.0*si::joule ),
		       0.0*si::mole );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluate( 0.0*si::joule ),
  		       1.0*si::mole );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluate( 0.5*si::joule ),
  		       1.0*si::mole );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluate( 1.0*si::joule ),
  		       1.0*si::mole );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluate( 2.0*si::joule ),
  		       0.0*si::mole );
}

//---------------------------------------------------------------------------//
// Check that the PDF can be evaluated
TEUCHOS_UNIT_TEST( UniformDistribution, evaluatePDF )
{
  TEST_EQUALITY_CONST( distribution->evaluatePDF( -2.0 ), 0.0 );
  TEST_EQUALITY_CONST( distribution->evaluatePDF( -1.0 ), 0.5 );
  TEST_EQUALITY_CONST( distribution->evaluatePDF( 0.0 ), 0.5 );
  TEST_EQUALITY_CONST( distribution->evaluatePDF( 1.0 ), 0.5 );
  TEST_EQUALITY_CONST( distribution->evaluatePDF( 2.0 ), 0.0 );
}

//---------------------------------------------------------------------------//
// Check that the unit aware PDF can be evaluated
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, evaluatePDF )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluatePDF( -1.0*si::joule ),
  		       0.0/si::joule );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluatePDF( 0.0*si::joule ),
  		       1.0/si::joule );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluatePDF( 0.5*si::joule ),
  		       1.0/si::joule );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluatePDF( 1.0*si::joule ),
  		       1.0/si::joule );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluatePDF( 2.0*si::joule ),
  		       0.0/si::joule );
}

//---------------------------------------------------------------------------//
// Check that the CDF can be evaluated
TEUCHOS_UNIT_TEST( UniformDistribution, evaluateCDF )
{
  TEST_EQUALITY_CONST( tab_distribution->evaluateCDF( -2.0 ), 0.0 );
  TEST_EQUALITY_CONST( tab_distribution->evaluateCDF( -1.0 ), 0.0 );
  TEST_EQUALITY_CONST( tab_distribution->evaluateCDF( 0.0 ), 0.5 );
  TEST_EQUALITY_CONST( tab_distribution->evaluateCDF( 1.0 ), 1.0 );
  TEST_EQUALITY_CONST( tab_distribution->evaluateCDF( 2.0 ), 1.0 );
}

//---------------------------------------------------------------------------//
// Check that the unit aware CDF can be evaluated
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, evaluateCDF )
{
  TEST_EQUALITY_CONST(unit_aware_tab_distribution->evaluateCDF(-1.0*si::joule),
		      0.0 );
  TEST_EQUALITY_CONST(unit_aware_tab_distribution->evaluateCDF(0.0*si::joule),
		      0.0 );
  TEST_EQUALITY_CONST(unit_aware_tab_distribution->evaluateCDF(0.5*si::joule),
		      0.5 );
  TEST_EQUALITY_CONST(unit_aware_tab_distribution->evaluateCDF(1.0*si::joule),
		      1.0 );
  TEST_EQUALITY_CONST(unit_aware_tab_distribution->evaluateCDF(2.0*si::joule),
		      1.0 );
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled
TEUCHOS_UNIT_TEST( UniformDistribution, sample )
{
  std::vector<double> fake_stream( 3 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.5;
  fake_stream[2] = 1.0 - 1e-15;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  double sample = distribution->sample();
  TEST_EQUALITY_CONST( sample, -1.0 );

  sample = distribution->sample();
  TEST_EQUALITY_CONST( sample, 0.0 );

  sample = distribution->sample();
  TEST_FLOATING_EQUALITY( sample, 1.0, 1e-14 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the unit aware distribution can be sampled
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, sample )
{
  std::vector<double> fake_stream( 3 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.5;
  fake_stream[2] = 1.0 - 1e-15;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  quantity<si::energy> sample = unit_aware_distribution->sample();
  TEST_EQUALITY_CONST( sample, 0.0*si::joule );

  sample = unit_aware_distribution->sample();
  TEST_EQUALITY_CONST( sample, 0.5*si::joule );

  sample = unit_aware_distribution->sample();
  UTILITY_TEST_FLOATING_EQUALITY( sample, 1.0*si::joule, 1e-14 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled w/o an instance
TEUCHOS_UNIT_TEST( UniformDistribution, sample_static )
{
  std::vector<double> fake_stream( 3 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.5;
  fake_stream[2] = 1.0 - 1e-15;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  double sample = Utility::UniformDistribution::sample( -1.0, 1.0 );
  TEST_EQUALITY_CONST( sample, -1.0 );

  sample = Utility::UniformDistribution::sample( -1.0, 1.0 );
  TEST_EQUALITY_CONST( sample, 0.0 );

  sample = Utility::UniformDistribution::sample( -1.0, 1.0 );
  TEST_FLOATING_EQUALITY( sample, 1.0, 1e-14 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be sampled w/o an instance
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, sample_static )
{
  std::vector<double> fake_stream( 3 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.5;
  fake_stream[2] = 1.0 - 1e-15;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  quantity<si::energy> sample =
    Utility::UnitAwareUniformDistribution<si::energy>::sample(
					        0.0*si::joule, 1.0*si::joule );
  TEST_EQUALITY_CONST( sample, 0.0*si::joule );

  sample =
    Utility::UnitAwareUniformDistribution<si::energy>::sample(
					        0.0*si::joule, 1.0*si::joule );
  TEST_EQUALITY_CONST( sample, 0.5*si::joule );

  sample =
    Utility::UnitAwareUniformDistribution<si::energy>::sample(
					        0.0*si::joule, 1.0*si::joule );
  UTILITY_TEST_FLOATING_EQUALITY( sample, 1.0*si::joule, 1e-14 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled
TEUCHOS_UNIT_TEST( UniformDistribution, sampleAndRecordTrials )
{
  std::vector<double> fake_stream( 3 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.5;
  fake_stream[2] = 1.0 - 1e-15;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  unsigned trials = 0;

  double sample = distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, -1.0 );
  TEST_EQUALITY_CONST( 1.0/trials, 1.0 );

  sample = distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, 0.0 );
  TEST_EQUALITY_CONST( 2.0/trials, 1.0 );

  sample = distribution->sampleAndRecordTrials( trials );
  TEST_FLOATING_EQUALITY( sample, 1.0, 1e-14 );
  TEST_EQUALITY_CONST( 3.0/trials, 1.0 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the unit aware distribution can be sampled
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, sampleAndRecordTrials )
{
  std::vector<double> fake_stream( 3 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.5;
  fake_stream[2] = 1.0 - 1e-15;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  unsigned trials = 0;

  quantity<si::energy> sample =
    unit_aware_distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, 0.0*si::joule );
  TEST_EQUALITY_CONST( 1.0/trials, 1.0 );

  sample = unit_aware_distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, 0.5*si::joule );
  TEST_EQUALITY_CONST( 2.0/trials, 1.0 );

  sample = unit_aware_distribution->sampleAndRecordTrials( trials );
  UTILITY_TEST_FLOATING_EQUALITY( sample, 1.0*si::joule, 1e-14 );
  TEST_EQUALITY_CONST( 3.0/trials, 1.0 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be sampled w/o an instance
TEUCHOS_UNIT_TEST( UniformDistribution, sampleAndRecordTrials_static )
{
  std::vector<double> fake_stream( 3 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.5;
  fake_stream[2] = 1.0 - 1e-15;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  unsigned trials = 0;

  double sample = Utility::UniformDistribution::sampleAndRecordTrials(
							   -1.0, 1.0, trials );
  TEST_EQUALITY_CONST( sample, -1.0 );
  TEST_EQUALITY_CONST( 1.0/trials, 1.0 );

  sample = Utility::UniformDistribution::sampleAndRecordTrials(
							   -1.0, 1.0, trials );
  TEST_EQUALITY_CONST( sample, 0.0 );
  TEST_EQUALITY_CONST( 2.0/trials, 1.0 );

  sample = Utility::UniformDistribution::sampleAndRecordTrials(
							   -1.0, 1.0, trials );
  TEST_FLOATING_EQUALITY( sample, 1.0, 1e-14 );
  TEST_EQUALITY_CONST( 3.0/trials, 1.0 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be sampled w/o an instance
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, sampleAndRecordTrials_static )
{
  std::vector<double> fake_stream( 3 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.5;
  fake_stream[2] = 1.0 - 1e-15;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  unsigned trials = 0;

  quantity<si::energy> sample = Utility::UnitAwareUniformDistribution<si::energy,si::amount>::sampleAndRecordTrials(
				        0.0*si::joule, 1.0*si::joule, trials );
  TEST_EQUALITY_CONST( sample, 0.0*si::joule );
  TEST_EQUALITY_CONST( 1.0/trials, 1.0 );

  sample =
    Utility::UnitAwareUniformDistribution<si::energy,si::amount>::sampleAndRecordTrials(
				        0.0*si::joule, 1.0*si::joule, trials );
  TEST_EQUALITY_CONST( sample, 0.5*si::joule );
  TEST_EQUALITY_CONST( 2.0/trials, 1.0 );

  sample =
    Utility::UnitAwareUniformDistribution<si::energy,si::amount>::sampleAndRecordTrials(
				        0.0*si::joule, 1.0*si::joule, trials );
  UTILITY_TEST_FLOATING_EQUALITY( sample, 1.0*si::joule, 1e-14 );
  TEST_EQUALITY_CONST( 3.0/trials, 1.0 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled
TEUCHOS_UNIT_TEST( UniformDistribution, sampleAndRecordBinIndex )
{
  std::vector<double> fake_stream( 3 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.5;
  fake_stream[2] = 1.0 - 1e-15;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  unsigned bin_index;

  double sample = tab_distribution->sampleAndRecordBinIndex( bin_index );
  TEST_EQUALITY_CONST( sample, -1.0 );
  TEST_EQUALITY_CONST( bin_index, 0.0 );

  sample = tab_distribution->sampleAndRecordBinIndex( bin_index );
  TEST_EQUALITY_CONST( sample, 0.0 );
  TEST_EQUALITY_CONST( bin_index, 0.0 );

  sample = tab_distribution->sampleAndRecordBinIndex( bin_index );
  TEST_FLOATING_EQUALITY( sample, 1.0, 1e-14 );
  TEST_EQUALITY_CONST( bin_index, 0.0 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be sampled
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, sampleAndRecordBinIndex )
{
  std::vector<double> fake_stream( 3 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.5;
  fake_stream[2] = 1.0 - 1e-15;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  unsigned bin_index;

  quantity<si::energy> sample =
    unit_aware_tab_distribution->sampleAndRecordBinIndex( bin_index );
  TEST_EQUALITY_CONST( sample, 0.0*si::joule );
  TEST_EQUALITY_CONST( bin_index, 0.0 );

  sample = unit_aware_tab_distribution->sampleAndRecordBinIndex( bin_index );
  TEST_EQUALITY_CONST( sample, 0.5*si::joule );
  TEST_EQUALITY_CONST( bin_index, 0.0 );

  sample = unit_aware_tab_distribution->sampleAndRecordBinIndex( bin_index );
  UTILITY_TEST_FLOATING_EQUALITY( sample, 1.0*si::joule, 1e-14 );
  TEST_EQUALITY_CONST( bin_index, 0.0 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled
TEUCHOS_UNIT_TEST( UniformDistribution, sampleWithRandomNumber )
{
  double sample = tab_distribution->sampleWithRandomNumber( 0.0 );
  TEST_EQUALITY_CONST( sample, -1.0 );

  sample = tab_distribution->sampleWithRandomNumber( 0.5 );
  TEST_EQUALITY_CONST( sample, 0.0 );

  sample = tab_distribution->sampleWithRandomNumber( 1.0 );
  TEST_FLOATING_EQUALITY( sample, 1.0, 1e-14 );
}

//---------------------------------------------------------------------------//
// Check that the unit aware distribution can be sampled
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, sampleWithRandomNumber )
{
  quantity<si::energy> sample =
    unit_aware_tab_distribution->sampleWithRandomNumber( 0.0 );
  TEST_EQUALITY_CONST( sample, 0.0*si::joule );

  sample = unit_aware_tab_distribution->sampleWithRandomNumber( 0.5 );
  TEST_EQUALITY_CONST( sample, 0.5*si::joule );

  sample = unit_aware_tab_distribution->sampleWithRandomNumber( 1.0 );
  UTILITY_TEST_FLOATING_EQUALITY( sample, 1.0*si::joule, 1e-14 );
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled w/o an instance
TEUCHOS_UNIT_TEST( UniformDistribution, sampleWithRandomNumber_static )
{
  double sample = Utility::UniformDistribution::sampleWithRandomNumber(
							      -1.0, 1.0, 0.0 );
  TEST_EQUALITY_CONST( sample, -1.0 );

  sample = Utility::UniformDistribution::sampleWithRandomNumber(
							      -1.0, 1.0, 0.5 );
  TEST_EQUALITY_CONST( sample, 0.0 );

  sample = Utility::UniformDistribution::sampleWithRandomNumber(
							      -1.0, 1.0, 1.0 );
  TEST_FLOATING_EQUALITY( sample, 1.0, 1e-14 );
}

//---------------------------------------------------------------------------//
// Check that the unit aware distribution can be sampled w/o an instance
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution,
		   sampleWithRandomNumber_static )
{
  quantity<si::energy> sample = Utility::UnitAwareUniformDistribution<si::energy,si::amount>::sampleWithRandomNumber(
					   0.0*si::joule, 1.0*si::joule, 0.0 );
  TEST_EQUALITY_CONST( sample, 0.0*si::joule );

  sample = Utility::UnitAwareUniformDistribution<si::energy,si::amount>::sampleWithRandomNumber(
					   0.0*si::joule, 1.0*si::joule, 0.5 );
  TEST_EQUALITY_CONST( sample, 0.5*si::joule );

  sample = Utility::UnitAwareUniformDistribution<si::energy,si::amount>::sampleWithRandomNumber(
					   0.0*si::joule, 1.0*si::joule, 1.0 );
  UTILITY_TEST_FLOATING_EQUALITY( sample, 1.0*si::joule, 1e-14 );
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled
TEUCHOS_UNIT_TEST( UniformDistribution, sampleInSubrange )
{
  std::vector<double> fake_stream( 3 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.5;
  fake_stream[2] = 1.0 - 1e-15;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  double sample = tab_distribution->sampleInSubrange( 0.0 );
  TEST_EQUALITY_CONST( sample, -1.0 );

  sample = tab_distribution->sampleInSubrange( 0.0 );
  TEST_EQUALITY_CONST( sample, -0.5 );

  sample = tab_distribution->sampleInSubrange( 0.0 );
  UTILITY_TEST_FLOATING_EQUALITY( sample, 0.0, 1e-14 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the unit aware distribution can be sampled
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, sampleInSubrange )
{
  std::vector<double> fake_stream( 3 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.5;
  fake_stream[2] = 1.0 - 1e-15;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  quantity<si::energy> sample =
    unit_aware_tab_distribution->sampleInSubrange( 0.5*si::joule );
  TEST_EQUALITY_CONST( sample, 0.0*si::joule );

  sample = unit_aware_tab_distribution->sampleInSubrange( 0.5*si::joule );
  TEST_EQUALITY_CONST( sample, 0.25*si::joule );

  sample = unit_aware_tab_distribution->sampleInSubrange( 0.5*si::joule );
  UTILITY_TEST_FLOATING_EQUALITY( sample, 0.5*si::joule, 1e-14 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled
TEUCHOS_UNIT_TEST( UniformDistribution, sampleWithRandomNumberInSubrange )
{
  double sample = tab_distribution->sampleWithRandomNumberInSubrange( 0.0, 0.0 );
  TEST_EQUALITY_CONST( sample, -1.0 );

  sample = tab_distribution->sampleWithRandomNumberInSubrange( 0.5, 0.0 );
  TEST_EQUALITY_CONST( sample, -0.5 );

  sample = tab_distribution->sampleWithRandomNumberInSubrange( 1.0, 0.0 );
  UTILITY_TEST_FLOATING_EQUALITY( sample, 0.0, 1e-14 );
}

//---------------------------------------------------------------------------//
// Check that the unit aware distribution can be sampled
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution,
		   sampleWithRandomNumberInSubrange )
{
  quantity<si::energy> sample =
    unit_aware_tab_distribution->sampleWithRandomNumberInSubrange(
							  0.0, 0.5*si::joule );
  TEST_EQUALITY_CONST( sample, 0.0*si::joule );

  sample = unit_aware_tab_distribution->sampleWithRandomNumberInSubrange(
							  0.5, 0.5*si::joule );
  TEST_EQUALITY_CONST( sample, 0.25*si::joule );

  sample = unit_aware_tab_distribution->sampleWithRandomNumberInSubrange(
							  1.0, 0.5*si::joule );
  UTILITY_TEST_FLOATING_EQUALITY( sample, 0.5*si::joule, 1e-14 );
}

//---------------------------------------------------------------------------//
// Check that the upper bound of the distribution independent variable can be
// returned
TEUCHOS_UNIT_TEST( UniformDistribution, getUpperBoundOfIndepVar )
{
  TEST_EQUALITY_CONST( distribution->getUpperBoundOfIndepVar(), 1.0 );
}

//---------------------------------------------------------------------------//
// Check that the upper bound of the unit-aware distribution independent
// variable can be returned
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, getUpperBoundOfIndepVar )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->getUpperBoundOfIndepVar(),
		       1.0*si::joule );
}

//---------------------------------------------------------------------------//
// Check that the lower bound of the distribution independent variable can be
// returned
TEUCHOS_UNIT_TEST( UniformDistribution, getLowerBoundOfIndepVar )
{
  TEST_EQUALITY_CONST( distribution->getLowerBoundOfIndepVar(), -1.0 );
}

//---------------------------------------------------------------------------//
// Check that the lower bound of the unit-aware distribution independent
// variable can be returned
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, getLowerBoundOfIndepVar )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->getLowerBoundOfIndepVar(),
		       0.0*si::joule );
}

//---------------------------------------------------------------------------//
// Check that the distribution type can be returned
TEUCHOS_UNIT_TEST( UniformDistribution, getDistributionType )
{
  TEST_EQUALITY_CONST( distribution->getDistributionType(),
		       Utility::UNIFORM_DISTRIBUTION );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution type can be returned
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, getDistributionType )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->getDistributionType(),
		       Utility::UNIFORM_DISTRIBUTION );
}

//---------------------------------------------------------------------------//
// Check if the distribution is tabular
TEUCHOS_UNIT_TEST( UniformDistribution, isTabular )
{
  TEST_ASSERT( distribution->isTabular() );
}

//---------------------------------------------------------------------------//
// Check if the unit-aware distribution is tabular
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, isTabular )
{
  TEST_ASSERT( unit_aware_distribution->isTabular() );
}

//---------------------------------------------------------------------------//
// Check if the distribution is continuous
TEUCHOS_UNIT_TEST( UniformDistribution, isContinuous )
{
  TEST_ASSERT( distribution->isContinuous() );
}

//---------------------------------------------------------------------------//
// Check if the unit-aware distribution is continuous
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, isContinuous )
{
  TEST_ASSERT( unit_aware_distribution->isContinuous() );
}

//---------------------------------------------------------------------------//
// Check that the distribution can be written to an xml file
TEUCHOS_UNIT_TEST( UniformDistribution, toParameterList )
{
  Teuchos::RCP<Utility::UniformDistribution> true_distribution =
   Teuchos::rcp_dynamic_cast<Utility::UniformDistribution>( distribution );

  Teuchos::ParameterList parameter_list;

  parameter_list.set<Utility::UniformDistribution>( "test distribution",
						     *true_distribution );

  Teuchos::writeParameterListToXmlFile( parameter_list,
					"uniform_dist_test_list.xml" );

  Teuchos::RCP<Teuchos::ParameterList> read_parameter_list =
    Teuchos::getParametersFromXmlFile( "uniform_dist_test_list.xml" );

  TEST_EQUALITY( parameter_list, *read_parameter_list );

  Teuchos::RCP<Utility::UniformDistribution>
    copy_distribution( new Utility::UniformDistribution );

  *copy_distribution =
    read_parameter_list->get<Utility::UniformDistribution>(
							  "test distribution");

  TEST_EQUALITY( *copy_distribution, *true_distribution );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be written to an xml file
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, toParameterList )
{
  typedef Utility::UnitAwareUniformDistribution<si::energy,si::amount> UnitAwareUniformDistribution;

  Teuchos::RCP<UnitAwareUniformDistribution> true_distribution =
    Teuchos::rcp_dynamic_cast<UnitAwareUniformDistribution>( unit_aware_distribution );

  Teuchos::ParameterList parameter_list;

  parameter_list.set<UnitAwareUniformDistribution>( "test distribution",
						    *true_distribution );

  Teuchos::writeParameterListToXmlFile( parameter_list,
					"unit_aware_uniform_dist_test_list.xml" );

  Teuchos::RCP<Teuchos::ParameterList> read_parameter_list =
    Teuchos::getParametersFromXmlFile( "unit_aware_uniform_dist_test_list.xml" );

  TEST_EQUALITY( parameter_list, *read_parameter_list );

  Teuchos::RCP<UnitAwareUniformDistribution>
    copy_distribution( new UnitAwareUniformDistribution );

  *copy_distribution =
    read_parameter_list->get<UnitAwareUniformDistribution>(
							  "test distribution");

  TEST_EQUALITY( *copy_distribution, *true_distribution );
}

//---------------------------------------------------------------------------//
// Check that the distribution can be read from an xml file
TEUCHOS_UNIT_TEST( UniformDistribution, fromParameterList )
{
  Utility::UniformDistribution xml_distribution =
    test_dists_list->get<Utility::UniformDistribution>( "Uniform Distribution A" );

  TEST_EQUALITY_CONST( xml_distribution.getLowerBoundOfIndepVar(), -1.0 );
  TEST_EQUALITY_CONST( xml_distribution.getUpperBoundOfIndepVar(), 1.0 );
  TEST_EQUALITY_CONST( xml_distribution.evaluate( 0.0 ), 2.0 );

  xml_distribution =
    test_dists_list->get<Utility::UniformDistribution>( "Uniform Distribution B" );

  TEST_EQUALITY_CONST( xml_distribution.getLowerBoundOfIndepVar(), 0.0 );
  TEST_EQUALITY_CONST( xml_distribution.getUpperBoundOfIndepVar(),
		       2*Utility::PhysicalConstants::pi );
  TEST_EQUALITY_CONST( xml_distribution.evaluate( 1.0 ), 1.0 );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be read from an xml file
TEUCHOS_UNIT_TEST( UnitAwareUniformDistribution, fromParameterList )
{
  typedef Utility::UnitAwareUniformDistribution<si::energy,si::amount> UnitAwareUniformDistribution;

  UnitAwareUniformDistribution xml_distribution =
    test_dists_list->get<UnitAwareUniformDistribution>( "Unit-Aware Uniform Distribution A" );

  TEST_EQUALITY_CONST( xml_distribution.getLowerBoundOfIndepVar(), 0.0*si::joule );
  TEST_EQUALITY_CONST( xml_distribution.getUpperBoundOfIndepVar(), 10.0*si::joule);
  TEST_EQUALITY_CONST( xml_distribution.evaluate( 5.0*si::joule ), 3.0*si::mole );

  xml_distribution =
    test_dists_list->get<UnitAwareUniformDistribution>( "Unit-Aware Uniform Distribution B" );

  TEST_EQUALITY_CONST( xml_distribution.getLowerBoundOfIndepVar(), 0.0*si::joule );
  TEST_EQUALITY_CONST( xml_distribution.getUpperBoundOfIndepVar(),
		       Utility::PhysicalConstants::pi*si::joule );
  TEST_EQUALITY_CONST( xml_distribution.evaluate( 1.0*si::joule ), 1.0*si::mole );
}

//---------------------------------------------------------------------------//
// Check that a unit-aware distribution can be scaled
TEUCHOS_UNIT_TEST_TEMPLATE_4_DECL( UnitAwareUniformDistribution,
				   explicit_conversion,
				   IndepUnitA,
				   DepUnitA,
				   IndepUnitB,
				   DepUnitB )
{
  typedef typename Utility::UnitTraits<IndepUnitA>::template GetQuantityType<double>::type IndepQuantityA;
  typedef typename Utility::UnitTraits<typename Utility::UnitTraits<IndepUnitA>::InverseUnit>::template GetQuantityType<double>::type InverseIndepQuantityA;

  typedef typename Utility::UnitTraits<IndepUnitB>::template GetQuantityType<double>::type IndepQuantityB;
  typedef typename Utility::UnitTraits<typename Utility::UnitTraits<IndepUnitB>::InverseUnit>::template GetQuantityType<double>::type InverseIndepQuantityB;

  typedef typename Utility::UnitTraits<DepUnitA>::template GetQuantityType<double>::type DepQuantityA;
  typedef typename Utility::UnitTraits<DepUnitB>::template GetQuantityType<double>::type DepQuantityB;

  // Copy from a unitless distribution to distribution type A (static method)
  Utility::UnitAwareUniformDistribution<IndepUnitA,DepUnitA>
    unit_aware_dist_a_copy = Utility::UnitAwareUniformDistribution<IndepUnitA,DepUnitA>::fromUnitlessDistribution( *Teuchos::rcp_dynamic_cast<Utility::UniformDistribution>( distribution ) );

  // Copy from distribution type A to distribution type B (explicit cast)
  Utility::UnitAwareUniformDistribution<IndepUnitB,DepUnitB>
    unit_aware_dist_b_copy( unit_aware_dist_a_copy );

  IndepQuantityA indep_quantity_a =
    Utility::QuantityTraits<IndepQuantityA>::initializeQuantity( -1.0 );
  InverseIndepQuantityA inv_indep_quantity_a =
    Utility::QuantityTraits<InverseIndepQuantityA>::initializeQuantity( 0.5 );
  DepQuantityA dep_quantity_a =
    Utility::QuantityTraits<DepQuantityA>::initializeQuantity( 2.0 );

  IndepQuantityB indep_quantity_b( indep_quantity_a );
  InverseIndepQuantityB inv_indep_quantity_b( inv_indep_quantity_a );
  DepQuantityB dep_quantity_b( dep_quantity_a );

  UTILITY_TEST_FLOATING_EQUALITY(
			   unit_aware_dist_a_copy.evaluate( indep_quantity_a ),
			   dep_quantity_a,
			   1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			unit_aware_dist_a_copy.evaluatePDF( indep_quantity_a ),
			inv_indep_quantity_a,
			1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			   unit_aware_dist_b_copy.evaluate( indep_quantity_b ),
			   dep_quantity_b,
			   1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			unit_aware_dist_b_copy.evaluatePDF( indep_quantity_b ),
			inv_indep_quantity_b,
			1e-15 );

  Utility::setQuantity( indep_quantity_a, 1.0 );
  Utility::setQuantity( inv_indep_quantity_a, 0.5 );
  Utility::setQuantity( dep_quantity_a, 2.0 );

  indep_quantity_b = IndepQuantityB( indep_quantity_a );
  inv_indep_quantity_b = InverseIndepQuantityB( inv_indep_quantity_a );
  dep_quantity_b = DepQuantityB( dep_quantity_a );

  UTILITY_TEST_FLOATING_EQUALITY(
			   unit_aware_dist_a_copy.evaluate( indep_quantity_a ),
			   dep_quantity_a,
			   1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			unit_aware_dist_a_copy.evaluatePDF( indep_quantity_a ),
			inv_indep_quantity_a,
			1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			   unit_aware_dist_b_copy.evaluate( indep_quantity_b ),
			   dep_quantity_b,
			   1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			unit_aware_dist_b_copy.evaluatePDF( indep_quantity_b ),
			inv_indep_quantity_b,
			1e-15 );
}

typedef si::energy si_energy;
typedef cgs::energy cgs_energy;
typedef si::length si_length;
typedef cgs::length cgs_length;
typedef si::mass si_mass;
typedef cgs::mass cgs_mass;
typedef si::amount si_amount;
typedef si::dimensionless si_dimensionless;
typedef cgs::dimensionless cgs_dimensionless;

TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      si_energy,
				      si_amount,
				      cgs_energy,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      cgs_energy,
				      si_amount,
				      si_energy,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      si_length,
				      si_energy,
				      cgs_length,
				      cgs_energy );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      cgs_length,
				      cgs_energy,
				      si_length,
				      si_energy );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      si_length,
				      si_mass,
				      cgs_length,
				      cgs_mass );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      cgs_length,
				      cgs_mass,
				      si_length,
				      si_mass );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      si_length,
				      void,
				      cgs_length,
				      void );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      cgs_length,
				      void,
				      si_length,
				      void );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      void,
				      si_mass,
				      void,
				      cgs_mass );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      void,
				      cgs_mass,
				      void,
				      si_mass );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      si_length,
				      si_dimensionless,
				      cgs_length,
				      cgs_dimensionless );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      cgs_length,
				      cgs_dimensionless,
				      si_length,
				      si_dimensionless );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      si_dimensionless,
				      si_length,
				      cgs_dimensionless,
				      cgs_length );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareUniformDistribution,
				      explicit_conversion,
				      cgs_dimensionless,
				      cgs_length,
				      si_dimensionless,
				      si_length );

//---------------------------------------------------------------------------//
// Custom main function
//---------------------------------------------------------------------------//
int main( int argc, char** argv )
{
  std::string test_dists_xml_file;

  Teuchos::CommandLineProcessor& clp = Teuchos::UnitTestRepository::getCLP();

  clp.setOption( "test_dists_xml_file",
		 &test_dists_xml_file,
		 "Test distributions xml file name" );

  const Teuchos::RCP<Teuchos::FancyOStream> out =
    Teuchos::VerboseObjectBase::getDefaultOStream();

  Teuchos::CommandLineProcessor::EParseCommandLineReturn parse_return =
    clp.parse(argc,argv);

  if ( parse_return != Teuchos::CommandLineProcessor::PARSE_SUCCESSFUL ) {
    *out << "\nEnd Result: TEST FAILED" << std::endl;
    return parse_return;
  }

  TEUCHOS_ADD_TYPE_CONVERTER( Utility::UniformDistribution );
  typedef Utility::UnitAwareUniformDistribution<si::energy,si::amount> UnitAwareUniformDistribution;
  TEUCHOS_ADD_TYPE_CONVERTER( UnitAwareUniformDistribution );

  test_dists_list = Teuchos::getParametersFromXmlFile( test_dists_xml_file );

  // Initialize the unit-aware distributions
  unit_aware_tab_distribution.reset(
	     new Utility::UnitAwareUniformDistribution<si::energy,si::amount>(
				      quantity<si::energy>( 0.0*si::joule ),
				      quantity<si::energy>( 1.0*si::joule ),
				      quantity<si::amount>( 1.0*si::mole ) ) );

  unit_aware_distribution = unit_aware_tab_distribution;

  // Initialize the random number generator
  Utility::RandomNumberGenerator::createStreams();

  // Run the unit tests
  Teuchos::GlobalMPISession mpiSession( &argc, &argv );

  const bool success = Teuchos::UnitTestRepository::runUnitTests(*out);

  if (success)
    *out << "\nEnd Result: TEST PASSED" << std::endl;
  else
    *out << "\nEnd Result: TEST FAILED" << std::endl;

  clp.printFinalTimerSummary(out.ptr());

  return (success ? 0 : 1);
}

//---------------------------------------------------------------------------//
// end tstUniformDistribution.cpp
//---------------------------------------------------------------------------//


/******************************************************************************
 * Copyright (C) 2014 by Jerome Maye                                          *
 * jerome.maye@gmail.com                                                      *
 ******************************************************************************/

/** \file ErrorTermPoseTest.cpp
    \brief This file tests the ErrorTermPose class.
  */

#include <cmath>

#include <boost/make_shared.hpp>

#include <gtest/gtest.h>

#include <sm/kinematics/EulerAnglesYawPitchRoll.hpp>

#include <aslam/backend/test/ErrorTermTestHarness.hpp>
#include <aslam/backend/EuclideanPoint.hpp>
#include <aslam/backend/RotationQuaternion.hpp>

#include <aslam/calibration/statistics/NormalDistribution.h>

#include "aslam/calibration/car/error-terms/ErrorTermPose.h"

using namespace aslam::backend;
using namespace aslam::calibration;
using namespace sm::kinematics;

TEST(AslamCalibrationTestSuite, testErrorTermPose) {

  // estimated pose, angles: yaw-pitch-roll
  const ErrorTermPose::Input xe((ErrorTermPose::Input() << 100, 200, 300,
    M_PI / 2.0, M_PI / 4.0, M_PI / 8.0).finished());

  // covariance matrix of the measurement
  ErrorTermPose::Covariance Q = ErrorTermPose::Covariance::Zero();
  Q(0, 0) = 1e-4;
  Q(1, 1) = 1e-4;
  Q(2, 2) = 1e-4;
  Q(3, 3) = 4e-4;
  Q(4, 4) = 4e-4;
  Q(5, 5) = 4e-4;

  // measured pose
  const ErrorTermPose::Input xm(NormalDistribution<6>(xe, Q).getSample());

  // encapsulate the estimated pose into a transformation expression
  auto translation = boost::make_shared<EuclideanPoint>(xe.head<3>());
  EulerAnglesYawPitchRoll conv;
  auto rotation = boost::make_shared<RotationQuaternion>(
    conv.parametersToRotationMatrix(xe.tail<3>()));
  TransformationExpression T(RotationExpression(rotation.get()),
    EuclideanExpression(translation.get()));

  // build the error term
  ErrorTermPose e1(T, xm, Q);

  // test the error term
  try {
    ErrorTermTestHarness<6> harness(&e1);
    harness.testAll();
  }
  catch (const std::exception& e) {
    FAIL() << e.what();
  }

  // accessors test
  e1.setInput(xm);
  e1.setCovariance(Q);
  ASSERT_EQ(e1.getInput(), xm);
  ASSERT_EQ(e1.getCovariance(), Q);
}


/* Copyright 2021-2022 Massachusetts Institute of Technology

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

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

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

#include <catch2/catch.hpp>

#include <gentl/util/randutils.h>
#include <gentl/modeling/sequence.h>
#include <gentl/inference/particle_filter.h>
#include <iostream>
#include <Eigen/Dense>

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

namespace sequence = gentl::modeling::sequence;

struct Parameters {
    VectorXd prior;
    MatrixXd emission_matrix;
    MatrixXd transition_matrix;
    Parameters(VectorXd prior_, MatrixXd emission_matrix_, MatrixXd transition_matrix_) :
        prior(prior_), emission_matrix(emission_matrix_), transition_matrix(transition_matrix_) {}
};

class KernelTrace;
class BoundKernel {
    long input_;
public:
    explicit BoundKernel(long input) : input_(input) { }
    template <typename RNGType>
    std::pair<std::unique_ptr<KernelTrace>,double> generate(
            RNGType& rng, Parameters& parameters, const long& observation,
            const GenerateOptions& options) const;
    template <typename RNGType>
    std::pair<std::unique_ptr<KernelTrace>,double> generate(
            RNGType& rng, Parameters& parameters, const sequence::constraints::Empty&,
            const GenerateOptions& options) const;
};
struct Kernel {
    BoundKernel operator()(long input) const { return BoundKernel(input); }
};
constexpr Kernel kernel;

struct KernelTrace {
    long hidden_state_;
    long observation_;
    KernelTrace(long hidden_state, long observation) : hidden_state_(hidden_state), observation_(observation) {}
    [[nodiscard]] long return_value() const {
        return hidden_state_;
    }
};

template <typename RNGType>
std::pair<std::unique_ptr<KernelTrace>,double> BoundKernel::generate(
        RNGType& rng, Parameters& parameters, const long& observation, const GenerateOptions& options) const {
    assert(!options.precompute_gradient());
    std::discrete_distribution<long> dynamics;
    if (input_ == -1) {
        // initial state
        dynamics = std::discrete_distribution<long>(parameters.prior.cbegin(), parameters.prior.cend());
    } else {
        auto column_vector = parameters.transition_matrix(all, input_);
        assert(std::abs(column_vector.sum() - 1.0) < 1e-16);
        dynamics = std::discrete_distribution<long>(column_vector.cbegin(), column_vector.cend());
    }
    long hidden_state = dynamics(rng);
    double log_weight = std::log(parameters.emission_matrix(observation, hidden_state));
    auto trace = std::make_unique<KernelTrace>(hidden_state, observation);
    return {std::move(trace), log_weight};
}

TEST_CASE("sequence hmm", "[sequence, particle filtering]") {
    using ValueType = long;
    using Model = sequence::Model<Kernel,KernelTrace,ValueType,Parameters>;
    using Trace = sequence::Trace<Kernel,KernelTrace,ValueType,Parameters>;
    using std::cerr, std::endl;

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

    // define parameters
    VectorXd prior {{0.2, 0.3, 0.5}};
    MatrixXd emission_matrix {
            {0.1, 0.2, 0.7},
            {0.2, 0.7, 0.1},
            {0.7, 0.2, 0.1}
    };
    emission_matrix.transposeInPlace();
    MatrixXd transition_matrix {
            {0.4, 0.4, 0.2},
            {0.2, 0.3, 0.5},
            {0.9, 0.05, 0.05}
    };
    transition_matrix.transposeInPlace();
    Parameters parameters(prior, emission_matrix, transition_matrix);

    // initial distribution
    ValueType init_value = -1;
    Model model(kernel, init_value);

    // test particle filter
    size_t num_particles = 10000;
    gentl::smc::ParticleSystem<Trace,std::mt19937> filter{num_particles, rng};
    std::vector<long> obs = {0, 0, 1, 2};
    filter.init_step(model, parameters, sequence::constraints::empty);
    auto it = obs.begin();
    while (it != obs.end()) {
        filter.step(sequence::modelchange::extend, sequence::constraints::NewStepOnly<long>(*it++));
        filter.resample();
    }
    double estimate = filter.log_marginal_likelihood_estimate();
    cerr << "actual: " << estimate;
//    << ", expected: " << expected << endl;
//    REQUIRE(std::abs(actual - expected) < 0.02);

}


#include <iostream>
#include <algorithm>
#include <string>

#include <cstdio>
#include <cstdlib>

#include <boost/program_options.hpp>

#include "graphblas/graphblas.hpp"
#include "test/test.hpp"

int main( int argc, char** argv )
{
  std::vector<graphblas::Index> row_indices;
  std::vector<graphblas::Index> col_indices;
  std::vector<float> values;
  graphblas::Index nrows, ncols, nvals;

  // Parse arguments
  namespace po = boost::program_options;
  po::variables_map vm;
  parseArgs(argc, argv, &vm);
  int TA, TB, NT, NUM_ITER, DEVICE;
  bool ROW_MAJOR, DEBUG, SPLIT;
  if( vm.count("ta") )
    TA       = vm["ta"].as<int>(); // default values of TA, TB, NT will be used
  if( vm.count("tb") )
    TB       = vm["tb"].as<int>();
  if( vm.count("nt") )
    NT       = vm["nt"].as<int>();
  if( vm.count("debug") )
    DEBUG    = vm["debug"].as<bool>();
  if( vm.count("split") )
    SPLIT    = vm["split"].as<bool>();
  if( vm.count("iter") )
    NUM_ITER = vm["iter"].as<int>();
  }
  // ROW_MAJOR == 1: means row major
  // ROW_MAJOR == 0: means col major
  // TA == 0 && TB == 0 && NT == 0: means cusparse
  if( vm.count("major") ) {
    std::string major = vm["major"].as<std::string>();
    ROW_MAJOR = (major=="row");
    if( major=="cusparse" ) {
      TA = 0; TB = 0; NT = 0;
    } else if( major=="cusparse2" ) {
      TA = 0; TB = 0; NT = 1;
    }
  }

  if( DEBUG ) {
    std::cout << "ta:    " << TA        << "\n";
    std::cout << "tb:    " << TB        << "\n";
    std::cout << "nt:    " << NT        << "\n";
    std::cout << "row:   " << ROW_MAJOR << "\n";
    std::cout << "debug: " << DEBUG     << "\n";
    std::cout << "split: " << SPLIT     << "\n";
  }

  // Read in sparse matrix
  if (argc < 2) {
    fprintf(stderr, "Usage: %s [matrix-market-filename]\n", argv[0]);
    exit(1);
  } else {
    readTsv( argv[argc-1], row_indices, col_indices, values, nrows, ncols,
    nvals, DEBUG );
  }

  // Matrix A
  graphblas::Matrix<float> a(nrows, ncols);
  a.build( row_indices, col_indices, values, nvals );
  a.nrows( nrows );
  a.ncols( ncols );
  a.nvals( nvals );
  if( DEBUG ) a.print();

  // Matrix B
  graphblas::Matrix<float> b(nrows, ncols);
  b.build( row_indices, col_indices, values, nvals );
  b.nrows( nrows );
  b.ncols( ncols );
  b.nvals( nvals );

  graphblas::Matrix<float> c(nrows, ncols);

  // Warmup
  CpuTimer warmup;
  warmup.Start();
  graphblas::mxm<float, float, float>( c, a, b );
  warmup.Stop();

  CpuTimer cpu_mxm;
  //cudaProfilerStart();
  cpu_mxm.Start();
  for( int i=0; i<NUM_ITER; i++ ) {
    if( SPLIT )
      graphblas::mxmCompute<float, float, float>( c, a, b );
    else
    graphblas::mxm<float, float, float>( c, a, b );
  }
  //cudaProfilerStop();
  cpu_mxm.Stop();

  float flop = 0;
  if( DEBUG ) std::cout << "warmup, " << warmup.ElapsedMillis() << ", " <<
    flop/warmup.ElapsedMillis()/1000000.0 << "\n";
  float elapsed_mxm = gpu_mxm.ElapsedMillis();
  std::cout << "spgemm, " << elapsed_mxm/NUM_ITER << "\n";

  if( DEBUG ) c.print();
  /*c.extractTuples( out_denseVal );
  for( int i=0; i<nvals; i++ ) {
    graphblas::Index row = row_indices[i];
    graphblas::Index col = col_indices[i];
    float            val = values[i];
    if( col<max_ncols ) {
      // Row major order
      if( ROW_MAJOR )
      //std::cout << row << " " << col << " " << val << " " << out_denseVal[row*max_ncols+col] << std::endl;
        BOOST_ASSERT( val==out_denseVal[row*max_ncols+col] );
      else
      // Column major order
      //std::cout << row << " " << col << " " << val << " " << out_denseVal[col*nrows+row] << std::endl;
        BOOST_ASSERT( val==out_denseVal[col*nrows+row] );
    }
  }*/
  return 0;
}


// Copyright 2018 The Simons Foundation, Inc. - All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef NETKET_NDM_SPIN_PHASE_HPP
#define NETKET_NDM_SPIN_PHASE_HPP

#include <Eigen/Dense>
#include <iostream>
#include <vector>
#include "Machine/rbm_spin.hpp"
#include "Utils/all_utils.hpp"
#include "Utils/lookup.hpp"
#include "abstract_density_matrix.hpp"

namespace netket {

/** Neural Density Matrix machine class with spin 1/2 hidden units.
This version has real-valued weights and two NDMs parameterizing phase and
amplitude
 *
 */
class NdmSpinPhase : public AbstractDensityMatrix {
  // number of visible units
  int nv_;

  // number of hidden units
  int nh_;

  // number of ancillary units
  int na_;

  // number of parameters
  int npar_;

  // visible units bias
  RealVectorType b1_;
  RealVectorType b2_;

  // hidden units bias
  RealVectorType h1_;
  RealVectorType h2_;

  // ancillary units bias
  RealVectorType d1_;

  // hidden unit weights
  RealMatrixType W1_;
  RealMatrixType W2_;

  // ancillary unit weights
  RealMatrixType U1_;
  RealMatrixType U2_;

  // Caches
  RealVectorType thetas_r1_;
  RealVectorType thetas_r2_;
  RealVectorType thetas_c1_;
  RealVectorType thetas_c2_;
  RealVectorType lnthetas_r1_;
  RealVectorType lnthetas_r2_;
  RealVectorType lnthetas_c1_;
  RealVectorType lnthetas_c2_;
  RealVectorType thetasnew_r1_;
  RealVectorType thetasnew_r2_;
  RealVectorType thetasnew_c1_;
  RealVectorType thetasnew_c2_;
  RealVectorType lnthetasnew_r1_;
  RealVectorType lnthetasnew_r2_;
  RealVectorType lnthetasnew_c1_;
  RealVectorType lnthetasnew_c2_;

  RealVectorType thetas_a1_;
  RealVectorType thetas_a2_;
  RealVectorType thetasnew_a1_;
  RealVectorType thetasnew_a2_;
  VectorType pi_;
  VectorType lnpi_;
  VectorType lnpinew_;

  bool useb_;
  bool useh_;
  bool used_;

  const Complex I_;

 public:
  explicit NdmSpinPhase(std::shared_ptr<const AbstractHilbert> hilbert,
                        int nhidden = 0, int nancilla = 0, int alpha = 0,
                        int beta = 0, bool useb = true, bool useh = true,
                        bool used = true)
      : AbstractDensityMatrix(hilbert),
        nv_(hilbert->Size()),
        useb_(useb),
        useh_(useh),
        used_(used),
        I_(0, 1) {
    nh_ = std::max(nhidden, alpha * nv_);
    na_ = std::max(nancilla, beta * nv_);
    Init();
  }

  void Init() {
    W1_.resize(nv_, nh_);
    U1_.resize(nv_, na_);
    b1_.resize(nv_);
    h1_.resize(nh_);
    d1_.resize(na_);

    W2_.resize(nv_, nh_);
    U2_.resize(nv_, na_);
    b2_.resize(nv_);
    h2_.resize(nh_);

    thetas_r1_.resize(nh_);
    thetas_r2_.resize(nh_);
    thetas_c1_.resize(nh_);
    thetas_c2_.resize(nh_);

    lnthetas_r1_.resize(nh_);
    lnthetas_r2_.resize(nh_);
    lnthetas_c1_.resize(nh_);
    lnthetas_c2_.resize(nh_);

    thetasnew_r1_.resize(nh_);
    thetasnew_r2_.resize(nh_);
    thetasnew_c1_.resize(nh_);
    thetasnew_c2_.resize(nh_);

    lnthetasnew_r1_.resize(nh_);
    lnthetasnew_r2_.resize(nh_);
    lnthetasnew_c1_.resize(nh_);
    lnthetasnew_c2_.resize(nh_);

    thetas_a1_.resize(na_);
    thetas_a2_.resize(na_);
    thetasnew_a1_.resize(na_);
    thetasnew_a2_.resize(na_);
    pi_.resize(na_);
    lnpi_.resize(na_);
    lnpinew_.resize(na_);

    npar_ = 2 * nv_ * (nh_ + na_);

    if (useb_) {
      npar_ += 2 * nv_;
    } else {
      b1_.setZero();
      b2_.setZero();
    }

    if (useh_) {
      npar_ += 2 * nh_;
    } else {
      h1_.setZero();
      h2_.setZero();
    }

    if (used_) {
      npar_ += na_;
    } else {
      d1_.setZero();
    }

    InfoMessage() << "Phase NDM Initizialized with nvisible = " << nv_
                  << " and nhidden  = " << nh_ << " and nancilla = " << na_
                  << std::endl;
    InfoMessage() << "Using visible   bias = " << useb_ << std::endl;
    InfoMessage() << "Using hidden    bias  = " << useh_ << std::endl;
    InfoMessage() << "Using ancillary bias  = " << used_ << std::endl;
  }

  int Nvisible() const override { return nv_; }

  int Nhidden() const { return nh_; }

  int Nancilla() const { return na_; }

  int Npar() const override { return npar_; }

  void InitRandomPars(int seed, double sigma) override {
    RealVectorType par(npar_);

    netket::RandomGaussian(par, seed, sigma);

    SetParameters(VectorType(par));
  }

  void InitLookup(VisibleConstType v, LookupType &lt) override {
    if (lt.VectorSize() == 0) {
      lt.AddVector(h1_.size());  // row 1
      lt.AddVector(h2_.size());  // row 2
      lt.AddVector(h1_.size());  // col 1
      lt.AddVector(h2_.size());  // col 2
      lt.AddVector(d1_.size());  // ancilla modulus
      lt.AddVector(d1_.size());  // ancilla phase
    }
    if (lt.V(0).size() != h1_.size()) {
      lt.V(0).resize(h1_.size());
    }
    if (lt.V(1).size() != h2_.size()) {
      lt.V(1).resize(h2_.size());
    }
    if (lt.V(2).size() != h1_.size()) {
      lt.V(2).resize(h1_.size());
    }
    if (lt.V(3).size() != h2_.size()) {
      lt.V(3).resize(h2_.size());
    }
    if (lt.V(4).size() != d1_.size()) {
      lt.V(4).resize(d1_.size());
    }
    if (lt.V(5).size() != d1_.size()) {
      lt.V(5).resize(d1_.size());
    }

    VisibleConstType vr = v.head(GetHilbertPhysical().Size());
    VisibleConstType vc = v.tail(GetHilbertPhysical().Size());

    lt.V(0) = (W1_.transpose() * vr + h1_);
    lt.V(1) = (W2_.transpose() * vr + h2_);
    lt.V(2) = (W1_.transpose() * vc + h1_);
    lt.V(3) = (W2_.transpose() * vc + h2_);

    lt.V(4) = (0.5 * U1_.transpose() * (vr + vc) + d1_);
    lt.V(5) = (0.5 * U2_.transpose() * (vr - vc));
  }

  void UpdateLookup(VisibleConstType v, const std::vector<int> &tochange,
                    const std::vector<double> &newconf,
                    LookupType &lt) override {
    VisibleConstType vr = v.head(GetHilbertPhysical().Size());
    VisibleConstType vc = v.tail(GetHilbertPhysical().Size());

    if (tochange.size() != 0) {
      for (std::size_t s = 0; s < tochange.size(); s++) {
        const int sf = tochange[s];
        if (sf < Nvisible()) {
          lt.V(0) += W1_.row(sf) * (newconf[s] - vr(sf));
          lt.V(1) += W2_.row(sf) * (newconf[s] - vr(sf));

          lt.V(4) += 0.5 * U1_.row(sf) * (newconf[s] - vr(sf));
          lt.V(5) += 0.5 * U2_.row(sf) * (newconf[s] - vr(sf));
        } else {
          const int sfc = sf - Nvisible();
          lt.V(2) += W1_.row(sfc) * (newconf[s] - vc(sfc));
          lt.V(3) += W2_.row(sfc) * (newconf[s] - vc(sfc));

          lt.V(4) += 0.5 * U1_.row(sfc) * (newconf[s] - vc(sfc));
          lt.V(5) -= 0.5 * U2_.row(sfc) * (newconf[s] - vc(sfc));
        }
      }
    }
  }

  VectorType DerLog(VisibleConstType v) override {
    LookupType ltnew;
    InitLookup(v, ltnew);
    return DerLog(v, ltnew);
  }

  VectorType DerLog(VisibleConstType v, const LookupType &lt) override {
    VisibleConstType vr = v.head(GetHilbertPhysical().Size());
    VisibleConstType vc = v.tail(GetHilbertPhysical().Size());

    VectorType der(npar_);

    const int impar = (npar_ + na_ * used_) / 2;

    if (useb_) {
      der.head(nv_) = 0.5 * (vr + vc);
      der.segment(impar, nv_) = I_ * 0.5 * (vr - vc);
    }

    RbmSpin::tanh(lt.V(0).real(), lnthetas_r1_);
    RbmSpin::tanh(lt.V(1).real(), lnthetas_r2_);
    RbmSpin::tanh(lt.V(2).real(), lnthetas_c1_);
    RbmSpin::tanh(lt.V(3).real(), lnthetas_c2_);

    if (useh_) {
      der.segment(useb_ * nv_, nh_) = 0.5 * (lnthetas_r1_ + lnthetas_c1_);
      der.segment(impar + useb_ * nv_, nh_) =
          I_ * 0.5 * (lnthetas_r2_ - lnthetas_c2_);
    }

    thetas_a1_ = 0.5 * U1_.transpose() * (vr + vc) + d1_;
    thetas_a2_ = 0.5 * U2_.transpose() * (vr - vc);
    RbmSpin::tanh(lt.V(4).real() + I_ * lt.V(5).real(), lnpi_);

    if (used_) {
      der.segment(useb_ * nv_ + useh_ * nh_, na_) = lnpi_;
    }

    const int initw_1 = nv_ * useb_ + nh_ * useh_ + na_ * used_;
    const int initw_2 = nv_ * useb_ + nh_ * useh_;

    MatrixType wder =
        0.5 * (vr * lnthetas_r1_.transpose() + vc * lnthetas_c1_.transpose());
    der.segment(initw_1, nv_ * nh_) =
        Eigen::Map<VectorType>(wder.data(), nv_ * nh_);

    wder = 0.5 * I_ *
           (vr * lnthetas_r2_.transpose() - vc * lnthetas_c2_.transpose());
    der.segment(impar + initw_2, nv_ * nh_) =
        Eigen::Map<VectorType>(wder.data(), nv_ * nh_);

    const int initu_1 = initw_1 + nv_ * nh_;
    const int initu_2 = initw_2 + nv_ * nh_;

    MatrixType uder = 0.5 * (vr + vc) * lnpi_.transpose();
    der.segment(initu_1, nv_ * na_) =
        Eigen::Map<VectorType>(uder.data(), nv_ * na_);

    uder = 0.5 * I_ * (vr - vc) * lnpi_.transpose();
    der.segment(impar + initu_2, nv_ * na_) =
        Eigen::Map<VectorType>(uder.data(), nv_ * na_);

    return der;
  }

  VectorType GetParameters() override {
    VectorType pars(npar_);

    const int impar = (npar_ + na_ * used_) / 2;

    if (useb_) {
      pars.head(nv_) = b1_;
      pars.segment(impar, nv_) = b2_;
    }

    if (useh_) {
      pars.segment(nv_ * useb_, nh_) = h1_;
      pars.segment(impar + nv_ * useb_, nh_) = h2_;
    }

    if (used_) {
      pars.segment(useb_ * nv_ + useh_ * nh_, na_) = d1_;
    }

    const int initw_1 = nv_ * useb_ + nh_ * useh_ + na_ * used_;
    const int initw_2 = nv_ * useb_ + nh_ * useh_;

    pars.segment(initw_1, nv_ * nh_) =
        Eigen::Map<RealVectorType>(W1_.data(), nv_ * nh_);
    pars.segment(impar + initw_2, nv_ * nh_) =
        Eigen::Map<RealVectorType>(W2_.data(), nv_ * nh_);

    const int initu_1 = initw_1 + nv_ * nh_;
    const int initu_2 = initw_2 + nv_ * nh_;

    pars.segment(initu_1, nv_ * na_) =
        Eigen::Map<RealVectorType>(U1_.data(), nv_ * na_);
    pars.segment(impar + initu_2, nv_ * na_) =
        Eigen::Map<RealVectorType>(U2_.data(), nv_ * na_);

    return pars;
  }

  void SetParameters(VectorConstRefType pars) override {
    const int impar = (npar_ + na_ * used_) / 2;

    if (useb_) {
      b1_ = pars.head(nv_).real();
      b2_ = pars.segment(impar, nv_).real();
    }

    if (useh_) {
      h1_ = pars.segment(useb_ * nv_, nh_).real();
      h2_ = pars.segment(impar + useb_ * nv_, nh_).real();
    }

    if (used_) {
      d1_ = pars.segment(useb_ * nv_ + useh_ * nh_, na_).real();
    }

    const int initw_1 = nv_ * useb_ + nh_ * useh_ + na_ * used_;
    const int initw_2 = nv_ * useb_ + nh_ * useh_;

    VectorType Wpars = pars.segment(initw_1, nv_ * nh_);
    W1_ = Eigen::Map<MatrixType>(Wpars.data(), nv_, nh_).real();

    Wpars = pars.segment(impar + initw_2, nv_ * nh_);
    W2_ = Eigen::Map<MatrixType>(Wpars.data(), nv_, nh_).real();

    const int initu_1 = initw_1 + nv_ * nh_;
    const int initu_2 = initw_2 + nv_ * nh_;

    VectorType Upars = pars.segment(initu_1, nv_ * na_);
    U1_ = Eigen::Map<MatrixType>(Upars.data(), nv_, na_).real();

    Upars = pars.segment(impar + initu_2, nv_ * na_);
    U2_ = Eigen::Map<MatrixType>(Upars.data(), nv_, na_).real();
  }

  // Value of the logarithm of the wave-function
  Complex LogVal(VisibleConstType v) override {
    VisibleConstType vr = v.head(GetHilbertPhysical().Size());
    VisibleConstType vc = v.tail(GetHilbertPhysical().Size());

    RbmSpin::lncosh(W1_.transpose() * vr + h1_, lnthetas_r1_);
    RbmSpin::lncosh(W2_.transpose() * vr + h2_, lnthetas_r2_);
    RbmSpin::lncosh(W1_.transpose() * vc + h1_, lnthetas_c1_);
    RbmSpin::lncosh(W2_.transpose() * vc + h2_, lnthetas_c2_);

    thetas_a1_ = 0.5 * U1_.transpose() * (vr + vc) + d1_;
    thetas_a2_ = 0.5 * U2_.transpose() * (vr - vc);
    RbmSpin::lncosh(thetas_a1_ + I_ * thetas_a2_, lnpi_);

    auto gamma_1 =
        0.5 * (lnthetas_r1_.sum() + lnthetas_c1_.sum() + (vr + vc).dot(b1_));

    auto gamma_2 =
        0.5 * (lnthetas_r2_.sum() - lnthetas_c2_.sum() + (vr - vc).dot(b2_));

    return gamma_1 + I_ * gamma_2 + lnpi_.sum();
  }

  // Value of the logarithm of the wave-function
  // using pre-computed look-up tables for efficiency
  Complex LogVal(VisibleConstType v, const LookupType &lt) override {
    VisibleConstType vr = v.head(GetHilbertPhysical().Size());
    VisibleConstType vc = v.tail(GetHilbertPhysical().Size());

    RbmSpin::lncosh(lt.V(0).real(), lnthetas_r1_);
    RbmSpin::lncosh(lt.V(1).real(), lnthetas_r2_);
    RbmSpin::lncosh(lt.V(2).real(), lnthetas_c1_);
    RbmSpin::lncosh(lt.V(3).real(), lnthetas_c2_);
    RbmSpin::lncosh(lt.V(4).real() + I_ * lt.V(5).real(), lnpi_);

    auto gamma_1 =
        0.5 * (lnthetas_r1_.sum() + lnthetas_c1_.sum() + (vr + vc).dot(b1_));
    auto gamma_2 =
        0.5 * (lnthetas_r2_.sum() - lnthetas_c2_.sum() + (vr - vc).dot(b2_));

    return gamma_1 + I_ * gamma_2 + lnpi_.sum();
  }

  // Difference between logarithms of values, when one or more visible variables
  // are being flipped
  VectorType LogValDiff(
      VisibleConstType v, const std::vector<std::vector<int>> &tochange,
      const std::vector<std::vector<double>> &newconf) override {
    VisibleConstType vr = v.head(GetHilbertPhysical().Size());
    VisibleConstType vc = v.tail(GetHilbertPhysical().Size());

    const std::size_t nconn = tochange.size();

    thetas_r1_ = (W1_.transpose() * vr + h1_);
    thetas_r2_ = (W2_.transpose() * vr + h2_);
    thetas_c1_ = (W1_.transpose() * vc + h1_);
    thetas_c2_ = (W2_.transpose() * vc + h2_);

    RbmSpin::lncosh(thetas_r1_, lnthetas_r1_);
    RbmSpin::lncosh(thetas_r2_, lnthetas_r2_);
    RbmSpin::lncosh(thetas_c1_, lnthetas_c1_);
    RbmSpin::lncosh(thetas_c2_, lnthetas_c2_);

    thetas_a1_ = 0.5 * U1_.transpose() * (vr + vc) + d1_;
    thetas_a2_ = 0.5 * U2_.transpose() * (vr - vc);
    RbmSpin::lncosh(thetas_a1_ + I_ * thetas_a2_, lnpi_);

    Complex logtsum = 0.5 * (lnthetas_r1_.sum() + lnthetas_c1_.sum()) +
                      0.5 * I_ * (lnthetas_r2_.sum() - lnthetas_c2_.sum()) +
                      lnpi_.sum();

    VectorType logvaldiffs = VectorType::Zero(nconn);
    for (std::size_t k = 0; k < nconn; k++) {
      if (tochange[k].size() != 0) {
        thetasnew_r1_ = thetas_r1_;
        thetasnew_r2_ = thetas_r2_;
        thetasnew_c1_ = thetas_c1_;
        thetasnew_c2_ = thetas_c2_;

        thetasnew_a1_ = thetas_a1_;
        thetasnew_a2_ = thetas_a2_;

        for (std::size_t s = 0; s < tochange[k].size(); s++) {
          const int sf = tochange[k][s];

          if (sf < Nvisible()) {
            logvaldiffs(k) += 0.5 * b1_(sf) * (newconf[k][s] - vr(sf));
            logvaldiffs(k) += 0.5 * I_ * b2_(sf) * (newconf[k][s] - vr(sf));

            thetasnew_r1_ += W1_.row(sf) * (newconf[k][s] - vr(sf));
            thetasnew_r2_ += W2_.row(sf) * (newconf[k][s] - vr(sf));
            thetasnew_a1_ += 0.5 * U1_.row(sf) * (newconf[k][s] - vr(sf));
            thetasnew_a2_ += 0.5 * U2_.row(sf) * (newconf[k][s] - vr(sf));
          } else {
            const int sfc = tochange[k][s] - Nvisible();
            logvaldiffs(k) += 0.5 * b1_(sfc) * (newconf[k][s] - vc(sfc));
            logvaldiffs(k) -= 0.5 * I_ * b2_(sfc) * (newconf[k][s] - vc(sfc));

            thetasnew_c1_ += W1_.row(sfc) * (newconf[k][s] - vc(sfc));
            thetasnew_c2_ += W2_.row(sfc) * (newconf[k][s] - vc(sfc));
            thetasnew_a1_ += 0.5 * U1_.row(sfc) * (newconf[k][s] - vc(sfc));
            thetasnew_a2_ -= 0.5 * U2_.row(sfc) * (newconf[k][s] - vc(sfc));
          }
        }
        RbmSpin::lncosh(thetasnew_r1_, lnthetasnew_r1_);
        RbmSpin::lncosh(thetasnew_r2_, lnthetasnew_r2_);
        RbmSpin::lncosh(thetasnew_c1_, lnthetasnew_c1_);
        RbmSpin::lncosh(thetasnew_c2_, lnthetasnew_c2_);
        RbmSpin::lncosh(thetasnew_a1_ + I_ * thetasnew_a2_, lnpinew_);

        logvaldiffs(k) +=
            0.5 * (lnthetasnew_r1_.sum() + lnthetasnew_c1_.sum()) +
            0.5 * I_ * (lnthetasnew_r2_.sum() - lnthetasnew_c2_.sum()) +
            lnpinew_.sum() - logtsum;
      }
    }
    return logvaldiffs;
  }

  // Difference between logarithms of values, when one or more visible variables
  // are being flipped Version using pre-computed look-up tables for efficiency
  // on a small number of spin flips
  Complex LogValDiff(VisibleConstType v, const std::vector<int> &tochange,
                     const std::vector<double> &newconf,
                     const LookupType &lt) override {
    VisibleConstType vr = v.head(GetHilbertPhysical().Size());
    VisibleConstType vc = v.tail(GetHilbertPhysical().Size());

    Complex logvaldiff = 0.;

    if (tochange.size() != 0) {
      RbmSpin::lncosh(lt.V(0).real(), lnthetas_r1_);
      RbmSpin::lncosh(lt.V(1).real(), lnthetas_r2_);
      RbmSpin::lncosh(lt.V(2).real(), lnthetas_c1_);
      RbmSpin::lncosh(lt.V(3).real(), lnthetas_c2_);
      RbmSpin::lncosh(lt.V(4).real() + I_ * lt.V(5).real(), lnpi_);

      thetasnew_r1_ = lt.V(0).real();
      thetasnew_r2_ = lt.V(1).real();
      thetasnew_c1_ = lt.V(2).real();
      thetasnew_c2_ = lt.V(3).real();
      thetasnew_a1_ = lt.V(4).real();
      thetasnew_a2_ = lt.V(5).real();

      for (std::size_t s = 0; s < tochange.size(); s++) {
        const int sf = tochange[s];

        if (sf < Nvisible()) {
          logvaldiff += 0.5 * b1_(sf) * (newconf[s] - vr(sf));
          logvaldiff += 0.5 * I_ * b2_(sf) * (newconf[s] - vr(sf));

          thetasnew_r1_ += W1_.row(sf) * (newconf[s] - vr(sf));
          thetasnew_r2_ += W2_.row(sf) * (newconf[s] - vr(sf));
          thetasnew_a1_ += 0.5 * U1_.row(sf) * (newconf[s] - vr(sf));
          thetasnew_a2_ += 0.5 * U2_.row(sf) * (newconf[s] - vr(sf));
        } else {
          const int sfc = tochange[s] - Nvisible();
          logvaldiff += 0.5 * b1_(sfc) * (newconf[s] - vc(sfc));
          logvaldiff -= 0.5 * I_ * b2_(sfc) * (newconf[s] - vc(sfc));

          thetasnew_c1_ += W1_.row(sfc) * (newconf[s] - vc(sfc));
          thetasnew_c2_ += W2_.row(sfc) * (newconf[s] - vc(sfc));
          thetasnew_a1_ += 0.5 * U1_.row(sfc) * (newconf[s] - vc(sfc));
          thetasnew_a2_ -= 0.5 * U2_.row(sfc) * (newconf[s] - vc(sfc));
        }
      }

      RbmSpin::lncosh(thetasnew_r1_, lnthetasnew_r1_);
      RbmSpin::lncosh(thetasnew_r2_, lnthetasnew_r2_);
      RbmSpin::lncosh(thetasnew_c1_, lnthetasnew_c1_);
      RbmSpin::lncosh(thetasnew_c2_, lnthetasnew_c2_);
      RbmSpin::lncosh(thetasnew_a1_ + I_ * thetasnew_a2_, lnpinew_);

      logvaldiff += 0.5 * (lnthetasnew_r1_.sum() + lnthetasnew_c1_.sum());
      logvaldiff -= 0.5 * (lnthetas_r1_.sum() + lnthetas_c1_.sum());
      logvaldiff += 0.5 * I_ * (lnthetasnew_r2_.sum() - lnthetasnew_c2_.sum());
      logvaldiff -= 0.5 * I_ * (lnthetas_r2_.sum() - lnthetas_c2_.sum());
      logvaldiff += (lnpinew_.sum() - lnpi_.sum());
    }
    return logvaldiff;
  }

  inline static double lncosh(double x) {
    const double xp = std::abs(x);
    if (xp <= 12.) {
      return std::log(std::cosh(xp));
    } else {
      const static double log2v = std::log(2.);
      return xp - log2v;
    }
  }

  void Save(const std::string &filename) const override {
    json state;
    state["Name"] = "NdmSpinPhase";
    state["Nvisible"] = nv_;
    state["Nhidden"] = nh_;
    state["Nancilla"] = na_;
    state["UseVisibleBias"] = useb_;
    state["UseHiddenBias"] = useh_;
    state["UseAncillaBias"] = used_;
    state["b1"] = b1_;
    state["h1"] = h1_;
    state["d1"] = d1_;
    state["W1"] = W1_;
    state["U1"] = U1_;

    state["b2"] = b2_;
    state["h2"] = h2_;
    state["W2"] = W2_;
    state["U2"] = U2_;
    WriteJsonToFile(state, filename);
  }

  void Load(const std::string &filename) override {
    auto pars = ReadJsonFromFile(filename);
    std::string name = FieldVal<std::string>(pars, "Name");
    if (name != "NdmSpinPhase") {
      throw InvalidInputError(
          "Error while constructing RbmSpinPhase from input parameters");
    }

    if (FieldExists(pars, "Nvisible")) {
      nv_ = FieldVal<int>(pars, "Nvisible");
    }
    if (nv_ != GetHilbertPhysical().Size()) {
      throw InvalidInputError(
          "Number of visible units is incompatible with given "
          "Hilbert space");
    }

    if (FieldExists(pars, "Nhidden")) {
      nh_ = FieldVal<int>(pars, "Nhidden");
    } else {
      nh_ = nv_ * double(FieldVal<double>(pars, "Alpha"));
    }

    if (FieldExists(pars, "Nancilla")) {
      na_ = FieldVal<int>(pars, "Nancilla");
    } else {
      na_ = nv_ * double(FieldVal<double>(pars, "Beta"));
    }

    useb_ = FieldOrDefaultVal(pars, "UseVisibleBias", true);
    useh_ = FieldOrDefaultVal(pars, "UseHiddenBias", true);
    used_ = FieldOrDefaultVal(pars, "UseAncillaBias", true);

    Init();

    // Loading parameters, if defined in the input
    if (FieldExists(pars, "b1")) {
      b1_ = FieldVal<RealVectorType>(pars, "b1");
      b2_ = FieldVal<RealVectorType>(pars, "b2");
    } else {
      b1_.setZero();
      b2_.setZero();
    }

    if (FieldExists(pars, "h1")) {
      h1_ = FieldVal<RealVectorType>(pars, "h1");
      h2_ = FieldVal<RealVectorType>(pars, "h2");
    } else {
      h1_.setZero();
      h2_.setZero();
    }

    if (FieldExists(pars, "d1")) {
      d1_ = FieldVal<RealVectorType>(pars, "d1");
    } else {
      d1_.setZero();
    }

    if (FieldExists(pars, "W1")) {
      W1_ = FieldVal<RealMatrixType>(pars, "W1");
      W2_ = FieldVal<RealMatrixType>(pars, "W2");
    }

    if (FieldExists(pars, "U1")) {
      U1_ = FieldVal<RealMatrixType>(pars, "U1");
      U2_ = FieldVal<RealMatrixType>(pars, "U2");
    }
  }

  bool IsHolomorphic() const noexcept override { return false; }
};

}  // namespace netket

#endif  // NETKET_NDM_SPIN_PHASE_HPP


#include <iostream>
#include <map>
#include <vector>
#include <utility>
#include <tuple>
#include <regex>

#include <boost/multiprecision/cpp_int.hpp>

using boost::multiprecision::uint128_t;

class Part2 {
public:
    Part2(int max_score);
    std::pair<uint128_t, uint128_t> solve(int p1, int p2, int s1, int s2);

private:
    int max_;

    std::vector<std::tuple<int, int, int>> throws;
    std::map<std::tuple<int, int, int, int>, std::pair<uint128_t, uint128_t>> cache_;
};

Part2::Part2(int max_score) : max_(max_score) {
    for (int d1 = 1; d1 <= 3; ++d1) {
        for (int d2 = 1; d2 <= 3; ++d2) {
            for (int d3 = 1; d3 <= 3; ++d3) {
                throws.emplace_back(d1,d2,d3);
            }
        }
    }
}

std::pair<uint128_t, uint128_t> Part2::solve(int p1, int p2, int s1, int s2) {
    if (s1 >= max_)
        return std::make_pair(1, 0);
    if (s2 >= max_)
        return std::make_pair(0, 1);

    auto q = std::make_tuple(p1, p2, s1, s2);
    auto a = cache_.find(q);
    if (a != cache_.end()) {
        return a->second;
    }
    auto score = std::make_pair<uint128_t,uint128_t>(0, 0);
    for (auto& a : throws) {
        auto [d1, d2, d3] = a;
        int new_p1 = p1 + d1 + d2 + d3;
        new_p1 %= 10;
        if (new_p1 == 0) new_p1 = 10;
        int new_s1 = s1 + new_p1;
        auto [p2_wins, p1_wins] = solve(p2, new_p1, s2, new_s1);
        score.first += p1_wins;
        score.second += p2_wins;
    }

    cache_[q] = score;
    return score;
}

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

    int player_1_start = 0;
    int player_2_start = 0;

    for (int i = 0; i < 2; ++i) {
        std::string input;
        getline(std::cin, input);
        const std::regex re(R"(Player \d starting position: (\d*))");
        std::smatch match;
        if (regex_search(input, match, re)) {
            if (i == 0)
                player_1_start = stoi(match.str(1));
            else
                player_2_start = stoi(match.str(1));
        }
    }

    int player_1 = 0;
    int player_2 = 0;

    int p1 = player_1_start;
    int p2 = player_2_start;

    int dice = 1;
    while (player_1 < 1000 && player_2 < 1000) {
        p1 += 3*dice + 3;
        dice += 3;
        p1 %= 10;
        if (p1 == 0) p1 = 10;
        player_1 += p1;
        if (player_1 >= 1000)
            break;
        p2 += 3*dice + 3;
        dice += 3;
        p2 %= 10;
        if (p2 == 0) p2 = 10;
        player_2 += p2;
    }

    int part_1 = std::min(player_1, player_2) * (dice-1);
    std::cout << part_1 << std::endl;

    Part2 part2(21);
    auto a = part2.solve(player_1_start, player_2_start, 0, 0);
    std::cout << std::max(a.first, a.second) << std::endl;

    return 0;
}