// Copyright Yamaha 2021
// MIT License
// https://github.com/yamaha-bps/cbr_control/blob/master/LICENSE

#ifndef CBR_CONTROL__MPC__DLTV_OCP_HPP_
#define CBR_CONTROL__MPC__DLTV_OCP_HPP_

#include <Eigen/Dense>

#include <cbr_utils/utils.hpp>

#include <utility>

#include "ocp_common.hpp"

namespace cbr
{

/* ---------------------------------------------------------------------------------------------- */
/*               Continuous Time Varying Linear Optimal Control Problem Discretizer               */
/* ---------------------------------------------------------------------------------------------- */

template<typename cltv_pb_t, std::size_t _nPts = 100, std::size_t exp_order = 4>
class DltvOcp
{
public:
  // Must be defined in ctlv_pb problem
  constexpr static std::size_t nx = cltv_pb_t::nx;
  constexpr static std::size_t nu = cltv_pb_t::nu;
  constexpr static std::size_t nPts = _nPts;

  using problem_t = cltv_pb_t;

  // Create some useful aliases
  using state_t = Eigen::Matrix<double, nx, 1>;
  using input_t = Eigen::Matrix<double, nu, 1>;
  using time_t = Eigen::Matrix<double, nPts, 1>;
  using A_t = Eigen::Matrix<double, nx, nx>;
  using B_t = Eigen::Matrix<double, nx, nu>;
  using Q_t = Eigen::Matrix<double, nx, nx>;
  using R_t = Eigen::Matrix<double, nu, nu>;

  // Get return type of problem functions
  using Ar_t = std::result_of_t<decltype(&cltv_pb_t::get_A)(cltv_pb_t, double)>;
  using Br_t = std::result_of_t<decltype(&cltv_pb_t::get_B)(cltv_pb_t, double)>;
  using Qr_t = std::result_of_t<decltype(&cltv_pb_t::get_Q)(cltv_pb_t, double)>;
  using QTr_t = std::result_of_t<decltype(&cltv_pb_t::get_QT)(cltv_pb_t)>;
  using Rr_t = std::result_of_t<decltype(&cltv_pb_t::get_R)(cltv_pb_t, double)>;

  /* -------------------------------------------------------------------------- */
  /*                                  Optionals                                 */
  /* -------------------------------------------------------------------------- */

  // Check existance of get_E function
  constexpr static bool has_E_approx = std::experimental::is_detected_v<
    ocp_detail::has_E_continuous, cltv_pb_t>;
  constexpr static bool has_E = std::experimental::is_detected_exact_v<
    state_t, ocp_detail::has_E_continuous, cltv_pb_t>||
    std::experimental::is_detected_exact_v<
    const state_t &, ocp_detail::has_E_continuous, cltv_pb_t>;
  using Er_t = std::experimental::detected_or_t<
    state_t, ocp_detail::has_E_continuous, cltv_pb_t>;
  static_assert(
    !(has_E_approx && !has_E),
    "Detected get_E function doesn't have a correct return type. "
    "It must be an nx*1 Eigen::Matrix (or a const reference to one)");

  // Check existance of get_q function
  constexpr static bool has_q_approx = std::experimental::is_detected_v<
    ocp_detail::has_q_continuous, cltv_pb_t>;
  constexpr static bool has_q = std::experimental::is_detected_exact_v<
    state_t, ocp_detail::has_q_continuous, cltv_pb_t>||
    std::experimental::is_detected_exact_v<
    const state_t &, ocp_detail::has_q_continuous, cltv_pb_t>;
  using qr_t = std::experimental::detected_or_t<
    state_t, ocp_detail::has_q_continuous, cltv_pb_t>;
  static_assert(
    !(has_q_approx && !has_q),
    "Detected get_q function doesn't have a correct return type. "
    "It must be an nx*1 Eigen::Matrix (or a const reference to one)");

  // Check existance of get_qT function
  constexpr static bool has_qT_approx = std::experimental::is_detected_v<
    ocp_detail::has_qT_continuous, cltv_pb_t>;
  constexpr static bool has_qT = std::experimental::is_detected_exact_v<
    state_t, ocp_detail::has_qT_continuous, cltv_pb_t>||
    std::experimental::is_detected_exact_v<
    const state_t &, ocp_detail::has_qT_continuous, cltv_pb_t>;

  using qTr_t = std::experimental::detected_or_t<
    state_t, ocp_detail::has_qT_continuous, cltv_pb_t>; \
  static_assert(
    !(has_qT_approx && !has_qT),
    "Detected get_qT function doesn't have a correct return type. "
    "It must be an nx*1 Eigen::Matrix (or a const reference to one)");

  // Check existance of get_r function
  constexpr static bool has_r_approx = std::experimental::is_detected_v<
    ocp_detail::has_r_continuous, cltv_pb_t>;
  constexpr static bool has_r = std::experimental::is_detected_exact_v<
    input_t, ocp_detail::has_r_continuous, cltv_pb_t>||
    std::experimental::is_detected_exact_v<
    const input_t &, ocp_detail::has_r_continuous, cltv_pb_t>;
  using rr_t = std::experimental::detected_or_t<
    input_t, ocp_detail::has_r_continuous, cltv_pb_t>;
  static_assert(
    !(has_r_approx && !has_r),
    "Detected get_r function doesn't have a correct return type. "
    "It must be an nu*1 Eigen::Matrix (or a const reference to one)");

  // Check problem dimensions
  static_assert(nx > 0, "Number of states must be > 0.");
  static_assert(nu > 0, "Number of inputs must be > 0.");
  static_assert(nPts > 1, "Number of trajectory points must be > 1.");
  static_assert(exp_order < 20, "Exponential order must be < 20.");

  // Check return type of problem functions
  static_assert(
    std::is_same_v<std::decay_t<Ar_t>, A_t>,
    "The get_A method of the problem must return an nx*nx Eigen::Matrix (or a reference to one).");
  static_assert(
    std::is_same_v<std::decay_t<Br_t>, B_t>,
    "The get_B method of the problem must return an nx*nu Eigen::Matrix (or a reference to one).");
  static_assert(
    std::is_same_v<std::decay_t<Qr_t>, Q_t>,
    "The get_Q method of the problem must return an nx*nx Eigen::Matrix (or a reference to one).");
  static_assert(
    std::is_same_v<std::decay_t<qr_t>, state_t>,
    "The get_q method of the problem must return an nx*1 Eigen::Matrix (or a reference to one).");
  static_assert(
    std::is_same_v<std::decay_t<QTr_t>, Q_t>,
    "The get_QT method of the problem must return an nx*nx Eigen::Matrix (or a reference to one).");
  static_assert(
    std::is_same_v<std::decay_t<qTr_t>, state_t>,
    "The get_qT method of the problem must return an nx*1 Eigen::Matrix (or a reference to one).");
  static_assert(
    std::is_same_v<std::decay_t<Rr_t>, R_t>,
    "The get_R method of the problem must return an nu*nu Eigen::Matrix (or a reference to one).");
  static_assert(
    std::is_same_v<std::decay_t<rr_t>, input_t>,
    "The get_r method of the problem must return an nu*1 Eigen::Matrix (or a reference to one).");

public:
  DltvOcp() = delete;
  DltvOcp(const DltvOcp &) = default;
  DltvOcp(DltvOcp &&) = default;
  DltvOcp & operator=(const DltvOcp &) = default;
  DltvOcp & operator=(DltvOcp &&) = default;

  explicit DltvOcp(const cltv_pb_t & pb)
  : cltv_pb_(pb),
    dt_{compute_dt()}
  {}

  explicit DltvOcp(cltv_pb_t && pb)
  : cltv_pb_(std::move(pb)),
    dt_{compute_dt()}
  {}

  template<typename T1>
  DltvOcp(T1 && pb)
  : cltv_pb_(std::forward<T1>(pb)),
    dt_{compute_dt()}
  {}

  void get_x0(Eigen::Ref<state_t> x0) const
  {
    cltv_pb_.get_x0(x0);
  }

  void get_state_lb(std::size_t k, Eigen::Ref<state_t> state_lb) const
  {
    cltv_pb_.get_state_lb(indexToTime(k), state_lb);
  }

  void get_state_ub(std::size_t k, Eigen::Ref<state_t> state_ub) const
  {
    cltv_pb_.get_state_ub(indexToTime(k), state_ub);
  }

  void get_input_lb(std::size_t k, Eigen::Ref<input_t> input_lb) const
  {
    cltv_pb_.get_input_lb(indexToTime(k), input_lb);
  }

  void get_input_ub(std::size_t k, Eigen::Ref<input_t> input_ub) const
  {
    cltv_pb_.get_input_ub(indexToTime(k), input_ub);
  }

  A_t get_A(std::size_t k) const
  {
    // define identity for order = 0
    A_t expA = A_t::Identity();

    if constexpr (exp_order > 0) {
      const A_t Adt = dt_ * cltv_pb_.get_A(indexToTime(k));
      expA += Adt;
      if constexpr (exp_order > 1) {
        double c = 1.;
        A_t Adtp = Adt;
        for (std::size_t i = 2; i <= exp_order; i++) {
          Adtp *= Adt;
          c /= static_cast<double>(i);
          expA += c * Adtp;
        }
      }
    }

    return expA;
  }

  B_t get_B(std::size_t k) const
  {
    // define identity for order = 0
    A_t expA = A_t::Identity();

    if constexpr (exp_order > 0) {
      const A_t Adt = dt_ * cltv_pb_.get_A(indexToTime(k));
      expA += Adt / 2.;
      if constexpr (exp_order > 1) {
        double c = 0.5;
        A_t Adtp = Adt / 2.;
        for (std::size_t i = 2; i <= exp_order; i++) {
          Adtp *= Adt;
          c /= static_cast<double>(i + 1);
          expA += c * Adtp;
        }
      }
    }

    B_t expB = expA * cltv_pb_.get_B(indexToTime(k)) * dt_;

    return expB;
  }

  template<typename T = std::size_t>
  state_t get_E(std::enable_if_t<has_E, T> k)
  {
    // define identity for order = 0
    A_t expA = A_t::Identity();

    if constexpr (exp_order > 0) {
      const A_t Adt = dt_ * cltv_pb_.get_A(indexToTime(k));
      expA += Adt / 2.;
      if constexpr (exp_order > 1) {
        double c = 0.5;
        A_t Adtp = Adt / 2.;
        for (std::size_t i = 2; i <= exp_order; i++) {
          Adtp *= Adt;
          c /= static_cast<double>(i + 1);
          expA += c * Adtp;
        }
      }
    }

    state_t expE = expA * cltv_pb_.get_E(indexToTime(k)) * dt_;
    return expE;
  }

  Q_t get_Q(std::size_t k) const
  {
    return dt_ * cltv_pb_.get_Q(indexToTime(k));
  }

  template<typename T = std::size_t>
  state_t get_q(std::enable_if_t<has_q, T> k) const
  {
    return dt_ * cltv_pb_.get_q(indexToTime(k));
  }

  R_t get_R(std::size_t k) const
  {
    return dt_ * cltv_pb_.get_R(indexToTime(k));
  }

  template<typename T = std::size_t>
  input_t get_r(std::enable_if_t<has_r, T> k) const
  {
    return dt_ * cltv_pb_.get_r(indexToTime(k));
  }

  Q_t get_QT() const
  {
    double T;
    cltv_pb_.get_T(T);
    return cltv_pb_.get_QT();
  }

  template<typename T = void *>
  state_t get_qT([[maybe_unused]] std::enable_if_t<has_qT, T> k = nullptr) const
  {
    double TT;
    cltv_pb_.get_T(TT);
    return cltv_pb_.get_qT();
  }

  cltv_pb_t & problem()
  {
    return cltv_pb_;
  }

  double indexToTime(std::size_t k) const
  {
    return static_cast<double>(k) * dt_;
  }

protected:
  double compute_dt()
  {
    double T;
    cltv_pb_.get_T(T);
    return T / static_cast<double>(nPts - 1);
  }

protected:
  cltv_pb_t cltv_pb_{};
  double dt_{};
};

// Class template argument deduction guides
template<typename T>
DltvOcp(T)->DltvOcp<T>;

template<typename T1, typename T2>
DltvOcp(T1, T2)->DltvOcp<T1, T2::nPts, T2::expOrder>;

}  // namespace cbr

#endif  // CBR_CONTROL__MPC__DLTV_OCP_HPP_


//          Copyright Gavin Band 2008 - 2012.
// Distributed under the Boost Software License, Version 1.0.
//    (See accompanying file LICENSE_1_0.txt or copy at
//          http://www.boost.org/LICENSE_1_0.txt)

#include <iostream>
#include <iomanip>
#include <boost/ptr_container/ptr_vector.hpp>
#include <boost/function.hpp>
#include "config/config.hpp"
#if HAVE_CBLAS
	#include "cblas.h"
#endif
#include "Eigen/Core"
#include "Eigen/Eigenvalues"
#include "genfile/VariantIdentifyingData.hpp"
#include "genfile/VariantDataReader.hpp"
#include "genfile/vcf/get_set_eigen.hpp"
#include "appcontext/get_current_time_as_string.hpp"
#include "components/RelatednessComponent/PCALoadingComputer.hpp"
#include "components/RelatednessComponent/LapackEigenDecomposition.hpp"
#include "components/RelatednessComponent/mean_centre_genotypes.hpp"

// #define DEBUG_PCA_LOADING_COMPUTER 1

namespace {
	template< typename Vector1, typename Vector2, typename NonMissingVector >
	double compute_correlation( Vector1 const& v1, Vector2 const& v2, NonMissingVector const& non_missingness_indicator ) {
		assert( v1.size() == v2.size() ) ;
		double non_missingness = non_missingness_indicator.sum() ;
		double mean1 = 0.0 ;
		double mean2 = 0.0 ;
		for( int i = 0; i < v1.size(); ++i ) {
			if( non_missingness_indicator( i )) {
				mean1 += v1(i) / non_missingness ;
				mean2 += v2(i) / non_missingness ;
			}
		}

		double covariance = 0.0 ;
		double variance1 = 0.0 ;
		double variance2 = 0.0 ;
		for( int i = 0; i < v1.size(); ++i ) {
			if( non_missingness_indicator( i )) {
				covariance += ( v1(i) - mean1 ) * ( v2(i) - mean2 ) ;
				variance1 += ( v1(i) - mean1 ) * ( v1(i) - mean1 ) ;
				variance2 += ( v2(i) - mean2 ) * ( v2(i) - mean2 ) ;
			}
		}

		// We should divide the covariance by N-1 and also
		// divide each variance by the same quantity.
		// But this washes out in the ratio.

		return covariance / std::sqrt( variance1 * variance2 ) ;
	}

	std::string eigenvector_column_names( std::size_t N, std::string const& string1, std::string const& string2, std::size_t i ) {
		if( i < N ) {
			return string1 + genfile::string_utils::to_string( i+1 ) ;
		}
		else {
			return string2 + genfile::string_utils::to_string( i+1 -N ) ;
		}
	}
}

PCALoadingComputer::PCALoadingComputer( int number_of_loadings ):
	m_number_of_loadings( number_of_loadings ),
	m_number_of_snps( 1 )
{}

void PCALoadingComputer::set_UDUT( std::size_t number_of_snps, Matrix const& udut ) {
	assert( udut.cols() == udut.rows() + 1 ) ;
	int n = std::min( int( m_number_of_loadings ), int( udut.rows() ) ) ;
	m_D = udut.block( 0, 0, n, 1 ) ;
	m_sqrt_D_inverse = 1 / m_D.array().sqrt() ;
	m_U = udut.block( 0, 1, udut.rows(), n ) ;
	m_number_of_snps = number_of_snps ;
}

void PCALoadingComputer::begin_processing_snps( std::size_t number_of_samples, genfile::SNPDataSource::Metadata const& ) {
	assert( number_of_samples = std::size_t( m_U.rows() )) ;
	m_genotype_calls.resize( number_of_samples ) ;
	m_non_missingness.resize( number_of_samples ) ;
}

void PCALoadingComputer::processed_snp( genfile::VariantIdentifyingData const& snp, genfile::VariantDataReader& data_reader ) {
	data_reader.get(
		":genotypes:",
		genfile::vcf::get_threshholded_calls( m_genotype_calls, m_non_missingness, 0.9, 0, 0, 1, 2 )
	) ;
	assert( m_genotype_calls.size() == m_U.rows() ) ;
	assert( m_non_missingness.size() == m_U.rows() ) ;
	// setup the storage
	m_loading_vectors.resize( 2 * m_D.rows() ) ;
	m_loading_vectors.setConstant( std::numeric_limits< double >::quiet_NaN() ) ;
	double const allele_frequency = m_genotype_calls.sum() / ( 2.0 * m_non_missingness.sum() ) ;
	if( m_non_missingness.sum() > 0 && allele_frequency > 0.001 ) {
		//std::cerr << "pre-mean genotypes are: " << m_genotype_calls.head( 20 ).transpose() << "...\n" ;
		pca::mean_centre_genotypes( &m_genotype_calls, m_non_missingness, allele_frequency ) ;
		m_genotype_calls /= std::sqrt( 2.0 * allele_frequency * ( 1.0 - allele_frequency ) ) ;

#if DEBUG_PCA_LOADING_COMPUTER
		//std::cerr << "                    SNP: " << snp << ", allele frequency = " << allele_frequency << ".\n" ;
		//std::cerr << std::resetiosflags( std::ios::floatfield ) << std::setprecision( 5 ) ;
		//std::cerr << "pre-scale genotypes are: " << m_genotype_calls.head( 20 ).transpose() << "...\n" ;
		std::cerr << "          genotypes are: " << m_genotype_calls.head( 20 ).transpose() << "...\n" ;
		std::cerr << "     non-missingness is: " << m_non_missingness.head( 20 ).transpose() << "...\n" ;
		std::cerr << "                   U is: " << m_U.block(0,0,10,10) << "...\n" ;
		std::cerr << "                   D is: " << m_D << "...\n" ;
#endif // DEBUG_PCA_LOADING_COMPUTER

		//
		// Let X  be the L\times n matrix (L SNPs, n samples) of (mean-centred, scaled) genotypes.  We want
		// to compute the row of the matrix S of unit eigenvectors of the variance-covariance matrix
		// (1/L) X X^t that corresponds to the current SNP.
		// The matrix S is given by
		//
		//       S = (1/√L) X U D^{-½}
		//
		// where
		//       (1/L) X^t X = U D U^t
		// is the eigenvalue decomposition of (1/L) X^t X that we are passed in via set_UDUT (and L is the number of SNPs).
		//
		// This is true since then
		//
		// S^t S = D^{-½} U^t (1/L) X^t X U D^{-½} = id
		//
		// (so columns of S are orthogonal) while
		//
		// (1/L X X^t) S = (1/L√L) X X^t X U D^{-½}
		//           = (1/√L) X U D U^t U D^{-½}
		//           = (1/√L) X U D^½
		//           = SD
		//
		// (so columns of S are eigenvectors with eigenvalues given by D.)
		//
#if 0
		m_loading_vectors.segment( 0, m_U.cols() ) =
			( m_genotype_calls.transpose() * m_U ) * m_D.array().sqrt().matrix().asDiagonal()
			/ ( ( m_D.transpose().array() * m_number_of_snps ).sqrt() ) ;
#else
		m_loading_vectors.segment( 0, m_U.cols() ) =
			( m_genotype_calls.transpose() * m_U ) * m_sqrt_D_inverse.asDiagonal() ;
		m_loading_vectors /= std::sqrt( m_number_of_snps ) ;
#endif
		// We also wish to compute the correlation between the SNP and the PCA component.
		// With S as above, the PCA components are the projections of columns of X onto columns of S.
		// If we want samples to correspond to columns, this is
		//   S^t X
		// which can be re-written
		//   sqrt(L) U D^{1/2}
		// i.e. we may as well compute the correlation with columns of U.
		if( m_non_missingness.sum() > 10 ) {
			for( int i = 0; i < m_U.cols(); ++i ) {
				m_loading_vectors( m_U.cols() + i ) = compute_correlation( m_genotype_calls, m_U.col( i ), m_non_missingness ) ;
			}
		}
	}
	send_results(
		snp,
		m_non_missingness.sum(),
		allele_frequency,
		m_loading_vectors,
		boost::bind(
			&eigenvector_column_names,
			m_U.cols(),
			"eigenvector_",
			"correlation_",
			_1
		)
	) ;
}

void PCALoadingComputer::send_results_to( ResultCallback callback ) {
	m_result_signal.connect( callback ) ;
}

void PCALoadingComputer::send_results( genfile::VariantIdentifyingData const& snp, double const N, double const frequency, Eigen::VectorXd const& data, GetNames get_names ) {
	m_result_signal( snp, N, frequency, data, get_names ) ;
}

std::string PCALoadingComputer::get_metadata() const {
	using namespace genfile::string_utils ;
	return "Number of SNPs: " + to_string( m_number_of_snps ) + "\n"
		+ "Number of samples: " + to_string( m_U.rows() ) + "\n"
		+ "These loadings represent unit eigenvectors of the variance-covariance matrix\n"
		+ "    1/L X X^t\n"
		+ "where X is the LxN matrix of genotypes at L SNPs and N samples (normalised across rows.)" ;
}

void PCALoadingComputer::end_processing_snps() {}


#include <iostream>
#include <string>
#include <Eigen/Dense>
#include <fstream>
#include <sstream>
#include <cstdlib>
#include <string>
#include <unistd.h>
#include <sys/wait.h>
#include "H5Cpp.h"
#include "rapidjson/document.h"
#include "rapidjson/writer.h"
#include "rapidjson/stringbuffer.h"

using std::cout;
using std::endl;

using namespace H5;
using namespace rapidjson;
using Eigen::tanh;
using Eigen::MatrixXd;

inline double min(double a, double b) { return(((a)<(b))?(a):(b));}
inline double max(double a, double b) { return(((a)>(b))?(a):(b));}

double hard_sigmoid(double x){
    return(max(0.0, min(1.0, x*0.2+0.5)));
}

double sigmoid(double x){
    return 1.0 / (1.0 + exp(-x));
}

inline MatrixXd Mult(MatrixXd input1, MatrixXd input2){
    return(input1.array() * input2.array());
}

char* read_model(const char *filename)
{
    char *arch_json;

    H5File file(filename, H5F_ACC_RDWR);
    Attribute attr(file.openAttribute("model_config"));
    DataType type(attr.getDataType());

    attr.read(type, &arch_json);
    return arch_json;
}

inline MatrixXd readCSV(const char* filename, int row_n, int col_n){

    std::ifstream file(filename);
    int col_flag = 0;

    MatrixXd m(row_n, col_n);

    if(row_n == 1){
        col_flag = 1;
        m.resize(col_n, row_n);
    }

    std::string line;

    int row = 0;
    int col = 0;

    if(file.is_open()){
        while(std::getline(file, line)){
            char *ptr = (char *)line.c_str();
            int len = line.length();

            col = 0;

            char *start = ptr;
            for(int i = 0; i < len; i++){
                if(ptr[i] == ','){
                    m(row, col++) = atof(start);
                    start = ptr + i + 1;
                }
            }
            m(row, col) = atof(start);

            row++;
        }
        file.close();
    }

    if(col_flag)
        return(m.transpose());

    return(m);
}

class Embedding{
    public:
        int row;
        int col;
        int out_size;
        MatrixXd W;

        Embedding(){

        }

        Embedding(const char* filename, int row, int col){
            this->row = row;
            this->col = col;
            this->out_size = col;

            W.resize(this->row, this->col);
            this->W = readCSV(filename, this->row, this->col);

        }

        MatrixXd operator()(MatrixXd input){
            MatrixXd Out(input.cols(), this->out_size);
            for(int t=0; t < Out.rows(); t++)
                Out.row(t) = W.row( input(0, t) );

            return(Out);
        }
};

class LSTM{
    public:
        int inp_size;
        int out_size;
        MatrixXd kernel;
        MatrixXd recurrent_kernel;
        MatrixXd bias;

        LSTM(){

        }

        LSTM(const char* kernel_filename, const char* recurrent_kernel_filename, const char* bias_filename, int inp_size, int out_size){
            this->inp_size = inp_size;
            this->out_size = out_size;

            kernel.resize(this->inp_size, this->out_size * 4);
            kernel = readCSV(kernel_filename, this->inp_size, this->out_size * 4);

            recurrent_kernel.resize(this->out_size, this->out_size * 4);
            recurrent_kernel = readCSV(recurrent_kernel_filename, this->out_size, this->out_size * 4);

            bias.resize(1, this->out_size * 4);
            bias = readCSV(bias_filename, 1, this->out_size * 4);

        }

        MatrixXd operator()(MatrixXd input){
            int LSTM_OUT = this->out_size;
            int MAXLEN   = input.rows();

            MatrixXd C_t = MatrixXd::Zero(1, this->out_size);
            MatrixXd h_t = MatrixXd::Zero(1, this->out_size);

            for(int t = 0; t < MAXLEN; t++){
                MatrixXd tmp_out(1, 4*LSTM_OUT);
                MatrixXd i_t(1, LSTM_OUT);
                MatrixXd f_t(1, LSTM_OUT);
                MatrixXd o_t(1, LSTM_OUT);
                MatrixXd g_t(1, LSTM_OUT);

                // IFCO
                tmp_out = input.row(t) * kernel + h_t * recurrent_kernel + bias;

                i_t = (tmp_out.block(0, 0*LSTM_OUT, 1, LSTM_OUT)).unaryExpr(&hard_sigmoid);
                f_t = (tmp_out.block(0, 1*LSTM_OUT, 1, LSTM_OUT)).unaryExpr(&hard_sigmoid);
                o_t = (tmp_out.block(0, 3*LSTM_OUT, 1, LSTM_OUT)).unaryExpr(&hard_sigmoid);
                g_t = tanh((tmp_out.block(0, 2*LSTM_OUT, 1, LSTM_OUT)).array());

                C_t = f_t.array() * C_t.array() + i_t.array() * g_t.array();
                h_t = o_t.array() * tanh(C_t.array());
            }
            return h_t;
        }
};

class GRU{
    public:
        int inp_size;
        int out_size;
        MatrixXd kernel;
        MatrixXd recurrent_kernel;
        MatrixXd bias;

        GRU(){

        }

        GRU(const char* kernel_filename, const char* recurrent_kernel_filename, const char* bias_filename, int inp_size, int out_size){
            this->inp_size = inp_size;
            this->out_size = out_size;

            kernel.resize(this->inp_size, this->out_size * 4);
            kernel = readCSV(kernel_filename, this->inp_size, this->out_size * 4);

            recurrent_kernel.resize(this->out_size, this->out_size * 4);
            recurrent_kernel = readCSV(recurrent_kernel_filename, this->out_size, this->out_size * 4);

            bias.resize(1, this->out_size * 4);
            bias = readCSV(bias_filename, 1, this->out_size * 4);

        }

        MatrixXd operator()(MatrixXd input){
            int LSTM_OUT = this->out_size;
            int MAXLEN   = input.rows();

            MatrixXd hh_t = MatrixXd::Zero(1,LSTM_OUT);
            MatrixXd h_t = MatrixXd::Zero(1,LSTM_OUT);
            MatrixXd one_arr = MatrixXd::Ones(1,LSTM_OUT);

            MatrixXd tmp_out(1, 3*LSTM_OUT);
            MatrixXd tmp_W(1, LSTM_OUT);
            MatrixXd tmp_U(LSTM_OUT, LSTM_OUT);
            MatrixXd z_t(1, LSTM_OUT);
            MatrixXd r_t(1, LSTM_OUT);

            tmp_U = recurrent_kernel.block(0, 2*LSTM_OUT, LSTM_OUT, LSTM_OUT);

            for(int t = 0; t < MAXLEN; t++){

                // IFCO

                tmp_W = input.row(t) * kernel + bias ;
                tmp_out = tmp_W + h_t * recurrent_kernel;

                z_t = (tmp_out.block(0, 0*LSTM_OUT, 1, LSTM_OUT)).unaryExpr(&hard_sigmoid);
                r_t = (tmp_out.block(0, 1*LSTM_OUT, 1, LSTM_OUT)).unaryExpr(&hard_sigmoid);
                r_t = r_t.array() * h_t.array();

                hh_t = tanh((r_t * tmp_U + tmp_W.block(0, 2*LSTM_OUT, 1, LSTM_OUT)).array() );
                h_t = (one_arr - z_t).array() * hh_t.array() + z_t.array() * h_t.array();

            }
            return h_t;
        }
};

class Dense{
    public:
        int inp_size;
        int out_size;
        string activation;
        MatrixXd kernel;
        MatrixXd bias;

        Dense(){

        }

        Dense(const char* kernel_filename, const char* bias_filename, int inp_size, int out_size, string activation="linear"){
            this->inp_size = inp_size;
            this->out_size = out_size;
            this->activation = activation;

            kernel.resize(this->inp_size, this->out_size);
            kernel = readCSV(kernel_filename, this->inp_size, this->out_size);

            bias.resize(1, this->out_size);
            bias = readCSV(bias_filename, 1, this->out_size);
        }

        MatrixXd operator()(MatrixXd input){

            if(this->activation == "linear")
                return( input * kernel + bias );
            else if(this->activation == "sigmoid")
                return( (input * kernel + bias).unaryExpr(&sigmoid) );
        }
};


#include <NTL/ZZX.h>

NTL_CLIENT

#define TIME_IT(t, action) \
do { \
   double _t0, _t1; \
   long _iter = 1; \
   long _cnt = 0; \
   do { \
      _t0 = GetTime(); \
      for (long _i = 0; _i < _iter; _i++) { action; _cnt++; } \
      _t1 = GetTime(); \
   } while ( _t1 - _t0 < 2 && (_iter *= 2)); \
   t = (_t1 - _t0)/_iter; \
} while(0)

void FillRandom(ZZX& f, long n, long k)
{
   long sw = RandomBnd(2);
   f.SetLength(n);
   for (long i = 0; i < n; i++) {
      if (sw) {
         long kk = 1 + RandomBnd(k);
         RandomBits(f[i], kk);
      }
      else {
        long kk = RandomBnd(k);
        SetBit(f[i], kk);
      }
      if (RandomBnd(2)) NTL::negate(f[i], f[i]);
   }
   f.normalize();
}

int main()
{

   for (long iter = 0; iter < 4000; iter++) {
     if (iter % 100 == 0) cerr << ".";
     long na, nb, k;

     long sw = RandomBnd(3);

     if (sw == 0) {
        na = RandomBnd(20) + 1;
        nb = RandomBnd(20) + 1;
        k = RandomBnd(20) + 1;
     }
     else if (sw == 1) {
        na = RandomBnd(200) + 10;
        nb = RandomBnd(200) + 10;
        k = RandomBnd(200) + 10;
     }
     else {
        na = RandomBnd(3000) + 100;
        nb = RandomBnd(3000) + 100;
        k = RandomBnd(3000) + 100;
     }

     ZZX a, b, c, c1;
     FillRandom(a, na, k);
     FillRandom(b, nb, k);

     if (RandomBnd(2)) {
        SSMul(c, a, b);
        KarMul(c1, a, b);
        if (c != c1) Error("oops");
     }
     else {
        SSSqr(c, a);
        KarSqr(c1, a);
        if (c != c1) Error("oops");
     }
   }

   cerr << "\n";
}


#ifndef STAN_MATH_PRIM_MAT_PROB_CATEGORICAL_LOGIT_LPMF_HPP
#define STAN_MATH_PRIM_MAT_PROB_CATEGORICAL_LOGIT_LPMF_HPP

#include <stan/math/prim/scal/err/check_bounded.hpp>
#include <stan/math/prim/scal/err/check_finite.hpp>
#include <stan/math/prim/arr/fun/log_sum_exp.hpp>
#include <stan/math/prim/mat/fun/log_softmax.hpp>
#include <stan/math/prim/mat/fun/log_sum_exp.hpp>
#include <stan/math/prim/mat/fun/sum.hpp>
#include <stan/math/prim/scal/meta/include_summand.hpp>
#include <boost/math/tools/promotion.hpp>
#include <vector>

namespace stan {
  namespace math {

    // CategoricalLog(n|theta)  [0 < n <= N, theta unconstrained], no checking
    template <bool propto,
              typename T_prob>
    typename boost::math::tools::promote_args<T_prob>::type
    categorical_logit_lpmf(int n,
                          const Eigen::Matrix<T_prob, Eigen::Dynamic, 1>&
                          beta) {
      static const char* function("categorical_logit_lpmf");

      check_bounded(function, "categorical outcome out of support", n,
                    1, beta.size());
      check_finite(function, "log odds parameter", beta);

      if (!include_summand<propto, T_prob>::value)
        return 0.0;

      // FIXME:  wasteful vs. creating term (n-1) if not vectorized
      return beta(n - 1) - log_sum_exp(beta);  // == log_softmax(beta)(n-1);
    }

    template <typename T_prob>
    inline
    typename boost::math::tools::promote_args<T_prob>::type
    categorical_logit_lpmf(int n,
                          const Eigen::Matrix<T_prob, Eigen::Dynamic, 1>&
                          beta) {
      return categorical_logit_lpmf<false>(n, beta);
    }

    template <bool propto,
              typename T_prob>
    typename boost::math::tools::promote_args<T_prob>::type
    categorical_logit_lpmf(const std::vector<int>& ns,
                          const Eigen::Matrix<T_prob, Eigen::Dynamic, 1>&
                          beta) {
      static const char* function("categorical_logit_lpmf");

      for (size_t k = 0; k < ns.size(); ++k)
        check_bounded(function, "categorical outcome out of support",
                      ns[k], 1, beta.size());
      check_finite(function, "log odds parameter", beta);

      if (!include_summand<propto, T_prob>::value)
        return 0.0;

      if (ns.size() == 0)
        return 0.0;

      Eigen::Matrix<T_prob, Eigen::Dynamic, 1> log_softmax_beta
        = log_softmax(beta);

      // FIXME:  replace with more efficient sum()
      Eigen::Matrix<typename boost::math::tools::promote_args<T_prob>::type,
                    Eigen::Dynamic, 1> results(ns.size());
      for (size_t i = 0; i < ns.size(); ++i)
        results[i] = log_softmax_beta(ns[i] - 1);
      return sum(results);
    }

    template <typename T_prob>
    inline
    typename boost::math::tools::promote_args<T_prob>::type
    categorical_logit_lpmf(const std::vector<int>& ns,
                          const Eigen::Matrix<T_prob, Eigen::Dynamic, 1>&
                          beta) {
      return categorical_logit_lpmf<false>(ns, beta);
    }

  }
}
#endif


/**
 * @file lsh_test.cpp
 *
 * Unit tests for the 'LSHSearch' class.
 */
#include <mlpack/core.hpp>
#include <mlpack/core/metrics/lmetric.hpp>
#include <boost/test/unit_test.hpp>
#include "old_boost_test_definitions.hpp"

#include <mlpack/methods/lsh/lsh_search.hpp>

using namespace std;
using namespace mlpack;
using namespace mlpack::neighbor;

BOOST_AUTO_TEST_SUITE(LSHTest);

BOOST_AUTO_TEST_CASE(LSHSearchTest)
{
  // Force to specific random seed for these results.
  math::RandomSeed(0);

  // Precomputed hash width value.
  const double hashWidth = 4.24777;

  arma::mat rdata(2, 10);
  rdata << 3 << 2 << 4 << 3 << 5 << 6 << 0 << 8 << 3 << 1 << arma::endr <<
           0 << 3 << 4 << 7 << 8 << 4 << 1 << 0 << 4 << 3 << arma::endr;

  arma::mat qdata(2, 3);
  qdata << 3 << 2 << 0 << arma::endr << 5 << 3 << 4 << arma::endr;

  // INPUT TO LSH:
  // Number of points: 10
  // Number of dimensions: 2
  // Number of projections per table: 'numProj' = 3
  // Number of hash tables: 'numTables' = 2
  // hashWidth (computed): 'hashWidth' = 4.24777
  // Second hash size: 'secondHashSize' = 11
  // Size of the bucket: 'bucketSize' = 3

  // Things obtained by random sampling listed in the sequences
  // as they will be obtained in the 'LSHSearch::BuildHash()' private function
  // in 'LSHSearch' class.
  //
  // 1. The weights of the second hash obtained as:
  //    secondHashWeights = arma::floor(arma::randu(3) * 11.0);
  //    COR.SOL.: secondHashWeights = [9, 4, 8];
  //
  // 2. The offsets for all the 3 projections in each of the 2 tables:
  //    offsets.randu(3, 2)
  //    COR.SOL.: [0.7984 0.3352; 0.9116 0.7682; 0.1976 0.2778]
  //    offsets *= hashWidth
  //    COR.SOL.: [3.3916 1.4240; 3.8725 3.2633; 0.8392 1.1799]
  //
  // 3. The  (2 x 3) projection matrices for the 2 tables:
  //    projMat.randn(2, 3)
  //    COR.SOL.: Proj. Mat 1: [2.7020 0.0187 0.4355; 1.3692 0.6933 0.0416]
  //    COR.SOL.: Proj. Mat 2: [-0.3961 -0.2666 1.1001; 0.3895 -1.5118 -1.3964]
  LSHSearch<> lsh_test(rdata, 3, 2, hashWidth, 11, 3);
//   LSHSearch<> lsh_test(rdata, qdata, 3, 2, 0.0, 11, 3);

  // Given this, the 'LSHSearch::bucketRowInHashTable' should be:
  // COR.SOL.: [2 11 4 7 6 3 11 0 5 1 8]
  //
  // The 'LSHSearch::bucketContentSize' should be:
  // COR.SOL.: [2 0 1 1 3 1 0 3 3 3 1]
  //
  // The final hash table 'LSHSearch::secondHashTable' should be
  // of size (3 x 9) with the following content:
  // COR.SOL.:
  // [0 2 4; 1 7 8; 3 9 10; 5 10 10; 6 10 10; 0 5 6; 1 2 8; 3 10 10; 4 10 10]

  arma::Mat<size_t> neighbors;
  arma::mat distances;

  lsh_test.Search(qdata, 2, neighbors, distances);

  // The private function 'LSHSearch::ReturnIndicesFromTable(0, refInds)'
  // should hash the query 0 into the following buckets:
  // COR.SOL.: Table 1 Bucket 7, Table 2 Bucket 0, refInds = [0 2 3 4 9]
  //
  // The private function 'LSHSearch::ReturnIndicesFromTable(1, refInds)'
  // should hash the query 1 into the following buckets:
  // COR.SOL.: Table 1 Bucket 9, Table 2 Bucket 4, refInds = [1 2 7 8]
  //
  // The private function 'LSHSearch::ReturnIndicesFromTable(2, refInds)'
  // should hash the query 2 into the following buckets:
  // COR.SOL.: Table 1 Bucket 0, Table 2 Bucket 7, refInds = [0 2 3 4 9]

  // After search
  // COR.SOL.: 'neighbors' = [2 1 9; 3 8 2]
  // COR.SOL.: 'distances' = [2 0 2; 4 2 16]

  arma::Mat<size_t> true_neighbors(2, 3);
  true_neighbors << 2 << 1 << 9 << arma::endr << 3 << 8 << 2 << arma::endr;
  arma::mat true_distances(2, 3);
  true_distances << 2 << 0 << 2 << arma::endr << 4 << 2 << 16 << arma::endr;

  for (size_t i = 0; i < 3; i++)
  {
    for (size_t j = 0; j < 2; j++)
    {
//      BOOST_REQUIRE_EQUAL(neighbors(j, i), true_neighbors(j, i));
//      BOOST_REQUIRE_CLOSE(distances(j, i), true_distances(j, i), 1e-5);
    }
  }
}

BOOST_AUTO_TEST_CASE(LSHTrainTest)
{
  // This is a not very good test that simply checks that the re-trained LSH
  // model operates on the correct dimensionality and returns the correct number
  // of results.
  arma::mat referenceData = arma::randu<arma::mat>(3, 100);
  arma::mat newReferenceData = arma::randu<arma::mat>(10, 400);
  arma::mat queryData = arma::randu<arma::mat>(10, 200);

  LSHSearch<> lsh(referenceData, 3, 2, 2.0, 11, 3);

  lsh.Train(newReferenceData, 4, 3, 3.0, 12, 4);

  arma::Mat<size_t> neighbors;
  arma::mat distances;

  lsh.Search(queryData, 3, neighbors, distances);

  BOOST_REQUIRE_EQUAL(neighbors.n_cols, 200);
  BOOST_REQUIRE_EQUAL(neighbors.n_rows, 3);
  BOOST_REQUIRE_EQUAL(distances.n_cols, 200);
  BOOST_REQUIRE_EQUAL(distances.n_rows, 3);
}

BOOST_AUTO_TEST_CASE(EmptyConstructorTest)
{
  // If we create an empty LSH model and then call Search(), it should throw an
  // exception.
  LSHSearch<> lsh;

  arma::mat dataset = arma::randu<arma::mat>(5, 50);
  arma::mat distances;
  arma::Mat<size_t> neighbors;
  BOOST_REQUIRE_THROW(lsh.Search(dataset, 2, neighbors, distances),
      std::invalid_argument);

  // Now, train.
  lsh.Train(dataset, 4, 3, 3.0, 12, 4);

  lsh.Search(dataset, 3, neighbors, distances);

  BOOST_REQUIRE_EQUAL(neighbors.n_cols, 50);
  BOOST_REQUIRE_EQUAL(neighbors.n_rows, 3);
  BOOST_REQUIRE_EQUAL(distances.n_cols, 50);
  BOOST_REQUIRE_EQUAL(distances.n_rows, 3);
}

BOOST_AUTO_TEST_SUITE_END();


#pragma once

#include <boost/mpl/identity.hpp>
#include <Eigen/Dense>
#include <algorithm>
#include <boost/multi_array.hpp>
#include <boost/multiprecision/mpfr.hpp>
#include <fstream>
#include <iomanip>
#include <stdexcept>

#include "spectral/hermiten.hpp"
#include "spectral/mpfr/import_std_math.hpp"

namespace boltzmann {

// ----------------------------------------------------------------------
namespace detail {
template <typename NUMERIC_T>
class sentry
{
 public:
  typedef NUMERIC_T numeric_t;

 private:
  typedef std::vector<numeric_t> vec_t;

 public:
  sentry(int n) { this->init(n); }

  sentry() { /* empty */}

  void init(int n)
  {
    n_ = n;
    factors_.resize(6);
    std::for_each(factors_.begin(), factors_.end(), [&](vec_t& v) { v.reserve(n + 1); });
  }

  NUMERIC_T coeff(int i, int j, int t);
  NUMERIC_T operator()(int i, int j, numeric_t x);

 private:
  int n_;
  std::vector<vec_t> factors_;
};

// ----------------------------------------------------------------------
template <typename NUMERIC_T>
NUMERIC_T
sentry<NUMERIC_T>::coeff(int i, int j, int t)
{
  const int maxij = std::max(i, j);
  const int minij = std::min(i, j);
  // (min(i,j) ... 1)
  auto& v0 = factors_[0];
  v0.resize(minij);
  for (int k = 0; k < minij; ++k) {
    v0[k] = std::min(i, j) - k;
  }

  //  sqrt(max(i,j) .. min(i,j)+1)
  auto& v1 = factors_[1];
  v1.resize(maxij - minij);
  for (int k = 0; k < maxij - minij; ++k) {
    v1[k] = ::math::sqrt(numeric_t(maxij - k));
  }

  // (i-t)!
  auto& v2 = factors_[2];
  v2.resize(i - t);
  for (int k = 0; k < i - t; ++k) {
    v2[k] = 1 / numeric_t(i - t - k);
  }

  // (j-t)!
  auto& v3 = factors_[3];
  v3.resize(j - t);
  for (int k = 0; k < j - t; ++k) {
    v3[k] = 1 / numeric_t(j - t - k);
  }

  // t!
  auto& v4 = factors_[4];
  v4.resize(t);
  for (int k = 0; k < t; ++k) {
    v4[k] = 1 / numeric_t(t - k);
  }

  // 2^(t-(i+j)/2)
  auto& v5 = factors_[5];
  int exp2 = t - (i + j) / 2;
  v5.resize(std::abs(exp2));
  if (exp2 > 0)
    for (unsigned int k = 0; k < v5.size(); ++k) {
      v5[k] = 2;
    }
  else
    for (unsigned int k = 0; k < v5.size(); ++k) {
      v5[k] = 1 / numeric_t(2);
    }
  // cout << "v5:\t";
  // for_each(v5.begin(), v5.end(), [](numeric_t v) { cout << v << "\t"; });
  // cout << endl;

  std::sort(factors_.begin(), factors_.end(), [](const vec_t& v1, const vec_t& v2) {
    return v1.size() < v2.size();
  });
  std::vector<int> lengths;
  std::for_each(
      factors_.begin(), factors_.end(), [&](const vec_t& v) { lengths.push_back(v.size()); });

  numeric_t f = (std::abs(i - t) % 2) ? -1 : 1;
  // multiply
  for (int jp = 0; jp < lengths[0]; ++jp) {
    numeric_t loc = 1;
    for (unsigned int fi = 0; fi < factors_.size(); ++fi) {
      loc *= factors_[fi][jp];
    }
    f *= loc;
  }
  for (unsigned int l = 1; l < lengths.size(); ++l) {
    for (int jp = lengths[l - 1]; jp < lengths[l]; ++jp) {  // loop over remaining positions
      numeric_t loc = 1;
      for (unsigned int fi = l; fi < factors_.size(); ++fi) {
        loc *= factors_[fi][jp];
      }
      f *= loc;
    }
  }

  // 2^(t-(i+j)/2) missing term
  if ((i + j) % 2 && exp2 <= 0)
    f /= ::math::sqrt(numeric_t(2));
  else if ((i + j) % 2 && exp2 > 0)
    f *= ::math::sqrt(numeric_t(2));

  return f;
}

// ----------------------------------------------------------------------
template <typename NUMERIC_T>
NUMERIC_T
sentry<NUMERIC_T>::operator()(int i, int j, numeric_t x)
{
  numeric_t sum = 0;

  // TOOD implement Horner's scheme
  for (int l = std::abs(i - j); l <= i + j; l += 2) {
    numeric_t loc = this->coeff(i, j, (i + j - l) / 2);

    sum += loc * ::math::pow(x, l);
  }

  return sum * ::math::exp(numeric_t(-x * x / 4));
  ;
}

/**
 * @brief polyval (Horner scheme)
 *
 * @param coeffs
 * @param x
 * @param N length of coeffs
 */
template <typename NUMERIC_T>
inline NUMERIC_T
polyval(NUMERIC_T* coeffs, NUMERIC_T x, int N)
{
  typedef NUMERIC_T numeric_t;
  numeric_t b = coeffs[N - 1];
  for (int i = 1; i < N - 1; ++i) {
    b = coeffs[N - 1 - i] + b * x;
  }
  return x * b + coeffs[0];
}

}  // end namespace detail

// ----------------------------------------------------------------------
/**
 * @brief Assemble shift matrix \f$ S^{\bar{x}}\f$.
 *
 * @tparam NUMERIC_T numeric type
 *
 */
template <typename NUMERIC_T>
class HShiftMatrix
{
 private:
  typedef NUMERIC_T numeric_t;
  typedef std::vector<numeric_t> vec_t;

 public:
  typedef Eigen::Matrix<numeric_t, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> matrix_t;

 public:
  /**
   *
   * @param N  max polynomial degree
   */
  HShiftMatrix(int N) { this->init(N); }

  HShiftMatrix() { /* empty */}

  /**
   * Compute polynomial coefficients of p(i,j;x)
   *
   */
  void init(int N);

  /**
   * Create linear Operator S^x by evaluating the polynomial
   *
   * @param x
   */
  void setx(numeric_t x);

  const matrix_t& get() const { return S_; }
  void dump(std::string fname) const;

 private:
  int size_;
  boost::multi_array<numeric_t, 3> coeffs_;
  /* compute coefficients of the S_-entry polynomial */
  detail::sentry<numeric_t> G_;
  matrix_t S_;
};

// ----------------------------------------------------------------------
template <typename NUMERIC_T>
void
HShiftMatrix<NUMERIC_T>::init(int N)
{
  size_ = N + 1;
  G_.init(N + 1);
  S_.resize(N + 1, N + 1);
  coeffs_.resize(boost::extents[size_][size_][2 * size_ - 1]);

  // std::fill(coeffs_.origin(), coeffs_.origin() + coeffs_.num_elements(), 0);
  for (int i = 0; i < size_; ++i) {
    for (int j = 0; j < size_; ++j) {
      for (int l = std::abs(i - j); l <= i + j; l += 2) {
        coeffs_[i][j][l] = G_.coeff(i, j, (i + j - l) / numeric_t(2));
      }
    }
  }
}

// ----------------------------------------------------------------------
template <typename NUMERIC_T>
void
HShiftMatrix<NUMERIC_T>::setx(numeric_t x)
{
  // Hint: polyval can be further optimized by using knowledge
  //       about which coefficients are zero.
  numeric_t f = ::math::exp(numeric_t(-x * x / 4));
  for (int i = 0; i < size_; ++i) {
    for (int j = 0; j < size_; ++j) {
      S_(i, j) = f * detail::polyval(coeffs_[i][j].origin(), x, coeffs_.shape()[2]);
    }
  }
}

// ----------------------------------------------------------------------
template <typename NUMERIC_T>
void
HShiftMatrix<NUMERIC_T>::dump(std::string fname) const
{
  std::ofstream fout(fname);
  fout << std::setprecision(10);
  fout << std::scientific;
  fout << S_;
  fout.close();
}

}  // end namespace boltzmann


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS
  @copyright 2016 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_DIVS_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_DIVS_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

 /*!

    @ingroup group-arithmetic
    Function object implementing divs capabilities

    Computes  the truncated saturated division of its parameters.

    @par semantic:
    For any given value @c x,  @c y of type @c T:

    @code
    T r = divs(x, y);
    @endcode

    The code is similar to:

    @code
    T r = trunc(x/y);
    @endcode

    for integral types, if y is @ref Zero, it returns @ref Valmax (resp. @ref Valmin)
    if x is positive (resp. negative) and @ref Zero if x is @ref Zero.

    Saturated means that for signed integer types,
    @c divs(Valmin,-1) returns @ref Valmax.

    @par Alias

    @c rdivide

    @see  divides, rec, divfloor, divceil, divround, divround2even, divfix

  **/
  const boost::dispatch::functor<tag::divs_> divs = {};
} }
#endif

#include <boost/simd/function/scalar/divs.hpp>
#include <boost/simd/function/simd/divs.hpp>

#endif


#include <boost/variant.hpp>
#include <vector>
#include <string>
#include <iostream>

// This typedefs and methods will be in header,
// that wraps around native SQL interface.
typedef boost::variant<int, float, std::string> cell_t;
typedef std::vector<cell_t> db_row_t;

// This is just an example, no actual work with database.
db_row_t get_row(const char* /*query*/) {
    // See recipe "Type 'reference to string'"
    // for a better type for 'query' parameter.
    db_row_t row;
    row.push_back(10);
    row.push_back(10.1f);
    row.push_back("hello again");
    return row;
}

// This is how code required to sum values
// We can provide no template parameter
// to boost::static_visitor<> if our visitor returns nothing.
struct db_sum_visitor: public boost::static_visitor<double> {
    double operator()(int value) const {
        return value;
    }
    double operator()(float value) const {
        return value;
    }
    double operator()(const std::string& /*value*/) const {
        return 0.0;
    }
};

int main() {
    db_row_t row = get_row("Query: Give me some row, please.");
    double res = 0.0;
    for (db_row_t::const_iterator it = row.begin(), end = row.end(); it != end; ++it) {
        res += boost::apply_visitor(db_sum_visitor(), *it);
    }
    std::cout << "Sum of arithmetic types in database row is: " << res << std::endl;
}


#include <aslam/backend/ScalarDesignVariable.hpp>
#include <aslam/backend/ErrorTermObservation.hpp>
#include <aslam/backend/ErrorTermMotion.hpp>
#include <aslam/backend/ErrorTermPrior.hpp>
#include <aslam/backend/OptimizationProblem.hpp>
#include <aslam/backend/Optimizer.hpp>
#include <iostream>
// Bring in some random number generation from Schweizer Messer.
#include <sm/random.hpp>
#include <vector>
#include <algorithm>
#include <boost/foreach.hpp>

int main(int argc, char ** argv)
{
  if(argc != 2)
    {
      std::cout << "Usage: example K\n";
      std::cout << "The argument K is the number of timesteps to include in the optimization\n";
      return 1;
    }

  const int K = atoi(argv[1]);

  try
    {
      // The true wall position
      const double true_w = -5.0;

      // The noise properties.
      const double sigma_n = 0.01;
      const double sigma_u = 0.1;
      const double sigma_x = 0.01;

      // Create random odometry
      std::vector<double> true_u_k(K);
      BOOST_FOREACH(double & u, true_u_k)
	{
	  u = sm::random::uniform();
	}

      // Create the noisy odometry
      std::vector<double> u_k(K);
      for(int k = 0; k < K; ++k)
	{
	  u_k[k] = true_u_k[k] + sigma_u * sm::random::normal();
	}

      // Create the states from noisy odometry.
      std::vector<double> x_k(K);
      std::vector<double> true_x_k(K);
      x_k[0] = 0.0;
      true_x_k[0] = 0.0;
      for(int k = 1; k < K; ++k)
	{
	  true_x_k[k] = true_x_k[k-1] + true_u_k[k];
	  x_k[k] = x_k[k-1] + u_k[k];
	}

      // Create the noisy measurments
      std::vector<double> y_k(K);
      for(int k = 0; k < K; ++k)
	{
	  y_k[k] = (1.0 / (true_w - true_x_k[k])) + sigma_n * sm::random::normal();
	}

      // Now we can build an optimization problem.
      boost::shared_ptr<aslam::backend::OptimizationProblem> problem( new aslam::backend::OptimizationProblem);

      // First, create a design variable for the wall position.
      boost::shared_ptr<aslam::backend::ScalarDesignVariable> dv_w(new aslam::backend::ScalarDesignVariable(true_w + sm::random::normal()));
      // Setting this active means we estimate it.
      dv_w->setActive(true);
      // Add it to the optimization problem.
      problem->addDesignVariable(dv_w);

      // Now we add the initial state.
      boost::shared_ptr<aslam::backend::ScalarDesignVariable> dv_x_km1(new aslam::backend::ScalarDesignVariable(x_k[0]));
      // Setting this active means we estimate it.
      dv_x_km1->setActive(true);
      // Add it to the optimization problem.
      problem->addDesignVariable(dv_x_km1);

      // Now create a prior for this initial state.
      boost::shared_ptr<aslam::backend::ErrorTermPrior> prior(new aslam::backend::ErrorTermPrior(dv_x_km1.get(), true_x_k[0], sigma_x * sigma_x));
      // and add it to the problem.
      problem->addErrorTerm(prior);

      // Now march through the states creating design variables,
      // odometry error terms and measurement error terms.
      for(int k = 1; k < K; ++k)
	{
	  boost::shared_ptr<aslam::backend::ScalarDesignVariable> dv_x_k(new aslam::backend::ScalarDesignVariable(x_k[k]));
	  dv_x_k->setActive(true);
	  problem->addDesignVariable(dv_x_k);

	  // Create odometry error
	  boost::shared_ptr<aslam::backend::ErrorTermMotion> em(new aslam::backend::ErrorTermMotion(dv_x_km1.get(), dv_x_k.get(), u_k[k], sigma_u * sigma_u));
	  problem->addErrorTerm(em);

	  // Create observation error
	  boost::shared_ptr<aslam::backend::ErrorTermObservation> eo(new aslam::backend::ErrorTermObservation(dv_x_k.get(), dv_w.get(), y_k[k], sigma_n * sigma_n));
	  problem->addErrorTerm(eo);

	  // Move this design variable to the x_{k-1} position for use in the next loop.
	  dv_x_km1 = dv_x_k;
	}

      // Now we have a valid optimization problem full of design variables and error terms.
      // Create some optimization options.
      aslam::backend::OptimizerOptions options;
      options.verbose = true;
      options.linearSolver = "cholmod";
      options.levenbergMarquardtLambdaInit = 10;
      options.doSchurComplement = false;
      options.doLevenbergMarquardt = true;
      // Force it to over-optimize
      options.convergenceDeltaX = 1e-12;
      options.convergenceDeltaJ = 1e-12;
      // Then create the optimizer and go!
      aslam::backend::Optimizer optimizer(options);
      optimizer.setProblem( problem );
      optimizer.optimize();

    }
  catch(const std::exception & e)
    {
      std::cout << "Exception during processing: " << e.what();
      return 1;
    }

  std::cout << "Processing completed successfully\n";
  return 0;
}


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_ARITHMETIC_FUNCTIONS_SCALAR_ISQRT_HPP_INCLUDED
#define BOOST_SIMD_ARITHMETIC_FUNCTIONS_SCALAR_ISQRT_HPP_INCLUDED

#include <boost/simd/arithmetic/functions/isqrt.hpp>
#include <boost/simd/include/functions/scalar/is_ltz.hpp>
#include <boost/simd/include/functions/scalar/sqrt.hpp>
#include <boost/simd/include/functions/scalar/toint.hpp>
#include <boost/simd/include/constants/zero.hpp>
#include <boost/dispatch/meta/as_integer.hpp>

namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::isqrt_, tag::cpu_
                            , (A0)
                            , (scalar_< arithmetic_<A0> >)
                            )
  {
    typedef typename  dispatch::meta::as_integer<A0>::type result_type;
    BOOST_SIMD_FUNCTOR_CALL(1)
    {
      return itrunc(boost::simd::sqrt(a0));
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::isqrt_, tag::cpu_
                            , (A0)
                            , (scalar_< uint_<A0> >)
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL(1)
    {
      return static_cast<A0>(boost::simd::sqrt(result_type(a0)));
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::isqrt_, tag::cpu_
                            , (A0)
                            , (scalar_< int_<A0> >)
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL(1)
    {
      return (is_ltz(a0)) ?  Zero<A0>() : static_cast<A0>(boost::simd::sqrt(result_type(a0)));
    }
  };
} } }

#endif


#include "Common.h"
#include "CPS3.h"
#include "PropertiesHolder/PropertiesHolder.h"
#include "Material.h"
#include <Eigen/Dense>

void CPS3::SetIndices(const std::vector<int>& indices)
{
	assert(indices.size() == 3);
	m_nodes[0] = indices[0];
	m_nodes[1] = indices[1];
	m_nodes[2] = indices[2];
}

std::vector<int> CPS3::GetIndices() const
{
	std::vector<int> indices(3);
	indices[0] = m_nodes[0];
	indices[1] = m_nodes[1];
	indices[2] = m_nodes[2];
	return indices;
}

std::vector<Eigen::Vector3f> CPS3::GetFunctionValuesAtNodes(const Eigen::VectorXf& deforms)const
{
	Eigen::Matrix<float, 6, 1> uv;
	for (int i = 0; i < 3; ++i)
	{
		uv[2 * i + 0] = deforms[2 * m_nodes[i] + 0];
		uv[2 * i + 1] = deforms[2 * m_nodes[i] + 1];
	}

	Eigen::Vector3f strain = m_B * uv;

	std::vector<Eigen::Vector3f> output;
	for (int i = 0; i < 3; i++)
	{
		output.push_back(strain);
	}
	return output;
}

void CPS3::CalcK(const StrideDataArray& nodes, const tfem::MaterialPtr mat, std::vector<Eigen::Triplet<float> >& tripletVector)
{
	m_mat = mat;
	Eigen::Vector3f x;
	Eigen::Vector3f y;
	for (int i = 0; i < 3; ++i)
	{
		x[i] = nodes(m_nodes[i], 0);
		y[i] = nodes(m_nodes[i], 1);
	}
	Eigen::Matrix3f C;
	C << Eigen::Vector3f(1.0f, 1.0f, 1.0f), x, y;

	float area = C.determinant() / 2.0f;

	Eigen::Matrix3f IC = C.inverse();

	for (int i = 0; i < 3; i++)
	{
		m_B(0, 2 * i + 0) = IC(1, i);
		m_B(0, 2 * i + 1) = 0.0f;
		m_B(1, 2 * i + 0) = 0.0f;
		m_B(1, 2 * i + 1) = IC(2, i);
		m_B(2, 2 * i + 0) = IC(2, i);
		m_B(2, 2 * i + 1) = IC(1, i);
	}
	Eigen::Matrix<float, 6, 6> K = m_B.transpose() * mat->GetElasticityMatrix(fem::PT_FlatStress) * m_B * area;
	GrabTriplets(K, tripletVector);
}

void CPS3::GrabTriplets(const Eigen::Matrix<float, 6, 6>& K, std::vector<Eigen::Triplet<float> >& tripletVector) const
{
	for (int i = 0; i < 3; i++)
	{
		for (int j = 0; j < 3; j++)
		{
			Eigen::Triplet<float> trplt11(2 * m_nodes[i] + 0, 2 * m_nodes[j] + 0, K(2 * i + 0, 2 * j + 0));
			Eigen::Triplet<float> trplt12(2 * m_nodes[i] + 0, 2 * m_nodes[j] + 1, K(2 * i + 0, 2 * j + 1));
			Eigen::Triplet<float> trplt21(2 * m_nodes[i] + 1, 2 * m_nodes[j] + 0, K(2 * i + 1, 2 * j + 0));
			Eigen::Triplet<float> trplt22(2 * m_nodes[i] + 1, 2 * m_nodes[j] + 1, K(2 * i + 1, 2 * j + 1));

			tripletVector.push_back(trplt11);
			tripletVector.push_back(trplt12);
			tripletVector.push_back(trplt21);
			tripletVector.push_back(trplt22);
		}
	}
}

tfem::Material* CPS3::GetMaterial()
{
	return m_mat.get();
}

IElement* CPS3::Create()
{
	return new CPS3;
}

CPS3::CPS3()
{

}


#include <engine/Solver_Kernels.hpp>
#include <utility/Constants.hpp>
#include <Eigen/Dense>
#include <engine/Backend_par.hpp>

using namespace Utility;
using Utility::Constants::Pi;

namespace Engine
{
namespace Solver_Kernels
{

    void sib_transform(const vectorfield & spins, const vectorfield & force, vectorfield & out)
    {
        int n = spins.size();

        auto s = spins.data();
        auto f = force.data();
        auto o = out.data();

        Backend::par::apply( n, [s,f,o] SPIRIT_LAMBDA (int idx) {
            Vector3 e1, a2, A;
            scalar detAi;
            e1 = s[idx];
            A = 0.5 * f[idx];

            // 1/determinant(A)
            detAi = 1.0 / (1 + pow(A.norm(), 2.0));

            // calculate equation witho the predictor?
            a2 = e1 - e1.cross(A);

            o[idx][0] = (a2[0] * (A[0] * A[0] + 1   ) + a2[1] * (A[0] * A[1] - A[2]) + a2[2] * (A[0] * A[2] + A[1])) * detAi;
            o[idx][1] = (a2[0] * (A[1] * A[0] + A[2]) + a2[1] * (A[1] * A[1] + 1   ) + a2[2] * (A[1] * A[2] - A[0])) * detAi;
            o[idx][2] = (a2[0] * (A[2] * A[0] - A[1]) + a2[1] * (A[2] * A[1] + A[0]) + a2[2] * (A[2] * A[2] + 1   )) * detAi;
        } );
    }

    void oso_calc_gradients(vectorfield & grad, const vectorfield & spins, const vectorfield & forces)
    {
        const Matrix3 t = ( Matrix3() << 0,0,1,0,-1,0,1,0,0 ).finished();

        auto g=grad.data();
        auto s=spins.data();
        auto f=forces.data();

        Backend::par::apply( spins.size(), [g,s,f,t] SPIRIT_LAMBDA (int idx)
            {
                g[idx] = t * (-s[idx].cross(f[idx]));
            }
        );
    }

    void oso_rotate( std::vector<std::shared_ptr<vectorfield>> & configurations, std::vector<vectorfield> & searchdir)
    {
        int noi = configurations.size();
        int nos = configurations[0]->size();
        for(int img=0; img<noi; ++img)
        {

            auto s  = configurations[img]->data();
            auto sd = searchdir[img].data();

            Backend::par::apply( nos, [s, sd] SPIRIT_LAMBDA (int idx)
                {
                    scalar theta = (sd[idx]).norm();
                    scalar q = cos(theta), w = 1-q,
                           x = -sd[idx][0]/theta, y = -sd[idx][1]/theta, z = -sd[idx][2]/theta,
                           s1 = -y*z*w, s2 = x*z*w, s3 = -x*y*w,
                           p1 = x*sin(theta), p2 = y*sin(theta), p3 = z*sin(theta);

                    scalar t1, t2, t3;
                    if(theta > 1.0e-20) // if theta is too small we do nothing
                    {
                        t1 = (q+z*z*w) * s[idx][0] + (s1+p1)   * s[idx][1] + (s2+p2)   * s[idx][2];
                        t2 = (s1-p1)   * s[idx][0] + (q+y*y*w) * s[idx][1] + (s3+p3)   * s[idx][2];
                        t3 = (s2-p2)   * s[idx][0] + (s3-p3)   * s[idx][1] + (q+x*x*w) * s[idx][2];
                        s[idx][0] = t1;
                        s[idx][1] = t2;
                        s[idx][2] = t3;
                    };
                }
            );
        }
    }

    scalar maximum_rotation(const vectorfield & searchdir, scalar maxmove)
    {
        int nos = searchdir.size();
        scalar theta_rms = 0;
        theta_rms = sqrt( Backend::par::reduce(searchdir, [] SPIRIT_LAMBDA (const Vector3 & v){ return v.squaredNorm(); }) / nos );
        scalar scaling = (theta_rms > maxmove) ? maxmove/theta_rms : 1.0;
        return scaling;
    }

    void atlas_rotate(std::vector<std::shared_ptr<vectorfield>> & configurations, const std::vector<scalarfield> & a3_coords, const std::vector<vector2field> & searchdir)
    {
        int noi = configurations.size();
        int nos = configurations[0]->size();
        for(int img=0; img<noi; img++ )
        {
            auto spins = configurations[img]->data();
            auto d     = searchdir[img].data();
            auto a3    = a3_coords[img].data();
            Backend::par::apply(nos, [nos, spins, d, a3] SPIRIT_LAMBDA (int idx) {
                const scalar gamma = (1 + spins[idx][2] * a3[idx]);
                const scalar denom = (spins[idx].head<2>().squaredNorm())/gamma + 2 * d[idx].dot( spins[idx].head<2>() ) + gamma * d[idx].squaredNorm();
                spins[idx].head<2>() = 2*(spins[idx].head<2>() + d[idx]*gamma);
                spins[idx][2] = a3[idx] * (gamma - denom);
                spins[idx] *= 1/(gamma + denom);
            } );
        }
    }

    void atlas_calc_gradients(vector2field & residuals, const vectorfield & spins, const vectorfield & forces, const scalarfield & a3_coords)
    {
        auto s = spins.data();
        auto a3 = a3_coords.data();
        auto g = residuals.data();
        auto f = forces.data();

        Backend::par::apply(spins.size(), [s, a3, g, f] SPIRIT_LAMBDA (int idx) {

            scalar J00 =  s[idx][1] * s[idx][1] + s[idx][2]*(s[idx][2] + a3[idx]);
            scalar J10 = -s[idx][0] * s[idx][1];
            scalar J01 = -s[idx][0] * s[idx][1];
            scalar J11 =  s[idx][0] * s[idx][0]  + s[idx][2]*(s[idx][2] + a3[idx]);
            scalar J02 = -s[idx][0] * (s[idx][2] + a3[idx]);
            scalar J12 = -s[idx][1] * (s[idx][2] + a3[idx]);

            g[idx][0] = -(J00 * f[idx][0] + J01 * f[idx][1] + J02 * f[idx][2]);
            g[idx][1] = -(J10 * f[idx][0] + J11 * f[idx][1] + J12 * f[idx][2]);
        });
    }

    bool ncg_atlas_check_coordinates(const std::vector<std::shared_ptr<vectorfield>> & spins, std::vector<scalarfield> & a3_coords, scalar tol)
    {
        int noi = spins.size();
        int nos = (*spins[0]).size();

        // We use `int` instead of `bool`, because somehow cuda does not like pointers to bool
        // TODO: fix in future
        field<int> result = field<int>(1, int(false));

        for(int img=0; img<noi; img++)
        {
            auto s = spins[0]->data();
            auto a3 = a3_coords[img].data();
            int *res = &result[0];

            Backend::par::apply( nos, [s, a3, tol, res] SPIRIT_LAMBDA (int idx) {
                    if (s[idx][2]*a3[idx] < tol && res[0] == int(false))
                        res[0] = int(true);
            } );
        }

        return bool(result[0]);
    }

    void lbfgs_atlas_transform_direction(std::vector<std::shared_ptr<vectorfield>> & configurations, std::vector<scalarfield> & a3_coords, std::vector<field<vector2field>> & atlas_updates, std::vector<field<vector2field>> & grad_updates, std::vector<vector2field> & searchdir, std::vector<vector2field> & grad_pr, scalarfield & rho)
    {
        int noi = configurations.size();
        int nos = configurations[0]->size();

        for(int n=0; n<atlas_updates[0].size(); n++)
        {
            rho[n] = 1/rho[n];
        }

        for(int img=0; img<noi; img++)
        {
            auto s = (*configurations[img]).data();
            auto a3 = a3_coords[img].data();
            auto sd = searchdir[img].data();
            auto g_pr = grad_pr[img].data();
            auto rh = rho.data();

            auto n_mem = atlas_updates[img].size();

            field<Vector2*> t1(n_mem), t2(n_mem);
            for(int n=0; n<n_mem; n++)
            {
                t1[n] = (atlas_updates[img][n].data());
                t2[n] = (grad_updates[img][n].data());
            }

            auto a_up = t1.data();
            auto g_up = t2.data();

            Backend::par::apply(nos, [s, a3, sd, g_pr, rh, a_up, g_up, n_mem] SPIRIT_LAMBDA (int idx) {
                scalar factor = 1;
                if( s[idx][2]*a3[idx] < 0 )
                {
                    // Transform coordinates to optimal map
                    a3[idx] = (s[idx][2] > 0) ? 1 : -1;
                    factor  = (1 - a3[idx] * s[idx][2]) / (1 + a3[idx] * s[idx][2]);
                    sd[idx]   *= factor;
                    g_pr[idx] *= factor;

                    for(int n=0; n<n_mem; n++)
                    {
                        rh[n] = rh[n] + (factor*factor-1) * a_up[n][idx].dot(g_up[n][idx]);
                        a_up[n][idx] *= factor;
                        g_up[n][idx] *= factor;
                    }
                }
            });
        }

        for(int n=0; n<atlas_updates[0].size(); n++)
        {
            rho[n] = 1/rho[n];
        }
    }
}
}


#include <NTL/GF2E.h>
#include <NTL/GF2XFactoring.h>

#include "LHLExtractor.h"

uint32_t LHLExtractor::getInputLength() const {
    return n_;
}

uint32_t LHLExtractor::getSeedLen() const {
    return n_;
}

uint32_t LHLExtractor::getOutputLen() const {
    return m_;
}

double LHLExtractor::getMinEntropy() const {
    return k_;
}

double LHLExtractor::getError() const {
    return error_;
}

Bitstring LHLExtractor::extract(const WeakSource &input, const Bitstring &seed) const {
    assert(input.getMinEntropy() >= k_);
    assert(seed.size() == n_);
    NTL::GF2EPush push;
    NTL::GF2E::init(NTL::BuildSparseIrred_GF2X(n_));
    auto x = input.getData().asGF2E();
    auto y = seed.asGF2E();
    return Bitstring((Bitstring(x) + Bitstring(y)).substr(0, m_));
}

LHLExtractor::LHLExtractor(uint32_t n, uint32_t k, double eps): n_(n), k_(k), error_(eps) {
    m_ = ceil(k_ + n_ - 2 * log2(1. / error_));
}


#include "advent.hpp"

#include <cstdint>
#include <fstream>
#include <gsl/gsl_util>
#include <iostream>
#include <scn/scn.h>
#include <Eigen/Dense>

auto day05(int argc, char** argv) -> int
{
    std::ifstream infile(argv[1]); // NOLINT
    std::string str;

    using point = std::tuple<int64_t, int64_t>;
    using line = std::pair<point, point>;
    std::vector<line> lines;

    while(std::getline(infile, str)) {
        std::vector<std::string> tokens;
        util::tokenize(str, ' ', tokens);
        int x = 0;
        int y = 0;
        scn::scan(tokens[0], "{},{}", x, y);
        point p1{x, y};
        scn::scan(tokens[2], "{},{}", x, y);
        point p2{x, y};
        lines.emplace_back(p1, p2);
    }

    int64_t xmax = std::numeric_limits<int64_t>::min();
    int64_t ymax = xmax;

    for(auto [p1, p2] : lines) {
        auto [x1, y1] = p1;
        auto [x2, y2] = p2;
        xmax = std::max({xmax, x1, x2});
        ymax = std::max({ymax, y1, y2});
    }

    Eigen::Array<int, -1, -1> map = decltype(map)::Zero(xmax + 1, ymax + 1);
    for (auto [p1, p2] : lines) {
        auto [x1, y1] = p1;
        auto [x2, y2] = p2;

        if (x1 == x2 || y1 == y2 || std::abs(x1-x2) == std::abs(y1-y2)) {
            auto dx = util::sgn(x2-x1);
            auto dy = util::sgn(y2-y1);
            for (auto x = x1, y = y1; (dx == 0 || x != x2) && (dy == 0 || y != y2); x += dx, y += dy) {
                map(x, y) += 1;
            }
            map(x2, y2) += 1;
        }
    }
    auto count = (map > 1).count();
    fmt::print("count: {}\n", count);

    return 0;
}


/**TODO:  Add copyright*/

#define BOOST_TEST_MODULE Preprocessing test suite
#include <boost/test/included/unit_test.hpp>
#include <EvoNet/core/Preprocessing.h>

using namespace EvoNet;
using namespace std;

BOOST_AUTO_TEST_SUITE(preprocessing)

BOOST_AUTO_TEST_CASE(P_selectRandomElement)
{
	// [TODO: make test; currently, combined with selectRandomNode1]
}

BOOST_AUTO_TEST_CASE(P_UnitScaleFunctor)
{
	Eigen::Tensor<float, 2> data(2, 2);
	data.setValues({{ 0, 2 }, { 3, 4 }});
	UnitScaleFunctor<float> unit_scale(data);
	BOOST_CHECK_CLOSE(unit_scale.getUnitScale(), 0.25, 1e-6);

	Eigen::Tensor<float, 2> data_test = data.unaryExpr(UnitScaleFunctor<float>(data));

	BOOST_CHECK_CLOSE(data_test(0, 0), 0.0, 1e-6);
	BOOST_CHECK_CLOSE(data_test(1, 1), 1.0, 1e-6);
}

BOOST_AUTO_TEST_CASE(P_LinearScaleFunctor)
{
	Eigen::Tensor<float, 2> data(2, 2);
	data.setValues({ { 0, 2 }, { 4, 8 } });

	Eigen::Tensor<float, 2> data_test = data.unaryExpr(LinearScaleFunctor<float>(0, 8, -1, 1));

	BOOST_CHECK_CLOSE(data_test(0, 0), -1.0, 1e-6);
	BOOST_CHECK_CLOSE(data_test(0, 1), -0.5, 1e-6);
	BOOST_CHECK_CLOSE(data_test(1, 0), 0.0, 1e-6);
	BOOST_CHECK_CLOSE(data_test(1, 1), 1.0, 1e-6);
}

BOOST_AUTO_TEST_CASE(P_LinearScale)
{
  Eigen::Tensor<float, 3> data(2, 2, 2);
  data.setValues({
    {{ 0, 2 }, { 4, 8 }},
    {{ 1, 1 }, { 3, 5 }}
    });

  // Test default initialization for the domain and setters
  LinearScale<float, 3> linearScale1(-1, 1);
  linearScale1.setDomain(0, 8);
  Eigen::Tensor<float, 3> data_test = linearScale1(data);

  BOOST_CHECK_CLOSE(data_test(0, 0, 0), -1.0, 1e-6);
  BOOST_CHECK_CLOSE(data_test(0, 0, 1), -0.5, 1e-6);
  BOOST_CHECK_CLOSE(data_test(0, 1, 0), 0.0, 1e-6);
  BOOST_CHECK_CLOSE(data_test(0, 1, 1), 1.0, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 0, 0), -0.75, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 0, 1), -0.75, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 1, 0), -0.25, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 1, 1), 0.25, 1e-6);

  // Test with manual domain and range initialization
  LinearScale<float, 3> linearScale(0, 8, -1, 1);
  data_test = linearScale(data);

  BOOST_CHECK_CLOSE(data_test(0, 0, 0), -1.0, 1e-6);
  BOOST_CHECK_CLOSE(data_test(0, 0, 1), -0.5, 1e-6);
  BOOST_CHECK_CLOSE(data_test(0, 1, 0), 0.0, 1e-6);
  BOOST_CHECK_CLOSE(data_test(0, 1, 1), 1.0, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 0, 0), -0.75, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 0, 1), -0.75, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 1, 0), -0.25, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 1, 1), 0.25, 1e-6);

  // Test with domain calculation and range initialization
  LinearScale<float, 3> linearScale2(data, -1, 1);
  data_test = linearScale2(data);

  BOOST_CHECK_CLOSE(data_test(0, 0, 0), -1.0, 1e-6);
  BOOST_CHECK_CLOSE(data_test(0, 0, 1), -0.5, 1e-6);
  BOOST_CHECK_CLOSE(data_test(0, 1, 0), 0.0, 1e-6);
  BOOST_CHECK_CLOSE(data_test(0, 1, 1), 1.0, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 0, 0), -0.75, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 0, 1), -0.75, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 1, 0), -0.25, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 1, 1), 0.25, 1e-6);
}

BOOST_AUTO_TEST_CASE(P_Standardize)
{
  Eigen::Tensor<float, 3> data(2, 2, 2);
  data.setValues({
    {{ 0, 2 }, { 4, 8 }},
    {{ 1, 3 }, { 3, 5 }}
    });

  // Test default initialization with setters and getters
  Standardize<float, 3> standardize1;
  standardize1.setMeanAndVar(1, 2);
  BOOST_CHECK_CLOSE(standardize1.getMean(), 1, 1e-6);
  BOOST_CHECK_CLOSE(standardize1.getVar(), 2, 1e-6);
  standardize1.setMeanAndVar(data);
  BOOST_CHECK_CLOSE(standardize1.getMean(), 3.25, 1e-6);
  BOOST_CHECK_CLOSE(standardize1.getVar(), 6.21428585, 1e-6);

  // Test with data initialization and getters
  Standardize<float, 3> standardize(data);
  BOOST_CHECK_CLOSE(standardize.getMean(), 3.25, 1e-6);
  BOOST_CHECK_CLOSE(standardize.getVar(), 6.21428585, 1e-6);

  // Test operator
  Eigen::Tensor<float, 3> data_test = standardize(data);
  BOOST_CHECK_CLOSE(data_test(0, 0, 0), -1.30373025, 1e-6);
  BOOST_CHECK_CLOSE(data_test(0, 0, 1), -0.501434684, 1e-6);
  BOOST_CHECK_CLOSE(data_test(0, 1, 0), 0.300860822, 1e-6);
  BOOST_CHECK_CLOSE(data_test(0, 1, 1), 1.90545189, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 0, 0), -0.902582467, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 0, 1), -0.100286946, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 1, 0), -0.100286946, 1e-6);
  BOOST_CHECK_CLOSE(data_test(1, 1, 1), 0.702008605, 1e-6);
}

BOOST_AUTO_TEST_CASE(P_MakeShuffleMatrix)
{
  const int shuffle_dim_size = 8;
  std::vector<int> indices = { 0, 1, 2, 3, 4, 5, 6, 7 };

  // Test default initialization with setters and getters
  MakeShuffleMatrix<float> shuffle1;
  shuffle1.setIndices(8);
  BOOST_CHECK(shuffle1.getIndices() != indices);
  for (int i = 0; i < shuffle_dim_size; ++i) {
    BOOST_CHECK_GE(shuffle1.getIndices().at(i), 0);
    BOOST_CHECK_LE(shuffle1.getIndices().at(i), 7);
  }
  shuffle1.setShuffleMatrix(true);
  //std::cout << "Shuffle_matrix\n" << shuffle1.getShuffleMatrix() << std::endl;
  for (int i = 0; i < shuffle_dim_size; ++i) {
    Eigen::Tensor<float, 0> row_sum = shuffle1.getShuffleMatrix().chip(i, 0).sum();
    BOOST_CHECK_EQUAL(row_sum(0), 1);
  }

  // Test initialization with dim size
  MakeShuffleMatrix<float> shuffle2(shuffle_dim_size, true);
  BOOST_CHECK(shuffle2.getIndices() != indices);
  for (int i = 0; i < shuffle_dim_size; ++i) {
    BOOST_CHECK_GE(shuffle2.getIndices().at(i), 0);
    BOOST_CHECK_LE(shuffle2.getIndices().at(i), 7);
  }

  // Test initialization with indices to use
  MakeShuffleMatrix<float> shuffle3(indices, true);
  BOOST_CHECK(shuffle3.getIndices() == indices);
  //std::cout << "Shuffle_matrix\n" << shuffle3.getShuffleMatrix() << std::endl;
  for (int i = 0; i < shuffle_dim_size; ++i) {
    BOOST_CHECK_EQUAL(shuffle3.getShuffleMatrix()(i, i), 1);
    Eigen::Tensor<float, 0> row_sum = shuffle3.getShuffleMatrix().chip(i, 0).sum();
    BOOST_CHECK_EQUAL(row_sum(0), 1);
  }

  // Test row/column shuffling on toy data
  Eigen::Tensor<float, 2> data(2, 3);
  data.setValues({ {1,2,3},{4,5,6} });
  MakeShuffleMatrix<float> shuffle_col(std::vector<int>({1,2,0}), true);
  Eigen::Tensor<float, 2> col_shuffle = data;
  shuffle_col(col_shuffle, true);
  BOOST_CHECK_EQUAL(col_shuffle(0, 0), 2);
  BOOST_CHECK_EQUAL(col_shuffle(0, 1), 3);
  BOOST_CHECK_EQUAL(col_shuffle(0, 2), 1);
  BOOST_CHECK_EQUAL(col_shuffle(1, 0), 5);
  BOOST_CHECK_EQUAL(col_shuffle(1, 1), 6);
  BOOST_CHECK_EQUAL(col_shuffle(1, 2), 4);
  MakeShuffleMatrix<float> shuffle_row(std::vector<int>({ 1,0 }), false);
  Eigen::Tensor<float, 2> row_shuffle = data;
  shuffle_row(row_shuffle, false);
  BOOST_CHECK_EQUAL(row_shuffle(0, 0), 4);
  BOOST_CHECK_EQUAL(row_shuffle(0, 1), 5);
  BOOST_CHECK_EQUAL(row_shuffle(0, 2), 6);
  BOOST_CHECK_EQUAL(row_shuffle(1, 0), 1);
  BOOST_CHECK_EQUAL(row_shuffle(1, 1), 2);
  BOOST_CHECK_EQUAL(row_shuffle(1, 2), 3);

  // Test row/column shuffling on toy data
  Eigen::Tensor<double, 2> data_db(2, 3);
  data_db.setValues({ {1,2,3},{4,5,6} });
  MakeShuffleMatrix<double> shuffle_col_db(std::vector<int>({ 1,2,0 }), true);
  Eigen::Tensor<double, 2> col_shuffle_db = data_db;
  shuffle_col_db(col_shuffle_db, true);
  BOOST_CHECK_EQUAL(col_shuffle_db(0, 0), 2);
  BOOST_CHECK_EQUAL(col_shuffle_db(0, 1), 3);
  BOOST_CHECK_EQUAL(col_shuffle_db(0, 2), 1);
  BOOST_CHECK_EQUAL(col_shuffle_db(1, 0), 5);
  BOOST_CHECK_EQUAL(col_shuffle_db(1, 1), 6);
  BOOST_CHECK_EQUAL(col_shuffle_db(1, 2), 4);
  MakeShuffleMatrix<double> shuffle_row_db(std::vector<int>({ 1,0 }), false);
  Eigen::Tensor<double, 2> row_shuffle_db = data_db;
  shuffle_row_db(row_shuffle_db, false);
  BOOST_CHECK_EQUAL(row_shuffle_db(0, 0), 4);
  BOOST_CHECK_EQUAL(row_shuffle_db(0, 1), 5);
  BOOST_CHECK_EQUAL(row_shuffle_db(0, 2), 6);
  BOOST_CHECK_EQUAL(row_shuffle_db(1, 0), 1);
  BOOST_CHECK_EQUAL(row_shuffle_db(1, 1), 2);
  BOOST_CHECK_EQUAL(row_shuffle_db(1, 2), 3);
}

BOOST_AUTO_TEST_CASE(P_LabelSmoother)
{
	Eigen::Tensor<float, 1> data(2);
	data.setValues({ 0, 1 });

	Eigen::Tensor<float, 1> data_test = data.unaryExpr(LabelSmoother<float>(0.1, 0.2));

	BOOST_CHECK_CLOSE(data_test(0), 0.1, 1e-4);
	BOOST_CHECK_CLOSE(data_test(1), 0.8, 1e-4);
}

BOOST_AUTO_TEST_CASE(P_OneHotEncoder)
{
	// TODO
}

BOOST_AUTO_TEST_CASE(SFcheckNan)
{
	Eigen::Tensor<float, 1> values(2);
	values.setConstant(5.0f);
	Eigen::Tensor<float, 1> test(2);

	// control
  test = values.unaryExpr([](float c) { return checkNan<float>(c); });
	BOOST_CHECK_CLOSE(test(0), 5.0, 1e-3);
	BOOST_CHECK_CLOSE(test(1), 5.0, 1e-3);

	// test
	values(0) = NAN; //NaN
	values(1) = INFINITY; //infinity
  test = values.unaryExpr([](float c) { return checkNan<float>(c); });
	BOOST_CHECK_CLOSE(test(0), NAN, 1e-3);
	BOOST_CHECK_CLOSE(test(1), INFINITY, 1e-3);
}

BOOST_AUTO_TEST_CASE(SFsubstituteNanInf)
{
	Eigen::Tensor<float, 1> values(3);
	values.setConstant(5.0f);
	Eigen::Tensor<float, 1> test(3);

	// control
  test = values.unaryExpr([](float c) { return substituteNanInf<float>(c); });
	BOOST_CHECK_CLOSE(test(0), 5.0, 1e-3);
	BOOST_CHECK_CLOSE(test(1), 5.0, 1e-3);

	// test
	values(0) = NAN; //NaN
	values(1) = INFINITY; //infinity
	values(2) = -INFINITY; //infinity
  test = values.unaryExpr([](float c) { return substituteNanInf<float>(c); });
	BOOST_CHECK_CLOSE(test(0), 0.0, 1e-3);
	BOOST_CHECK_CLOSE(test(1), 1e9, 1e-3);
	BOOST_CHECK_CLOSE(test(2), -1e9, 1e-3);
}

BOOST_AUTO_TEST_CASE(SFClipOp)
{
	Eigen::Tensor<float, 1> net_input(3);
	net_input.setValues({ 0.0f, 1.0f, 0.5f });

	// test input
	Eigen::Tensor<float, 1> result = net_input.unaryExpr(ClipOp<float>(0.1f, 0.0f, 1.0f));
	BOOST_CHECK_CLOSE(result(0), 0.1, 1e-3);
	BOOST_CHECK_CLOSE(result(1), 0.9, 1e-3);
	BOOST_CHECK_CLOSE(result(2), 0.5, 1e-3);
}

BOOST_AUTO_TEST_CASE(SFRandBinaryOp)
{
	Eigen::Tensor<float, 1> net_input(2);
	net_input.setValues({ 2.0f, 2.0f });
	Eigen::Tensor<float, 1> result;

	// test input
	result = net_input.unaryExpr(RandBinaryOp<float>(0.0f));
	BOOST_CHECK_CLOSE(result(0), 2.0, 1e-3);
	BOOST_CHECK_CLOSE(result(1), 2.0, 1e-3);
	result = net_input.unaryExpr(RandBinaryOp<float>(1.0f));
	BOOST_CHECK_CLOSE(result(0), 0.0, 1e-3);
	BOOST_CHECK_CLOSE(result(1), 0.0, 1e-3);
}

BOOST_AUTO_TEST_CASE(assertClose)
{
	BOOST_CHECK(!assert_close<float>(1.1, 1.2, 1e-4, 1e-4));
	BOOST_CHECK(assert_close<float>(1.1, 1.2, 1, 1));
}

BOOST_AUTO_TEST_CASE(P_GaussianMixture)
{
	// TODO
}

BOOST_AUTO_TEST_CASE(P_SwissRoll)
{
	// TODO
}

BOOST_AUTO_TEST_CASE(P_GumbelSampler)
{
	Eigen::Tensor<float, 2> gumbel_samples = GumbelSampler<float>(2, 3);
	BOOST_CHECK_LE(gumbel_samples(0, 0), 10);
	BOOST_CHECK_GE(gumbel_samples(0, 0), -10);
	BOOST_CHECK_LE(gumbel_samples(1, 2), 10);
	BOOST_CHECK_GE(gumbel_samples(1, 2), -10);
	std::cout << gumbel_samples << std::endl;
}

BOOST_AUTO_TEST_CASE(P_GaussianSampler)
{
	Eigen::Tensor<float, 2> gaussian_samples = GaussianSampler<float>(2, 3);
	BOOST_CHECK_LE(gaussian_samples(0, 0), 2);
	BOOST_CHECK_GE(gaussian_samples(0, 0), -2);
	BOOST_CHECK_LE(gaussian_samples(1, 2), 2);
	BOOST_CHECK_GE(gaussian_samples(1, 2), -2);
	std::cout << gaussian_samples << std::endl;
}

BOOST_AUTO_TEST_SUITE_END()


/* Boost test/mul.cpp
 * test multiplication, division, square and square root on some intervals
 *
 * Copyright 2002-2003 Guillaume Melquiond
 *
 * 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 <boost/numeric/interval.hpp>
#include <boost/test/minimal.hpp>
#include "bugs.hpp"

typedef boost::numeric::interval<double> I;

static double min BOOST_PREVENT_MACRO_SUBSTITUTION (double a, double b, double c, double d) {
  return (std::min)((std::min)(a, b), (std::min)(c, d));
}

static double max BOOST_PREVENT_MACRO_SUBSTITUTION (double a, double b, double c, double d) {
  return (std::max)((std::max)(a, b), (std::max)(c, d));
}

static bool test_mul(double al, double au, double bl, double bu) {
  I a(al, au), b(bl, bu);
  I c = a * b;
  return c.lower() == (min)(al*bl, al*bu, au*bl, au*bu)
      && c.upper() == (max)(al*bl, al*bu, au*bl, au*bu);
}

static bool test_mul1(double ac, double bl, double bu) {
  I a(ac), b(bl, bu);
  I c = ac * b;
  I d = b * ac;
  I e = a * b;
  return equal(c, d) && equal(d, e);
}

static bool test_div(double al, double au, double bl, double bu) {
  I a(al, au), b(bl, bu);
  I c = a / b;
  return c.lower() == (min)(al/bl, al/bu, au/bl, au/bu)
      && c.upper() == (max)(al/bl, al/bu, au/bl, au/bu);
}

static bool test_div1(double al, double au, double bc) {
  I a(al, au), b(bc);
  I c = a / bc;
  I d = a / b;
  return equal(c, d);
}

static bool test_div2(double ac, double bl, double bu) {
  I a(ac), b(bl, bu);
  I c = ac / b;
  I d = a / b;
  return equal(c, d);
}

static bool test_square(double al, double au) {
  I a(al, au);
  I b = square(a);
  I c = a * a;
  return b.upper() == c.upper() &&
         (b.lower() == c.lower() || (c.lower() <= 0 && b.lower() == 0));
}

static bool test_sqrt(double al, double au) {
  I a(al, au);
  I b = square(sqrt(a));
  return subset(abs(a), b);
}

int test_main(int, char*[]) {
  BOOST_CHECK(test_mul(2, 3, 5, 7));
  BOOST_CHECK(test_mul(2, 3, -5, 7));
  BOOST_CHECK(test_mul(2, 3, -7, -5));
  BOOST_CHECK(test_mul(-2, 3, 5, 7));
  BOOST_CHECK(test_mul(-2, 3, -5, 7));
  BOOST_CHECK(test_mul(-2, 3, -7, -5));
  BOOST_CHECK(test_mul(-3, -2, 5, 7));
  BOOST_CHECK(test_mul(-3, -2, -5, 7));
  BOOST_CHECK(test_mul(-3, -2, -7, -5));

  BOOST_CHECK(test_mul1(3, 5, 7));
  BOOST_CHECK(test_mul1(3, -5, 7));
  BOOST_CHECK(test_mul1(3, -7, -5));
  BOOST_CHECK(test_mul1(-3, 5, 7));
  BOOST_CHECK(test_mul1(-3, -5, 7));
  BOOST_CHECK(test_mul1(-3, -7, -5));

  BOOST_CHECK(test_div(30, 42, 2, 3));
  BOOST_CHECK(test_div(30, 42, -3, -2));
  BOOST_CHECK(test_div(-30, 42, 2, 3));
  BOOST_CHECK(test_div(-30, 42, -3, -2));
  BOOST_CHECK(test_div(-42, -30, 2, 3));
  BOOST_CHECK(test_div(-42, -30, -3, -2));

  BOOST_CHECK(test_div1(30, 42, 3));
  BOOST_CHECK(test_div1(30, 42, -3));
  BOOST_CHECK(test_div1(-30, 42, 3));
  BOOST_CHECK(test_div1(-30, 42, -3));
  BOOST_CHECK(test_div1(-42, -30, 3));
  BOOST_CHECK(test_div1(-42, -30, -3));

  BOOST_CHECK(test_div2(30, 2, 3));
  BOOST_CHECK(test_div2(30, -3, -2));
  BOOST_CHECK(test_div2(-30, 2, 3));
  BOOST_CHECK(test_div2(-30, -3, -2));

  BOOST_CHECK(test_square(2, 3));
  BOOST_CHECK(test_square(-2, 3));
  BOOST_CHECK(test_square(-3, 2));

  BOOST_CHECK(test_sqrt(2, 3));
  BOOST_CHECK(test_sqrt(5, 7));
  BOOST_CHECK(test_sqrt(-1, 2));

# ifdef __BORLANDC__
  ::detail::ignore_warnings();
# endif
  return 0;
}


#pragma once

#include "craam/RMDP.hpp"

#include <Eigen/Dense>
#include <rm/range.hpp>

namespace craam{namespace algorithms{

using namespace std;
using namespace Eigen;

/// Internal helper functions
namespace internal{

    /// Helper function to deal with variable indexing
    template<class SType>
    inline Transition mean_transition_state(const SType& state, long index, const pair<indvec,vector<numvec>>& policies){
        return state.mean_transition(policies.first[index], policies.second[index]);
    }

    /// Helper function to deal with variable indexing
    template<class SType>
    inline Transition mean_transition_state(const SType& state, long index, const indvec& policy){
        return state.mean_transition(policy[index]);
    }

    /// Helper function to deal with variable indexing
    template<class SType>
    inline prec_t mean_reward_state(const SType& state, long index, const pair<indvec,vector<numvec>>& policies){
        return state.mean_reward(policies.first[index], policies.second[index]);
    }

    /// Helper function to deal with variable indexing
    template<class SType>
    inline prec_t mean_reward_state(const SType& state, long index, const indvec& policy){
        return state.mean_reward(policy[index]);
    }
}

/**
Constructs the transition (or its transpose) matrix for the policy.

\tparam SType Type of the state in the MDP (regular vs robust)
\tparam Policy Type of the policy. Either a single policy for
                the standard MDP evaluation, or a pair of a deterministic
                policy and a randomized policy of the nature
\param rmdp Regular or robust MDP
\param policies The policy (indvec) or the pair of the policy and the policy
        of nature (pair<indvec,vector<numvec> >). The nature is typically
        a randomized policy
\param transpose (optional, false) Whether to return the transpose of the transition matrix.
        This is useful for computing occupancy frequencies
*/
template<typename SType, typename Policies>
inline MatrixXd transition_mat(const GRMDP<SType>& rmdp, const Policies& policies, bool transpose = false) {
    const size_t n = rmdp.state_count();
    MatrixXd result = MatrixXd::Zero(n,n);

    const auto& states = rmdp.get_states();
    #pragma omp parallel for
    for(size_t s = 0; s < n; s++){
        const Transition&& t = internal::mean_transition_state(states[s], s, policies);

        const auto& indexes = t.get_indices();
        const auto& probabilities = t.get_probabilities();

        if(!transpose){
            for(size_t j=0; j < t.size(); j++)
                result(s,indexes[j]) = probabilities[j];
        }else{
            for(size_t j=0; j < t.size(); j++)
                result(indexes[j],s) = probabilities[j];
        }
    }
    return result;
}

/**
Constructs the rewards vector for each state for the RMDP.

\tparam Policy Type of the policy. Either a single policy for
                the standard MDP evaluation, or a pair of a deterministic
                policy and a randomized policy of the nature
\param rmdp Regular or robust MDP
\param policies The policy (indvec) or the pair of the policy and the policy
        of nature (pair<indvec,vector<numvec> >). The nature is typically
        a randomized policy
 */
template<typename SType, typename Policy>
inline numvec rewards_vec(const GRMDP<SType>& rmdp, const Policy& policies){

    const auto n = rmdp.state_count();
    numvec rewards(n);

    #pragma omp parallel for
    for(size_t s=0; s < n; s++){
        const SType& state = rmdp[s];
        if(state.is_terminal())
            rewards[s] = 0;
        else
            rewards[s] = internal::mean_reward_state(state, s, policies);
    }
    return rewards;
}

/**
Computes occupancy frequencies using matrix representation of transition
probabilities. This method may not scale well

\tparam SType Type of the state in the MDP (regular vs robust)
\tparam Policy Type of the policy. Either a single policy for
                the standard MDP evaluation, or a pair of a deterministic
                policy and a randomized policy of the nature
\param init Initial distribution (alpha)
\param discount Discount factor (gamma)
\param policies The policy (indvec) or the pair of the policy and the policy
        of nature (pair<indvec,vector<numvec> >). The nature is typically
        a randomized policy
*/
template<typename SType, typename Policies>
inline numvec
occfreq_mat(const GRMDP<SType>& rmdp, const Transition& init, prec_t discount,
                 const Policies& policies) {
    const auto n = rmdp.state_count();

    // initial distribution
    const numvec& ivec = init.probabilities_vector(n);
    const VectorXd initial_vec = Map<const VectorXd,Unaligned>(ivec.data(),ivec.size());

    // get transition matrix and construct (I - gamma * P^T)
    MatrixXd t_mat = MatrixXd::Identity(n,n)  - discount * transition_mat(rmdp, policies, true);

    // solve set of linear equations
    numvec result(n,0);
    Map<VectorXd,Unaligned>(result.data(),result.size()) = HouseholderQR<MatrixXd>(t_mat).solve(initial_vec);

    return result;
}

}}


////////////////////////////////////////////////////////////////////////////////
/// Copyright 2018-present Xinyan DAI<xinyan.dai@outlook.com>
///
/// permission is hereby granted, free of charge, to any person obtaining a copy
/// of this software and associated documentation files (the "Software"), to
/// deal in the Software without restriction, including without limitation the
/// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
/// sell copies of the Software, and to permit persons to whom the Software is
/// furnished to do so, subject to the following conditions:
///
/// The above copyright notice and this permission notice shall be included in
/// all copies or substantial portions ofthe Software.
///
/// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
/// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
/// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
/// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
/// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
/// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
/// IN THE SOFTWARE.

/// @version 0.1
/// @author  Xinyan DAI
/// @contact xinyan.dai@outlook.com
//////////////////////////////////////////////////////////////////////////////
#pragma once

#include <eigen3/Eigen/Dense>

#include <map>
#include <vector>
#include <random>
#include <iostream>
#include <functional>
#include <boost/progress.hpp>

#include "index.hpp"

namespace ss {
    template<typename DataType>
    class Node {
    protected:
        DataType                            _medians;
        vector<DataType >                   _projector;
        vector<int >                        _idx;
        Node                                *_left;
        Node                                *_right;
        const u_int                         _depth;
        const u_int                         _dim;
        const u_int                         _max_depth;

    public:
        explicit Node(u_int dim, u_int depth, u_int max_depth) :
            _depth(depth), _max_depth(max_depth), _dim(dim), _projector(dim) {
        }

        bool IsLeaf() {
            return this->_depth == this->_max_depth;
        }

        void _GenerateProjector(const Matrix<DataType> &data, const vector<int> &idx) {
            // TODO(Xinyan) choose the eigenvector as random projection vector
            std::default_random_engine generator;
            std::normal_distribution<DataType > distribution(0.0, 1.0);

            for (int j=0; j<_dim; j++) {
                this->_projector[j] = distribution(generator);
            }
        }

        void MakeTree(const Matrix<DataType> &data, const vector<int > & idx) {
            if (IsLeaf()) {
                _idx = idx;
                return;
            }
            int N = idx.size();
            this->_GenerateProjector(data, idx);
            vector <DataType> projected_value(idx.size());
            for (int i = 0; i < idx.size(); ++i) {
                projected_value[i] = ss::InnerProduct(data[idx[i]], _projector.data(), _projector.size());
            }
            vector<size_t > sorted_idx = ss::SortIndexes(projected_value);
            vector<size_t > sorted_left_inx = vector<size_t>(sorted_idx.begin(), sorted_idx.begin() + N/2);
            vector<size_t > sorted_right_inx = vector<size_t >(sorted_idx.begin() + N/2, sorted_idx.end());

            this->_medians = (projected_value[sorted_idx[N/2-1]]+ projected_value[sorted_idx[N/2+1]]) / 2.0;

            this->_left = new Node(_dim, _depth + 1, _max_depth);
            this->_right = new Node(_dim, _depth + 1, _max_depth);

            this->_left->MakeTree(data, ss::FancyIndex(idx, sorted_left_inx));
            this->_right->MakeTree(data, ss::FancyIndex(idx, sorted_right_inx));
        }

        const vector<int>& ProbeLeaf(const DataType * query) {
            if (IsLeaf()) {
                return this->_idx;
            }
            DataType projected_value= ss::InnerProduct(query, _projector.data(), _projector.size());
            if (projected_value < _medians) {
                return this->_left->ProbeLeaf(query);
            } else {
                return this->_right->ProbeLeaf(query);
            }
        }

        void LeafIdx(vector<vector<int > * > & leafIdx) {
            if (IsLeaf()) {
                leafIdx.push_back(&this->_idx);
            } else {

                this->_left->LeafIdx(leafIdx);
                this->_right->LeafIdx(leafIdx);
            }
        }

    };

    template<typename DataType>
    class RPTIndex: public Index<DataType > {
    protected:
        Node<DataType >             _root;
    public:
        explicit RPTIndex(const parameter & para) :
                Index <DataType >(para),
                _root(para.dim, 0, para.num_bit) {
        }

        ~RPTIndex() {}

        void Train(const Matrix<DataType> & data) override {
            this->_root.MakeTree(data, ss::Range<int>(0, data.getSize()));

        }

        void Add(const Matrix<DataType> &data) override {
        }

        vector<vector<int > * > LeafIdx() {
            vector<vector<int > * > leafIdx;
            leafIdx.reserve(1 << this->_para.num_bit);
            this->_root.LeafIdx(leafIdx);
            return leafIdx;
        }

        void Search(const DataType* query, const std::function<void (int)>& prober) override {
            const vector<int >& idx = this->_root.ProbeLeaf(query);
            for (int id : idx) {
                prober(id);
            }
        }
    };
} // namespace ss

// ------------------------- implementation -------------------------


/*
 For more information, please see: http://software.sci.utah.edu

 The MIT License

 Copyright (c) 2015 Scientific Computing and Imaging Institute,
 University of Utah.

 Permission is hereby granted, free of charge, to any person obtaining a
 copy of this software and associated documentation files (the "Software"),
 to deal in the Software without restriction, including without limitation
 the rights to use, copy, modify, merge, publish, distribute, sublicense,
 and/or sell copies of the Software, and to permit persons to whom the
 Software is furnished to do so, subject to the following conditions:

 The above copyright notice and this permission notice shall be included
 in all copies or substantial portions of the Software.

 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
 THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 DEALINGS IN THE SOFTWARE.
*/

#include <gtest/gtest.h>
#include <boost/test/floating_point_comparison.hpp>

#include <Testing/Util/MatrixTestUtilities.h>

#include <Core/Math/sci_lapack.h>
#include <Core/Datatypes/MatrixFwd.h>
#include <Core/Datatypes/DenseMatrix.h>
#include <Core/Datatypes/ColumnMatrix.h>
#include <Core/Datatypes/MatrixOperations.h>
#include <Core/Datatypes/MatrixTypeConverter.h>

#include <Core/Exceptions/LapackError.h>

using namespace boost::assign;
using namespace SCIRun;
using namespace SCIRun::TestUtils;

TEST(InvertMatrixTest, InvertZeroException)
{
  const int dim = 3;
  DenseMatrix zero_matrix(dim, dim);
  zero_matrix.zero();

  EXPECT_THROW(lapackinvert(zero_matrix.get_data_pointer(), dim), LapackError);
}

TEST(InvertMatrixTest, InvertArgException)
{
  const int dim = 5;
  const int bad_dim = 4;
  DenseMatrixHandle identity(DenseMatrix::identity(dim));
  DenseMatrixHandle identityOriginal(DenseMatrix::identity(dim));

  EXPECT_THROW(lapackinvert(identity->get_data_pointer(), bad_dim), LapackError);

  EXPECT_FALSE(compare_exactly(*identity, *identityOriginal));
}

TEST(SolveLinSysWithLapackTest, SolveLinSysZeroException)
{
  const int dim = 6;
  DenseMatrix zero_matrix(dim, dim);
  zero_matrix.zero();
  ColumnMatrix rhs = MAKE_COLUMN_MATRIX((107) (60) (71) (43) (82));

  EXPECT_THROW(lapacksolvelinearsystem(zero_matrix.get_raw_2D_pointer(), dim, dim,
                                       rhs.get_data_pointer(),  dim, 1),
               LapackError);
}

TEST(InvertMatrixTest, CanInvertIdentity)
{
  const int dim = 3;
  DenseMatrixHandle identity(DenseMatrix::identity(dim));
  DenseMatrixHandle identityOriginal(DenseMatrix::identity(dim));

  EXPECT_NO_THROW(lapackinvert(identity->get_data_pointer(), dim));

  EXPECT_TRUE(compare_exactly(*identity, *identityOriginal));
}

TEST(SolveLinSysWithLapackTest, SolvingSimpleCase)
{
  DenseMatrix M = MAKE_DENSE_MATRIX(
    (2, 4, 7, 9, 8)
    (6, 9, 2, 5, 2)
    (6, 3, 5, 1, 8)
    (1, 5, 6, 1, 2)
    (1, 2, 8, 2, 9));

   ColumnMatrix rhs = MAKE_COLUMN_MATRIX((107) (60) (71) (43) (82));

   EXPECT_NO_THROW(lapacksolvelinearsystem(M.get_raw_2D_pointer(), 5, 5, rhs.get_data_pointer(),  5, 1));

   EXPECT_COLUMN_MATRIX_EQ_TO(rhs,
    (1.0)
    (2.0)
    (3.0)
    (4.0)
    (5.0));

}

TEST(InvertMatrixTest, CanInvertWithMemberFunction)
{
  const int rows = 3, cols = 3;

  DenseMatrix m = MAKE_DENSE_MATRIX(
    (1, 0, 1)
    (0, 2, 0)
    (0, 0, -1));

  const MatrixHandle original(m.clone());
  //std::cout << "Matrix:" << std::endl;
  //std::cout << matrix_to_string(m) << std::endl;

  EXPECT_TRUE(m.invert());
  const MatrixHandle inverseFromMethod(m.clone());
  //std::cout << "Inverse from method:" << std::endl;
  //std::cout << to_string(inverseFromMethod) << std::endl;

  EXPECT_TRUE(m.invert());
  //std::cout << "Back to original matrix:" << std::endl;
  //std::cout << matrix_to_string(m) << std::endl;

  //std::cout << "Inversion via direct call to lapack:" << std::endl;

  EXPECT_NO_THROW(lapackinvert(m.get_data_pointer(), rows));

  //std::cout << matrix_to_string(m) << std::endl;
  {
    const MatrixHandle inverseFromDirectLapack(m.clone());

    //std::cout << "Difference matrix:" << std::endl;
    MatrixHandle diff = inverseFromDirectLapack - inverseFromMethod;
    //std::cout << to_string(diff) << std::endl;

    DenseMatrixHandle zero(DenseMatrix::zero_matrix(rows, cols));

    EXPECT_TRUE(compare_exactly(*diff, *zero));
  }

  MatrixHandle id3(DenseMatrix::identity(rows));
  MatrixHandle product = inverseFromMethod * original;
  EXPECT_TRUE(compare_exactly(*id3, *product));
}

//Note: this test can be a template for further lapack function testing.
TEST(SVDTest, ExampleFromWikiPage)
{
  const int rows = 4, cols = 5;

  DenseMatrix m = MAKE_DENSE_MATRIX(
    (1, 0, 0, 0, 2)
    (0, 0, 3, 0, 0)
    (0, 0, 0, 0, 0)
    (0, 4, 0, 0, 0));

  DenseMatrix u(rows, rows);
  ColumnMatrix s(rows);
  DenseMatrix v_transpose(cols, cols);

  EXPECT_NO_THROW(lapacksvd(m.get_raw_2D_pointer(), rows, cols,
                            s.get_data_pointer(),
                            u.get_raw_2D_pointer(),
                            v_transpose.get_raw_2D_pointer()));

  EXPECT_MATRIX_EQ_TO(u,
    (0,0,1,0)
    (0,1,0,0)
    (0,0,0,-1)
    (1,0,0,0));

  EXPECT_COLUMN_MATRIX_EQ_TO(s,
    (4.0)
    (3.0)
    (2.23606798)
    (0));

  EXPECT_MATRIX_EQ_TO(v_transpose,
    (0.0,1.0,0.0,0.0,0.0)
    (0,  0,  1,  0,  0)
    (sqrt(0.2), 0, 0, 0, sqrt(0.8))
    (0, 0, 0, 1, 0)
    (-sqrt(0.8), 0, 0, 0, sqrt(0.2)));

  MatrixHandle U(u.clone());
  MatrixHandle fullS(DenseMatrix::make_diagonal_from_column(s, rows, cols));
  MatrixHandle V_transpose(v_transpose.clone());

  EXPECT_MATRIX_EQ(m, *(U * fullS * V_transpose));
}