#include <cstdio>
#include <cmath>
#include <cfloat>
#include <cstdarg>
#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <map>
#include <Eigen/Dense>
#include <unsupported/Eigen/CXX11/Tensor>
#include <boost/filesystem.hpp>
#include <omp.h>
#include <H5Cpp.h>

#include "timer/timer.h"

/** \brief Read all files in a directory matching the given extension.
 * \param[in] directory path to directory
 * \param[out] files read file paths
 * \param[in] extension extension to filter for
 */
void read_directory(const boost::filesystem::path directory, std::map<int, boost::filesystem::path>& files, const std::string extension = ".off") {
  files.clear();
  boost::filesystem::directory_iterator end;

  for (boost::filesystem::directory_iterator it(directory); it != end; ++it) {
    if (it->path().extension().string() == extension) {
      if (!boost::filesystem::is_empty(it->path()) && !it->path().empty() && it->path().filename().string() != "") {
        int number = std::stoi(it->path().filename().string());
        files.insert(std::pair<int, boost::filesystem::path>(number, it->path()));
      }
    }
  }
}

/** \brief Just encapsulating vertices and faces. */
class Mesh {
public:
  /** \brief Empty constructor. */
  Mesh() {

  }

  /** \brief Add a vertex.
   * \param[in] vertex vertex to add
   */
  void add_vertex(Eigen::Vector3f& vertex) {
    this->vertices.push_back(vertex);
  }

  /** \brief Get the number of vertices.
   * \return number of vertices
   */
  int num_vertices() const {
    return static_cast<int>(this->vertices.size());
  }

  Eigen::Vector3f vertex(int v) {
    assert(v >= 0 && v < this->vertices.size());
    return this->vertices[v];
  }

  /** \brief Add a face.
   * \param[in] face face to add
   */
  void add_face(Eigen::Vector3i& face) {
    this->faces.push_back(face);
  }

  /** \brief Get the number of faces.
   * \return number of faces
   */
  int num_faces() const {
    return static_cast<int>(this->faces.size());
  }

  Eigen::Vector3i face(int f) {
    assert(f >= 0 && f < this->num_faces());
    return this->faces[f];
  }

  /** \brief Rotate the point cloud around the origin.
   * \param[in] rotation rotation matrix
   */
  void rotate(const Eigen::Matrix3f &rotation) {
    for (int v = 0; v < this->num_vertices(); ++v) {
      this->vertices[v] = rotation*this->vertices[v];
    }
  }

  /** \brief Translate the mesh.
   * \param[in] translation translation vector
   */
  void translate(const Eigen::Vector3f& translation) {
    for (int v = 0; v < this->num_vertices(); ++v) {
      for (int i = 0; i < 3; ++i) {
        this->vertices[v](i) += translation(i);
      }
    }
  }

  /** \brief Scale the mesh.
   * \param[in] scale scale vector
   */
  void scale(const Eigen::Vector3f& scale) {
    for (int v = 0; v < this->num_vertices(); ++v) {
      for (int i = 0; i < 3; ++i) {
        this->vertices[v](i) *= scale(i);
      }
    }
  }

  /** \brief Sample points from the mesh
   * \param[in] mesh mesh to sample from
   * \param[in] n batch index in points
   * \param[in] points pre-initialized tensor holding points
   */
  bool sample(const int N, const int k, Eigen::Tensor<float, 3, Eigen::RowMajor> &points) {

    // Stores the areas of faces.
    std::vector<float> areas(this->num_faces());
    float sum = 0;

    // Build a probability distribution over faces.
    for (int f = 0; f < this->num_faces(); f++) {
      Eigen::Vector3f a = this->vertices[this->faces[f][0]];
      Eigen::Vector3f b = this->vertices[this->faces[f][1]];
      Eigen::Vector3f c = this->vertices[this->faces[f][2]];

      // Angle between a->b and a->c.
      Eigen::Vector3f ab = b - a;
      Eigen::Vector3f ac = c - a;
      float cos_angle = ab.dot(ac)/(ab.norm()*ac.norm());
      float angle = std::acos(cos_angle);

      // Compute triangle area.
      float area = std::max(0., 0.5*ab.norm()*ac.norm()*std::sin(angle));
      //std::cout << area << " " << std::pow(area, 1./4.) << " " << angle << " " << ab.norm() << " " << ac.norm() << " " << std::sin(angle) << std::endl;

      // Accumulate.
      //area = std::sqrt(area);
      areas[f] = area;
      sum += area;
      //areas.push_back(1);
      //sum += 1;
    }

    //std::cout << sum << std::endl;
    if (sum < 1e-6) {
      std::cout << "[Error] face area sum of " << sum << std::endl;
      return false;
    }

    for (int f = 0; f < this->num_faces(); f++) {
      //std::cout << areas[f] << " ";
      areas[f] /= sum;
      //std::cout << areas[f] << std::endl;
    }

    std::vector<float> cum_areas(areas.size());
    cum_areas[0] = areas[0];

    for (int f = 1; f < this->num_faces(); f++) {
      cum_areas[f] = areas[f] + cum_areas[f - 1];
    }

    int j = 0;
    for (int f = 0; f < this->num_faces(); f++) {
      int n = std::max(static_cast<int>(areas[f]*N), 1);

      for (int i = 0; i < n; i++) {
        float r1 = 0;
        float r2 = 0;
        do {
          r1 = static_cast<float>(std::rand())/static_cast<float>(RAND_MAX);
          r2 = static_cast<float>(std::rand())/static_cast<float>(RAND_MAX);
        }
        while (r1 + r2 > 1.f);

        int s = std::rand()%3;
        //std::cout << face << " " << areas[face] << std::endl;

        Eigen::Vector3f a = this->vertices[this->faces[f](s)];
        Eigen::Vector3f b = this->vertices[this->faces[f]((s + 1)%3)];
        Eigen::Vector3f c = this->vertices[this->faces[f]((s + 2)%3)];

        Eigen::Vector3f ab = b - a;
        Eigen::Vector3f ac = c - a;

        Eigen::Vector3f point = a + r1*ab + r2*ac;
        points(k, j, 0) = point(0);
        points(k, j, 1) = point(1);
        points(k, j, 2) = point(2);
        j++;
      }
    }

    return true;
  }

  /** \brief Reading an off file and returning the vertices x, y, z coordinates and the
   * face indices.
   * \param[in] filepath path to the OFF file
   * \param[out] mesh read mesh with vertices and faces
   * \return success
   */
  static bool from_off(const std::string filepath, Mesh& mesh) {

    std::ifstream* file = new std::ifstream(filepath.c_str());
    std::string line;
    std::stringstream ss;
    int line_nb = 0;

    std::getline(*file, line);
    ++line_nb;

    if (line != "off" && line != "OFF") {
      std::cout << "[Error] Invalid header: \"" << line << "\", " << filepath << std::endl;
      return false;
    }

    size_t n_edges;
    std::getline(*file, line);
    ++line_nb;

    int n_vertices;
    int n_faces;
    ss << line;
    ss >> n_vertices;
    ss >> n_faces;
    ss >> n_edges;

    for (size_t v = 0; v < n_vertices; ++v) {
      std::getline(*file, line);
      ++line_nb;

      ss.clear();
      ss.str("");

      Eigen::Vector3f vertex;
      ss << line;
      ss >> vertex(0);
      ss >> vertex(1);
      ss >> vertex(2);

      mesh.add_vertex(vertex);
    }

    size_t n;
    for (size_t f = 0; f < n_faces; ++f) {
      std::getline(*file, line);
      ++line_nb;

      ss.clear();
      ss.str("");

      size_t n;
      ss << line;
      ss >> n;

      if(n != 3) {
        std::cout << "[Error] Not a triangle (" << n << " points) at " << (line_nb - 1) << std::endl;
        return false;
      }

      Eigen::Vector3i face;
      ss >> face(0);
      ss >> face(1);
      ss >> face(2);

      mesh.add_face(face);
    }

    if (n_vertices != mesh.num_vertices()) {
      std::cout << "[Error] Number of vertices in header differs from actual number of vertices." << std::endl;
      return false;
    }

    if (n_faces != mesh.num_faces()) {
      std::cout << "[Error] Number of faces in header differs from actual number of faces." << std::endl;
      return false;
    }

    file->close();
    delete file;

    return true;
  }

  /** \brief Write mesh to OFF file.
   * \param[in] filepath path to OFF file to write
   * \return success
   */
  bool to_off(const std::string filepath) {
    std::ofstream* out = new std::ofstream(filepath, std::ofstream::out);
    if (!static_cast<bool>(*out)) {
      return false;
    }

    (*out) << "OFF" << std::endl;
    (*out) << this->vertices.size() << " " << this->num_faces() << " 0" << std::endl;

    for (unsigned int v = 0; v < this->vertices.size(); v++) {
      (*out) << this->vertices[v](0) << " " << this->vertices[v](1) << " " << this->vertices[v](2) << std::endl;
    }

    for (unsigned int f = 0; f < this->num_faces(); f++) {
      (*out) << "3 " << this->faces[f](0) << " " << this->faces[f](1) << " " << this->faces[f](2) << std::endl;
    }

    out->close();
    delete out;

    return true;
  }

  /** \brief Write mesh to obj file.
   * \param[in] filepath
   * \param[in] mtl_lib
   * \param[in] materials
   * \return success
   */
  bool to_obj(const std::string filepath) {

    std::ofstream* out = new std::ofstream(filepath, std::ofstream::out);
    if (!static_cast<bool>(out)) {
      return false;
    }

    for (unsigned int v = 0; v < this->vertices.size(); v++) {
      (*out) << "v " << this->vertices[v](0) << " " << this->vertices[v](1) << " " << this->vertices[v](2) << std::endl;
    }

    for (unsigned int f = 0; f < this->num_faces(); f++) {
      (*out) << "f " << this->faces[f](0) + 1 << " " << this->faces[f](1) + 1 << " " << this->faces[f](2) + 1 << std::endl;
    }

    out->close();
    delete out;

    return true;
  }

private:

  /** \brief Vertices as (x,y,z)-vectors. */
  std::vector<Eigen::Vector3f> vertices;

  /** \brief Faces as list of vertex indices. */
  std::vector<Eigen::Vector3i> faces;
};

/** \brief Write the given set of volumes to h5 file.
 * \param[in] filepath h5 file to write
 * \param[in] n number of volumes
 * \param[in] height height of volumes
 * \param[in] width width of volumes
 * \param[in] depth depth of volumes
 * \param[in] dense volume data
 */
bool write_hdf5(const std::string filepath, Eigen::Tensor<float, 3, Eigen::RowMajor>& dense) {

  try {

    /*
     * Turn off the auto-printing when failure occurs so that we can
     * handle the errors appropriately
     */
    H5::Exception::dontPrint();

    /*
     * Create a new file using H5F_ACC_TRUNC access,
     * default file creation properties, and default file
     * access properties.
     */
    H5::H5File file(filepath, H5F_ACC_TRUNC);

    /*
     * Define the size of the array and create the data space for fixed
     * size dataset.
     */
    hsize_t rank = 3;
    hsize_t dimsf[rank];
    dimsf[0] = dense.dimension(0);
    dimsf[1] = dense.dimension(1);
    dimsf[2] = dense.dimension(2);
    H5::DataSpace dataspace(rank, dimsf);

    /*
     * Define datatype for the data in the file.
     * We will store little endian INT numbers.
     */
    H5::IntType datatype(H5::PredType::NATIVE_FLOAT);
    datatype.setOrder(H5T_ORDER_LE);

    /*
     * Create a new dataset within the file using defined dataspace and
     * datatype and default dataset creation properties.
     */
    H5::DataSet dataset = file.createDataSet("tensor", datatype, dataspace);

    /*
     * Write the data to the dataset using default memory space, file
     * space, and transfer properties.
     */
    float* data = static_cast<float*>(dense.data());
    dataset.write(data, H5::PredType::NATIVE_FLOAT);
  }  // end of try block

  // catch failure caused by the H5File operations
  catch(H5::FileIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSet operations
  catch(H5::DataSetIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSpace operations
  catch(H5::DataSpaceIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSpace operations
  catch(H5::DataTypeIException error) {
    error.printError();
    return false;
  }

  return true;
}

/** \brief Main entrance point of script.
 * Expects one argument, the path to the config file.
 */
int main(int argc, char** argv) {
  if (argc < 2) {
    std::cout << "[Error] Usage: sample off_directory h5_file" << std::endl;
    exit(1);
  }

  boost::filesystem::path off_directory(argv[1]);
  boost::filesystem::path h5_file(argv[2]);

  if (!boost::filesystem::is_directory(off_directory)) {
    std::cout << "[Error] directory " << off_directory << " not found" << std::endl;
    exit(1);
  }

  std::map<int, boost::filesystem::path> off_files;
  read_directory(off_directory, off_files);
  std::cout << "[ICP] found " << off_files.size() << " files" << std::endl;

  std::vector<int> indices;
  for (std::map<int, boost::filesystem::path>::iterator it = off_files.begin(); it != off_files.end(); it++) {
    indices.push_back(it->first);
  }

  int N_points = 1000000;
  Eigen::Tensor<float, 3, Eigen::RowMajor> points(indices.size(), N_points, 3);
  points.setZero();

  float total = 0;

  omp_set_num_threads(16);
  #pragma omp parallel
  {
    #pragma omp for
    for (unsigned int i = 0; i < indices.size(); i++) {
      int n = indices[i];
      std::string off_file = off_files[n].string();

      Mesh mesh;
      Mesh::from_off(off_file, mesh);

      Timer timer;
      timer.start();
      mesh.sample(N_points, i, points);
      timer.stop();

      float elapsed = timer.getElapsedTimeInMilliSec();
      std::cout << "[Sample] sampled " << off_file << " (" << elapsed << "ms)" << std::endl;

      #pragma omp critical
      {
        total += elapsed;
      }
    }
  }

  write_hdf5(h5_file.string(), points);
  std::cout << "[Sample] wrote " << h5_file.string() << std::endl;
  std::cout << "[Sample] took  on average " << total/indices.size() << "ms" << std::endl;
  exit(0);
}


// Copyright (c) 2020 Kent Hu

/// @author kent hu

#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <exception>
#include <deque>
#include <functional>
#include <iostream>
#include <limits>
#include <random>
#include <string>
#include <thread>
#include <utility>
#include <vector>

#include <boost/lexical_cast.hpp>
#include <boost/optional/optional.hpp>
#include <comma/application/command_line_options.h>
#include <comma/base/exception.h>
#include <comma/base/types.h>
#include <comma/csv/ascii.h>
#include <comma/csv/options.h>
#include <comma/csv/stream.h>
#include <comma/string/split.h>
#include <tbb/blocked_range.h>

#define TBB_PREVIEW_GLOBAL_CONTROL 1 // required to #include <tbb/global_control.h>

#include <tbb/global_control.h>
#include <tbb/parallel_for.h>
#include <tbb/parallel_reduce.h>
#include "../../visiting/traits.h"

#ifdef SNARK_USE_CUDA

#include "math-k-means/device.h"
#include "math-k-means/math_k_means.h"

#endif

void usage( const bool verbose )
{
    std::cerr << std::endl;
    std::cerr << "run k-means on input data" << std::endl;
    std::cerr << std::endl;
    std::cerr << "usage: cat sample.csv | math-k-means [<options>]" << std::endl;
    std::cerr << std::endl;
    std::cerr << "fields" << std::endl;
    std::cerr << "    block: block number; output k-means centroid id for each" << std::endl;
    std::cerr << "           contiguous block of samples with the same block id" << std::endl;
    std::cerr << "    data: vector of size <size>" << std::endl;
    std::cerr << "    default: data" << std::endl;
    std::cerr << "output" << std::endl;
    std::cerr << "    appended fields: centroid/id,centroid/data" << std::endl;
    std::cerr << "    centroids only: if block field present: block,centroid/id,centroid/data" << std::endl;
    std::cerr << "                    if no block field present: centroid/id,centroid/data" << std::endl;
    std::cerr << "    binary format: 64-bit floating-point for centroid/data, 32-bit unsigned integer for centroid/id" << std::endl;
    std::cerr << "options" << std::endl;
#ifdef SNARK_USE_CUDA
    std::cerr << "    --cuda: use gpu; ATTENTION: only 32-bit float precision currently implemented; if you need 64-bit float precision, run without --cuda" << std::endl;
    std::cerr << "    --cuda-use-pitched,--use-pitched,--pitched: use pitched memory for 2d array (padded 2d array for memory coalescing; faster but uses more memory)" << std::endl;
#endif
    std::cerr << "    --help,-h: show this help; --help --verbose: more help" << std::endl;
    std::cerr << "    --ignore-tolerance: ignore tolerance value used for early exit" << std::endl;
    std::cerr << "    --max-iterations,--iterations=<n>: number of iterations for Lloyd's algorithm; default: 300" << std::endl;
    std::cerr << "    --max-threads,--threads=<n>: maximum number of threads to run, if 0, set to number of cores in system (" << std::thread::hardware_concurrency() << ')' << "; default: 0" << std::endl;
    std::cerr << "    --number-of-clusters,--clusters=<n>: number of k-means cluster per block/all" << std::endl;
    std::cerr << "    --number-of-runs,--runs=<n>: number of times to run k-means, best run is one with lowest inertia (sum of distance of cluster points to cluster centroid); default: 10" << std::endl;
    std::cerr << "    --output-centroids,--centroids: output centroids only" << std::endl;
    std::cerr << "    --seed=[<seed>]: use seed for random initialization of centroids" << std::endl;
    std::cerr << "    --size=[<n>]: a hint of number of elements in the data vector; ignored, if data indices specified, e.g. data[0],data[1],data[2]" << std::endl;
    std::cerr << "    --tolerance=<distance>: difference between two consecutive iteration centroid l2 norms to declare convergence and stop iterating in LLoyd's algorithm; default: 1.0e-4" << std::endl;
    if( verbose ) { std::cerr << std::endl << "csv options" << std::endl << comma::csv::options::usage() << std::endl; }
    std::cerr << std::endl;
    std::cerr << "examples" << std::endl;
    std::cerr << "    cat sample.csv | math-k-means --number-of-clusters=3" << std::endl;
    std::cerr << "    cat sample.csv | math-k-means --number-of-clusters=3 --fields=data --size=3" << std::endl;
    std::cerr << "    cat sample.csv | math-k-means --number-of-clusters=3 --fields=,,,data[0],,,,data[1],,data[2]" << std::endl;
    std::cerr << "    cat sample.bin | math-k-means --number-of-clusters=3 --fields=block,data --size=3 --binary=ui,3d" << std::endl;
    std::cerr << std::endl;
    exit( 0 );
}

static comma::csv::options csv;
static boost::optional< unsigned int > size = boost::none;
static bool output_centroids = false;
static boost::optional< std::mt19937::result_type > seed = boost::none;
static bool use_block = false;
#ifdef SNARK_USE_CUDA
static bool use_pitched = false;
#endif
static bool verbose = false;

namespace snark {

#ifdef SNARK_USE_CUDA
namespace cuda { namespace k_means {

struct input_t
{
    std::vector< float > data;
    comma::uint32 block = 0;

    input_t() : data( *size ), block( 0 ) {}
};

class k_means
{
public:
    k_means( float tolerance, unsigned int max_iterations, unsigned int number_of_runs,
             comma::uint32 number_of_clusters, unsigned int size ) noexcept : tolerance_( tolerance ),
                                                                              max_iterations_( max_iterations ),
                                                                              number_of_runs_( number_of_runs ),
                                                                              number_of_clusters_( number_of_clusters ),
                                                                              ncols_( size ) {}

    template < template < typename > class Matrix >
    std::tuple< std::vector< float >, std::vector< comma::uint32 > >
    run_on_block( const std::vector< float >& h_matrix ) const
    {
        const size_t nrows = h_matrix.size() / ncols_;

        // device arrays/matrices
        Matrix< float > d_matrix( nrows, ncols_, h_matrix );
        Matrix< float > d_means( number_of_clusters_, ncols_ );
        Matrix< float > d_sums( number_of_clusters_, ncols_ );
        CUDA_CHECK_ERRORS( d_sums.fill( 0.0f ) )

        device::array< comma::uint32 > d_assignments( nrows );
        device::array< unsigned int > d_counts( number_of_clusters_ );
        CUDA_CHECK_ERRORS( d_counts.fill( 0 ) )

        // host vectors
        std::vector< float > h_means( number_of_clusters_* ncols_ );
        std::vector< comma::uint32 > h_assignments( nrows );

        // save runs to choose best centroid
        std::vector< float > all_scores;
        all_scores.reserve( number_of_runs_ );
        std::vector< decltype( h_means ) > all_centroids;
        all_centroids.reserve( number_of_runs_ );
        std::vector< decltype( h_assignments ) > all_centroid_assignments;
        all_centroid_assignments.reserve( number_of_runs_ );

        cudaDeviceProp prop{};
        CUDA_CHECK_ERRORS( cudaGetDeviceProperties( &prop, 0 ) )
        for( unsigned int run = 0; run < number_of_runs_; ++run )
        {
            CUDA_CHECK_ERRORS( d_means.to_device( initialize_centroids( h_matrix ) ) )
            float difference = std::numeric_limits< float >::max();
            for( unsigned int iteration = 0; iteration < max_iterations_; ++iteration )
            {
                assign_centroids( d_matrix, d_means, d_assignments, d_sums, d_counts, number_of_clusters_, prop.maxThreadsPerBlock );
                update_centroids( d_means, d_sums, d_counts, number_of_clusters_ );
                if( tolerance_ < 0 ) { continue; }
                std::vector< float > new_means( number_of_clusters_* ncols_ ); // temp buffer for calculating tolerance difference
                CUDA_CHECK_ERRORS( d_means.to_host( new_means ) )
                difference = tbb::parallel_reduce( tbb::blocked_range< comma::uint32 >( 0, number_of_clusters_ ), 0.0f,
                                                   [&]( const tbb::blocked_range< comma::uint32 > chunk, float difference ) -> float
                                                   {
                                                       for( comma::uint32 i = chunk.begin(); i < chunk.end(); ++i )
                                                       {
                                                           difference += squared_euclidean_distance( ncols_, &h_means[i * ncols_], &new_means[i * ncols_] );
                                                       }
                                                       return difference;
                                                   },
                                                   std::plus< float >()
                );
                h_means = std::move( new_means );
                if( difference < tolerance_ * tolerance_ ) { break; }
            }
            CUDA_CHECK_ERRORS( d_assignments.to_host( h_assignments ) )
            if( tolerance_ < 0 || difference > tolerance_ * tolerance_ ) { CUDA_CHECK_ERRORS( d_means.to_host( h_means ) ) }
            const float score = tbb::parallel_reduce( tbb::blocked_range< size_t >( 0, nrows ), 0.0f,
                                                      [&]( const tbb::blocked_range< size_t > chunk, float score ) -> float
                                                      {
                                                          for( size_t point = chunk.begin(); point < chunk.end(); ++point )
                                                          {
                                                              const auto cluster_assignment = h_assignments[point];
                                                              score += std::sqrt( squared_euclidean_distance( ncols_, &h_matrix[point * ncols_], &h_means[cluster_assignment * ncols_] ) );
                                                          }
                                                          return score;
                                                      },
                                                      std::plus< float >()
            );
            all_scores.emplace_back( score );
            all_centroids.emplace_back( h_means );
            all_centroid_assignments.emplace_back( h_assignments );
        }
        const auto& best_index = std::distance( all_scores.begin(), std::min_element( all_scores.begin(), all_scores.end() ) );
        return std::make_tuple( all_centroids[best_index], all_centroid_assignments[best_index] );
    }

private:
    std::vector< float > initialize_centroids( const std::vector< float >& h_matrix ) const
    {
        static std::mt19937 generator( seed.is_initialized() ? *seed : std::random_device{}() );
        static std::vector< float > init_centroids( number_of_clusters_* ncols_ );
        static std::uniform_int_distribution< size_t > indices( 0, h_matrix.size() / ncols_ - 1 );
        for( size_t cluster = 0; cluster < number_of_clusters_; ++cluster )
        {
            for( size_t col = 0; col < ncols_; ++col )
            {
                init_centroids[cluster * ncols_ + col] = h_matrix[indices( generator ) * ncols_ + col];
            }
        }
        return init_centroids;
    }

    const float tolerance_;
    const unsigned int max_iterations_;
    const unsigned int number_of_runs_;
    const unsigned int number_of_clusters_;
    const unsigned int ncols_;
};

} } // namespace cuda { namespace k_means {

#endif

namespace k_means {

static double square( const double value ) { return value * value; }

static double squared_euclidean_distance( const std::vector< double >& first, const std::vector< double >& second )
{
    double ret = 0.0;
    for( size_t i = 0; i < first.size(); ++i ) { ret += square( first[i] - second[i] ); }
    return ret;
}

struct input_t
{
    std::vector< double > data;
    comma::uint32 block;

    input_t() : data( *size ), block( 0 ) {}
};

struct k_means
{
    const double tolerance;
    const unsigned int max_iterations;
    const unsigned int number_of_runs;
    const comma::uint32 number_of_clusters;

    k_means( double tolerance, unsigned int max_iterations, unsigned int number_of_runs, comma::uint32 number_of_clusters ) noexcept :
            tolerance( tolerance ),
            max_iterations( max_iterations ),
            number_of_runs( number_of_runs ),
            number_of_clusters( number_of_clusters ) {}

    std::vector< std::vector< double > > initialize_centroids( const std::deque< std::vector< double > >& dataframe ) const
    {
        static std::mt19937 generator( seed.is_initialized() ? *seed : std::random_device{}() );
        std::vector< std::vector< double > > centroids( number_of_clusters );
        std::uniform_int_distribution< size_t > indices( 0, dataframe.size() - 1 );
        std::generate( centroids.begin(), centroids.end(), [&]() { return dataframe[indices( generator )]; } );
        return centroids;
    }

    std::vector< comma::uint32 >
    assign_centroids( const std::vector< std::vector< double > >& centroids, const std::deque< std::vector< double > >& dataframe ) const
    {
        std::vector< comma::uint32 > centroid_assignments( dataframe.size() );
        tbb::parallel_for( tbb::blocked_range< size_t >( 0, centroid_assignments.size() ), [&]( const tbb::blocked_range< size_t >& chunk )
        {
            for( size_t point = chunk.begin(); point < chunk.end(); ++point )
            {
                auto best_distance = std::numeric_limits< double >::max();
                comma::uint32 best_centroid = 0;
                for( comma::uint32 centroid_i = 0; centroid_i < number_of_clusters; ++centroid_i )
                {
                    const double distance = squared_euclidean_distance( dataframe[point], centroids[centroid_i] );
                    if( distance < best_distance ) { best_distance = distance; best_centroid = centroid_i; }
                }
                centroid_assignments[point] = best_centroid;
            }
        } );
        return centroid_assignments;
    }

    std::vector< std::vector< double > >
    update_centroids( const std::vector< comma::uint32 >& centroid_assignments, const std::deque< std::vector< double > >& dataframe ) const
    {
        // for each cluster, calculate sum of data vectors and number of data vectors
        std::vector< size_t > size_of_centroid( number_of_clusters );
        std::vector< std::vector< double > > sum_of_points_in_centroid( number_of_clusters, std::vector< double >( *size, 0.0 ) );
        for( size_t point = 0; point < dataframe.size(); ++point )
        {
            const auto centroid_i = centroid_assignments[point];
            for( size_t point_i = 0;
                 point_i < dataframe[point].size(); ++point_i ) { sum_of_points_in_centroid[centroid_i][point_i] += dataframe[point][point_i]; }
            ++size_of_centroid[centroid_i];
        }
        // calculate new centroid means
        std::vector< std::vector< double > > new_centroids( number_of_clusters );
        tbb::parallel_for( tbb::blocked_range< comma::uint32 >( 0, number_of_clusters ), [&]( const tbb::blocked_range< comma::uint32 >& chunk )
        {
            for( comma::uint32 centroid_i = chunk.begin(); centroid_i < chunk.end(); ++centroid_i )
            {
                new_centroids[centroid_i].reserve( *size );
                const auto centroid_size = std::max< size_t >( 1, size_of_centroid[centroid_i] ); // minimum size is at least 1
                for( size_t point_i = 0; point_i < sum_of_points_in_centroid[centroid_i].size(); ++point_i )
                {
                    new_centroids[centroid_i].emplace_back( sum_of_points_in_centroid[centroid_i][point_i] / centroid_size );
                }
            }
        } );
        return new_centroids;
    }

    // refer to http://www.goldsborough.me/c++/python/cuda/2017/09/10/20-32-46-exploring_k-means_in_python,_c++_and_cuda/
    std::tuple< std::vector< std::vector< double > >, std::vector< comma::uint32 > > run_on_block( const std::deque< std::vector< double > >& dataframe ) const
    {
        std::vector< double > all_scores( number_of_runs );
        std::vector< std::vector< std::vector< double > > > all_centroids( number_of_runs );
        std::vector< std::vector< comma::uint32 > > all_centroid_assignments( number_of_runs );
        tbb::parallel_for( tbb::blocked_range< unsigned int >( 0, number_of_runs ), [&]( const tbb::blocked_range< unsigned int >& chunk )
        {
            for( unsigned int run = chunk.begin(); run < chunk.end(); ++run )
            {
                std::vector< comma::uint32 > centroid_assignments;
                std::vector< std::vector< double > > run_centroids = initialize_centroids( dataframe );
                for( unsigned int iteration = 0; iteration < max_iterations; ++iteration )
                {
                    centroid_assignments = assign_centroids( run_centroids, dataframe );
                    std::vector< std::vector< double > > new_centroids = update_centroids( centroid_assignments, dataframe );
                    if( tolerance < 0 ) { run_centroids = std::move( new_centroids ); continue; }
                    double summed_difference = tbb::parallel_reduce( tbb::blocked_range< comma::uint32 >( 0, number_of_clusters ), 0.0,
                            [&]( const tbb::blocked_range< comma::uint32 > chunk, double difference ) -> double
                            {
                                for( comma::uint32 centroid_i = chunk.begin(); centroid_i < chunk.end(); ++centroid_i )
                                {
                                    difference += squared_euclidean_distance( run_centroids[centroid_i], new_centroids[centroid_i] );
                                }
                                return difference;
                            },
                            std::plus< double >()
                    );
                    run_centroids = std::move( new_centroids );
                    if( summed_difference < tolerance * tolerance ) { break; }
                }
                double run_score = tbb::parallel_reduce( tbb::blocked_range< size_t >( 0, dataframe.size() ), 0.0,
                                                         [&]( const tbb::blocked_range< size_t > chunk, double score ) -> double
                                                         {
                                                             for( size_t point = chunk.begin(); point < chunk.end(); ++point )
                                                             {
                                                                 score += std::sqrt( squared_euclidean_distance( dataframe[point], run_centroids[centroid_assignments[point]] ) );
                                                             }
                                                             return score;
                                                         },
                                                         std::plus< double >()
                );
                all_scores[run] = run_score;
                all_centroids[run] = std::move( run_centroids );
                all_centroid_assignments[run] = std::move( centroid_assignments );
            }
        } );
        const std::vector< double >::iterator& min_element = std::min_element( all_scores.begin(), all_scores.end() );
        const auto best_index = std::distance( all_scores.begin(), min_element );
        return std::make_pair( all_centroids[best_index], all_centroid_assignments[best_index] );
    }
};

} // namespace k_means {

} // namespace snark {

static std::string get_size()
{
    std::string first;
    const std::vector< std::string >& fields = comma::split( csv.fields, ',' );
    if( csv.has_field( "data" ) )
    {
        unsigned int count;
        if( csv.binary() )
        {
            count = csv.format().count();
        }
        else
        {
            while( std::cin.good() && first.empty() ) { std::getline( std::cin, first ); }
            count = comma::split( first, csv.delimiter ).size(); // quick and dirty, wasteful
        }
        size = count - fields.size() + 1;
    }
    else
    {
        unsigned int max = 0;
        for( const auto& field : fields )
        {
            if( field.substr( 0, 5 ) == "data[" && *field.rbegin() == ']' )
            {
                const auto k = boost::lexical_cast< unsigned int >( field.substr( 5, field.size() - 6 ) ) + 1;
                if( k > max ) { max = k; }
            }
        }
        if( max == 0 ) { COMMA_THROW( comma::exception, "please specify valid data fields" ) }
        size = max;
    }
    return first;
}

#ifdef SNARK_USE_CUDA

static int run_cuda_( const float tolerance, const unsigned int max_iterations, const unsigned int number_of_runs, const comma::uint32 number_of_clusters )
{
    using namespace snark::cuda::k_means;
    comma::uint32 block = 0;
    std::vector< float > dataframe;
    std::vector< std::string > input_lines;
    auto set_input_values_ = [&]( const input_t& p, const std::string& line )
    {
        block = p.block;
        input_lines.emplace_back( line );
        dataframe.insert( end( dataframe ), begin( p.data ), end( p.data ) );
    };
    if( !size )
    {
        const auto& first_line = get_size();
        if( !first_line.empty() ) { set_input_values_( comma::csv::ascii< input_t >( csv ).get( first_line ), first_line ); }
    }
    auto write_centroids_only_ = [number_of_clusters]( const std::vector< float >& centroids, const comma::uint32 block )
    {
        for( comma::uint32 i = 0; i < number_of_clusters; ++i )
        {
            if( csv.binary() )
            {
                if( use_block ) { std::cout.write( reinterpret_cast< const char* >( &block ), sizeof( block ) ); }
                std::cout.write( reinterpret_cast< const char* >( &i ), sizeof( i ) );
                for( size_t col = 0; col < *size; ++col )
                {
                    const auto point = static_cast< double >( centroids[i * *size + col] );
                    std::cout.write( reinterpret_cast< const char* >( &point ), sizeof( point ) );
                }
            }
            else
            {
                if( use_block ) { std::cout << block << csv.delimiter; }
                std::cout << i;
                for( size_t col = 0; col < *size; ++col ) { std::cout << csv.delimiter << static_cast< double >( centroids[i * *size + col] ); }
                std::cout << std::endl;
            }
        }
    };
    auto write_lines_ = []( const std::vector< std::string >& input_lines, const std::vector< float >& centroids,
                            const std::vector< comma::uint32 >& centroid_assignments )
    {
        for( size_t i = 0; i < centroid_assignments.size(); ++i )
        {
            const comma::uint32 centroid_assignment = centroid_assignments[i];
            std::cout << input_lines[i];
            if( csv.binary() )
            {
                std::cout.write( reinterpret_cast< const char* >( &centroid_assignment ), sizeof( centroid_assignment ) );
                for( size_t col = 0; col < *size; ++col )
                {
                    const auto point = static_cast< double >( centroids[centroid_assignment * *size + col] );
                    std::cout.write( reinterpret_cast< const char* >( &point ), sizeof( point ) );
                }
            }
            else
            {
                std::cout << csv.delimiter << centroid_assignment;
                for( size_t col = 0; col < *size; ++col ) { std::cout << csv.delimiter << static_cast< double >( centroids[centroid_assignment * *size + col] ); }
                std::cout << std::endl;
            }
        }
    };
    std::vector< float > centroids;
    std::vector< comma::uint32 > centroid_assignments;
    k_means operation{ tolerance, max_iterations, number_of_runs, number_of_clusters, *size };
    comma::csv::input_stream< input_t > istream( std::cin, csv );
    while( istream.ready() || std::cin.good() )
    {
        const input_t* p = istream.read();
        if( !dataframe.empty() && ( !p || block != p->block ) )
        {
            std::tie( centroids, centroid_assignments ) = use_pitched ? operation.run_on_block< snark::cuda::device::pitched_matrix >( dataframe ) : operation.run_on_block< snark::cuda::device::matrix >( dataframe );
            output_centroids ? write_centroids_only_( centroids, block ) : write_lines_( input_lines, centroids, centroid_assignments );
            if( csv.flush ) { std::cout.flush(); }
            dataframe.clear();
            input_lines.clear();
        }
        if( !p ) { break; }
        set_input_values_( *p, istream.last() );
    }
    return 0;
}

#endif

static int run_( const double tolerance, const unsigned int max_iterations, const unsigned int number_of_runs, const comma::uint32 number_of_clusters )
{
    using namespace snark::k_means;
    std::deque< std::vector< double > > dataframe;
    std::deque< std::string > input_lines;
    comma::uint32 block = 0;
    auto set_input_values_ = [&]( const input_t& p, const std::string& line )
    {
        block = p.block;
        input_lines.emplace_back( line );
        dataframe.emplace_back( p.data );
    };
    if( !size )
    {
        const auto& first_line = get_size();
        if( !first_line.empty()) { set_input_values_( comma::csv::ascii< input_t >( csv ).get( first_line ), first_line ); }
    }
    auto write_centroids_only_ = [number_of_clusters]( const std::vector< std::vector< double > >& centroids, const comma::uint32 block )
    {
        for( comma::uint32 i = 0; i < number_of_clusters; ++i )
        {
            if( csv.binary() )
            {
                if( use_block ) { std::cout.write( reinterpret_cast< const char* >( &block ), sizeof( block ) ); }
                std::cout.write( reinterpret_cast< const char* >( &i ), sizeof( i ) );
                for( const auto point : centroids[i] ) { std::cout.write( reinterpret_cast< const char* >( &point ), sizeof( point ) ); }
            }
            else
            {
                if( use_block ) { std::cout << block << csv.delimiter; }
                std::cout << i;
                for( const auto point : centroids[i] ) { std::cout << csv.delimiter << point; }
                std::cout << std::endl;
            }
        }
    };
    auto write_lines_ = []( const std::deque< std::string >& input_lines,
                            const std::vector< std::vector< double > >& centroids,
                            const std::vector< comma::uint32 >& centroid_assignments )
    {
        for( size_t i = 0; i < centroid_assignments.size(); ++i )
        {
            std::cout << input_lines[i];
            const auto centroid_assignment = centroid_assignments[i];
            const std::vector< double >& centroid = centroids[centroid_assignment];
            if( csv.binary() )
            {
                std::cout.write( reinterpret_cast< const char* >( &centroid_assignment ), sizeof( centroid_assignment ) );
                for( const auto point : centroid ) { std::cout.write( reinterpret_cast< const char* >( &point ), sizeof( point ) ); }
            }
            else
            {
                std::cout << csv.delimiter << centroid_assignment;
                for( const auto point : centroid ) { std::cout << csv.delimiter << point; }
                std::cout << std::endl;
            }
        }
    };
    std::vector< std::vector< double > > centroids;
    std::vector< comma::uint32 > centroid_assignments;
    k_means operation{ tolerance, max_iterations, number_of_runs, number_of_clusters };
    comma::csv::input_stream< input_t > istream( std::cin, csv );
    while( istream.ready() || std::cin.good() )
    {
        const input_t* p = istream.read();
        if( !dataframe.empty() && ( !p || block != p->block ) )
        {
            std::tie( centroids, centroid_assignments ) = operation.run_on_block( dataframe );
            if( output_centroids ) { write_centroids_only_( centroids, block ); }
            else { write_lines_( input_lines, centroids, centroid_assignments ); }
            if( csv.flush ) { std::cout.flush(); }
            dataframe.clear();
            input_lines.clear();
        }
        if( !p ) { break; }
        set_input_values_( *p, istream.last() );
    }
    return 0;
}

static int run( const comma::command_line_options& options )
{
    output_centroids = options.exists( "--output-centroids,--centroids" );
    csv = comma::csv::options( options, "data" );
    std::cout.precision( csv.precision );
    use_block = csv.has_field( "block" );
    seed = options.optional< std::mt19937::result_type >( "--seed" );
    const auto max_iterations = options.value< unsigned int >( "--max-iterations,--iterations", 300 );
    if( max_iterations == 0 ) { std::cerr << "math-k-means: got --max-iterations=0, --max-iterations should be at least 1" << std::endl; return 1; }
    const auto number_of_clusters = options.value< comma::uint32 >( "--number-of-clusters,--clusters" );
    if( number_of_clusters == 0 ) { std::cerr << "math-k-means: got --number-of-clusters=0, --number-of-clusters should be at least 1" << std::endl; return 1; }
    const auto number_of_runs = options.value< unsigned int >( "--number-of-runs,--runs", 10 );
    if( number_of_runs == 0 ) { std::cerr << "math-k-means: got --number-of-runs=0, --number-of-runs should be at least 1" << std::endl; return 1; }
    size = options.optional< unsigned int >( "--size" );
    auto tolerance = options.value< double >( "--tolerance", 1.0e-4 );
    if( tolerance <= 0 ) { std::cerr << "math-k-means: got --tolerance=" << tolerance << ", --tolerance should be greater than 0" << std::endl; return 1; }
    if( options.exists( "--ignore-tolerance" ) ) { tolerance = -1.0; }
#ifdef SNARK_USE_CUDA
    use_pitched = options.exists( "--cuda-use-pitched,--use-pitched,--pitched" );
    return options.exists( "--cuda" ) ? run_cuda_( static_cast< float >( tolerance ), max_iterations, number_of_runs, number_of_clusters ) : run_( tolerance, max_iterations, number_of_runs, number_of_clusters );
#else
    if( options.exists( "--cuda" ) ) { std::cerr << "math-k-means: given --cuda, but built without cuda support; run anyway, but may be slower than you expect" << std::endl; }
    return run_( tolerance, max_iterations, number_of_runs, number_of_clusters );
#endif
}

namespace comma { namespace visiting {

template <>
struct traits< snark::k_means::input_t >
{
    template < typename K, typename V >
    static void visit( const K&, snark::k_means::input_t& p, V& v )
    {
        v.apply( "data", p.data );
        v.apply( "block", p.block );
    }

    template < typename K, typename V >
    static void visit( const K&, const snark::k_means::input_t& p, V& v )
    {
        v.apply( "data", p.data );
        v.apply( "block", p.block );
    }
};

#ifdef SNARK_USE_CUDA
template <>
struct traits< snark::cuda::k_means::input_t >
{
    template < typename K, typename V >
    static void visit( const K&, snark::cuda::k_means::input_t& p, V& v )
    {
        v.apply( "data", p.data );
        v.apply( "block", p.block );
    }

    template < typename K, typename V >
    static void visit( const K&, const snark::cuda::k_means::input_t& p, V& v )
    {
        v.apply( "data", p.data );
        v.apply( "block", p.block );
    }
};
#endif

} } // namespace comma { namespace visiting {

int main( int argc, char** argv )
{
    try
    {
        comma::command_line_options options( argc, argv, usage );
        if( options.exists( "--input-fields" ) ) { std::cout << "data" << std::endl; return 0; }
        if( options.exists( "--output-fields" ) ) { std::cout << "data,centroid/id,centroid/data" << std::endl; return 0; }
        verbose = options.exists( "--verbose,-v" );
        auto max_threads = options.value< unsigned int >( "--max-threads,--threads", std::thread::hardware_concurrency() );
        if( max_threads <= 1 )
        {
            std::cerr << "math-k-means: warning: you set or std::thread::hardware_concurrency returned " << max_threads << "; will not explicitly call tbb::global_control::max_allowed_parallelism(); everything still will work, but multithreading may not be set to what you expected" << std::endl;
            return run( options );
        }
        tbb::global_control gc( tbb::global_control::max_allowed_parallelism, max_threads );
        return run( options );
    }
    catch( std::exception& ex ) { std::cerr << "math-k-means: " << ex.what() << std::endl; }
    catch( ... ) { std::cerr << "math-k-means: unknown exception" << std::endl; }
    return 1;
}


#include "SimpleTriMesh.h"

#include <Eigen/StdVector>

#include <map>

void SimpleTriMesh::addVertex(double x, double y, double z)
{
	m_vertices.push_back(Eigen::Vector3d(x, y, z));
}
void SimpleTriMesh::addTri(int i1, int i2, int i3)
{
	m_triangles.push_back(Eigen::Vector3i(i1, i2, i3));
}

void SimpleTriMesh::setFromMatrices(const Eigen::MatrixXd &vertices, const Eigen::MatrixXi &tris)
{
	m_vertices.resize(vertices.cols());
	m_triangles.resize(tris.cols());

	for (int i = 0; i < (int)m_vertices.size(); i++)
	{
		m_vertices[i] = vertices.col(i);
	}
	for (int i = 0; i < (int)m_triangles.size(); i++)
	{
		m_triangles[i] = tris.col(i);
	}
}
void SimpleTriMesh::toMatrices(Eigen::MatrixXd &vertices, Eigen::MatrixXi &tris) const
{
	vertices.resize(3, m_vertices.size());
	tris.resize(3, m_triangles.size());

	for (int i = 0; i < (int)m_vertices.size(); i++)
	{
		vertices.col(i) = m_vertices[i];
	}
	for (int i = 0; i < (int)m_triangles.size(); i++)
	{
		tris.col(i) = m_triangles[i];
	}
}
class TriBaryCoords
{
public:
	TriBaryCoords(int triIdx, const Eigen::Vector2d &u)
		:m_triIdx(triIdx),
		m_u(u)
	{
	}

	int getTriIdx() const
	{
		return m_triIdx;
	}

	Eigen::Vector3d getN() const
	{
		Eigen::Vector3d N;
		N << 1.0 - m_u.sum(), m_u;
		return N;
	}

private:
	Eigen::Vector2d m_u;
	int m_triIdx;
public:
	EIGEN_MAKE_ALIGNED_OPERATOR_NEW
};

void SimpleTriMesh::upsample(int nAdditionalPointsPerEdge)
{
	//the following implementation removes non used vertices

	int numPointsPerEdge = nAdditionalPointsPerEdge + 2;
	int numTrisPerOldFace = (numPointsPerEdge - 1)*(numPointsPerEdge - 1);
	int numTris = numTrisPerOldFace * this->m_triangles.size();

	std::vector<TriBaryCoords, Eigen::aligned_allocator<TriBaryCoords> > baryCoords;

	std::vector<int> oldVertexToBaryCoords(this->m_vertices.size(), -1);
	std::vector<std::map<int, int> > edgeToBaryCoords(this->m_vertices.size());

	auto getOldVertexIdx = [&oldVertexToBaryCoords, &baryCoords](int oldTriIdx, int v, const Eigen::Vector2d &u) {
		if (oldVertexToBaryCoords[v] == -1)
		{
			baryCoords.push_back(TriBaryCoords(oldTriIdx, u));
			oldVertexToBaryCoords[v] = baryCoords.size() - 1;
		}
		return oldVertexToBaryCoords[v];
	};
	auto getVertexIdxOnEdge = [numPointsPerEdge, &getOldVertexIdx, &baryCoords, &edgeToBaryCoords](
		int oldTriIdx, int v0, const Eigen::Vector2d &u0, int v1, const Eigen::Vector2d &u1, int col) {
		if (col == 0) return getOldVertexIdx(oldTriIdx, v0, u0);
		if (col == numPointsPerEdge - 1) return getOldVertexIdx(oldTriIdx, v1, u1);

		Eigen::Vector2d uCopy0 = u0;
		Eigen::Vector2d uCopy1 = u1;

		if (v0 > v1)
		{
			std::swap(v0, v1);
			Eigen::Vector2d ut = uCopy0;
			uCopy0 = uCopy1;
			uCopy1 = ut;
			col = numPointsPerEdge - col - 1;
		}

		if (edgeToBaryCoords[v0].find(v1) == edgeToBaryCoords[v0].end())
		{
			double lambda = col / double(numPointsPerEdge - 1);
			Eigen::Vector2d u = (1.0 - lambda) * uCopy0 + lambda * uCopy1;
			baryCoords.push_back(TriBaryCoords(oldTriIdx, u));
			edgeToBaryCoords[v0].emplace(v1, baryCoords.size() - 1);
		}

		return edgeToBaryCoords[v0].find(v1)->second;
	};

	auto getVertexIdx = [this, &getVertexIdxOnEdge, &baryCoords, numPointsPerEdge](int oldTriIdx, int row, int col) {
		const Eigen::Vector3i &tri = this->m_triangles[oldTriIdx];

		Eigen::Vector2d u0(0.0, 0.0);
		Eigen::Vector2d u1(1.0, 0.0);
		Eigen::Vector2d u2(0.0, 1.0);

		if (row == 0)
		{
			return getVertexIdxOnEdge(oldTriIdx, tri[0], u0, tri[1], u1, col);
		}
		if (col == 0)
		{
			return getVertexIdxOnEdge(oldTriIdx, tri[0], u0, tri[2], u2, row);
		}
		if (row - col == 0)
		{
			return getVertexIdxOnEdge(oldTriIdx, tri[1], u1, tri[2], u2, row);
		}

		int numPointsInRow = numPointsPerEdge - row;

		double lambda = row / double(numPointsPerEdge - 1);
		double lambda2 = col / double(numPointsInRow - 1);

		//Eigen::Vector2d p1 = (1.0 - lambda) * u0 + lambda * u2;
		//Eigen::Vector2d p2 = (1.0 - lambda) * u1 + lambda * u2;

		//Eigen::Vector2d u = (1.0 - lambda2) * p1 + lambda2 * p2;

		Eigen::Vector2d u(lambda2 * (1.0 - lambda), lambda);

		//if (numPointsInRow == 1) u = Eigen::Vector2d(0.0, lambda);

		baryCoords.push_back(TriBaryCoords(oldTriIdx, u));

		return (int)baryCoords.size() - 1;
	};

	Eigen::MatrixXi newFaces(3, numTris);

	for (int oldTriIdx = 0; oldTriIdx < this->m_triangles.size(); oldTriIdx++)
	{
		int ciTri = 0;
		for (int row = 0; row < numPointsPerEdge - 1; row++)
		{
			int numTrisInRow = numPointsPerEdge - 2 * row;
			for (int k = 0, baseVertex = 0; k < numTrisInRow; k++, baseVertex++)
			{
				Eigen::Vector3i newTri;
				newTri[0] = getVertexIdx(oldTriIdx, row, baseVertex);
				newTri[1] = getVertexIdx(oldTriIdx, row, baseVertex + 1);
				newTri[2] = getVertexIdx(oldTriIdx, row + 1, baseVertex);
				newFaces.col(oldTriIdx * numTrisPerOldFace + ciTri) = newTri;
				ciTri += 1;

				k++;
				if (k >= numTrisInRow) break;

				newTri[0] = getVertexIdx(oldTriIdx, row, baseVertex + 1);
				newTri[1] = getVertexIdx(oldTriIdx, row + 1, baseVertex + 1);
				newTri[2] = getVertexIdx(oldTriIdx, row + 1, baseVertex);
				newFaces.col(oldTriIdx * numTrisPerOldFace + ciTri) = newTri;
				ciTri += 1;
			}
		}
		assert(ciTri == numTrisPerOldFace);
	}

	Eigen::MatrixXd newVertices(3, baryCoords.size());
	for (int i = 0; i < baryCoords.size(); i++)
	{
		const Eigen::Vector3i &tri = this->m_triangles[baryCoords[i].getTriIdx()];
		Eigen::Vector3d N = baryCoords[i].getN();
		Eigen::Vector3d v = Eigen::Vector3d::Zero();
		for (int j = 0; j < 3; j++)
		{
			v += N[j] * this->m_vertices[tri[j]];
		}
		newVertices.col(i) = v;
	}

	setFromMatrices(newVertices, newFaces);
}

SimpleTriMesh SimpleTriMesh::icosphere()
{
	double t = (1.0 + std::sqrt(5.0)) / 2.0;
	double normalizationFactor = 1.0 / (1.0 + t*t);
	double u = t * normalizationFactor;
	double o = 1.0 * normalizationFactor;
	SimpleTriMesh mesh;

	mesh.addVertex(-o, u, 0);
	mesh.addVertex(o, u, 0);
	mesh.addVertex(-o, -u, 0);
	mesh.addVertex(o, -u, 0);

	mesh.addVertex(0, -o, u);
	mesh.addVertex(0, o, u);
	mesh.addVertex(0, -o, -u);
	mesh.addVertex(0, o, -u);

	mesh.addVertex(u, 0, -o);
	mesh.addVertex(u, 0, o);
	mesh.addVertex(-u, 0, -o);
	mesh.addVertex(-u, 0, o);

	mesh.addTri(0, 5, 1);
	mesh.addTri(0, 11, 5);
	mesh.addTri(0, 1, 7);
	mesh.addTri(0, 10, 11);
	mesh.addTri(0, 7, 10);

	mesh.addTri(3, 9, 4);
	mesh.addTri(3, 4, 2);
	mesh.addTri(3, 2, 6);
	mesh.addTri(3, 6, 8);
	mesh.addTri(3, 8, 9);

	mesh.addTri(1, 5, 9);
	mesh.addTri(7, 1, 8);
	mesh.addTri(10, 7, 6);
	mesh.addTri(5, 11, 4);
	mesh.addTri(11, 10, 2);

	mesh.addTri(4, 9, 5);
	mesh.addTri(9, 8, 1);
	mesh.addTri(6, 2, 10);
	mesh.addTri(2, 4, 11);
	mesh.addTri(8, 6, 7);

	return mesh;
}
SimpleTriMesh SimpleTriMesh::icosphere(int nSubdivisions)
{
	SimpleTriMesh mesh = SimpleTriMesh::icosphere();
	for (int i = 0; i < nSubdivisions; i++)
	{
		mesh.upsample(1);
		for (auto &vertex : mesh.vertices())
		{
			vertex.normalize();
		}
	}
	return mesh;
}


#ifndef __PARTICLE_FILTER_HPP
#define __PARTICLE_FILTER_HPP

#include <boost/random/discrete_distribution.hpp>

#include "particle_filter_base.hpp"

template <size_t Size>
ParticleFilter<Size>::ParticleFilter(const PredictionModel<Size> &pm,
				     const LikelihoodModel<Size> &lm,
				     unsigned int nParticles,
				     const double *initState):
  pm(pm),
  lm(lm),
  particles(std::vector<Particle<Size> >(nParticles)),
  epoch(0)
{
  if (initState)
  {
    pm.init(initState, particles);
  }
  else
  {
    double *_initState = new double[Size];
    std::fill(_initState, _initState+Size, 0);
    pm.init(initState, particles);
    delete []_initState;
  }
}

template <size_t Size>
const std::vector<Particle<Size> > &ParticleFilter<Size>::getParticles() const
{
  return particles;
}

template <size_t Size>
const Particle<Size> &ParticleFilter<Size>::getBestParticle() const
{

  double bestLikelihood = particles.at(0).likelihood;
  int bestID = 0;

  for (typename std::vector<Particle<Size> >::const_iterator it=particles.begin()+1;
       it!=particles.end(); ++it)
  {
    const Particle<Size> &p = *it;
    if (p.likelihood>bestLikelihood)
    {
      bestLikelihood = p.likelihood;
      bestID = it-particles.begin();
    }
  }

  return particles.at(bestID);
}

template <size_t Size>
void ParticleFilter<Size>::step(double deltaT)
{
  // Create a copy of the current particles
  std::vector<Particle<Size> > tempParticles(particles.size());
  std::copy(particles.begin(), particles.end(), tempParticles.begin());

  // Particle filter iteration
  // - update the particles state using the provided state-update model
  pm.update(tempParticles, deltaT);

  // - compute the importance weights using the provided likelihood model
  lm.eval(tempParticles);

  // Normalize the particles weight (and build the list of normalized weights
  // used in the next selection step)
  double wSum = 0.;
  for (typename std::vector<Particle<Size> >::iterator it=tempParticles.begin();
       it!=tempParticles.end(); ++it)
  {
    const Particle<Size> &p = *it;
    wSum += p.likelihood;
  }

  std::vector<double> normW;
  for (typename std::vector<Particle<Size> >::iterator it=tempParticles.begin();
       it!=tempParticles.end(); ++it)
  {
    Particle<Size> &p = *it;
    p.likelihood /= wSum;
    normW.push_back(p.likelihood);
  }

  // - selection step
  // TODO: parameterize at which epoch execute selection
  //if (!(epoch%10))
  if (true)
  {
    boost::random::discrete_distribution<> D(normW.begin(), normW.end());

    std::vector<int> idx;
    for (int i=0; i<tempParticles.size(); i++) idx.push_back(D(rng));

    // Copy back the selected particles
    for (int i=0; i<idx.size(); i++) particles.at(i) = tempParticles.at(idx.at(i));
  }
  else
  {
    std::copy(tempParticles.begin(), tempParticles.end(), particles.begin());
  }

  epoch++;

  // Done
}

#endif // __PARTICLE_FILTER_HPP


/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
/*                                                                           */
/*  This file is part of the library KASKADE 7                               */
/*    see http://www.zib.de/projects/kaskade7-finite-element-toolbox         */
/*                                                                           */
/*  Copyright (C) 2014 Zuse Institute Berlin                                 */
/*                                                                           */
/*  KASKADE 7 is distributed under the terms of the ZIB Academic License.    */
/*    see $KASKADE/academic.txt                                              */
/*                                                                           */
/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */

#include <numeric>

#include <boost/timer/timer.hpp>

#include <dune/common/fmatrix.hh>
#include "dune/grid/config.h"

#include "linalg/threadedMatrix.hh"

using namespace Kaskade;

// Create a NxN lower triangular matrix with 1x2 entries
// as a mockup of the divergence discretization in Stokes problems
template <class Index>
void divergence(Index const N)
{
  std::cout << "running divergence with N=" << N << "\n";

  // First, sparsity pattern creation. We proceed row by row.
  NumaCRSPatternCreator<Index> creator(N,N,false);

  boost::timer::cpu_timer timer;
  std::vector<std::vector<Index>> cols(N), rows(N);
  for (Index k=0; k<cols.size(); ++k)
  {
    cols[k].resize(k+1);
    rows[k].resize(1);
    std::iota(begin(cols[k]),end(cols[k]),0);
    rows[k][0] = k;
  }
  std::cerr << "en bloc filling of indices: " << timer.format() << "\n";
  timer.start();
  creator.addElements(rows,cols);
  std::cerr << "en bloc adding elements:    " << timer.format() << "\n";

  // Next, create a sparsity pattern from the creator and a matrix.
  timer.start();
  creator.balance();
  std::cerr << "balance:                    " << timer.format() << "\n";
  timer.start();
  std::shared_ptr<NumaCRSPattern<Index>> pattern(new NumaCRSPattern<Index>(creator));
  std::cerr << "pattern creation:           " << timer.format() << "\n";
  std::cout << "nonzero entries: " << pattern->nonzeroes() << "\n";
  std::cout << "storage size:    " << pattern->storage() << "\n";
  for (int i=0; i<pattern->nodes(); ++i)
  {
    auto const& cp = *pattern->pattern(i);
    std::cout << "chunk " << i << ": [" << cp.first() << "," << cp.last() << "[  nnz=" << cp.nonzeroes() << "  size=" << cp.storage() << "\n";
  }
  timer.start();
  NumaBCRSMatrix<Dune::FieldMatrix<double,1,2>,Index> matrix(pattern);
  std::cerr << "matrix creation:            " << timer.format() << "\n";

  // Fill the matrix with ones.
  for (auto r=matrix.begin(); r!=matrix.end(); ++r)
    for (auto c=r->begin(); c!=r->end(); ++c)
      *c = 1.0;

  Dune::BlockVector<Dune::FieldVector<double,2>> x(N);
  Dune::BlockVector<Dune::FieldVector<double,1>> y(N);
  for (int i=0; i<N; ++i)
  {
    x[i] = 1;
    y[i] = 0;
  }

  timer.start();
  // Compute 1000 matrix-vector products
  for (int i=0; i<1000; ++i)
    matrix.umv(x,y);
  std::cout << "1000 Ax: " << timer.format() << "\n";

  timer.start();
  for (int i=0; i<1000; ++i)
    matrix.umtv(x,y);
  std::cout << "1000 A^Tx: " << timer.format() << "\n";
  std::cout << "-----------------------------\n\n";
}

// Create a NxN Toeplitz matrix of bandwidth 5
template <class Index>
void toeplitz(Index const N, bool symmetric)
{
  std::cout << "running toeplitz with symmetric=" << symmetric << "\n";

  // First, sparsity pattern creation. The pattern is a superposition of
  // 5x5 blocks on the diagonal. To begin with, we add a couple of those blocks
  // individually.
  NumaCRSPatternCreator<Index> creator(N,N,symmetric,14);

  std::vector<int> idx(5);
  for (int k=0; k<100; ++k)
  {
    std::iota(begin(idx),end(idx),k);
    creator.addElements(begin(idx),end(idx),begin(idx),end(idx));
  }

  // The rest is filled en bloc.
  boost::timer::cpu_timer timer;
  std::vector<std::vector<Index>> cols(N-104), rows(N-104);
  for (Index k=0; k<cols.size(); ++k)
  {
    cols[k].resize(5);
    rows[k].resize(5);
    std::iota(begin(cols[k]),end(cols[k]),k+100);
    std::iota(begin(rows[k]),end(rows[k]),k+100);
  }
  std::cerr << "en bloc filling of indices: " << timer.format() << "\n";
  timer.start();
  creator.addElements(rows,cols);
  std::cerr << "en bloc adding elements:    " << timer.format() << "\n";

  // Next, create a sparsity pattern from the creator and a matrix.
  timer.start();
  creator.balance();
  std::cerr << "balance:                    " << timer.format() << "\n";
  timer.start();
  std::shared_ptr<NumaCRSPattern<Index>> pattern(new NumaCRSPattern<Index>(creator));
  std::cerr << "pattern creation:           " << timer.format() << "\n";
  std::cout << "nonzero entries: " << pattern->nonzeroes() << "\n";
  std::cout << "storage size:    " << pattern->storage() << "\n";
  for (int i=0; i<pattern->nodes(); ++i)
  {
    auto const& cp = *pattern->pattern(i);
    std::cout << "chunk " << i << ": [" << cp.first() << "," << cp.last() << "[  nnz=" << cp.nonzeroes() << "  size=" << cp.storage() << "\n";
  }
  timer.start();
  NumaBCRSMatrix<Dune::FieldMatrix<double,1,1>,Index> matrix(pattern);
  std::cerr << "matrix creation:            " << timer.format() << "\n";

  // Fill the matrix with ones.
  for (auto r=matrix.begin(); r!=matrix.end(); ++r)
    for (auto c=r->begin(); c!=r->end(); ++c)
      *c = 1.0;

  Dune::BlockVector<Dune::FieldVector<double,1>> x(N), y(N);
  for (int i=0; i<N; ++i)
    x[i] = 1;

  timer.start();
  // Compute maximum eigenvalue by power method
  double mx = 0;
  double dp;
  for (int i=0; i<2000; ++i)
  {
    dp = matrix.smv(.125,x,y);
    if (i%100==0 || i>1995)
    {
      mx = 0;
      for (int i=0; i<N; i+=16)
        mx += std::abs(y[i][0]);
      for (int i=0; i<N; ++i)
        x[i] = y[i]/mx;
    }
  }
  std::cout << "2000 iterations power method: " << timer.format() << "\n";
  std::cout.precision(8);
  std::cout << "max eigenvalue: " << 8*mx << "\n";
  std::cout << "dp: " << dp << "\n";

  matrix = 1.0;

  if (!symmetric)
  {
    for (int i=0; i<N; ++i)
      x[i] = 1;
    timer.start();
    for (int i=0; i<2000; ++i)
    {
      matrix.smtv(.125,x,y);
      if (i%100==0 || i>1995)
      {
        mx = 0;
        for (int i=0; i<N; i+=16)
          mx += std::abs(y[i][0]);
        for (int i=0; i<N; ++i)
          x[i] = y[i]/mx;
      }
    }
    std::cout << "2000 iterations power method: " << timer.format() << "\n";
    std::cout.precision(8);
    std::cout << "max eigenvalue: " << 8*mx << "\n";
  }
  std::cout << "-----------------------------\n\n";
}

int main(void)
{
  int N = 64000;
  toeplitz(N,false);
  toeplitz(static_cast<size_t>(N),true);
  divergence(2000);
  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.
//
// by S. Efthymiou, October 2018

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

#ifndef NETKET_MPS_PERIODIC_HPP
#define NETKET_MPS_PERIODIC_HPP

namespace netket {

template <typename T, bool diag>
class MPSPeriodic : public AbstractMachine<T> {
  using VectorType = typename AbstractMachine<T>::VectorType;
  using MatrixType = typename AbstractMachine<T>::MatrixType;
  using VectorRefType = typename AbstractMachine<T>::VectorRefType;
  using VectorConstRefType = typename AbstractMachine<T>::VectorConstRefType;
  using VisibleConstType = typename AbstractMachine<T>::VisibleConstType;

  const AbstractHilbert &hilbert_;

  // Number of sites
  int N_;
  // Physical dimension
  int d_;
  // Bond dimension
  int D_;
  // Second matrix dimension (D for normal MPS, 1 for diagonal)
  int Dsec_;
  // D squared
  int Dsq_;
  // Number of variational parameters
  int npar_;
  // Period of translational symmetry (has to be a divisor of N)
  int symperiod_;

  // Used for tree look up
  int Nleaves_;
  // Map from site that changes to the corresponding "leaves"
  // Shape (N_, nr of leaves for the corresponding site)
  std::vector<std::vector<int>> leaves_of_site_;
  // Contractions needed to produce each leaf
  // Shape (total leaves, 2)
  std::vector<std::vector<int>> leaf_contractions_;

  // MPS Matrices (stored as [symperiod, d, D, D] or [symperiod, d, D, 1])
  std::vector<std::vector<MatrixType>> W_;

  // Map from Hilbert states to MPS indices
  std::map<double, int> confindex_;
  // Identity Matrix
  MatrixType identity_mat_;

 public:
  using StateType = T;
  using LookupType = Lookup<T>;

  explicit MPSPeriodic(const AbstractHilbert &hilbert, double bond_dim,
                       int symperiod = -1)
      : hilbert_(hilbert),
        N_(hilbert.Size()),
        d_(hilbert.LocalSize()),
        D_(bond_dim),
        symperiod_(symperiod) {
    if (symperiod_ == -1) {
      symperiod_ = N_;
    }

    Init();
  }

  inline MatrixType prod(const MatrixType &m1, const MatrixType &m2) const {
    if (diag) {
      return m1.cwiseProduct(m2);
    }

    return m1 * m2;
  }

  inline T trace(const MatrixType &m) const {
    if (diag) {
      return m.sum();
    }
    return m.trace();
  }

  inline void setparamsident(MatrixType &m, VectorConstRefType pars) const {
    if (diag) {
      for (int i = 0; i < D_; i++) {
        m(i, 0) = T(1, 0) + pars(i);
      }
    } else {
      for (int i = 0; i < D_; i++) {
        for (int j = 0; j < D_; j++) {
          m(i, j) = pars(i * D_ + j);
          if (i == j) {
            m(i, j) += T(1, 0);
          }
        }
      }
    }
  }

  // Auxiliary function that defines the matrices
  void Init() {
    // Initialize parameters
    std::vector<MatrixType> pushback_vec;
    if (diag) {
      Dsec_ = 1;
      identity_mat_ = MatrixType::Ones(D_, 1);
    } else {
      Dsec_ = D_;
      identity_mat_ = MatrixType::Identity(D_, D_);
    }
    MatrixType init_mat = MatrixType::Zero(D_, Dsec_);
    Dsq_ = D_ * Dsec_;
    npar_ = symperiod_ * d_ * Dsq_;

    for (int site = 0; site < symperiod_; site++) {
      W_.push_back(pushback_vec);
      for (int spin = 0; spin < d_; spin++) {
        W_[site].push_back(init_mat);
      }
    }

    // Initialize tree parameters
    InitTree();

    // Machine creation messages
    if (diag) {
      InfoMessage() << "Periodic diagonal MPS machine with " << N_
                    << " sites created" << std::endl;
    } else {
      InfoMessage() << "Periodic MPS machine with " << N_ << " sites created"
                    << std::endl;
    }
    InfoMessage() << "Physical dimension d = " << d_
                  << " and bond dimension D = " << D_ << std::endl;
    if (symperiod_ < N_) {
      InfoMessage() << "Translation invariance is used. Number of "
                       "variational parameters is "
                    << npar_ << " instead of " << npar_ * N_ / symperiod_
                    << std::endl;
    } else {
      InfoMessage() << "Number of variational parameters is " << npar_
                    << std::endl;
    }
    // Initialize map from Hilbert space states to MPS indices
    auto localstates = hilbert_.LocalStates();
    for (int i = 0; i < d_; i++) {
      confindex_[localstates[i]] = i;
    }
  }

  void InitTree() {
    // Initializes vectors used for tree look up tables
    // leaves_of_site_ and leaf_contractions_

    std::vector<int> two_vector(2), empty_vector, level_start;
    std::vector<std::vector<int>> above;
    int level = 1, available_ind = 0, available_level = 0;
    bool available = false;

    level_start.push_back(0);
    level_start.push_back(N_);

    int length = level_start[level] - level_start[level - 1];
    while (length > 1 or available) {
      above.push_back(empty_vector);
      // Iterate level-1
      for (int i = 0; i < length - 1; i += 2) {
        // Construct above for level-1
        above.back().push_back(level_start[level] + i / 2);
        above.back().push_back(level_start[level] + i / 2);
        // Construct leaf_contractions_ for level
        two_vector[0] = level_start[level - 1] + i;
        two_vector[1] = level_start[level - 1] + i + 1;
        leaf_contractions_.push_back(two_vector);
      }
      if (length % 2 == 1) {
        if (available) {
          // Connect the two odd leaves
          above.back().push_back(level_start[level] + (length - 1) / 2);
          above[available_level].push_back(level_start[level] +
                                           (length - 1) / 2);

          two_vector[0] = level_start[level] - 1;
          two_vector[1] = available_ind;
          leaf_contractions_.push_back(two_vector);

          level_start.push_back(level_start.back() + length / 2 + 1);
          available = false;
        } else {
          available = true;
          available_ind = level_start[level] - 1;
          available_level = level - 1;
          level_start.push_back(level_start.back() + length / 2);
        }
      } else {
        level_start.push_back(level_start.back() + length / 2);
      }
      level++;
      length = level_start[level] - level_start[level - 1];
    }
    Nleaves_ = level_start.back();

    // Flatten above vector
    std::vector<int> flat_above;
    for (std::size_t l = 0; l < above.size(); l++) {
      for (std::size_t k = 0; k < above[l].size(); k++) {
        flat_above.push_back(above[l][k]);
      }
    }

    // Create leaves_of_site_ from above vector
    for (int i = 0; i < N_; i++) {
      std::vector<int> leaves;
      leaves.push_back(flat_above[i]);
      while (flat_above[leaves.back()] < Nleaves_ - 1) {
        leaves.push_back(flat_above[leaves.back()]);
      }
      leaves.push_back(Nleaves_ - 1);
      leaves_of_site_.push_back(leaves);
      leaves.clear();
    }
    Nleaves_ += -N_;
  }

  int Npar() const override { return npar_; }

  VectorType GetParameters() override {
    int k = 0;
    VectorType pars(npar_);

    for (int site = 0; site < symperiod_; site++) {
      for (int spin = 0; spin < d_; spin++) {
        for (int i = 0; i < D_; i++) {
          for (int j = 0; j < Dsec_; j++) {
            pars(k) = W_[site][spin](i, j);
            k++;
          }
        }
      }
    }
    return pars;
  }

  void SetParameters(VectorConstRefType pars) override {
    int k = 0;

    for (int site = 0; site < symperiod_; site++) {
      for (int spin = 0; spin < d_; spin++) {
        for (int i = 0; i < D_; i++) {
          for (int j = 0; j < Dsec_; j++) {
            W_[site][spin](i, j) = pars(k);
            k++;
          }
        }
      }
    }
  }

  // Auxiliary function used for setting initial random parameters and adding
  // identities in every matrix
  void SetParametersIdentity(VectorConstRefType pars) {
    int k = 0;
    for (int site = 0; site < symperiod_; site++) {
      for (int spin = 0; spin < d_; spin++) {
        setparamsident(W_[site][spin], pars.segment(k, Dsq_));
        k += Dsq_;
      }
    }
  }

  void InitRandomPars(int seed, double sigma) override {
    VectorType pars(npar_);

    netket::RandomGaussian(pars, seed, sigma);
    SetParametersIdentity(pars);
  }

  int Nvisible() const override { return N_; }

  void InitLookup(VisibleConstType v, LookupType &lt) override {
    for (int k = 0; k < Nleaves_; k++) {
      _InitLookup_check(lt, k);
      if (leaf_contractions_[k][0] < N_) {
        if (leaf_contractions_[k][1] < N_) {
          lt.M(k) = prod(W_[leaf_contractions_[k][0] % symperiod_]
                           [confindex_[v(leaf_contractions_[k][0])]],
                         W_[leaf_contractions_[k][1] % symperiod_]
                           [confindex_[v(leaf_contractions_[k][1])]]);
        } else {
          lt.M(k) = prod(W_[leaf_contractions_[k][0] % symperiod_]
                           [confindex_[v(leaf_contractions_[k][0])]],
                         lt.M(leaf_contractions_[k][1] - N_));
        }

      } else {
        if (leaf_contractions_[k][1] < N_) {
          lt.M(k) = prod(lt.M(leaf_contractions_[k][0] - N_),
                         W_[leaf_contractions_[k][1] % symperiod_]
                           [confindex_[v(leaf_contractions_[k][1])]]);
        } else {
          lt.M(k) = prod(lt.M(leaf_contractions_[k][0] - N_),
                         lt.M(leaf_contractions_[k][1] - N_));
        }
      }
    }
  }

  // Auxiliary function
  inline void _InitLookup_check(LookupType &lt, int i) {
    if (lt.MatrixSize() == i) {
      lt.AddMatrix(D_, Dsec_);
    } else {
      lt.M(i).resize(D_, Dsec_);
    }
  }

  // Auxiliary function for sorting indeces
  // (copied from stackexchange - original answer by Lukasz Wiklendt)
  inline std::vector<std::size_t> sort_indeces(const std::vector<int> &v) {
    // initialize original index locations
    std::vector<std::size_t> idx(v.size());
    std::iota(idx.begin(), idx.end(), 0);
    // sort indexes based on comparing values in v
    std::sort(idx.begin(), idx.end(),
              [&v](std::size_t i1, std::size_t i2) { return v[i1] < v[i2]; });
    return idx;
  }

  void UpdateLookup(VisibleConstType v, const std::vector<int> &tochange,
                    const std::vector<double> &newconf,
                    LookupType &lt) override {
    std::size_t nchange = tochange.size();
    if (nchange <= 0) {
      return;
    }

    MatrixType empty_matrix = MatrixType::Zero(D_, Dsec_);
    std::vector<std::size_t> sorted_ind = sort_indeces(tochange);

    std::set<int> leaves2update;
    std::size_t set_len = 0;
    std::map<int, MatrixType *> base_mat;

    for (std::size_t k = 0; k < nchange; k++) {
      int site = tochange[sorted_ind[k]];
      // Add changed visible matrices to map
      base_mat[site] =
          &(W_[site % symperiod_][confindex_[newconf[sorted_ind[k]]]]);
      // Add the rest of matrices that affects
      for (std::size_t l = 0; l < leaves_of_site_[site].size(); l++) {
        int leaf = leaves_of_site_[site][l];
        leaves2update.insert(leaf);
        if (leaves2update.size() > set_len) {
          set_len++;
          for (int i = 0; i < 2; i++) {
            int lc = leaf_contractions_[leaf - N_][i];
            if (lc < N_) {
              if (base_mat.count(lc) == 0) {
                base_mat[lc] = &(W_[lc % symperiod_][confindex_[v(lc)]]);
              }
            } else {
              base_mat[lc] = &(lt.M(lc - N_));
            }
          }
        }
      }
    }
    for (auto leaf : leaves2update) {
      lt.M(leaf - N_) = prod(*(base_mat[leaf_contractions_[leaf - N_][0]]),
                             *(base_mat[leaf_contractions_[leaf - N_][1]]));
    }
  }

  // Auxiliary function that calculates contractions from site1 to site2
  inline MatrixType mps_contraction(VisibleConstType v, const int &site1,
                                    const int &site2) {
    MatrixType c = identity_mat_;
    for (int site = site1; site < site2; site++) {
      c = prod(c, W_[site % symperiod_][confindex_[v(site)]]);
    }
    return c;
  }

  T LogVal(VisibleConstType v) override {
    return std::log(trace(mps_contraction(v, 0, N_)));
  }

  T LogVal(VisibleConstType /* v */, const LookupType &lt) override {
    return std::log(trace(lt.M(Nleaves_ - 1)));
  }

  VectorType LogValDiff(
      VisibleConstType v, const std::vector<std::vector<int>> &tochange,
      const std::vector<std::vector<double>> &newconf) override {
    const std::size_t nconn = tochange.size();

    std::vector<std::size_t> sorted_ind;
    VectorType logvaldiffs = VectorType::Zero(nconn);
    StateType current_psi = trace(mps_contraction(v, 0, N_));
    MatrixType new_prods(D_, Dsec_);

    for (std::size_t k = 0; k < nconn; k++) {
      std::size_t nchange = tochange[k].size();
      if (nchange > 0) {
        sorted_ind = sort_indeces(tochange[k]);
        int site = tochange[k][sorted_ind[0]];

        if (site == 0) {
          new_prods = W_[0][confindex_[newconf[k][sorted_ind[0]]]];
        } else {
          new_prods = prod(
              mps_contraction(v, 0, site),
              W_[site % symperiod_][confindex_[newconf[k][sorted_ind[0]]]]);
        }

        for (std::size_t i = 1; i < nchange; i++) {
          site = tochange[k][sorted_ind[i]];
          new_prods = prod(
              new_prods,
              prod(mps_contraction(v, tochange[k][sorted_ind[i - 1]] + 1, site),
                   W_[site % symperiod_]
                     [confindex_[newconf[k][sorted_ind[i]]]]));
        }
        site = tochange[k][sorted_ind[nchange - 1]];
        if (site < N_ - 1) {
          new_prods = prod(new_prods, mps_contraction(v, site + 1, N_));
        }
        logvaldiffs(k) = std::log(trace(new_prods) / current_psi);
      }
    }
    return logvaldiffs;
  }

  T LogValDiff(VisibleConstType v, const std::vector<int> &toflip,
               const std::vector<double> &newconf,
               const LookupType &lt) override {
    // Assumes number of levels > 1 (?)
    std::size_t nflip = toflip.size();
    if (nflip <= 0) {
      return T(0, 0);
    }

    MatrixType empty_matrix = MatrixType::Zero(D_, Dsec_);
    std::vector<std::size_t> sorted_ind = sort_indeces(toflip);

    std::set<int> leaves2update;
    std::map<int, MatrixType> ltpM;
    std::size_t set_len = 0;

    std::vector<bool> two_vector(2, false);
    std::map<int, std::vector<bool>> contractions_change;

    for (std::size_t k = 0; k < nflip; k++) {
      int site = toflip[sorted_ind[k]];
      // Add changed visible matrices to map
      ltpM[site] = W_[site % symperiod_][confindex_[newconf[sorted_ind[k]]]];

      // Add the rest of matrices that affects
      for (std::size_t l = 0; l < leaves_of_site_[site].size(); l++) {
        int leaf = leaves_of_site_[site][l];
        leaves2update.insert(leaf);
        if (leaves2update.size() > set_len) {
          set_len++;
          ltpM[leaf] = empty_matrix;
        }
        // Check the contractions of the new matrix (if they change)
        for (int i = 0; i < 2; i++) {
          if (std::find(toflip.begin(), toflip.end(),
                        leaf_contractions_[leaf - N_][i]) != toflip.end()) {
            two_vector[i] = true;
          } else if (leaves2update.find(leaf_contractions_[leaf - N_][i]) !=
                     leaves2update.end()) {
            two_vector[i] = true;
          }
        }
        contractions_change[leaf] = two_vector;
        two_vector[0] = false;
        two_vector[1] = false;
      }
    }

    // Calculate products
    for (auto leaf : leaves2update) {
      std::vector<MatrixType> m(2);
      for (int i = 0; i < 2; i++) {
        int lc = leaf_contractions_[leaf - N_][i];
        if (contractions_change[leaf][i]) {
          m[i] = ltpM[lc];
        } else {
          if (lc < N_) {
            m[i] = W_[lc % symperiod_][confindex_[v(lc)]];
          } else {
            m[i] = lt.M(lc - N_);
          }
        }
      }
      ltpM[leaf] = prod(m[0], m[1]);
    }
    return std::log(trace(ltpM[Nleaves_ + N_ - 1]) / trace(lt.M(Nleaves_ - 1)));
  }

  // Derivative with full calculation
  VectorType DerLog(VisibleConstType v) override {
    MatrixType temp_product(D_, Dsec_);
    std::vector<MatrixType> left_prods, right_prods;
    VectorType der = VectorType::Zero(npar_);

    // Calculate products
    left_prods.push_back(W_[0][confindex_[v(0)]]);
    right_prods.push_back(W_[(N_ - 1) % symperiod_][confindex_[v(N_ - 1)]]);
    for (int site = 1; site < N_ - 1; site++) {
      left_prods.push_back(prod(left_prods[site - 1],
                                W_[site % symperiod_][confindex_[v(site)]]));
      right_prods.push_back(
          prod(W_[(N_ - 1 - site) % symperiod_][confindex_[v(N_ - 1 - site)]],
               right_prods[site - 1]));
    }
    left_prods.push_back(prod(
        left_prods[N_ - 2], W_[(N_ - 1) % symperiod_][confindex_[v(N_ - 1)]]));
    right_prods.push_back(prod(W_[0][confindex_[v(0)]], right_prods[N_ - 2]));

    der.segment(confindex_[v(0)] * Dsq_, Dsq_) +=
        Eigen::Map<VectorType>(right_prods[N_ - 2].transpose().data(), Dsq_);
    for (int site = 1; site < N_ - 1; site++) {
      temp_product = prod(right_prods[N_ - site - 2], left_prods[site - 1]);
      der.segment((d_ * (site % symperiod_) + confindex_[v(site)]) * Dsq_,
                  Dsq_) +=
          Eigen::Map<VectorType>(temp_product.transpose().data(), Dsq_);
    }
    der.segment((d_ * ((N_ - 1) % symperiod_) + confindex_[v(N_ - 1)]) * Dsq_,
                Dsq_) +=
        Eigen::Map<VectorType>(left_prods[N_ - 2].transpose().data(), Dsq_);

    return der / trace(left_prods[N_ - 1]);
  }

  const AbstractHilbert &GetHilbert() const noexcept override {
    return hilbert_;
  }

  // Json functions
  void to_json(json &j) const override {
    j["Name"] = "MPSperiodic";
    j["Length"] = N_;
    j["BondDim"] = D_;
    j["PhysDim"] = d_;
    j["Diagonal"] = diag;
    j["SymmetryPeriod"] = symperiod_;
    for (int i = 0; i < symperiod_; i++) {
      for (int k = 0; k < d_; k++) {
        j["W" + std::to_string(d_ * i + k)] = W_[i][k];
      }
    }
  }

  void from_json(const json &pars) override {
    if (pars.at("Name") != "MPSperiodic") {
      throw InvalidInputError("Error while constructing MPS from Json input");
    }

    if (FieldExists(pars, "Length")) {
      N_ = pars["Length"];
    }
    if (N_ != hilbert_.Size()) {
      throw InvalidInputError(
          "Number of spins is incompatible with given Hilbert space");
    }

    if (FieldExists(pars, "PhysDim")) {
      d_ = pars["PhysDim"];
    }
    if (d_ != hilbert_.LocalSize()) {
      throw InvalidInputError(
          "Number of spins is incompatible with given Hilbert space");
    }

    if (FieldExists(pars, "BondDim")) {
      D_ = pars["BondDim"];
    } else {
      throw InvalidInputError("Unspecified bond dimension");
    }

    if (FieldExists(pars, "SymmetryPeriod")) {
      symperiod_ = pars["SymmetryPeriod"];
    } else {
      // Default is symperiod = N, resp. no translational symmetry
      symperiod_ = N_;
    }

    Init();

    // Loading parameters, if defined in the input
    from_jsonWeights(pars);
  }

  inline void from_jsonWeights(const json &pars) {
    for (int i = 0; i < symperiod_; i++) {
      for (int k = 0; k < d_; k++) {
        if (FieldExists(pars, "W" + std::to_string(d_ * i + k))) {
          W_[i][k] = pars["W" + std::to_string(d_ * i + k)];
        }
      }
    }
  }
};

}  // namespace netket

#endif