//
// Original authors: Andrew Dowsey <andrew.dowsey@manchester.ac.uk>
//
// Copyright 2014 CADET Bioinformatics Laboratory
//                University of Manchester, UK
//
// 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.
//

#pragma once

#include <iostream>
#include <sstream>
#include <math.h>
#include <vector>
#include <boost/filesystem.hpp>
#include <boost/gil/gil_all.hpp>
#include <boost/gil/extension/io/png_dynamic_io.hpp>
#include <Eigen/Core>
#include <Eigen/SparseCore>

using namespace std;
using namespace boost;
using namespace Eigen;

// fp is the desired floating point precision
typedef float fp;
//typedef double fp;

enum ClrMap  {simpleJet, matlabJet, cubeHelix};

template <class fp>
class Triplets
{
protected:
	vector< Triplet<fp> > trips;

public:
	Triplets() {}

	void push_back(int i, int j, fp v)
	{
		trips.push_back(Triplet<fp>(i,j,v));
	}

	void init(SparseMatrix<fp>& mat)
	{
		mat.setFromTriplets(trips.begin(), trips.end());
	}
};

template <class fp>
class PNGWriter
{
public:
	PNGWriter() {}

	void write(const filesystem::path& path, Matrix<fp,Dynamic,Dynamic>& mat, fp max_counts, bool show_sparsity = false);
	void write(const filesystem::path& path, Matrix<fp,Dynamic,Dynamic>& mat, bool show_sparsity = false);

	void write(const filesystem::path& path, SparseMatrix<fp>& mat, fp max_counts, bool show_sparsity = false);
	void write(const filesystem::path& path, SparseMatrix<fp>& mat, bool show_sparsity = false);
};

///////////////////////////////////////////////////////////////////////////////

typedef struct {
    fp r,g,b;
} COLOUR;

COLOUR GetColour(fp v,fp vmin,fp vmax);

template <class fp>
void
PNGWriter<fp>::
write(const filesystem::path& filename, Matrix<fp,Dynamic,Dynamic>& mat, fp max_counts, bool show_sparsity)
{
	fp min_intensity = log(2.0f*sqrt(3.0f/8.0f));
	fp max_intensity = log(2.0f*sqrt(max_counts+3.0f/8.0f));

	gil::rgb8_image_t png_img(mat.cols(), mat.rows());
	gil::rgb8_view_t png_view = gil::view(png_img);
	gil::rgb8_pixel_t bg(0, 0, show_sparsity ? 0 : 255);
    gil::fill_pixels(png_view, bg);

	for (int k = 0; k < mat.outerSize(); ++k)
	for (typename Matrix<fp,Dynamic,Dynamic>::InnerIterator it(mat,k); it; ++it)
	{
		COLOUR c = GetColour(log(2.0f*sqrt(it.value()+3.0f/8.0f)), min_intensity, max_intensity);
		png_view(it.col(), it.row()) = gil::rgb8_pixel_t((unsigned char) (c.r*255.9999f), (unsigned char) (c.g*255.9999f), (unsigned char) (c.b*255.9999f));
	}
    gil::png_write_view(filename.string(), gil::const_view(png_img));
}

template <class fp>
void
PNGWriter<fp>::
write(const filesystem::path& filename, Matrix<fp,Dynamic,Dynamic>& mat, bool show_sparsity)
{
	fp max_counts = 0.0;
	for (int k = 0; k < mat.outerSize(); ++k)
	for (typename Matrix<fp,Dynamic,Dynamic>::InnerIterator it(mat,k); it; ++it)
	{
		max_counts = max_counts > it.value() ? max_counts : it.value();
	}
	write(filename, mat, max_counts, show_sparsity);
}

template <class fp>
void
PNGWriter<fp>::
write(const filesystem::path& filename, SparseMatrix<fp>& mat, fp max_counts, bool show_sparsity)
{
	fp max_intensity = log(2.0f*sqrt(max_counts+3.0f/8.0f));

	gil::rgb8_image_t png_img(mat.cols(), mat.rows());
	gil::rgb8_view_t png_view = gil::view(png_img);
	gil::rgb8_pixel_t bg(0, 0, show_sparsity ? 0 : 255);
    gil::fill_pixels(png_view, bg);

	for (int k = 0; k < mat.outerSize(); ++k)
	for (typename SparseMatrix<fp>::InnerIterator it(mat,k); it; ++it)
	{
		COLOUR c = GetColour(log(2.0f*sqrt(it.value()+3.0f/8.0f)), 0.0, max_intensity);
		png_view(it.col(), it.row()) = gil::rgb8_pixel_t((unsigned char) (c.r*255.9999f), (unsigned char) (c.g*255.9999f), (unsigned char) (c.b*255.9999f));
	}
    gil::png_write_view(filename.string(), gil::const_view(png_img));
}

template <class fp>
void
PNGWriter<fp>::
write(const filesystem::path& filename, SparseMatrix<fp>& mat, bool show_sparsity)
{
	fp max_counts = 0.0;
	for (int k = 0; k < mat.outerSize(); ++k)
	for (typename SparseMatrix<fp>::InnerIterator it(mat,k); it; ++it)
	{
		max_counts = max_counts > it.value() ? max_counts : it.value();
	}
	write(filename, mat, max_counts, show_sparsity);
}


#pragma once

#include "../value.hpp"

#include <boost/fusion/adapted/struct/adapt_struct.hpp>
#include <boost/optional.hpp>

#include <string>
#include <cstdint>

namespace WretchedCss { namespace ValueTypes
{
    // Source: http://www.w3schools.com/cssref/css_units.asp
    // WARNING: MODIFY SOURCE IF YOU DO CHANGES TO THIS! AND PARSER
    enum class Unit
    {
        em, ex, percent, px, cm, mm, in, pt, pc,
        ch,
        rem,
        vh, vw,
        vmin,
        vmax
    };

    /**
     *  Holds values like 10px, 2em, 12.5%
     */
    struct NumericValue : public Value
                        , public JSON::Stringifiable <NumericValue>
                        , public JSON::Parsable <NumericValue>
    {
        Unit unit;
        double value;

        std::string getTypeString() const override;
        std::string toString() const override;
        NumericValue* clone() const override;

        NumericValue(Unit unit = Unit::px, double value = 0.);
    };

    /**
     *  Tries to convert the string into a NumericValue value object.
     */
    boost::optional <NumericValue> tryGetNumericValueFromString(std::string const& str);

    /**
     *  Converts a string into an enum value.
     *
     *  @param unit A unit as a string, such as "px".
     *  @return Returns a unit as the enum value.
     */
    Unit getUnitFromString(std::string const& unit);

    /**
     *  Converts a unit enum value into a string.
     *
     *  @param unit A unit as an enum value, such as "px".
     *  @return Returns a unit as a string.
     */
    std::string unitToString(Unit unit);

} // namespace ValueTypes
} // namespace WretchedCss

BOOST_FUSION_ADAPT_STRUCT
(
    WretchedCss::ValueTypes::NumericValue,
	(WretchedCss::ValueTypes::Unit, unit)
	(double, value)
)


#include "experiments/util/bitmap_db.hpp"
#include "experiments/util/threading.hpp"

#include <dtl/bitmap/bah.hpp>
#include <dtl/bitmap/bbc.hpp>
#include <dtl/bitmap/concise.hpp>
#include <dtl/bitmap/dynamic_bitmap.hpp>
#include <dtl/bitmap/dynamic_roaring_bitmap.hpp>
#include <dtl/bitmap/dynamic_wah.hpp>
#include <dtl/bitmap/teb_wrapper.hpp>
#include <dtl/bitmap/uah.hpp>
#include <dtl/bitmap/util/convert.hpp>
#include <dtl/bitmap/util/random.hpp>
#include <dtl/dtl.hpp>
#include <dtl/env.hpp>

#include <boost/algorithm/string.hpp>
#include <boost/filesystem.hpp>

#include <atomic>
#include <cstddef>
#include <fstream>
#include <iostream>
#include <limits>
#include <ostream>
#include <sstream>
#include <string>
#include <vector>
//===----------------------------------------------------------------------===//
namespace fs = boost::filesystem;
//===----------------------------------------------------------------------===//
template<typename T>
std::size_t
get_compressed_size(dtl::bitmap& bm) {
  T enc(bm);
  std::size_t size = enc.size_in_bytes();
  // Validation
  const auto dec = dtl::to_bitmap_using_iterator(enc);
  if ((bm & dec) != bm) {
    std::cerr << "Validation failed for type " << enc.name() << std::endl;
    std::exit(1);
  }
  return size;
}
// Disable validation for WAH64, since the iterator on that type is not
// implemented.
template<>
std::size_t
get_compressed_size<dtl::dynamic_wah64>(dtl::bitmap& bm) {
  dtl::dynamic_wah64 enc(bm);
  std::size_t size = enc.size_in_bytes();
  return size;
}
//===----------------------------------------------------------------------===//
void run(const std::string& dir, std::ostream& result_out,
    std::ostream& result_csv_out) {
  result_out
      << "//===----------------------------------------------------------------------===//"
      << "\ndata set: '" << dir << "'"
      << std::endl;

  std::size_t min_val = std::numeric_limits<std::size_t>::max();
  std::size_t max_val = 0;
  std::size_t file_cnt = 0;

  std::vector<std::string> filenames;
  fs::directory_iterator end_it;
  for (fs::directory_iterator dir_it(dir); dir_it != end_it; ++dir_it) {
    const auto file = dir_it->path().string();
    filenames.push_back(file);
    ++file_cnt;
    std::cout << "reading file: " << file << std::endl;
    std::ifstream is(file);
    std::string token;
    while (std::getline(is, token, ',')) {
      const auto val = std::stoull(token);
      if (val < min_val) min_val = val;
      if (val > max_val) max_val = val;
    }
    is.close();
  }
  std::sort(filenames.begin(), filenames.end());

  const auto n = max_val + 1;
  const auto n_pow2 = dtl::next_power_of_two(n);
  const auto c = file_cnt;

  result_out << "min_val: " << min_val
             << " max_val: " << max_val
             << " bitmap_cnt: " << file_cnt
             << " n: " << n
             << std::endl;

  $f64 min_d = std::numeric_limits<f64>::max();
  $f64 max_d = std::numeric_limits<f64>::min();
  $f64 sum_d = 0.0;

  $f64 min_f = std::numeric_limits<f64>::max();
  $f64 max_f = std::numeric_limits<f64>::min();
  $f64 sum_f = 0.0;

  std::atomic<std::size_t> bytes_roaring { 0 };
  std::atomic<std::size_t> bytes_wah { 0 };
  std::atomic<std::size_t> bytes_wah64 { 0 };
  std::atomic<std::size_t> bytes_teb { 0 };
  std::atomic<std::size_t> bytes_uah8 { 0 };
  std::atomic<std::size_t> bytes_uah16 { 0 };
  std::atomic<std::size_t> bytes_uah32 { 0 };
  std::atomic<std::size_t> bytes_bah { 0 };
  std::atomic<std::size_t> bytes_concise { 0 };
  std::atomic<std::size_t> bytes_bbc { 0 };

  std::vector<dtl::bitmap> bitmaps;
  std::vector<dtl::bitmap> bitmaps_pow2;

  std::size_t total_bit_cnt = 0;
  for (auto& file : filenames) {
    std::ifstream is(file);
    std::string token;

    dtl::bitmap bm(n);

    while (std::getline(is, token, ',')) {
      const auto val = std::stoull(token);
      bm[val] = true;
    }
    is.close();

    bitmaps.push_back(bm);

    const auto d = dtl::determine_bit_density(bm);
    min_d = std::min(d, min_d);
    max_d = std::max(d, max_d);
    sum_d += d;

    const auto f = dtl::determine_clustering_factor(bm);
    min_f = std::min(f, min_f);
    max_f = std::max(f, max_f);
    sum_f += f;

    total_bit_cnt += bm.count();
  }
  result_out << "cardinality: " << total_bit_cnt << std::endl;
  result_out << "d (min/avg/max): "
             << min_d << " / "
             << (sum_d / file_cnt) << " / "
             << max_d << std::endl;
  result_out << "f (min/avg/max): "
             << min_f << " / "
             << (sum_f / file_cnt) << " / "
             << max_f << std::endl;

  auto thread_fn = [&](const std::size_t bid, std::ostream& os) {
    auto& bm = bitmaps[bid];
    auto& bm_pow2 = bitmaps_pow2[bid];

    std::size_t r = 0;
    std::size_t t = 0;
    std::size_t w = 0;
    std::size_t w64 = 0;
    std::size_t uah8 = 0;
    std::size_t uah16 = 0;
    std::size_t uah32 = 0;
    std::size_t bah = 0;
    std::size_t concise = 0;
    std::size_t bbc = 0;
    {
      using T = dtl::dynamic_roaring_bitmap;
      r = get_compressed_size<T>(bm);
    }
    {
      using T = dtl::dynamic_wah32;
      w = get_compressed_size<T>(bm);
    }
    {
      using T = dtl::dynamic_wah64;
      w64 = get_compressed_size<T>(bm);
    }
    {
      using T = dtl::teb_wrapper;
      t = get_compressed_size<T>(bm);
    }
    {
      using T = dtl::uah8;
      uah8 = get_compressed_size<T>(bm);
    }
    {
      using T = dtl::uah16;
      uah16 = get_compressed_size<T>(bm);
    }
    {
      using T = dtl::uah32;
      uah32 = get_compressed_size<T>(bm);
    }
    {
      using T = dtl::bah;
      bah = get_compressed_size<T>(bm);
    }
    {
      using T = dtl::concise;
      concise = get_compressed_size<T>(bm);
    }
    {
      using T = dtl::bbc;
      bbc = get_compressed_size<T>(bm);
//      dtl::bbc x(bm);
//      bbc = x.size_in_bytes();
//      const auto dec = dtl::to_bitmap_using_iterator(x);
//      if (bm != dec) {
//        std::cerr << "BBC: Validation failed." << std::endl;
//        std::cout << "n=" << bm.size() << std::endl;
//        std::cout << "expected: ";
//        auto it_a = dtl::plain_bitmap_iter<dtl::bitmap>(bm);
//        while (!it_a.end()) {
//          std::cout << "[" << it_a.pos() << "," << (it_a.pos() + it_a.length())
//            << "),";
//          it_a.next();
//        }
//        std::cout << std::endl;
//        std::cout << "but got:  ";
//        auto it_b = x.it();
//        while (!it_b.end()) {
//          std::cout << "[" << it_b.pos() << "," << (it_b.pos() + it_b.length())
//              << "),";
//          it_b.next();
//        }
//        std::cout << std::endl;
//        std::cout << x << std::endl;
//        std::exit(1);
//      }
    }
    // clang-format off
// TEB (lossy compressed)
//    {
//      const auto fpr = 0.0001;
//      dtl::teb_wrapper teb(bm_pow2, fpr);
//      t = teb.size_in_bytes();
//      const auto dec = dtl::to_bitmap_using_iterator(teb);
//      if ((bm_pow2 & dec) != bm_pow2) {
//        std::cerr << "Validation failed." << std::endl;
//        std::exit(1);
//      }
//      u64 max_fp_cnt = static_cast<u64>(teb.size() * fpr);
//      u64 fp_cnt = (bm_pow2 ^ dec).count();
//      std::cout << "fp_cnt=" << fp_cnt << std::endl;
//      if (fp_cnt > max_fp_cnt) {
//        std::cerr << "Validation failed. Max FP count exceeded." << std::endl;
//        std::exit(1);
//      }
//      os << teb.info() << std::endl;
//    }
    // clang-format on

    bytes_roaring += r;
    bytes_wah += w;
    bytes_wah64 += w64;
    bytes_teb += t;
    bytes_uah8 += uah8;
    bytes_uah16 += uah16;
    bytes_uah32 += uah32;
    bytes_bah += bah;
    bytes_concise += concise;
    bytes_bbc += bbc;
  };

  dispatch(0, bitmaps.size(), thread_fn);

  result_out << "roaring: " << std::setw(15) << bytes_roaring << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_roaring * 8.0) / total_bit_cnt) << " bits/int" << std::endl;
  result_out << "teb:     " << std::setw(15) << bytes_teb << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_teb * 8.0) / total_bit_cnt) << " bits/int" << std::endl;
  result_out << "wah:     " << std::setw(15) << bytes_wah << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_wah * 8.0) / total_bit_cnt) << " bits/int" << std::endl;
  result_out << "wah64:   " << std::setw(15) << bytes_wah64 << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_wah64 * 8.0) / total_bit_cnt) << " bits/int" << std::endl;
  result_out << "uah8 :   " << std::setw(15) << bytes_uah8 << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_uah8 * 8.0) / total_bit_cnt) << " bits/int" << std::endl;
  result_out << "uah16 :  " << std::setw(15) << bytes_uah16 << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_uah16 * 8.0) / total_bit_cnt) << " bits/int" << std::endl;
  result_out << "uah32 :  " << std::setw(15) << bytes_uah32 << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_uah32 * 8.0) / total_bit_cnt) << " bits/int" << std::endl;
  result_out << "bah :    " << std::setw(15) << bytes_bah << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_bah * 8.0) / total_bit_cnt) << " bits/int" << std::endl;
  result_out << "concise: " << std::setw(15) << bytes_concise << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_concise * 8.0) / total_bit_cnt) << " bits/int" << std::endl;
  result_out << "bbc:     " << std::setw(15) << bytes_bbc << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_bbc * 8.0) / total_bit_cnt) << " bits/int" << std::endl;

  result_csv_out << "\"" << dir << "\"";
  result_csv_out << "," << std::setprecision(4) << ((bytes_roaring * 8.0) / total_bit_cnt);
  result_csv_out << "," << std::setprecision(4) << ((bytes_teb * 8.0) / total_bit_cnt);
  result_csv_out << "," << std::setprecision(4) << ((bytes_wah * 8.0) / total_bit_cnt);
  result_csv_out << "," << std::setprecision(4) << ((bytes_bah * 8.0) / total_bit_cnt);
  result_csv_out << "," << std::setprecision(4) << ((bytes_concise * 8.0) / total_bit_cnt);
  result_csv_out << "," << std::setprecision(4) << ((bytes_bbc * 8.0) / total_bit_cnt);
  result_csv_out << std::endl;
}
//===----------------------------------------------------------------------===//
$i32 main() {
  const auto basedir = dtl::env<std::string>::get("DIR",
      "./real-roaring-dataset");
  std::vector<std::string> dirs {
    basedir + "/census1881/",
    basedir + "/census1881_srt/",
    basedir + "/census-income/",
    basedir + "/census-income_srt/",
    basedir + "/weather_sept_85/",
    basedir + "/weather_sept_85_srt/",
    basedir + "/wikileaks-noquotes/",
    basedir + "/wikileaks-noquotes_srt/",
  };

  std::stringstream results;
  std::stringstream results_csv;
  for (auto& dir : dirs) {
    run(dir, results, results_csv);
  }
  std::cout << "Results:" << std::endl;
  std::cout << results.str() << std::endl;
  std::cout << std::endl;
  std::cout << "Results [bits per value]:" << std::endl;
  std::cout << "=========================" << std::endl;
  std::cout << std::endl;
  std::cout << "data_set,roaring,teb,wah,bah,concise,bbc" << std::endl;
  std::cout << results_csv.str() << std::endl;
}
//===----------------------------------------------------------------------===//


/* Copyright (c) 2015, Julian Straub <jstraub@csail.mit.edu>, Randi Cabezas <rcabezas@csail.mit.edu>
 * Licensed under the MIT license. See the license file LICENSE.
 */

#include <iostream>

#include <boost/program_options.hpp>
#include <boost/algorithm/string.hpp> //case insensitve string comparison

#include <dpMM/dirMultiNaiveBayes.hpp>
#include <dpMM/niwBaseMeasure.hpp>
#include <dpMM/niwSphere.hpp>
#include <dpMM/dirBaseMeasure.hpp>
#include <dpMM/normalSphere.hpp> //for sampling points on the sphere
#include <dpMM/iw.hpp>
#include <dpMM/typedef.h>
#include <dpMM/timer.hpp>

namespace po = boost::program_options;
using boost::iequals;

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

  // Declare the supported options.
  po::options_description desc("Allowed options");
  desc.add_options()
    ("help,h", "produce help message")
	("ompMaxThreads,m", po::value<int>(), "max number of OMP threads to use [default: num cores]")
	("K,K", po::value<int>(), "number of initial clusters")
	("T,T", po::value<int>(), "iterations")
	("v,v", po::value<bool>(), "verbose output")
    ("params,p", po::value<string>(),
      "path to file containing parameters (see class definition)")
	("alpha,a", po::value<double>(), "alpha value")
	("nu,n", po::value<std::vector<double> >()->multitoken(), "NIW mean prior strength to use for distributions")
	("deltaOffset,d", po::value<std::vector<double> >()->multitoken(), "NIW sigma prior strength to use for distributions")
    ("input,i", po::value<string>(),
      "path to input dataset .csv file (rows: dimensions; cols: different datapoints)")
    ("output,o", po::value<string>(),
      "path to output labels .csv file (rows: time; cols: different datapoints)")
	("b,b", po::value<std::vector<string> >()->multitoken(),
	  "base class to use for components (must be the same size as M in the data). "
	  "valid values: NiwSampled, NiwSphere, DirSampled [default: NiwSampled].");

    po::variables_map vm;
	po::store(po::parse_command_line(argc, argv, desc), vm);
	po::notify(vm);

	if (vm.count("help")) {
		cout << desc << "\n";
		return 1;
	}

	bool verbose = false;

	uint NumObs = 1; //num observations (number of components of multi-dimention data)
	uint K=2; //num clusters
	uint T=10; //iterations
	uint M=2; //num docs
	double alphaInput = 10; //weighting of the alpha

	vector<uint> N(NumObs, 100) ; //num data points (total)
	vector<uint> D(NumObs, 2);  //dimention of data
	vector<double> nuIn; //nu
	vector<double> deltaOffsetIn; //delta offset

	//use the params provided
	string paramsF="";
	if (vm.count("params"))
		paramsF = vm["params"].as<string>();

	if (vm.count("K"))
		K = vm["K"].as<int>();
	if (vm.count("T"))
		T = vm["T"].as<int>();
	if (vm.count("v"))
		verbose = vm["v"].as<bool>();
	if (vm.count("a"))
		alphaInput = vm["a"].as<double>();

	//handle ompMaxNumber of threads
	int Mproc = 1000000;
	if(vm.count("ompMaxThreads"))
		Mproc = vm["ompMaxThreads"].as<int>();

	#if defined(_OPENMP)
		//numCores to use is min(Mproc,#cores)
		Mproc = (Mproc <= omp_get_max_threads()) ? Mproc : omp_get_max_threads();
		omp_set_num_threads(Mproc);
	#endif

	string pathIn ="";
	string pathOut ="";
	if(vm.count("input"))
		pathIn = vm["input"].as<string>();
	if(vm.count("output"))
		pathOut= vm["output"].as<string>();
	if(vm.count("nu"))
		nuIn = vm["nu"].as<vector<double> >();
	if(vm.count("deltaOffset"))
		deltaOffsetIn = vm["deltaOffset"].as<vector<double> >();

	vector<vector< Matrix<double, Dynamic, Dynamic> > > x;
	x.reserve(NumObs);
	if (!pathIn.compare(""))
	{
		cout<<"making some data up" <<endl;
		for (uint n=0; n<NumObs; ++n)
		{
			if(verbose)
				cout << "Nobs " << n << ": "<< endl;
			vector<Matrix<double, Dynamic, Dynamic> > tempDoc;
			tempDoc.reserve(M);
			uint Ndoc=M;
			uint Nword=int(N[n]/M);
			for(uint i=0; i<Ndoc; ++i)
			{
				if(verbose)
					cout << "\tdoc " << i << ": ";
				/*MatrixXd  xdoc(D[n],Nword);
				for(uint w=0; w<Nword; ++w)
				{
					if(i<Ndoc/2)
						xdoc.col(w) <<  (NumObs*(n%2))*VectorXd::Ones(D[n]);
					else
						xdoc.col(w) <<  (NumObs*(n%2)+1)*VectorXd::Ones(D[n]);
				} */

				//sampleClustersOnSphere(xdoc,K);
				/*MatrixXd xdoc;
				if(i%2==0) {
					xdoc = (NumObs*(i%2))*MatrixXd::Ones(D[n],Nword);
				} else {
					xdoc = (NumObs*(i%2)+99)*MatrixXd::Ones(D[n],Nword);
				}
				*/

				//MatrixXd  xdoc(D[n],Nword);
				//VectorXd fill = VectorXd::Zero(D[n]);
				//for(uint w=0; w<Nword; ++w) {
				//	VectorXd tempFill = fill;
				//	int ind = int(w<Nword/2);
				//	tempFill[ind] = 1;
				//	xdoc.col(w) <<  tempFill;
				//}

				MatrixXd  xdoc = MatrixXd::Zero(D[n],Nword);
				if(i<Ndoc/2) {
					xdoc.row(0) = MatrixXd::Ones(1,Nword);
				} else {
					xdoc.row(1) = MatrixXd::Ones(1,Nword);
				}

				tempDoc.push_back(xdoc);
				if(verbose)
					cout << xdoc << endl;
			}

			x.push_back(tempDoc);
		}

	}else{
		cout<<"loading data from "<<pathIn<<endl;
		std::ifstream fin(pathIn.data(),std::ifstream::in);

		if (!fin.good()) {
			cout << "could not open file...returning" << endl;
			return(-1);
		}

		//read data [sufficient statistics]
		fin>>NumObs;
		N.clear(); N.reserve(NumObs);
		D.clear(); D.reserve(NumObs);

		for(uint32_t n=0; n<NumObs; ++n)
		{
			vector< Matrix<double, Dynamic, Dynamic> > xiter;
			uint Niter, Diter;

			fin>>Niter;
			fin>>M;
			fin>>Diter;

			N.push_back(Niter);
			D.push_back(Diter);

			MatrixXd data(Diter,Niter);
			VectorXu words(M);

			for (uint j=0; j<M; ++j)
				fin>>words(j);

			for (uint j=1; j<(Diter+1); ++j)
				for (uint i=0; i<Niter; ++i)
					fin>>data(j-1,i);

			uint count = 0;
			for (uint j=0; j<M; ++j)
			{
				xiter.push_back(data.middleCols(count,words[j]));
				count+=words[j];
			}
			x.push_back(xiter);
		}
		fin.close();
	}

	if(uint(nuIn.size())!=NumObs && !nuIn.empty()) {
		cerr << "Error specified number nu parameters does not equal number of data components" << endl;
		cerr << "#N nu=" << nuIn.size() << ", #data components=" << NumObs <<  "." << endl;
		cerr << "exiting" << endl;
		return(-1);
	}
	if(uint(deltaOffsetIn.size())!=NumObs && !deltaOffsetIn.empty()) {
		cerr << "Error specified number delta offset parameters does not equal number of data components" << endl;
		cerr << "#N nu=" << deltaOffsetIn.size() << ", #data components=" << NumObs <<  "." << endl;
		cerr << "exiting" << endl;
		return(-1);
	}

	boost::mt19937 rndGen(9191);
	DirMultiNaiveBayes<double> *naive_samp;
	vector<string> useBaseDist(NumObs, " ");

	if(paramsF.compare("")){
		//initialize from file
		std::ifstream fin(paramsF.data(),std::ifstream::in);
			naive_samp = new DirMultiNaiveBayes<double>(fin,&rndGen);
			naive_samp->loadData(x);
		fin.close();
	} else {
		//initialize here

		vector<string> baseDist; //base distributions for each component

		if(vm.count("b")) {
			baseDist.clear();
			baseDist = vm["b"].as<vector<string> >();

			if(uint(baseDist.size())!=NumObs) {
				cerr << "Error specified number of base distributions not equal number of data components" << endl;
				cerr << "#bases =" << baseDist.size() << ", #data components=" << NumObs <<  "." << endl;
				cerr << "exiting" << endl;
				return(-1);
			}
		} else {
			baseDist.clear();
			baseDist = vector<string>(NumObs, "NiwSampled");
		}

		useBaseDist = baseDist;

		VectorXd alpha = (alphaInput)*VectorXd::Ones(K);

		Dir<Catd,double> dir(alpha,&rndGen);
		vector<vector<boost::shared_ptr<BaseMeasure<double> > > > thetas;
		thetas.reserve(NumObs);

		//creates thetas
		for(uint m=0;m<NumObs ; ++m)
		{
			if(iequals(baseDist[m], "NiwSampled")) { //sampled normal inversed wishart
				//sufficient statistics size:(1+D+D^2)x1
				//[N, sum(x_i) outerProduct(x_i)_row_major_form ]
				vector<boost::shared_ptr<BaseMeasure<double> > > niwSampled;

				//input data will be size given above (but this is SS size and we want size data)
				//solve size from above and you get this
				int dNorm = int( (sqrt(D[m]*4+1)-1)/2);

				double nu;
				if(nuIn.empty()) {
					nu = dNorm+1;
				}else {
					nu = nuIn[m];
				}
				double kappa = dNorm+1;

				MatrixXd Delta = MatrixXd::Identity(dNorm,dNorm);
				if(deltaOffsetIn.empty()) {
					Delta *= 0.1;
				} else {
					Delta *= deltaOffsetIn[m];
				}
				Delta *= nu;

				VectorXd theta = VectorXd::Zero(dNorm);

				NIW<double> niw(Delta,theta,nu,kappa,&rndGen);
				boost::shared_ptr<NiwSampled<double> > tempBase( new NiwSampled<double>(niw));

				//simple copy
				for(int k=0; k<int(K);++k) {
					niwSampled.push_back(boost::shared_ptr<BaseMeasure<double> >(tempBase->copy()));
				}
				thetas.push_back(niwSampled);

			} else if(iequals(baseDist[m], "NiwSphere")) {
				//sufficient statistics size:(1+D+(D-1)^2)x1
				//[N, Karcher mean, zeros]
				//IW is one dimention smaller (makes sense since it's on tangent space)
				//double nu = D[m];

				//input data will be size given above (but this is SS size and we want size data)
				//solve size from above and you get this
				int dNorm = int( (sqrt(D[m]*4-7)+1)/2 );

				vector<boost::shared_ptr<BaseMeasure<double> > > niwTangent;

				double nu;
				if(nuIn.empty()) {
					nu = dNorm+1;
				}else {
					nu = nuIn[m];
				}
//				double kappa = dNorm;

				MatrixXd Delta = MatrixXd::Identity(dNorm-1,dNorm-1);
				if(deltaOffsetIn.empty()) {
					Delta *= pow(5.0*M_PI/180.0,2);
				} else {
					Delta *= deltaOffsetIn[m];
				}
				Delta *= nu;

				IW<double> iw(Delta,nu,&rndGen);
				boost::shared_ptr<NiwSphere<double> > tempBase( new NiwSphere<double>(iw,&rndGen));

				//simple copycd
				for(int k=0; k<int(K);++k) {
					niwTangent.push_back(boost::shared_ptr<BaseMeasure<double> >(tempBase->copy()));
				}
				thetas.push_back(niwTangent);

			} else if(iequals(baseDist[m], "DirSampled")) {
				//sufficient statistics size:(D)x1
				//[counts]

				vector<boost::shared_ptr<BaseMeasure<double> > > dirSampled;

				for(int k=0; k<int(K); ++k) {
					VectorXd alpha = VectorXd::Ones(D[m]);
					int highProbPos = k%D[m];
					if(nuIn.empty()) {
						alpha *= 10;
						alpha(highProbPos) = 10/2;
					} else {
						alpha *= nuIn[m];
						alpha(highProbPos) = nuIn[m]/2;
					}

					Dir<Catd,double> dirBase(alpha,&rndGen);

					boost::shared_ptr<DirSampled<Catd, double> > tempBase( new DirSampled<Catd, double>(dirBase));
					dirSampled.push_back(boost::shared_ptr<BaseMeasure<double> >(tempBase));
				}

				thetas.push_back(dirSampled);

			} else if(iequals(baseDist[m], "MF")) {
				//sufficient statistics size:(D)x1
				//[counts]
				vector<boost::shared_ptr<BaseMeasure<double> > > mfBases;
        for(uint32_t k=0; k<K; ++k)
        {
          uint32_t nIter = 30; // iterations of internal MF sampler
          // mixture over the 6 different MF directions
          Matrix<double,Dynamic,1> alpha = Matrix<double,Dynamic,1>::Ones(6)*10;
          Dir<Cat<double>, double> dir(alpha,&rndGen);
          // tangent space IW priors
          double nu = 4;
          Matrix<double,Dynamic,Dynamic> Delta =
            Matrix<double,Dynamic,Dynamic>::Identity(2,2);
          Delta *= pow(5.0*M_PI/180.,2)*nu;
          IWd iw0(Delta,nu,&rndGen);
          std::vector<shared_ptr<BaseMeasure<double> > > iwTs;
          for(uint32_t j=0; j<6; ++j)
          {
            iwTs.push_back(shared_ptr<IwTangent<double> >(
                  new IwTangent<double>(iw0,&rndGen)));
          }
          DirMM<double> dirMM(dir,iwTs);
          MfPrior<double> mfPrior(dirMM, nIter);
//          MF<T> mf(R,pi,TGs);
          mfBases.push_back(boost::shared_ptr<BaseMeasure<double> >(
                new MfBase<double>(mfPrior)));
        }
        //set
				thetas.push_back(mfBases);

			} else {
				cerr << "error with base distributions (check help) ... returning." << endl;
				return(-1);
			}
		}

		naive_samp = new DirMultiNaiveBayes<double>(dir,thetas);
		naive_samp->initialize( x );
	}

	Timer tlocal;
	tlocal.tic();

	cout << "multiObsNaiveBayesian Clustering:" << endl;
	cout << "Ndocs=" << M << endl;
	cout << "NumComp=" << NumObs << endl;
	cout << "NumData,Dim,baseDist= ";
	for(uint n=0; n<NumObs; ++n)
		cout << "[" << N[n] << ", " << D[n] << "," << useBaseDist[n] << "]; ";
	cout << endl;

	cout << "Num Cluster = " << K << ", (" << T << " iterations)." << endl;

	naive_samp->inferAll(T,verbose);

	if (pathOut.compare(""))
	{
		std::ofstream fout(pathOut.data(),std::ofstream::out);
		naive_samp->dump_clean(fout);
		fout.close();
	}

	tlocal.displayElapsedTimeAuto();

	delete naive_samp;

	return(0);

};


/////////////////////
//// halo_array.h
////
//// Declaration of the halo_array class
////
////////////////////

#ifndef ____HALO_ARRAY__
#define ____HALO_ARRAY__
#include "constants.hpp"
#include <armadillo>

class halo_vec{

  arma::vec array;
  int Nx;
  int rulepts;
  int fdpts;

public:
  /// Constructor
  halo_vec(const int _Nx, const int _rulepts):
    Nx{_Nx},rulepts{_rulepts},
    array(_Nx+_rulepts-1,arma::fill::zeros){fdpts = (rulepts-1)*0.5;}
  /// Destructor
  ~halo_vec(){}

  // Overloaded operators
  inline double& operator()(const int i) {return array(i+fdpts);}
  inline halo_vec& operator=(const arma::vec& eqarr)
  {array = eqarr; return *this;}
  inline void zeros() {array.zeros();}
};

class halo1D_mat{

  arma::mat array;
  int Nx, Ny;
  int rulepts;
  int fdpts;

public:
  /// Constructor
  halo1D_mat(const int _Nx, const int _Ny, const int _rulepts):
    Nx{_Nx},Ny{_Ny},rulepts{_rulepts},
    array(_Nx+_rulepts-1,_Ny,arma::fill::zeros){fdpts = (rulepts-1)*0.5;}
  /// Destructor
  ~halo1D_mat(){}

  // Overloaded operators
  inline double& operator()(const int i, const int j)
  {return array(i+fdpts,j);}
  inline halo1D_mat& operator=(const arma::mat& eqarr)
  {array = eqarr; return *this;}
  inline void zeros() {array.zeros();}
};

class halo_mat{

  arma::mat array;
  int Nx, Ny;
  int xrulepts, yrulepts;
  int xfdpts, yfdpts;

public:
  /// Constructor
  halo_mat(const int _Nx, const int _Ny,
	   const int _xrulepts, const int _yrulepts):
    Nx{_Nx}, Ny{_Ny},
    xrulepts{_xrulepts}, yrulepts{_yrulepts},
    array(_Nx+_xrulepts-1,_Ny+_yrulepts-1,arma::fill::zeros)
  {xfdpts = (xrulepts-1)*0.5; yfdpts = (yrulepts-1)*0.5;}
  /// Destructor
  ~halo_mat(){}

  // Overloaded operators
  inline double& operator()(const int i, const int j)
  {return array(i+xfdpts,j+yfdpts);}
  inline halo_mat& operator=(const arma::mat& eqmat)
  {array = eqmat; return *this;}
  inline void zeros() {array.zeros();}
};

class halo_cx_mat{

  arma::cx_mat array;
  int Nx, Ny;
  int xrulepts, yrulepts;
  int xfdpts, yfdpts;

public:
  /// Constructor
  halo_cx_mat(const int _Nx, const int _Ny,
	      const int _xrulepts, const int _yrulepts):
    Nx{_Nx}, Ny{_Ny},
    xrulepts{_xrulepts}, yrulepts{_yrulepts},
    array(_Nx+_xrulepts-1,_Ny+_yrulepts-1,arma::fill::zeros)
  {xfdpts = (xrulepts-1)*0.5; yfdpts = (yrulepts-1)*0.5;}
  /// Destructor
  ~halo_cx_mat(){}

  // Overloaded operators
  inline dcomplex& operator()(const int i, const int j)
  {return array(i+xfdpts,j+yfdpts);}
  inline halo_cx_mat& operator=(const arma::cx_mat& eqcmat)
  {array = eqcmat; return *this;}
  inline halo_cx_mat& operator=(const dcomplex& eqcmat)
  {array = eqcmat; return *this;}

  inline void zeros() {array.zeros();}
  inline void set_value(const int i, const int j, const dcomplex value)
  {array(i+xfdpts, j+yfdpts) = value;}
};

#endif // ____HALO_ARRAY__ //


/* Copyright 2015,2016 Tao Xu
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

#include "TreeRegressor.h"

#include <cstdlib>
#include <limits>
#include <boost/random/uniform_real.hpp>

#include "Tree.h"
#include "GbmFun.h"
#include "DataSet.h"
#include "gflags/gflags.h"
#include "glog/logging.h"

DEFINE_int32(min_leaf_examples, 256,
             "minimum number of data points in the leaf");

namespace boosting {

using namespace std;

// Return true with approximately the desired probability;
// actual probability may differ by ~ 1/RAND_MAX
// Not thread safe. But:
// (1) we're only using this to sample some examples and features,
//     so we don't care (?)
// (2) currently TreeRegressor isn't multithreading anyway
inline bool biasedCoinFlip(double probabilityOfTrue) {
  return (rand() < probabilityOfTrue * RAND_MAX);
}

TreeRegressor::SplitNode::SplitNode(const vector<int>* st):
  subset(st), fid(-1), fv(0), gain(0), selected(false),
  left(NULL), right(NULL) {
}

TreeRegressor::TreeRegressor(
  const DataSet& ds,
  const boost::scoped_array<double>& y,
  const GbmFun& fun) : ds_(ds), y_(y), fun_(fun) {
}

TreeRegressor::~TreeRegressor() {
  for (SplitNode* split : allSplits_) {
    delete split;
  }
}

void TreeRegressor::splitExamples(
  const SplitNode& split,
  vector<int>* left,
  vector<int>* right) {

  const int fid = split.fid;
  const uint16_t fv = split.fv;

  auto &f = ds_.features_[fid];

  if (f.encoding == BYTE) {
    boosting::split<uint8_t>(*(split.subset), left, right, *(f.bvec), fv);
  } else {
    CHECK(f.encoding == SHORT);
    boosting::split<uint16_t>(*(split.subset), left, right, *(f.svec), fv);
  }
}

void TreeRegressor::getBestSplitFromHistogram(
  const TreeRegressor::Histogram& hist,
  int* idx,
  double* gain) {

  // The loss function should really be
  //   (sum of (y - y_mean)^2 for observations left of idx)
  //     + (sum of (y - y_mean)^2 for observations right of idx)
  // By math, this equals
  //   (sum of squares of all y-values)
  //     - (sum of y-values on left)^2 / (number of observations on left)
  //     - (sum of y-values on right)^2 / (number of observations on right)
  // Since the first term (sum of squares of all y-values) is independent
  // of our choice of where to split, it makes no difference, so we ignore it
  // in calculating loss.

  // loss function if we don't split at all
  double lossBefore = -1.0 * hist.totalSum * hist.totalSum / hist.totalCnt;

  int cntLeft = 0;       // number of observations on or to left of idx
  double sumLeft = 0.0;  // number of observations strictly to right of idx

  double bestGain = 0.0;
  int bestIdx = -1;      // everything strictly to right of idx

  CHECK(hist.num >= 1);

  for (int i = 0; i < hist.num - 1; i++) {

    cntLeft += hist.cnt[i];
    sumLeft += hist.sumy[i];

    double sumRight = hist.totalSum - sumLeft;
    int cntRight = hist.totalCnt - cntLeft;

    if (cntLeft < FLAGS_min_leaf_examples) {
      continue;
    }
    if (cntRight < FLAGS_min_leaf_examples) {
      break;
    }

    double lossAfter =
      -1.0 * sumLeft * sumLeft / cntLeft
      - 1.0 * sumRight * sumRight / cntRight;

    double gain = lossBefore - lossAfter;
    if (gain > bestGain) {
      bestGain = gain;
      bestIdx = i;
    }
  }

  *idx = bestIdx;
  *gain = bestGain;
}

TreeRegressor::SplitNode*
TreeRegressor::getBestSplit(const vector<int>* subset,
                            double featureSamplingRate,
                            bool terminal) {

  SplitNode* split = new SplitNode(subset);
  if (terminal) {
    allSplits_.push_back(split);
    return split;
  }

  int bestFid = -1;       // which feature to split on, -1 is invalid
  int bestFv = 0;         // critical value of that feature

  // gain in prediction accuracy from that split:
  // initialize to 0 instead of std::numeric_limits<double>::lowest() because,
  // if no split results in a positive gain, we would rather report that, than
  // return a valid but degenerate split
  double bestGain = 0.0;

  double totalSum = 0.0;  // sum of all target values

  for (auto& id : *subset) {
    totalSum += y_[id];
  }

  // For each of a random sampling of features, see if splitting on that
  // feature results in the biggest improvement so far.
  // TODO(tiankai): The various fid's can be processed in parallel.
  for (int fid = 0; fid < ds_.numFeatures_; fid++) {
    const auto& f = ds_.features_[fid];

    if (f.encoding == EMPTY || !biasedCoinFlip(featureSamplingRate)) {
      continue;
    }

    Histogram hist(f.transitions.size() + 1, subset->size(), totalSum);

    if (f.encoding == BYTE) {
      buildHistogram<uint8_t>(*subset, *(f.bvec), hist);
    } else {
      CHECK(f.encoding == SHORT);
      buildHistogram<uint16_t>(*subset, *(f.svec), hist);
    }

    int fv;
    double gain;
    getBestSplitFromHistogram(hist, &fv, &gain);

    if (gain > bestGain) {
      bestFid = fid;
      bestFv = fv;
      bestGain = gain;
    }
  }
  split->fid = bestFid;
  split->fv = bestFv;
  split->gain = bestGain;

  frontiers_.push_back(split);
  allSplits_.push_back(split);
  return split;
}

TreeNode<uint16_t>* TreeRegressor::getTree(
  const int numLeaves,
  const double exampleSamplingRate,
  const double featureSamplingRate,
  double fimps[]) {

  // randomly sample data in ds_
  vector<int>* subset = new vector<int>();
  for (int i = 0; i < ds_.getNumExamples(); i++) {
    if (biasedCoinFlip(exampleSamplingRate)) {
      subset->push_back(i);
    }
  }
  CHECK(subset->size() >= FLAGS_min_leaf_examples * numLeaves);

  // compute the decision tree in SplitNode's
  SplitNode* root = getBestSplits(subset, numLeaves - 1, featureSamplingRate);

  // convert the decision tree to PartitionNode's and LeafNode's
  return getTreeHelper(root, fimps);
}

TreeNode<uint16_t>* TreeRegressor::getTreeHelper(
  SplitNode* split,
  double fimps[]) {

  if (split == NULL) {
    return NULL;
  } else if (!split->selected) {
    // leaf of decision tree
    double fvote = fun_.getLeafVal(*(split->subset), y_);
    LOG(INFO) << "leaf:  " << fvote << ", #examples:"
              << split->subset->size();
    CHECK(split->subset->size() >= FLAGS_min_leaf_examples);

    return new LeafNode<uint16_t>(fvote);
  } else {
    // internal node of decision tree
    LOG(INFO) << "select split: " << split->fid << ":" << split->fv
              << " gain: " << split->gain << ", #examples:"
              << split->subset->size() << ", min partition: "
              << std::min(split->left->subset->size(), split->right->subset->size());

    fimps[split->fid] += split->gain;
    double fvote = fun_.getLeafVal(*(split->subset), y_);
    PartitionNode<uint16_t>* node = new PartitionNode<uint16_t>(split->fid, split->fv);
    node->setLeft(getTreeHelper(split->left, fimps));
    node->setRight(getTreeHelper(split->right, fimps));
    node->setVote(fvote);

    return node;
  }
}

TreeRegressor::SplitNode* TreeRegressor::getBestSplits(
  const vector<int>* subset, const int numSplits, double featureSamplingRate) {

  CHECK(subset != NULL);

  // Compute the root of the decision tree.
  SplitNode* firstSplit = getBestSplit(subset, featureSamplingRate, false);

  int numSelected = 0;
  do {
    // frontiers_.size() = #leaves = #internal nodes + 1 = numSelected + 1
    CHECK(frontiers_.size() == numSelected+1);

    // Do a linear search over the leaves to find the next split with the most
    // gain.
    double bestGain = 0.0;
    vector<SplitNode*>::iterator best_it = frontiers_.end();
    for (auto it = frontiers_.begin(); it != frontiers_.end(); it++) {
      if ((*it)->gain > bestGain) {
        bestGain = (*it)->gain;
        best_it = it;
      }
    }

    if (best_it == frontiers_.end()) {
      // no gain from any split
      break;
    }

    CHECK(bestGain > 0.0);

    (*best_it)->selected = true;
    numSelected++;
    SplitNode* bestSplit = *best_it;
    frontiers_.erase(best_it);

    // Now that we've selected bestSplit, expand its left and right children.
    vector<int>* left = new vector<int>();
    vector<int>* right = new vector<int>();

    splitExamples(*bestSplit, left, right);
    bool terminal = (numSelected == numSplits);

    bestSplit->left = getBestSplit(left, featureSamplingRate, terminal);
    bestSplit->right = getBestSplit(right, featureSamplingRate, terminal);
  } while (numSelected < numSplits);

  return firstSplit;
}

}


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2007-2011 Barend Gehrels, Amsterdam, the Netherlands.
// Copyright (c) 2008-2011 Bruno Lalande, Paris, France.
// Copyright (c) 2009-2011 Mateusz Loskot, London, UK.

// Parts of Boost.Geometry are redesigned from Geodan's Geographic Library
// (geolib/GGL), copyright (c) 1995-2010 Geodan, Amsterdam, the Netherlands.

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

#ifndef BOOST_GEOMETRY_MULTI_STRATEGIES_CARTESIAN_CENTROID_AVERAGE_HPP
#define BOOST_GEOMETRY_MULTI_STRATEGIES_CARTESIAN_CENTROID_AVERAGE_HPP

#include <boost/numeric/conversion/cast.hpp>
#include <boost/type_traits.hpp>

#include <boost/geometry/core/coordinate_type.hpp>
#include <boost/geometry/core/point_type.hpp>
#include <boost/geometry/arithmetic/arithmetic.hpp>
#include <boost/geometry/strategies/centroid.hpp>

namespace boost { namespace geometry
{

namespace strategy { namespace centroid
{

/*!
\brief Centroid calculation taking average of points
\ingroup strategies
*/
template
<
    typename PointCentroid,
    typename Point = PointCentroid
>
class average
{
private :

    /*! subclass to keep state */
    class sum
    {
        friend class average;
        int count;
        PointCentroid centroid;

    public :
        inline sum()
            : count(0)
        {
            assign_zero(centroid);
        }
    };

public :
    typedef sum state_type;
    typedef PointCentroid centroid_point_type;
    typedef Point point_type;

    static inline void apply(Point const& p, sum& state)
    {
        add_point(state.centroid, p);
        state.count++;
    }

    static inline void result(sum const& state, PointCentroid& centroid)
    {
        centroid = state.centroid;
        divide_value(centroid, state.count);
    }

};

#ifndef DOXYGEN_NO_STRATEGY_SPECIALIZATIONS

namespace services
{

template <typename Point, std::size_t DimensionCount, typename Geometry>
struct default_strategy
<
    cartesian_tag,
    pointlike_tag,
    DimensionCount,
    Point,
    Geometry
>
{
    typedef average
        <
            Point,
            typename point_type<Geometry>::type
        > type;
};

} // namespace services

#endif

}} // namespace strategy::centroid

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_MULTI_STRATEGIES_CARTESIAN_CENTROID_AVERAGE_HPP


// Copyright (C) 2017 Vicente J. Botet Escriba
//
//  Distributed under the Boost Software License, Version 1.0. (See accompanying
//  file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#include "flat_non_zero_integer.hpp"

#include <boost/detail/lightweight_test.hpp>
#include <iostream>

int main()
{
#ifdef COMPILE_FAILS
    {
        non_zero_integer a = 1; // compile fails
    }
#endif

#if __cplusplus > 201402L && defined __clang__
    {
        opt_non_zero_integer onzi = non_zero_integer::make(5);
        BOOST_TEST(bool(onzi));
        BOOST_TEST(onzi->underlying() == 5);
        non_zero_integer nzi = *onzi;
        std::cout << nzi << '\n';
    }
    {
        opt_non_zero_integer onzi = non_zero_integer::make(0);
        BOOST_TEST( ! bool(onzi) );
    }
#endif
    {
        non_zero_integer a = non_zero_integer::cast(5);
        BOOST_TEST(a == a);
    }
    {
        int i = 5;
        // while the following operation is not safe, we know the value of i
        non_zero_integer a = non_zero_integer::cast(i);
        BOOST_TEST(a == 5);
    }
    {
        non_zero_integer a = non_zero_integer::cast(5);
        BOOST_TEST(a == 5);
    }
#ifdef RUN_FAILS
    {
        non_zero_integer a = non_zero_integer::cast(0);
        static_cast<void>(a);
    }
#endif
    {
        non_zero_integer a = non_zero_integer::cast(5);
        static_assert(std::is_same<decltype(-a), non_zero_integer>::value, "");
        BOOST_TEST(a == -(-a));
        BOOST_TEST(-a == -5);
    }
    {
        non_zero_integer a = non_zero_integer::cast(5);
        static_assert(std::is_same<decltype(+a), non_zero_integer>::value, "");
        BOOST_TEST(a == +a);
        BOOST_TEST(a == 5);
    }
    {
        non_zero_integer a = non_zero_integer::cast(5);
        BOOST_TEST(a+1 == 1+a);
        static_assert(std::is_same<decltype(a+1), int>::value, "");
        static_assert(std::is_same<decltype(1+a), int>::value, "");
    }
    {
        non_zero_integer a = non_zero_integer::cast(5);
        non_zero_integer b = a;
        BOOST_TEST(a == b);
        BOOST_TEST(a+b == b+a);
        static_assert(std::is_same<decltype(a+b), non_zero_integer>::value, "");
        static_assert(std::is_same<decltype(b+a), non_zero_integer>::value, "");
        BOOST_TEST(a+b == 10);
    }
    {
        non_zero_integer a = non_zero_integer::cast(5);
        non_zero_integer b = a;
        a += b;
        BOOST_TEST(a == 10);
        BOOST_TEST(b == 5);
    }
    {
        non_zero_integer a = non_zero_integer::cast(5);
        BOOST_TEST(a-1 == 4);
        BOOST_TEST(1-a == -4);
        BOOST_TEST(a-1 == -(1-a));
        static_assert(std::is_same<decltype(a-1), int>::value, "");
        static_assert(std::is_same<decltype(1-a), int>::value, "");
    }
    {
        non_zero_integer a = non_zero_integer::cast(5);
        non_zero_integer b = a;
        BOOST_TEST(a-b == 0);
        static_assert(std::is_same<decltype(a-b), int>::value, "");
    }

#ifdef COMPILE_FAILS
    {
        non_zero_integer a = non_zero_integer::cast(5);
        non_zero_integer b = a;
        a -= b; // compile fails
    }
#endif
    {
        non_zero_integer a = non_zero_integer::cast(5);
        BOOST_TEST(a*1 == 1*a);
        BOOST_TEST(a*1 == a);
        static_assert(std::is_same<decltype(a*1), int>::value, "");
        static_assert(std::is_same<decltype(1*a), int>::value, "");
    }
    {
        non_zero_integer a = non_zero_integer::cast(5);
        non_zero_integer b = a;
        BOOST_TEST(a == b);
        BOOST_TEST(a*b == b*a);
        static_assert(std::is_same<decltype(a*b), non_zero_integer>::value, "");
        static_assert(std::is_same<decltype(b*a), non_zero_integer>::value, "");
        BOOST_TEST(a*b == 25);
    }
    {
        non_zero_integer a = non_zero_integer::cast(5);
        non_zero_integer b = a;
        a *= b;
        BOOST_TEST(a == b*b);
        BOOST_TEST(a == 25);
        BOOST_TEST(b == 5);
    }

    return ::boost::report_errors();
}


#include <hre/config.h>

#include <hre/user.h>
#include <dm/dm.h>
#include <dm/dm_boost.h>
#include <ltsmin-lib/ltsmin-standard.h>

#include <popt.h>

#include <vector>
#include <iostream>

#include <boost/config.hpp>
#include <boost/graph/graph_traits.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/cuthill_mckee_ordering.hpp>
#include <boost/graph/properties.hpp>
#include <boost/graph/bandwidth.hpp>
#include <boost/graph/profile.hpp>
#include <boost/graph/wavefront.hpp>
#include <boost/graph/sloan_ordering.hpp>
#include <boost/graph/minimum_degree_ordering.hpp>
#include <boost/graph/king_ordering.hpp>

static int sloan_w1 = 1;
static int sloan_w2 = 16;

struct poptOption boost_options[] = {
    { "sloan-w1", 0, POPT_ARG_INT|POPT_ARGFLAG_SHOW_DEFAULT, &sloan_w1, 0, "use <W1> as weight 1 for the Sloan algorithm", "<W1>"},
    { "sloan-w2", 0, POPT_ARG_INT|POPT_ARGFLAG_SHOW_DEFAULT, &sloan_w2, 0, "use <W2> as weight 2 for the Sloan algorithm", "<W2>"},
    POPT_TABLEEND
};

using namespace boost;
using namespace std;

template<typename DMGraph>
static DMGraph
create_graph(const matrix_t* m, const int total)
{
    Warning(infoLong, "Creating bipartite graph: G");

    // size of the bipartite graph of m
    int size = dm_nrows(m) + dm_ncols(m);

    if (total) {
        /* Compute number of extra vertices.
         * This will be the number of non-zeros in the matrix. */
        for (int i = 0; i < dm_nrows(m); i++) {
            size += dm_ones_in_row((matrix_t*)m, i);
        }
    }

    // create the adjacency graph
    DMGraph g = DMGraph(size);

    // add edges for the bipartite graph
    for (int i = 0; i < dm_nrows(m); i++) {
        for (int j = 0; j < dm_ncols(m); j++) {
            if (dm_is_set(m, i, j)) add_edge(i, dm_nrows(m) + j, g);
        }
    }

    if (total) {
        Warning(infoLong, "Creating total graph of G: T(G)");

        // add edges for every adjacent vertex in G
        int offset = dm_nrows(m) + dm_ncols(m);
        for (int i = 0; i < dm_nrows(m); i++) {
            for (int j = 0; j < dm_ncols(m); j++) {
                if (dm_is_set(m, i, j)) {
                    add_edge(i, offset, g);
                    add_edge(dm_nrows(m) + j, offset, g);
                    offset++;
                }
            }
        }

        // add the edges for incident edges in G, by iterating over rows.
        offset = dm_nrows(m) + dm_ncols(m);
        for (int i = 0; i <  dm_nrows(m); i++) {
            for (int j = 0, k = -1; j < dm_ncols(m); j++) {
                if (dm_is_set(m, i, j)) {
                    if (k > -1) add_edge(k, offset, g);
                    k = offset;
                    offset++;
                }
            }
        }

        /* Now we need to add edges by iterating over columns.
         * But this is a bit harder, because we don't know the
         * vertex numbers on-the-fly. So we precompute them
         * and put them in a map. */
        std::map<int, int> offsets;
        for (int i = 0, c = 0; i < dm_nrows(m); i++) {
            for (int j = 0; j < dm_ncols(m); j++) {
                if (dm_is_set(m, i, j)) {
                    offsets[i * dm_ncols(m) + j] = c++ + dm_nrows(m) + dm_ncols(m);
                }
            }
        }

        // add the edges for incident edges in G, by iterating over columns.
        for (int i = 0; i < dm_ncols(m); i++) {
            for (int j = 0, k = -1; j < dm_nrows(m); j++) {
                if (dm_is_set(m, j, i)) {
                    if (k > -1) add_edge(offsets[k * dm_ncols(m) + i], offsets[j * dm_ncols(m) + i], g);
                    k = j;
                }
            }
        }
    }

    return g;
}

template<typename Vertex, typename Graph>
static void
parse_ordering(std::vector<Vertex> inv_perm, typename property_map<Graph, vertex_index_t>::type index_map,
               int* row_perm, int* col_perm, const matrix_t* m)
{
    Warning(infoLong, "Parsing ordering");
    int r = 0, c = 0;
    for (typename std::vector<Vertex>::const_iterator i=inv_perm.begin(); i != inv_perm.end(); ++i) {
        if (index_map[*i] < (unsigned int) dm_nrows(m) + dm_ncols(m)) {
            if (index_map[*i] < (unsigned int) dm_nrows(m)) row_perm[r++] = index_map[*i];
            else col_perm[c++] = index_map[*i] - dm_nrows(m);
        }
    }
    if (r != dm_nrows(m) || c != dm_ncols(m)) {
        Warning(error, "Got invalid permutation from boost. "
                "This may happen when the graph representation of the matrix "
                "has disconnected sub graphs. This needs to be fixed in boost. "
                "A bug report has been filed here: "
                "https://svn.boost.org/trac/boost/ticket/11355.");
        HREexit(LTSMIN_EXIT_FAILURE);
    }
}

template<typename Graph, typename VertexIndexMap>
static typename graph_traits<Graph>::vertices_size_type
span(Graph g, VertexIndexMap index)
{
    typedef typename graph_traits<Graph>::vertices_size_type vertices_size_type;
    vertices_size_type span = 0;
    typename graph_traits<Graph>::vertex_iterator i, end;
    for (boost::tie(i, end) = vertices(g); i != end; ++i){
        if (out_degree(*i, g) == 0) continue;
        int min = num_vertices(g);
        int max = 0;
        typename graph_traits<Graph>::out_edge_iterator e, end;
        for (boost::tie(e, end) = out_edges(*i, g); e != end; ++e) {
            int i = get(index, target(*e, g));
            if (i < min) min = i;
            if (i > max) max = i;
        }
        span += ((max - min) + 1);
    }
    return span;
}

template<typename Vertex, typename Graph>
static void
graph_stats(std::vector<Vertex> inv_perm, Graph g, typename property_map<Graph, vertex_index_t>::type index_map, const int metrics)
{
    if (log_active(infoLong) || metrics) {

        log_t l = infoLong;
        if (metrics) l = info;

        Warning(infoLong, "Computing graph stats");
        std::vector<int> perm(num_vertices(g));

        for (unsigned int c = 0; c != num_vertices(g); ++c) {
            perm[index_map[inv_perm[c]]] = c;
        }

        Warning(l, "bandwidth: %zu", bandwidth(g, make_iterator_property_map(&perm[0], index_map, perm[0])));
        Warning(l, "profile: %zu", profile(g, make_iterator_property_map(&perm[0], index_map, perm[0])));
        Warning(l, "span: %zu", span(g, make_iterator_property_map(&perm[0], index_map, perm[0])));
        Warning(l, "average wavefront: %.*g", (int) std::ceil(std::log10(num_vertices(g) * num_vertices(g))), aver_wavefront(g, make_iterator_property_map(&perm[0], index_map, perm[0])));
        Warning(l, "RMS wavefront: %.*g", (int) std::ceil(std::log10(num_vertices(g) * num_vertices(g))), rms_wavefront(g, make_iterator_property_map(&perm[0], index_map, perm[0])));
    }
}

void
boost_ordering(const matrix_t* m, int* row_perm, int* col_perm, const boost_reorder_t alg, const int total, const int metrics)
{
    switch(alg) {
        case BOOST_NONE: {
            typedef adjacency_list<vecS, vecS, undirectedS,
                property<vertex_color_t, default_color_type,
                    property<vertex_degree_t,int> > > Graph;

            typedef graph_traits<Graph>::vertex_descriptor Vertex;

            Graph g = create_graph<Graph>(m, total);

            property_map<Graph, vertex_index_t>::type index_map = get(vertex_index, g);

            std::vector<Vertex> inv_perm(num_vertices(g));

            for (unsigned int i = 0; i < num_vertices(g); i++) inv_perm[i] = i;

            graph_stats(inv_perm, g, index_map, metrics);
            parse_ordering<Vertex, Graph>(inv_perm, index_map, row_perm, col_perm, m);
        } break;
        case BOOST_CM: {
            typedef adjacency_list<vecS, vecS, undirectedS,
                property<vertex_color_t, default_color_type,
                    property<vertex_degree_t,int> > > Graph;

            typedef graph_traits<Graph>::vertex_descriptor Vertex;

            Graph g = create_graph<Graph>(m, total);

            property_map<Graph, vertex_index_t>::type index_map = get(vertex_index, g);

            std::vector<Vertex> inv_perm(num_vertices(g));

            Warning(infoLong, "Computing Cuthill McKee ordering");
            cuthill_mckee_ordering(g, inv_perm.rbegin(), get(vertex_color, g), make_degree_map(g));
            graph_stats(inv_perm, g, index_map, metrics);

            parse_ordering<Vertex, Graph>(inv_perm, index_map, row_perm, col_perm, m);
        } break;
        case BOOST_SLOAN: {
            typedef adjacency_list<setS, vecS, undirectedS,
                    property<vertex_color_t, default_color_type,
                        property<vertex_degree_t, int,
                            property< vertex_priority_t, double > > > > Graph;

            typedef graph_traits<Graph>::vertex_descriptor Vertex;

            Graph g = create_graph<Graph>(m, total);

            property_map<Graph, vertex_index_t>::type index_map = get(vertex_index, g);

            std::vector<Vertex> inv_perm(num_vertices(g));

            Warning(infoLong, "Computing Sloan ordering");
            sloan_ordering(g, inv_perm.begin(), get(vertex_color, g), make_degree_map(g), get(vertex_priority, g), sloan_w1, sloan_w2);
            graph_stats(inv_perm, g, index_map, metrics);

            parse_ordering<Vertex, Graph>(inv_perm, index_map, row_perm, col_perm, m);
        } break;
        case BOOST_KING: {
            typedef adjacency_list<vecS, vecS, undirectedS,
                property<vertex_color_t, default_color_type,
                    property<vertex_degree_t,int> > > Graph;

            typedef graph_traits<Graph>::vertex_descriptor Vertex;

            Graph g = create_graph<Graph>(m, total);

            property_map<Graph, vertex_index_t>::type index_map = get(vertex_index, g);

            std::vector<Vertex> inv_perm(num_vertices(g));

            Warning(infoLong, "Computing King ordering");
            king_ordering(g, inv_perm.rbegin(), get(vertex_color, g), make_degree_map(g), index_map);
            graph_stats(inv_perm, g, index_map, metrics);

            parse_ordering<Vertex, Graph>(inv_perm, index_map, row_perm, col_perm, m);
        } break;
        default: {
            Warning(error, "Unsupported Boost ordering");
            HREexit(LTSMIN_EXIT_FAILURE);
        }
    }
}


// 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 G. Mazzola, May-Aug 2018

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

#ifndef NETKET_JAS_SYMM_HPP
#define NETKET_JAS_SYMM_HPP

namespace netket {

// Jastrow with permutation symmetries
class JastrowSymm : public AbstractMachine {
  const AbstractHilbert &hilbert_;
  const AbstractGraph &graph_;

  std::vector<std::vector<int>> permtable_;
  int permsize_;

  // number of visible units
  int nv_;

  // number of parameters
  int npar_;

  // number of parameters without symmetries
  int nbarepar_;

  // weights
  MatrixType W_;

  // weights with symmetries
  MatrixType Wsymm_;

  VectorType thetas_;
  VectorType thetasnew_;

  Eigen::MatrixXd DerMatSymm_;
  Eigen::MatrixXi Wtemp_;

 public:
  // constructor
  explicit JastrowSymm(const AbstractHilbert &hilbert)
      : hilbert_(hilbert), graph_(hilbert.GetGraph()), nv_(hilbert.Size()) {
    Init(graph_);

    SetBareParameters();
  }

  void Init(const AbstractGraph &graph) {
    if (nv_ < 2) {
      throw InvalidInputError(
          "Cannot construct Jastrow states with less than two visible units");
    }

    permtable_ = graph.SymmetryTable();
    permsize_ = permtable_.size();

    for (int i = 0; i < permsize_; i++) {
      assert(int(permtable_[i].size()) == nv_);
    }

    W_.resize(nv_, nv_);
    W_.setZero();
    thetas_.resize(nv_);
    thetasnew_.resize(nv_);

    nbarepar_ = (nv_ * (nv_ - 1)) / 2;

    // Constructing the matrix that maps the bare derivatives to the symmetric
    // ones

    Wtemp_ = Eigen::MatrixXi::Zero(nv_, nv_);

    std::map<int, int> params;
    int k = 0;

    for (int i = 0; i < nv_; i++) {
      for (int j = i + 1; j < nv_; j++) {
        for (int l = 0; l < permsize_; l++) {
          int isymm = permtable_[l][i];
          int jsymm = permtable_[l][j];

          if (isymm < 0 || isymm >= nv_ || jsymm < 0 || jsymm >= nv_) {
            std::cerr << "Error in JastrowSymm" << std::endl;
            std::abort();
          }
          Wtemp_(isymm, jsymm) = k;
          Wtemp_(jsymm, isymm) = k;
        }  // l
        k++;
      }  // j
    }    // i

    int nk_unique = 0;

    for (int i = 0; i < nv_; i++) {
      for (int j = i + 1; j < nv_; j++) {
        k = Wtemp_(i, j);
        if (params.count(k) == 0) {
          nk_unique++;
          params.insert(std::pair<int, int>(k, nk_unique));
        }
      }
    }

    npar_ = params.size();

    for (int i = 0; i < nv_; i++) {
      for (int j = i + 1; j < nv_; j++) {
        if (params.count(Wtemp_(i, j))) {
          Wtemp_(i, j) = params.find(Wtemp_(i, j))->second;
        } else {
          std::cerr << "Error in JastrowSymm" << std::endl;
          std::abort();
        }
        Wtemp_(j, i) = Wtemp_(i, j);
      }
    }

    DerMatSymm_ = Eigen::MatrixXd::Zero(npar_, nbarepar_);
    Wsymm_.resize(npar_, 1);  // used to stay close to RbmSpinSymm class

    int kbare = 0;
    for (int i = 0; i < nv_; i++) {
      for (int j = i + 1; j < nv_; j++) {
        int ksymm = Wtemp_(i, j);
        if (ksymm < 1 || ksymm - 1 >= npar_) {
          std::cerr << "Error in JastrowSymm" << std::endl;
          std::abort();
        }
        DerMatSymm_(ksymm - 1, kbare) = 1;
        kbare++;
      }
    }

    InfoMessage() << "Jastrow WF Initizialized with nvisible = " << nv_
                  << std::endl;
    InfoMessage() << "Symmetries are being used : " << npar_
                  << " parameters left, instead of " << nbarepar_ << std::endl;
  }

  int Nvisible() const override { return nv_; }

  // int Nhidden() const { return nh_; }

  int Npar() const override { return npar_; }

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

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

    SetParameters(par);
  }

  void InitLookup(VisibleConstType v, LookupType &lt) override {
    if (lt.VectorSize() == 0) {
      lt.AddVector(v.size());
    }
    if (lt.V(0).size() != v.size()) {
      lt.V(0).resize(v.size());
    }
    lt.V(0) = (W_.transpose() * v);  // does not matter the transpose W is symm
  }

  // same as RBM
  void UpdateLookup(VisibleConstType v, const std::vector<int> &tochange,
                    const std::vector<double> &newconf,
                    LookupType &lt) override {
    if (tochange.size() != 0) {
      for (std::size_t s = 0; s < tochange.size(); s++) {
        const int sf = tochange[s];
        lt.V(0) += W_.row(sf) * (newconf[s] - v(sf));
      }
    }
  }

  VectorType BareDerLog(VisibleConstType v) {
    VectorType der(nbarepar_);

    int k = 0;
    for (int i = 0; i < nv_; i++) {
      for (int j = i + 1; j < nv_; j++) {
        der(k) = v(i) * v(j);
        k++;
      }
    }

    return der;
  }

  // now unchanged w.r.t. RBM spin symm
  VectorType DerLog(VisibleConstType v) override {
    return DerMatSymm_ * BareDerLog(v);
  }

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

    int k = 0;

    for (int i = 0; i < npar_; i++) {
      pars(k) = Wsymm_(i, 0);
      k++;
    }
    return pars;
  }

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

    for (int i = 0; i < npar_; i++) {
      Wsymm_(i, 0) = pars(k);
      k++;
    }

    SetBareParameters();
  }

  void SetBareParameters() {
    for (int i = 0; i < nv_; i++) {
      for (int j = i + 1; j < nv_; j++) {
        W_(i, j) = Wsymm_(Wtemp_(i, j) - 1, 0);
        W_(j, i) = W_(i, j);  // create the lover triangle
        W_(i, i) = Complex(0);
      }
    }
  }

  Complex LogVal(VisibleConstType v) override { return 0.5 * v.dot(W_ * v); }

  // Value of the logarithm of the wave-function
  // using pre-computed look-up tables for efficiency
  Complex LogVal(VisibleConstType v, const LookupType &lt) override {
    return 0.5 * v.dot(lt.V(0));
  }

  // 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 {
    const std::size_t nconn = tochange.size();
    VectorType logvaldiffs = VectorType::Zero(nconn);

    thetas_ = (W_.transpose() * v);
    Complex logtsum = 0.5 * v.dot(thetas_);

    for (std::size_t k = 0; k < nconn; k++) {
      if (tochange[k].size() != 0) {
        thetasnew_ = thetas_;
        Eigen::VectorXd vnew(v);

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

          thetasnew_ += W_.row(sf) * (newconf[k][s] - v(sf));
          vnew(sf) = newconf[k][s];
        }

        logvaldiffs(k) = 0.5 * vnew.dot(thetasnew_) - logtsum;
      }
    }

    return logvaldiffs;
  }

  Complex LogValDiff(VisibleConstType v, const std::vector<int> &tochange,
                     const std::vector<double> &newconf,
                     const LookupType &lt) override {
    Complex logvaldiff = 0.;

    if (tochange.size() != 0) {
      Complex logtsum = 0.5 * v.dot(lt.V(0));
      thetasnew_ = lt.V(0);
      Eigen::VectorXd vnew(v);

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

        thetasnew_ += W_.row(sf) * (newconf[s] - v(sf));
        vnew(sf) = newconf[s];
      }

      logvaldiff = 0.5 * vnew.dot(thetasnew_) - logtsum;
    }

    return logvaldiff;
  }

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

  void to_json(json &j) const override {
    j["Name"] = "JastrowSymm";
    j["Nvisible"] = nv_;
    j["Wsymm"] = Wsymm_;
  }

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

    if (FieldExists(pars, "Nvisible")) {
      nv_ = pars["Nvisible"];
    }
    if (nv_ != hilbert_.Size()) {
      throw InvalidInputError(
          "Number of visible units is incompatible with given "
          "Hilbert space");
    }

    Init(graph_);

    if (FieldExists(pars, "Wsymm")) {
      Wsymm_ = pars["Wsymm"];
    }

    SetBareParameters();
  }
};

}  // namespace netket

#endif


// Distributed under the MIT License.
// See LICENSE.txt for details.

#pragma once

#include <array>
#include <boost/optional.hpp>
#include <cstddef>
#include <limits>
#include <string>
#include <unordered_map>

#include "DataStructures/Tensor/TypeAliases.hpp"
#include "Utilities/TypeTraits.hpp"
/// \cond
class FunctionOfTime;
namespace PUP {
class er;
}  // namespace PUP
/// \endcond

namespace CoordMapsTimeDependent {
/*!
 * \ingroup CoordMapsTimeDependentGroup
 * \brief CubicScale map defined by \f$x = a(t)*\xi+(b(t)-a(t))\xi^3/X^2\f$,
 *  where \f$X\f$ is the outer boundary.
 *
 * The map scales the coordinates \f$\xi\f$ near the center by a factor
 * \f$a(t)\f$, while the coordinates near the outer boundary \f$X\f$, are scaled
 * by a factor \f$b(t)\f$. Here \f$a(t)\f$ and \f$b(t)\f$ are FunctionsOfTime.
 *
 * Currently, only a 1-D implementation.
 */
class CubicScale {
 public:
  static constexpr size_t dim = 1;

  explicit CubicScale(double outer_boundary) noexcept;
  CubicScale() = default;

  template <typename T>
  std::array<tt::remove_cvref_wrap_t<T>, 1> operator()(
      const std::array<T, 1>& source_coords, double time,
      const std::unordered_map<std::string, FunctionOfTime&>& map_list) const
      noexcept;

  /// Returns boost::none if the point is outside the range of the map.
  template <typename T>
  boost::optional<std::array<tt::remove_cvref_wrap_t<T>, 1>> inverse(
      const std::array<T, 1>& target_coords, double time,
      const std::unordered_map<std::string, FunctionOfTime&>& map_list) const
      noexcept;

  template <typename T>
  std::array<tt::remove_cvref_wrap_t<T>, 1> frame_velocity(
      const std::array<T, 1>& source_coords, double time,
      const std::unordered_map<std::string, FunctionOfTime&>& map_list) const
      noexcept;

  template <typename T>
  tnsr::Ij<tt::remove_cvref_wrap_t<T>, 1, Frame::NoFrame> inv_jacobian(
      const std::array<T, 1>& source_coords, double time,
      const std::unordered_map<std::string, FunctionOfTime&>& map_list) const
      noexcept;

  template <typename T>
  tnsr::Ij<tt::remove_cvref_wrap_t<T>, 1, Frame::NoFrame> jacobian(
      const std::array<T, 1>& source_coords, double time,
      const std::unordered_map<std::string, FunctionOfTime&>& map_list) const
      noexcept;

  template <typename T>
  tnsr::Iaa<tt::remove_cvref_wrap_t<T>, 1, Frame::NoFrame> hessian(
      const std::array<T, 1>& source_coords, double time,
      const std::unordered_map<std::string, FunctionOfTime&>& map_list) const
      noexcept;

  // clang-tidy: google-runtime-references
  void pup(PUP::er& p) noexcept;  // NOLINT

 private:
  friend bool operator==(const CubicScale& lhs, const CubicScale& rhs) noexcept;

  std::string f_of_t_a_ = "expansion_a";
  std::string f_of_t_b_ = "expansion_b";
  double outer_boundary_{std::numeric_limits<double>::signaling_NaN()};
};

inline bool operator!=(const CoordMapsTimeDependent::CubicScale& lhs,
                       const CoordMapsTimeDependent::CubicScale& rhs) noexcept {
  return not(lhs == rhs);
}

}  // namespace CoordMapsTimeDependent


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG (haftungsbeschränkt), www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also license.mtl.txt in the distribution.

#ifndef MTL_SQUARED_ABS_INCLUDE
#define MTL_SQUARED_ABS_INCLUDE

#include <cmath>
#include <complex>
#include <boost/mpl/or.hpp>
#include <boost/utility/enable_if.hpp>
#include <boost/type_traits/is_floating_point.hpp>
#include <boost/type_traits/is_integral.hpp>
#include <boost/numeric/mtl/interface/vpt.hpp>

namespace mtl {

/// When squaring magnitudes of intrinsic types the abs can be omitted
template <typename T>
typename boost::enable_if<boost::mpl::or_<boost::is_integral<T>, boost::is_floating_point<T> >, T>::type
inline squared_abs(const T& x)
{
    vampir_trace<10> tracer;
    return x * x;
}

/// Squaring complex numbers can be computed without square root
template <typename T>
T inline squared_abs(const std::complex<T>& z)
{
    vampir_trace<11> tracer;
    using std::norm;
    return norm(z);
}

/// When squaring magnitudes of non-intrinsic and non-complex types we use abs for security
template <typename T>
typename boost::disable_if<boost::mpl::or_<boost::is_integral<T>, boost::is_floating_point<T> >, T>::type
inline squared_abs(const T& x)
{
    vampir_trace<12> tracer;
    using std::abs;
    T a= abs(x);
    return a * a;
}

} // namespace mtl

#endif // MTL_SQUARED_ABS_INCLUDE


#include <iostream>
#include <Eigen/Dense>
#include <vector>
#include <cmath>
#include <thread>
#include <chrono>

#include "Open3D/Open3D.h"

using namespace open3d;
using namespace std;

using Eigen::Vector3d;

static vector<Vector3d> vv;
static float t_start;
static auto cloud_ptr = make_shared<geometry::PointCloud>();

void step(){
  Vector3d v;
  vv.clear();

  for (int t = 0; t < 2000; t++)
  {
    float T = t_start + t;
    v(2) = (float)t * 0.1;
    v(0) = -10 * sin(T / 1000.0 * 5 * M_PI);
    v(1) =  10 * cos(T/ 1000.0 * 5 * M_PI);
    vv.push_back(v);
  }
  t_start += 50;
  cloud_ptr->points_ = vv;
}

int main()
{

  t_start = 0;
  step();

  visualization::DrawGeometriesWithAnimationCallback(
      {cloud_ptr},
      [&](visualization::Visualizer *vis) {
        step();
        vis->UpdateGeometry();
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
        return false;
      },
      "CloudAnim", 640, 480);
}


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

#include <Eigen/Eigenvalues>

#include <tvm/task_dynamics/ProportionalDerivative.h>
#include <tvm/function/abstract/Function.h>

namespace tvm
{

namespace task_dynamics
{

  ProportionalDerivative::ProportionalDerivative(double kp, double kv)
    : kp_(kp)
    , kv_(kv)
  {
  }

  ProportionalDerivative::ProportionalDerivative(const Eigen::VectorXd& kp, const Eigen::VectorXd& kv)
    : kp_(kp)
    , kv_(kv)
  {
  }

  ProportionalDerivative::ProportionalDerivative(const Eigen::MatrixXd& kp, const Eigen::MatrixXd& kv)
    : kp_(kp)
    , kv_(kv)
  {
  }

  ProportionalDerivative::ProportionalDerivative(double kp, const Eigen::VectorXd& kv)
    : kp_(kp)
    , kv_(kv)
  {
  }

  ProportionalDerivative::ProportionalDerivative(double kp, const Eigen::MatrixXd& kv)
    : kp_(kp)
    , kv_(kv)
  {
  }

  ProportionalDerivative::ProportionalDerivative(const Eigen::VectorXd& kp, double kv)
    : kp_(kp)
    , kv_(kv)
  {
  }

  ProportionalDerivative::ProportionalDerivative(const Eigen::VectorXd& kp, const Eigen::MatrixXd& kv)
    : kp_(kp)
    , kv_(kv)
  {
  }

  ProportionalDerivative::ProportionalDerivative(const Eigen::MatrixXd& kp, double kv)
    : kp_(kp)
    , kv_(kv)
  {
  }

  ProportionalDerivative::ProportionalDerivative(const Eigen::MatrixXd& kp, const Eigen::VectorXd& kv)
    : kp_(kp)
    , kv_(kv)
  {
  }

  ProportionalDerivative::ProportionalDerivative(double kp)
    : ProportionalDerivative(kp, 2 * std::sqrt(kp))
  {
  }

  ProportionalDerivative::ProportionalDerivative(const Eigen::VectorXd& kp)
    : kp_(kp)
    , kv_(mpark::in_place_index<1>, 2*kp.cwiseSqrt())
  {
  }

  ProportionalDerivative::ProportionalDerivative(const Eigen::MatrixXd& kp)
    : kp_(kp)
    , kv_(Eigen::MatrixXd(kp.rows(), kp.cols()))
  {
    Eigen::RealSchur<Eigen::MatrixXd> dec(kp);
    assert(dec.matrixT().isDiagonal(1e-8) && "kp is not symmetric.");
    assert((dec.matrixT().diagonal().array() >= 0).all() && "kp is undefinite.");
    kv_.emplace<2>(2 * dec.matrixU() * dec.matrixT().diagonal().asDiagonal() * dec.matrixU().transpose());
  }

  std::unique_ptr<abstract::TaskDynamicsImpl> ProportionalDerivative::impl_(FunctionPtr f, constraint::Type t, const Eigen::VectorXd& rhs) const
  {
    return std::make_unique<Impl>(f, t, rhs, kp_, kv_);
  }

  Order ProportionalDerivative::order_() const
  {
    return Order::Two;
  }

  ProportionalDerivative::Impl::Impl(FunctionPtr f, constraint::Type t, const Eigen::VectorXd& rhs, const Gain& kp, const Gain& kv)
    : TaskDynamicsImpl(Order::Two, f, t, rhs)
    , kp_(kp)
    , kv_(kv)
  {
    assert((kp.index() == 0                                           // Scalar gain
        || (kp.index() == 1 && mpark::get<Eigen::VectorXd>(kp).size() == f->size())   // Diagonal gain
        || (kp.index() == 2 && mpark::get<Eigen::MatrixXd>(kp).cols() == f->size()    // Matrix gain
                            && mpark::get<Eigen::MatrixXd>(kp).rows() == f->size()))
      && "Gain kp and function have incompatible sizes");

    assert((kv.index() == 0                                           // Scalar gain
        || (kv.index() == 1 && mpark::get<Eigen::VectorXd>(kv).size() == f->size())   // Diagonal gain
        || (kv.index() == 2 && mpark::get<Eigen::MatrixXd>(kv).cols() == f->size()    // Matrix gain
                            && mpark::get<Eigen::MatrixXd>(kv).rows() == f->size()))
      && "Gain kv and function have incompatible sizes");
  }

  void ProportionalDerivative::Impl::updateValue()
  {
    switch (kv_.index())
    {
    case 0: value_ = -mpark::get<double>(kv_) * function().velocity(); break;
    case 1: value_.noalias() = -(mpark::get<Eigen::VectorXd>(kv_).asDiagonal() * function().velocity()); break;
    case 2: value_.noalias() = -mpark::get<Eigen::MatrixXd>(kv_) * function().velocity(); break;
    default: assert(false);
    }
    switch (kp_.index())
    {
    case 0: value_ -= mpark::get<double>(kp_) * (function().value() - rhs()); break;
    case 1: value_.noalias() -= mpark::get<Eigen::VectorXd>(kp_).asDiagonal() * (function().value() - rhs()); break;
    case 2: value_.noalias() -= mpark::get<Eigen::MatrixXd>(kp_) * (function().value() - rhs()); break;
    default: assert(false);
    }
  }

  std::pair<const PD::Gain&, const PD::Gain&> ProportionalDerivative::Impl::gains() const
  {
    return {kp_, kv_};
  }

  void ProportionalDerivative::Impl::gains(double kp, double kv)
  {
    gains_(kp, kv);
  }

  void ProportionalDerivative::Impl::gains(const Eigen::VectorXd& kp, const Eigen::VectorXd& kv)
  {
    gains_(kp,kv);
  }

  void ProportionalDerivative::Impl::gains(const Eigen::MatrixXd& kp, const Eigen::MatrixXd& kv)
  {
    gains_(kp,kv);
  }

  void ProportionalDerivative::Impl::gains(double kp, const Eigen::VectorXd& kv)
  {
    gains_(kp,kv);
  }

  void ProportionalDerivative::Impl::gains(double kp, const Eigen::MatrixXd& kv)
  {
    gains_(kp,kv);
  }

  void ProportionalDerivative::Impl::gains(const Eigen::VectorXd& kp, double kv)
  {
    gains_(kp,kv);
  }

  void ProportionalDerivative::Impl::gains(const Eigen::VectorXd& kp, const Eigen::MatrixXd& kv)
  {
    gains_(kp,kv);
  }

  void ProportionalDerivative::Impl::gains(const Eigen::MatrixXd& kp, double kv)
  {
    gains_(kp,kv);
  }

  void ProportionalDerivative::Impl::gains(const Eigen::MatrixXd& kp, const Eigen::VectorXd& kv)
  {
    gains_(kp,kv);
  }

  void ProportionalDerivative::Impl::gains(double kp)
  {
    kp_ = kp;
    kv_ = 2 * std::sqrt(kp);
  }

  void ProportionalDerivative::Impl::gains(const Eigen::VectorXd& kp)
  {
    checkGainSize(kp);
    kp_ = kp;
    kv_.emplace<1>(2 * kp.cwiseSqrt());
  }

  void ProportionalDerivative::Impl::gains(const Eigen::MatrixXd& kp)
  {
    checkGainSize(kp);
    Eigen::RealSchur<Eigen::MatrixXd> dec(kp);
    assert(dec.matrixT().isDiagonal(1e-8) && "kp is not symmetric.");
    assert((dec.matrixT().diagonal().array() >= 0).all() && "kp is undefinite.");
    kv_.emplace<2>(2*dec.matrixU() * dec.matrixT().diagonal().asDiagonal() * dec.matrixU().transpose());
  }

  void ProportionalDerivative::Impl::checkGainSize(double k) const
  {
    //do nothing
  }

  void ProportionalDerivative::Impl::checkGainSize(const Eigen::VectorXd& k) const
  {
    if (k.size() != function().size())
    {
      throw std::runtime_error("[task_dynamics::ProportionalDerivative::Impl::gains] Gain and function have incompatible sizes.");
    }
  }

  void ProportionalDerivative::Impl::checkGainSize(const Eigen::MatrixXd& k) const
  {
    if (k.rows() != function().size() || k.cols() != function().size())
    {
      throw std::runtime_error("[task_dynamics::ProportionalDerivative::Impl::gains] Gain and function have incompatible sizes.");
    }
  }

}  // namespace task_dynamics

}  // namespace tvm


//
// Created by david on 2018-08-06.
//

#include "class_eigsolver_armadillo.h"
#include <armadillo>

void class_eigsolver_armadillo::subtract_phase(std::complex<double>  *eigvecs, int L) {
        using namespace std::complex_literals;
        auto uL = static_cast<unsigned int> (L);
        arma::cx_mat arma_eigvecs(eigvecs,uL,uL,false,false);
        for (int i = 0; i < L; i++) {
            auto begin = arma_eigvecs.begin() + i * L;
            auto end = begin + L;
            std::complex<double> inv_phase = -1.0i * std::arg(arma_eigvecs[i * L]);
            std::complex<double> exp_inv_phase = std::exp(inv_phase);
            std::transform(begin, end, begin,
                           [exp_inv_phase](std::complex<double> num) -> std::complex<double>
                           {
                               auto result =  num * exp_inv_phase;
                               if (std::abs(result.imag()) < 1e-15){result = result.real();}
                               return result;
                           });
        }
}

void class_eigsolver_armadillo::eig_sym(int L, std::complex<double> *hermitian_matrix, double *eigvals, std::complex<double> *eigvecs, bool remove_phase){
    // for matrices with complex elements
    auto uL = static_cast<unsigned int> (L);
    arma::cx_mat arma_matrix(hermitian_matrix, uL, uL ,false, false);
    arma::vec    arma_eigvals(eigvals,uL,false,false);
    arma::cx_mat arma_eigvecs(eigvecs,uL,uL,false,false);
    arma::eig_sym(arma_eigvals, arma_eigvecs, arma_matrix, "dc");
    if (remove_phase){
        subtract_phase(eigvecs, L);
    }
}

void class_eigsolver_armadillo::eig_sym(int L, double *symmetric_matrix, double *eigvals, double *eigvecs, bool remove_phase){
    // for matrices with complex elements
    auto uL = static_cast<unsigned int> (L);
    arma::mat arma_matrix(symmetric_matrix, uL, uL ,false, false);
    arma::vec    arma_eigvals(eigvals,uL,false,false);
    arma::mat arma_eigvecs(eigvecs,uL,uL,false,false);
    arma::eig_sym(arma_eigvals, arma_eigvecs, arma_matrix, "dc");
}


#if GOOGLE_CUDA

#define EIGEN_USE_GPU

#include <stdio.h>
#include <cfloat>
#include <vector>
#include <ctime>
#include <cstdlib>
#include <thrust/device_vector.h>
#include <Eigen/Geometry>
#include "average_distance_loss_op_gpu.h"

#define POSE_CHANNELS 4

#define CUDA_1D_KERNEL_LOOP(i, n)                            \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \
       i += blockDim.x * gridDim.x)

// namespace tensorflow {
using namespace tensorflow;

inline
cudaError_t checkCuda(cudaError_t result)
{
  if (result != cudaSuccess) {
    fprintf(stderr, "CUDA Runtime Error: %s\n",
            cudaGetErrorString(result));
    assert(result == cudaSuccess);
  }
  return result;
}

template <typename Dtype>
__global__ void AveragedistanceForward(const int nthreads, const Dtype* prediction, const Dtype* target,
    const Dtype* weight, const Dtype* point, const Dtype* symmetry, const int batch_size, const int num_classes,
    const int num_points, const float margin, Dtype* rotations, Dtype* losses, Dtype* diffs)
{
  CUDA_1D_KERNEL_LOOP(index_thread, nthreads)
  {
    // batch index
    int n = index_thread / num_points;
    int p = index_thread % num_points;

    // find the class label and pose of this object
    int index_cls = -1, ind;
    Dtype s, u, v, w;
    for (int i = 0; i < POSE_CHANNELS * num_classes; i += POSE_CHANNELS)
    {
      int index = n * POSE_CHANNELS * num_classes + i;
      if (weight[index] > 0)
      {
        index_cls = i / POSE_CHANNELS;

        // gt quaternion
        s = target[index + 0];
        u = target[index + 1];
        v = target[index + 2];
        w = target[index + 3];

        // gt rotation matrix
        ind = n * num_points * 6 * 9 + p * 6 * 9;
        rotations[ind + 0] = s * s + u * u - v * v - w * w;
        rotations[ind + 1] = 2 * (u * v - s * w);
        rotations[ind + 2] = 2 * (u * w + s * v);
        rotations[ind + 3] = 2 * (u * v + s * w);
        rotations[ind + 4] = s * s - u * u + v * v - w * w;
        rotations[ind + 5] = 2 * (v * w - s * u);
        rotations[ind + 6] = 2 * (u * w - s * v);
        rotations[ind + 7] = 2 * (v * w + s * u);
        rotations[ind + 8] = s * s - u * u - v * v + w * w;

        // predicted quaternion
        s = prediction[index + 0];
        u = prediction[index + 1];
        v = prediction[index + 2];
        w = prediction[index + 3];

        // predicted rotation matrix
        ind = n * num_points * 6 * 9 + p * 6 * 9 + 9;
        rotations[ind + 0] = s * s + u * u - v * v - w * w;
        rotations[ind + 1] = 2 * (u * v - s * w);
        rotations[ind + 2] = 2 * (u * w + s * v);
        rotations[ind + 3] = 2 * (u * v + s * w);
        rotations[ind + 4] = s * s - u * u + v * v - w * w;
        rotations[ind + 5] = 2 * (v * w - s * u);
        rotations[ind + 6] = 2 * (u * w - s * v);
        rotations[ind + 7] = 2 * (v * w + s * u);
        rotations[ind + 8] = s * s - u * u - v * v + w * w;
        break;
      }
    }
    if (index_cls == -1)
      continue;

    // derivatives of Ru to quaternion
    ind = n * num_points * 6 * 9 + p * 6 * 9 + 18;
    rotations[ind + 0] = 2 * s;
    rotations[ind + 1] = -2 * w;
    rotations[ind + 2] = 2 * v;
    rotations[ind + 3] = 2 * w;
    rotations[ind + 4] = 2 * s;
    rotations[ind + 5] = -2 * u;
    rotations[ind + 6] = -2 * v;
    rotations[ind + 7] = 2 * u;
    rotations[ind + 8] = 2 * s;

    ind = n * num_points * 6 * 9 + p * 6 * 9 + 27;
    rotations[ind + 0] = 2 * u;
    rotations[ind + 1] = 2 * v;
    rotations[ind + 2] = 2 * w;
    rotations[ind + 3] = 2 * v;
    rotations[ind + 4] = -2 * u;
    rotations[ind + 5] = -2 * s;
    rotations[ind + 6] = 2 * w;
    rotations[ind + 7] = 2 * s;
    rotations[ind + 8] = -2 * u;

    ind = n * num_points * 6 * 9 + p * 6 * 9 + 36;
    rotations[ind + 0] = -2 * v;
    rotations[ind + 1] = 2 * u;
    rotations[ind + 2] = 2 * s;
    rotations[ind + 3] = 2 * u;
    rotations[ind + 4] = 2 * v;
    rotations[ind + 5] = 2 * w;
    rotations[ind + 6] = -2 * s;
    rotations[ind + 7] = 2 * w;
    rotations[ind + 8] = -2 * v;

    ind = n * num_points * 6 * 9 + p * 6 * 9 + 45;
    rotations[ind + 0] = -2 * w;
    rotations[ind + 1] = -2 * s;
    rotations[ind + 2] = 2 * u;
    rotations[ind + 3] = 2 * s;
    rotations[ind + 4] = -2 * w;
    rotations[ind + 5] = 2 * v;
    rotations[ind + 6] = 2 * u;
    rotations[ind + 7] = 2 * v;
    rotations[ind + 8] = 2 * w;

    // for the point
    int index = index_cls * num_points * 3 + p * 3;
    ind = n * num_points * 6 * 9 + p * 6 * 9;

    // rotate the first point
    Dtype x1 = rotations[ind + 9 + 0] * point[index + 0] + rotations[ind + 9 + 1] * point[index + 1] + rotations[ind + 9 + 2] * point[index + 2];
    Dtype y1 = rotations[ind + 9 + 3] * point[index + 0] + rotations[ind + 9 + 4] * point[index + 1] + rotations[ind + 9 + 5] * point[index + 2];
    Dtype z1 = rotations[ind + 9 + 6] * point[index + 0] + rotations[ind + 9 + 7] * point[index + 1] + rotations[ind + 9 + 8] * point[index + 2];

    int index_min;
    Dtype x2, y2, z2;
    if (symmetry[index_cls] > 0)
    {
      // find the closet point for symmetry object
      Dtype dmin = FLT_MAX;
      for (int i = 0; i < num_points; i++)
      {
        int index2 = index_cls * num_points * 3 + i * 3;
        x2 = rotations[ind + 0] * point[index2 + 0] + rotations[ind + 1] * point[index2 + 1] + rotations[ind + 2] * point[index2 + 2];
        y2 = rotations[ind + 3] * point[index2 + 0] + rotations[ind + 4] * point[index2 + 1] + rotations[ind + 5] * point[index2 + 2];
        z2 = rotations[ind + 6] * point[index2 + 0] + rotations[ind + 7] * point[index2 + 1] + rotations[ind + 8] * point[index2 + 2];
        Dtype distance = (x1 - x2) * (x1 - x2) + (y1 - y2) * (y1 - y2) + (z1 - z2) * (z1 - z2);
        if (distance < dmin)
        {
          dmin = distance;
          index_min = index2;
        }
      }
    }
    else
      index_min = index;

    x2 = rotations[ind + 0] * point[index_min + 0] + rotations[ind + 1] * point[index_min + 1] + rotations[ind + 2] * point[index_min + 2];
    y2 = rotations[ind + 3] * point[index_min + 0] + rotations[ind + 4] * point[index_min + 1] + rotations[ind + 5] * point[index_min + 2];
    z2 = rotations[ind + 6] * point[index_min + 0] + rotations[ind + 7] * point[index_min + 1] + rotations[ind + 8] * point[index_min + 2];

    Dtype distance = ((x1 - x2) * (x1 - x2) + (y1 - y2) * (y1 - y2) + (z1 - z2) * (z1 - z2));
    if (distance < margin)
      continue;

    losses[index_thread] = (distance - margin) / (2.0 * batch_size * num_points);

    int index_diff = n * num_points * POSE_CHANNELS * num_classes + p * POSE_CHANNELS * num_classes + POSE_CHANNELS * index_cls;
    for (int j = 0; j < 3; j++)
    {
      Dtype diff;
      if (j == 0)
        diff = x1 - x2;
      else if (j == 1)
        diff = y1 - y2;
      else
        diff = z1 - z2;
      for (int k = 0; k < 3; k++)
      {
        ind = n * num_points * 6 * 9 + p * 6 * 9 + 18;
        diffs[index_diff + 0] += diff * point[index + k] * rotations[ind + j * 3 + k] / (batch_size * num_points);
        ind = n * num_points * 6 * 9 + p * 6 * 9 + 27;
        diffs[index_diff + 1] += diff * point[index + k] * rotations[ind + j * 3 + k] / (batch_size * num_points);
        ind = n * num_points * 6 * 9 + p * 6 * 9 + 36;
        diffs[index_diff + 2] += diff * point[index + k] * rotations[ind + j * 3 + k] / (batch_size * num_points);
        ind = n * num_points * 6 * 9 + p * 6 * 9 + 45;
        diffs[index_diff + 3] += diff * point[index + k] * rotations[ind + j * 3 + k] / (batch_size * num_points);
      }
    }
  }
}

template <typename Dtype>
__global__ void sum_losses_gradients(const int nthreads, const Dtype* losses, const Dtype* diffs, const int batch_size,
    const int num_classes, const int num_points, Dtype* loss_batch, Dtype* bottom_diff)
{
  CUDA_1D_KERNEL_LOOP(index, nthreads)
  {
    int n = index / (POSE_CHANNELS * num_classes);
    int c = index % (POSE_CHANNELS * num_classes);
/*
    // find the most violated point
    Dtype lmax = -FLT_MAX;
    int pmax;
    for (int p = 0; p < num_points; p++)
    {
      if (losses[n * num_points + p] > lmax)
      {
        lmax = losses[n * num_points + p];
        pmax = p;
      }
    }

    int index_diff = n * num_points * POSE_CHANNELS * num_classes + pmax * POSE_CHANNELS * num_classes + c;
    bottom_diff[index] = diffs[index_diff] * num_points;

    if (c == 0)
      loss_batch[n] = lmax * num_points;
*/

    bottom_diff[index] = 0;
    for (int p = 0; p < num_points; p++)
    {
      int index_diff = n * num_points * POSE_CHANNELS * num_classes + p * POSE_CHANNELS * num_classes + c;
      bottom_diff[index] += diffs[index_diff];
    }

    if (c == 0)
    {
      loss_batch[n] = 0;
      for (int p = 0; p < num_points; p++)
        loss_batch[n] += losses[n * num_points + p];
    }

  }
}

// bottom_data: (batch_size, 4 * num_classes)
void AveragedistanceForwardLaucher(OpKernelContext* context,
    const float* bottom_prediction, const float* bottom_target, const float* bottom_weight, const float* bottom_point,
    const float* bottom_symmetry, const int batch_size, const int num_classes, const int num_points, const float margin,
    float* top_data, float* bottom_diff, const Eigen::GpuDevice& d)
{
  // run kernels
  cudaError_t err;
  const int kThreadsPerBlock = 512;
  int output_size;

  // temp losses
  int dims[2];
  dims[0] = batch_size;
  dims[1] = num_points;
  TensorShape output_shape_losses;
  TensorShapeUtils::MakeShape(dims, 2, &output_shape_losses);
  Tensor losses_tensor;
  OP_REQUIRES_OK(context, context->allocate_temp(DT_FLOAT, output_shape_losses, &losses_tensor));
  float* losses = losses_tensor.flat<float>().data();
  checkCuda(cudaMemset(losses, 0, batch_size * num_points * sizeof(float)));

  TensorShape output_shape_loss_batch;
  TensorShapeUtils::MakeShape(&batch_size, 1, &output_shape_loss_batch);
  Tensor loss_batch_tensor;
  OP_REQUIRES_OK(context, context->allocate_temp(DT_FLOAT, output_shape_loss_batch, &loss_batch_tensor));
  float* loss_batch = loss_batch_tensor.flat<float>().data();
  checkCuda(cudaMemset(loss_batch, 0, batch_size * sizeof(float)));

  // temp diffs
  int dims_diff[3];
  dims_diff[0] = batch_size;
  dims_diff[1] = num_points;
  dims_diff[2] = POSE_CHANNELS * num_classes;
  TensorShape output_shape_diff;
  TensorShapeUtils::MakeShape(dims_diff, 3, &output_shape_diff);
  Tensor diffs_tensor;
  OP_REQUIRES_OK(context, context->allocate_temp(DT_FLOAT, output_shape_diff, &diffs_tensor));
  float* diffs = diffs_tensor.flat<float>().data();
  checkCuda(cudaMemset(diffs, 0, batch_size * num_points * POSE_CHANNELS * num_classes * sizeof(float)));

  // temp rotations
  int dims_rot[3];
  dims_rot[0] = batch_size;
  dims_rot[1] = num_points;
  dims_rot[2] = 6 * 9;
  TensorShape output_shape_rotations;
  TensorShapeUtils::MakeShape(dims_rot, 3, &output_shape_rotations);
  Tensor rotations_tensor;
  OP_REQUIRES_OK(context, context->allocate_temp(DT_FLOAT, output_shape_rotations, &rotations_tensor));
  float* rotations = rotations_tensor.flat<float>().data();
  checkCuda(cudaMemset(rotations, 0, batch_size * num_points * 6 * 9 * sizeof(float)));

  // compute the losses and gradients
  output_size = batch_size * num_points;
  AveragedistanceForward<<<(output_size + kThreadsPerBlock - 1) / kThreadsPerBlock,
                       kThreadsPerBlock, 0, d.stream()>>>(
      output_size, bottom_prediction, bottom_target, bottom_weight, bottom_point, bottom_symmetry,
      batch_size, num_classes, num_points, margin, rotations, losses, diffs);
  cudaDeviceSynchronize();

  err = cudaGetLastError();
  if(cudaSuccess != err)
  {
    fprintf( stderr, "cudaCheckError() failed: %s\n", cudaGetErrorString( err ) );
    exit( -1 );
  }

  // sum the diffs
  checkCuda(cudaMemset(bottom_diff, 0, batch_size * POSE_CHANNELS * num_classes * sizeof(float)));
  output_size = batch_size * POSE_CHANNELS * num_classes;
  sum_losses_gradients<<<(output_size + kThreadsPerBlock - 1) / kThreadsPerBlock,
                       kThreadsPerBlock, 0, d.stream()>>>(
      output_size, losses, diffs, batch_size, num_classes, num_points, loss_batch, bottom_diff);
  cudaDeviceSynchronize();

  // sum the loss
  checkCuda(cudaMemset(top_data, 0, sizeof(float)));
  thrust::device_ptr<float> losses_ptr(loss_batch);
  float loss = thrust::reduce(losses_ptr, losses_ptr + batch_size);
  cudaMemcpy(top_data, &loss, sizeof(float), cudaMemcpyHostToDevice);

  err = cudaGetLastError();
  if(cudaSuccess != err)
  {
    fprintf( stderr, "cudaCheckError() failed: %s\n", cudaGetErrorString( err ) );
    exit( -1 );
  }
}

template <typename Dtype>
__global__ void AveragedistanceBackward(const int nthreads, const Dtype* top_diff,
    const Dtype* bottom_diff, Dtype* output)
{
  CUDA_1D_KERNEL_LOOP(index, nthreads)
  {
    output[index] = top_diff[0] * bottom_diff[index];
  }
}

bool AveragedistanceBackwardLaucher(const float* top_diff, const float* bottom_diff, const int batch_size,
    const int channels, float* output, const Eigen::GpuDevice& d)
{
  const int kThreadsPerBlock = 512;
  const int output_size = batch_size * channels;
  cudaError_t err;

  AveragedistanceBackward<<<(output_size + kThreadsPerBlock - 1) / kThreadsPerBlock,
                       kThreadsPerBlock, 0, d.stream()>>>(
      output_size, top_diff, bottom_diff, output);

  cudaDeviceSynchronize();
  err = cudaGetLastError();
  if(cudaSuccess != err)
  {
    fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) );
    exit( -1 );
  }

  return d.ok();
}

// }  // namespace tensorflow

#endif  // GOOGLE_CUDA