#include "../include/cte_bits/hamiltonian.hpp"
#include "../include/cte_bits/random.hpp"
#include <algorithm>
#include <boost/math/tools/roots.hpp>
#include <cmath>
#include <cstdlib>
#include <ibs>
#include <iostream>
#include <map>
#include <math.h>
#include <string>

namespace cte_distributions {
std::vector<double> BiGaussian4D(const double betax, const double ex,
                                 const double betay, const double ey,
                                 int seed) {
  std::vector<double> out;

  // std::printf("%-30s %12.6e \n", "betx", betax);
  // std::printf("%-30s %12.6e \n", "bety", betay);
  // std::printf("%-30s %12.6e \n", "ex", ex);
  // std::printf("%-30s %12.6e \n", "ey", ey);
  static double ampx, ampy, amp, r1, r2, facc;
  static double x, px, y, py;

  // 1 sigma rms beam sizes using average ring betas
  ampx = sqrt(betax * ex);
  ampy = sqrt(betay * ey);

  // generate bi-gaussian distribution in the x-px phase-space
  do {
    r1 = 2 * cte_random::ran3(&seed) - 1;
    r2 = 2 * cte_random::ran3(&seed) - 1;
    amp = r1 * r1 + r2 * r2;
  } while (amp >= 1);

  facc =
      sqrt(-2 * log(amp) /
           amp); // transforming [-1,1] uniform to gaussian - inverse transform

  x = ampx * r1 * facc;  // scaling the gaussian
  px = ampx * r2 * facc; // scaling the gaussian

  // generate bi-gaussian distribution in the y-py phase-space
  do {
    r1 = 2 * cte_random::ran3(&seed) - 1;
    r2 = 2 * cte_random::ran3(&seed) - 1;
    amp = r1 * r1 + r2 * r2;
  } while (amp >= 1);

  facc =
      sqrt(-2 * log(amp) /
           amp); // transforming [-1,1] uniform to gaussian - inverse transform

  y = ampy * r1 * facc;  // scaling the gaussian
  py = ampy * r2 * facc; // scaling the gaussian

  // std::printf("%12.6e %12.6e %12.6e %12.6e\n", x, px, y, py);
  out.push_back(x);
  out.push_back(px);
  out.push_back(y);
  out.push_back(py);

  return out;
}

std::vector<double>
BiGaussian6DLongMatched(double betax, double ex, double betay, double ey,
                        std::vector<double> &h, std::vector<double> &v,
                        std::map<std::string, double> &twheader,
                        std::map<std::string, double> &longparam, int seed) {

  double h0 = h[0];
  double tauhat = twheader["tauhat"];
  double omega = twheader["omega"];
  double ampt = longparam["sigs"] / clight;
  // Max value Hamiltonian that is stable
  // is, with the sign convention used, left of the ham contour
  // at 180-phis (The Ham rises lin to the right.)
  double hom = (h0 * omega);
  double ts = longparam["phis"] / hom;
  int npi = int(longparam["phis"] / (2.0 * pi));
  double tperiod = 2.0 * pi / hom;
  double ts2 = longparam["phisNext"] / hom + double(npi) * tperiod;
  double delta = longparam["sige"];

  std::vector<double> out;
  out = BiGaussian4D(betax, ex, betay, ey, seed);

  // adding two zeros
  out.push_back(0.0);
  out.push_back(0.0);

  double r1, r2, amp, facc, tc, pc, ham, hammin;
  tc = (omega * twheader["eta"] * h0);
  pc = (omega * twheader["CHARGE"]) /
       (2.0 * pi * twheader["PC"] * 1.0e9 * twheader["betar"]);

  // std::printf("%-20s %16.8e\n", "tc", tc);
  // std::printf("%-20s %16.8e\n", "pc", pc);
  // max Hamiltonian
  double hammax =
      cte_hamiltonian::Hamiltonian(twheader, longparam, h, v, tc, ts, 0.0);
  /*
  std::printf("%-20s %16.8e\n", "hammax", hammax);
  std::printf("%-20s %16.8e\n", "ts", ts);
  std::printf("%-20s %16.8e\n", "ts2", ts2);
  */
  // select valid t values
  do {
    // looper++;
    // std::printf("%i\n", looper);
    r1 = 2 * cte_random::ran3(&seed) - 1;
    r2 = 2 * cte_random::ran3(&seed) - 1;
    amp = r1 * r1 + r2 * r2;
    if (amp >= 1)
      continue;

    facc = sqrt(-2 * log(amp) / amp);
    // std::printf("%-20s %16.8e\n", "out4", out[4]);
    out[4] = ts2 + ampt * r1 * facc;
    /*
        std::printf("%-20s %16.8e\n", "ts", ts);
        std::printf("%-20s %16.8e\n", "ampt", ampt);
        std::printf("%-20s %16.8e\n", "r1", r1);
        std::printf("%-20s %16.8e\n", "facc", facc);
        std::printf("%-20s %16.8e\n", "out4", out[4]);
        std::printf("%-20s %16.8e\n", "out4-ts", out[4] - ts);
        std::printf("%-20s %16.8e\n", "tauhat", tauhat);
    */
    if (abs(out[4] - ts2) >= abs(ts - ts2))
      continue;

    // min Hamiltonian
    hammin = cte_hamiltonian::Hamiltonian(twheader, longparam, h, v, tc, out[4],
                                          0.0);
    // std::printf("%-20s %16.8e\n", "hammin", hammin);
    // std::printf("%-20s %16.8e\n", "hammax", hammax);

  } while ((hammin > hammax) || (abs(out[4] - ts2) >= abs(ts - ts2)));

  // select matched deltas
  do {
    // looper++;
    // std::printf("%i\n", looper);
    r1 = 2 * cte_random::ran3(&seed) - 1;
    r2 = 2 * cte_random::ran3(&seed) - 1;
    amp = r1 * r1 + r2 * r2;

    if (amp >= 1)
      continue;

    facc = sqrt(-2 * log(amp) / amp);
    out[5] = longparam["sige"] * r2 * facc;
    // std::printf("%-20s %16.8e\n", "delta", out[5]);

    ham = cte_hamiltonian::Hamiltonian(twheader, longparam, h, v, tc, out[4],
                                       out[5]);
    // std::printf("%-20s %16.8e\n", "hammin", hammin);
    // std::printf("%-20s %16.8e\n", "ham", ham);
    // std::printf("%-20s %16.8e\n", "hammax", hammax);
  } while ((ham < hammin) || (ham > hammax));

  return out;
}

std::vector<std::vector<double>> GenerateDistributionMatched(
    int nMacro, double betax, double ex, double betay, double ey,
    std::vector<double> &h, std::vector<double> &v,
    std::map<std::string, double> &twheader,
    std::map<std::string, double> &longparam, int seed) {
  std::vector<std::vector<double>> out;

  for (int i = 0; i < nMacro; i++) {
    out.push_back(BiGaussian6DLongMatched(betax, ex, betay, ey, h, v, twheader,
                                          longparam, seed));
  }
  return out;
}

} // namespace cte_distributions


//===-- parameter_tests - Parameter Unit Tests ---------------------------------===//
//
// Part of the SAM Project
// Created by Amir Masoud Abdol on 14/01/2021
//
//===----------------------------------------------------------------------===//
///
/// @file
/// This file contains some tests for Parameter<T>
///
//===----------------------------------------------------------------------===//

#define BOOST_TEST_DYN_LINK

#define BOOST_TEST_MODULE Parameter Tests

#include "nlohmann/json.hpp"
#include "Parameter.h"
#include <boost/test/unit_test.hpp>
#include <boost/mpl/list.hpp>
#include <boost/utility/identity_type.hpp>

using namespace sam;
using json = nlohmann::ordered_json;

typedef boost::mpl::list<int, float, float> test_types;
typedef boost::mpl::list<float, float> dist_test_types;

BOOST_AUTO_TEST_CASE_TEMPLATE( parameter_type_test, T, test_types) {

  Parameter<T> parameter{1};

  BOOST_TEST(parameter.n_elem == 1);
  BOOST_TEST(sizeof(parameter.at(0)) == sizeof(T));
}

BOOST_AUTO_TEST_CASE_TEMPLATE( scalar_parameter, T, test_types) {

  Parameter<T> parameter{1};

  T b = static_cast<T>(parameter);

  // If the Parameter is not a distirbution, then, it'll not be randomized
  BOOST_TEST(static_cast<T>(parameter()) == b);

  json config = R"(
    {
      "param": 1
    }
  )"_json;

  Parameter<T> jparameter(config["param"], 1);
  BOOST_TEST(jparameter.at(0) == 1);

  jparameter = {2};
  BOOST_TEST(jparameter.at(0) == 2);
}

BOOST_AUTO_TEST_CASE_TEMPLATE( array_of_scalars_parameter, T, test_types) {

  Parameter<T> parameter{1, 2, 3};

  arma::Row<T> b = parameter;

  // If the Parameter is not a distirbution, then, it'll not be randomized
  BOOST_TEST(arma::approx_equal(parameter, b, "absdiff", 0.001));

  json config = R"(
    {
    "param": [1, 2, 3]
    }
  )"_json;

  Parameter<T> jparameter(config["param"], 3);

  arma::Row<T> a{1, 2, 3};
  BOOST_TEST(arma::approx_equal(jparameter,
                                a, "absdiff", 0.001));

  a = arma::Row<T>({4, 5, 6});
  BOOST_TEST(!arma::approx_equal(jparameter,
                                a, "absdiff", 0.001));

}

BOOST_AUTO_TEST_CASE_TEMPLATE( univariate_parameter, T, dist_test_types) {

  json config = R"(
    {
      "param": {
        "dist": "normal_distribution",
        "mean": 0,
        "stddev": 1
      }
    }
  )"_json;

  Parameter<T> parameter(config["param"], 1);

  BOOST_TEST(parameter.n_elem == 1);

  T a = static_cast<T>(parameter());
  T b = static_cast<T>(parameter());
  BOOST_TEST(sizeof(a) == sizeof(b));
  BOOST_TEST(a != b);
}

BOOST_AUTO_TEST_CASE_TEMPLATE( multivariate_parameter, T, dist_test_types) {

  json config = R"(
  {
    "param": {
      "dist": "mvnorm_distribution",
      "means": [0, 0, 0, 0],
      "covs": 0,
      "stddevs": 1
    }
  }
  )"_json;

  Parameter<T> parameter(config["param"], 4);

  BOOST_TEST(parameter.n_elem == 4);

  arma::Row<T> a = parameter();
  arma::Row<T> b = parameter();

  BOOST_TEST(!arma::approx_equal(a, b, "absdiff", 0.001));
}


// Boost.Bimap
//
// Copyright (c) 2006-2007 Matias Capeletto
//
// 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)

// Boost.Bimap Example
//-----------------------------------------------------------------------------
// Hashed indices can be used as an alternative to ordered indices when fast
// lookup is needed and sorting information is of no interest. The example
// features a word counter where duplicate entries are checked by means of a
// hashed index.

#include <boost/config.hpp>

//[ code_mi_to_b_path_hashed_indices

#include <iostream>
#include <iomanip>

#include <boost/tokenizer.hpp>

#include <boost/bimap/bimap.hpp>
#include <boost/bimap/unordered_set_of.hpp>
#include <boost/bimap/multiset_of.hpp>
#include <boost/bimap/support/lambda.hpp>

using namespace boost::bimaps;

struct word        {};
struct occurrences {};

typedef bimap
<

     multiset_of< tagged<unsigned int,occurrences>, std::greater<unsigned int> >,
unordered_set_of< tagged< std::string,       word>                             >

> word_counter;

typedef boost::tokenizer<boost::char_separator<char> > text_tokenizer;

int main()
{

    std::string text=
        "Relations between data in the STL are represented with maps."
        "A map is a directed relation, by using it you are representing "
        "a mapping. In this directed relation, the first type is related to "
        "the second type but it is not true that the inverse relationship "
        "holds. This is useful in a lot of situations, but there are some "
        "relationships that are bidirectional by nature.";

    // feed the text into the container

    word_counter   wc;
    text_tokenizer tok(text,boost::char_separator<char>(" \t\n.,;:!?'\"-"));
    unsigned int   total_occurrences = 0;

    for( text_tokenizer::const_iterator it = tok.begin(), it_end = tok.end();
         it != it_end ; ++it )
    {
        ++total_occurrences;

        word_counter::map_by<occurrences>::iterator wit =
            wc.by<occurrences>().insert(
                 word_counter::map_by<occurrences>::value_type(0,*it)
            ).first;

        wc.by<occurrences>().modify_key( wit, ++_key);
    }

    // list words by frequency of appearance

    std::cout << std::fixed << std::setprecision(2);

    for( word_counter::map_by<occurrences>::const_iterator
            wit     = wc.by<occurrences>().begin(),
            wit_end = wc.by<occurrences>().end();

         wit != wit_end; ++wit )
    {
        std::cout << std::setw(15) << wit->get<word>() << ": "
                  << std::setw(5)
                  << 100.0 * wit->get<occurrences>() / total_occurrences << "%"
                  << std::endl;
    }

    return 0;
}
//]


/* Copyright (c) 2015, Julian Straub <jstraub@csail.mit.edu>
 * Licensed under the MIT license. See the license file LICENSE.
 */
#pragma once
#include <iostream>
#include <stdint.h>
#include <vector>
#include <Eigen/Dense>

#include <boost/shared_ptr.hpp>

#include <dpMM/dpMM.hpp>
#include <dpMM/cat.hpp>
#include <dpMM/dir.hpp>
#include <dpMM/niw.hpp>
#include <dpMM/sampler.hpp>
#include <dpMM/basemeasure.hpp>

using namespace Eigen;
using std::endl; using std::cout;
using boost::shared_ptr;

  // algorithm 2 of Neal
  // http://www.stat.columbia.edu/npbayes/papers/neal_sampling.pdf
template<typename T>
class CrpMM : public DpMM<T>
{
public:
  CrpMM(const T alpha, const shared_ptr<BaseMeasure<T> >& theta, uint32_t K0,
      boost::mt19937* pRndGen);
  virtual ~CrpMM()
  {};

  virtual void initialize(const shared_ptr<Matrix<T,Dynamic,Dynamic> >& spx);
  virtual void initialize(const Matrix<T,Dynamic,Dynamic>& x)
  {this->spx_=shared_ptr<Matrix<T,Dynamic,Dynamic> >(new
      Matrix<T,Dynamic,Dynamic>(x)); initialize(this->spx_);};
  virtual void sampleLabels();
  virtual void sampleParameters();
  virtual void proposeSplits() {};
  virtual void proposeMerges() {};
  // call this right before sampleParameters()
  virtual void resampleFromBase(uint32_t Kmax) {};
  virtual const VectorXu & getLabels(){return z_;};
  virtual const VectorXu & labels(){return z_;};
  virtual Matrix<T,Dynamic,1> getCounts();
  virtual uint32_t getK() const { return K_;};
  virtual double logJoint();

//  virtual MatrixXu mostLikelyInds(uint32_t n)
//{ return MatrixXu::Zero(n,1);};
  virtual MatrixXu mostLikelyInds(uint32_t n, Matrix<T,Dynamic,Dynamic>& logLikes)
  { return MatrixXu::Zero(n,1);};

private:
  void removeEmptyClusters();

  uint32_t K_;
  T alpha_; // dp concentration paramter

  shared_ptr<BaseMeasure<T> > theta0_;
  vector<shared_ptr<BaseMeasure<T> > > thetas_;

  VectorXu z_;

  boost::mt19937* pRndGen_;
};

// ------------------------- impl -----------------------------------
template<typename T>
CrpMM<T>::CrpMM(const T alpha, const shared_ptr<BaseMeasure<T> >& theta,
    uint32_t K0,  boost::mt19937* pRndGen)
  : K_(K0), alpha_(alpha), theta0_(theta), pRndGen_(pRndGen)
{};

template <typename T>
Matrix<T,Dynamic,1> CrpMM<T>::getCounts()
{
  return counts<T,uint32_t>(z_,K_);
};

template<typename T>
void CrpMM<T>::initialize(const shared_ptr<Matrix<T,Dynamic,Dynamic> >& spx)
{
  cout<<"init"<<endl;
  this->spx_ = spx;
  // randomly init labels from prior
  z_.setZero(spx->cols());

  Matrix<T,Dynamic,1> alpha = (alpha_/K_)*Matrix<T,Dynamic,1>::Ones(K_);
  Cat<T> pi = Dir<Cat<T>,T>(alpha,pRndGen_).sample();
  for(uint32_t i=0; i<z_.size(); ++i)
  {
    z_(i) = pi.sample();
  }

  cout<<"init counts "<<this->getCounts().transpose()<<endl;

  // init the parameters
  if(thetas_.size() == 0)
  {
    cout<<"creating thetas"<<endl;
    for (uint32_t k=0; k<K_; ++k)
      thetas_.push_back(shared_ptr<BaseMeasure<T> >(theta0_->copy()));
  }
};

template<typename T>
void CrpMM<T>::sampleLabels()
{
  Matrix<T,Dynamic,1> Nk(K_+1);
  Nk.topRows(K_) = counts<T,uint32_t>(z_,K_);
  Nk(K_) = alpha_;
//  cout<<"Nk="<<Nk.transpose()<<" alpha="<<alpha_<<endl;
  for(uint32_t i=0; i<z_.size(); ++i)
  {
    Nk(z_(i)) --; // take the data-point out of that cluster
    Matrix<T,Dynamic,1> pdf = Nk.array().log().matrix();
    pdf = pdf.array() - log(Nk.sum());
//    cout<<pdf.transpose()<< " -> ";
    // add the log posteriors of the individual clusters
    for(uint32_t k=0; k<K_; ++k)
      pdf(k) += thetas_[k]->logLikelihood(this->spx_->col(i));
    // add the log likelihood of data under the hyperparameters
    pdf(K_) += theta0_->logPdfUnderPriorMarginalized(this->spx_->col(i));
//    cout<<pdf.transpose()<<" exp: ";
    // normalize and exponentiate pdf
    pdf = (pdf.array() - logSumExp(pdf)).exp().matrix();
//    cout<<pdf.transpose()<<endl;
    z_(i) = Catd(pdf,pRndGen_).sample();
    Nk(z_(i)) ++; // add data-point into new cluster

    // add a new cluster from the base measure
    if(z_(i)==K_)
    {
      thetas_.push_back(shared_ptr<BaseMeasure<T> >(theta0_->copy()));
      // TODO might want to add z_i to the SS of the new cluster?
      thetas_[z_(i)]->posterior(*this->spx_,z_,K_); // TODO slow
      ++K_;
      Nk.conservativeResize(K_+1);
      Nk(z_(i)) -= alpha_;
      Nk(K_) = alpha_;
    }
    if(z_.size()>10000 && i %(z_.size()/100) == 0) cout<<" CrpMM<T>::sampleLabel: "<<(i/(z_.size()/100))<<"% done"<<endl;
  }
  this->removeEmptyClusters();
};

template<typename T>
void CrpMM<T>::sampleParameters()
{
//#pragma omp parallel for
  for(uint32_t k=0; k<K_; ++k)
  {
    thetas_[k]->posterior(*this->spx_,z_,k);
    cout<<"k:"<<k<<" ";
    thetas_[k]->print();
  }
};

template <typename T>
void CrpMM<T>::removeEmptyClusters()
{
//  cout<<"K="<<K_<<" z="<<z_.transpose()<<endl;
  std::vector<bool> toDelete(K_,true);
//#pragma omp parallel for
  for(uint32_t k=0; k<K_; ++k)
  {
    toDelete[k] = true;
    for(uint32_t i=0; i<z_.size(); ++i)
      toDelete[k] = toDelete[k] && (z_(i)!=k);
  }

  std::vector<uint32_t> labelMap(K_,0);
  {
    uint32_t k=0;
    for(k=0; k<K_; ++k) if(!toDelete[k]) break;
    for(k=k+1; k<K_; ++k)
      if(toDelete[k])
        labelMap[k] = labelMap[k-1];
      else
        labelMap[k] = labelMap[k-1]+1;
//    for(k=0; k<K_; ++k) cout<<toDelete[k]<<" "; cout<<endl;
//    for(k=0; k<K_; ++k) cout<<labelMap[k]<<" "; cout<<endl;
  }
#pragma omp parallel for
  for(uint32_t i=0; i<z_.size(); ++i)
    z_(i) = labelMap[z_(i)];
//  cout<<z_.transpose()<<endl;

  for(int32_t k=K_-1; k>=0; --k)
    if (toDelete[k])
      thetas_.erase(thetas_.begin()+k);
  K_ = labelMap[K_-1]+1;
}

template <typename T>
double CrpMM<T>::logJoint()
{
  double logP = 0.;
  for (uint32_t i=0; i<z_.size(); ++i)
  {
    logP += thetas_[z_(i)]->logLikelihood(this->spx_->col(i));
  }
  return logP;
}


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

int main(){
    auto t1 = std::chrono::high_resolution_clock::now();

    int mat_size = 5;
    Eigen::MatrixXd kern_mat = Eigen::MatrixXd::Zero(mat_size, mat_size);
    Eigen::VectorXd vec_curr(5);

    #pragma omp parallel for
    for (int i = 0; i < mat_size; i ++){
        std::this_thread::sleep_for(std::chrono::seconds(1));
        vec_curr << i, i+1, i+2, i+3, i+4;
        kern_mat.col(i) = vec_curr;
    }

    auto t2 = std::chrono::high_resolution_clock::now();
    auto tot_time =
        std::chrono::duration_cast<std::chrono::milliseconds>(t2-t1).count();
    std::cout << tot_time << std::endl;
    std::cout << kern_mat << std::endl;

}


//---------------------------------------------------------------------------//
// Copyright (c) 2013 Kyle Lutz <kyle.r.lutz@gmail.com>
//
// 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
//
// See http://kylelutz.github.com/compute for more information.
//---------------------------------------------------------------------------//

#define BOOST_TEST_MODULE TestIota
#include <boost/test/unit_test.hpp>

#include <boost/compute/system.hpp>
#include <boost/compute/context.hpp>
#include <boost/compute/algorithm/iota.hpp>
#include <boost/compute/container/vector.hpp>
#include <boost/compute/iterator/permutation_iterator.hpp>

#include "check_macros.hpp"
#include "context_setup.hpp"

namespace bc = boost::compute;

BOOST_AUTO_TEST_CASE(iota_int)
{
    bc::vector<int> vector(4);
    bc::iota(vector.begin(), vector.end(), 0);
    CHECK_RANGE_EQUAL(int, 4, vector, (0, 1, 2, 3));

    bc::iota(vector.begin(), vector.end(), 10);
    CHECK_RANGE_EQUAL(int, 4, vector, (10, 11, 12, 13));

    bc::iota(vector.begin() + 2, vector.end(), -5);
    CHECK_RANGE_EQUAL(int, 4, vector, (10, 11, -5, -4));

    bc::iota(vector.begin(), vector.end() - 2, 4);
    CHECK_RANGE_EQUAL(int, 4, vector, (4, 5, -5, -4));
}

BOOST_AUTO_TEST_CASE(iota_doctest)
{
    boost::compute::vector<int> vec(3, context);

//! [iota]
boost::compute::iota(vec.begin(), vec.end(), 0, queue);
//! [iota]

    CHECK_RANGE_EQUAL(int, 3, vec, (0, 1, 2));
}

BOOST_AUTO_TEST_CASE(iota_permutation_iterator)
{
    bc::vector<int> output(5);
    bc::fill(output.begin(), output.end(), 0);

    int map_data[] = { 2, 0, 1, 4, 3 };
    bc::vector<int> map(map_data, map_data + 5);

    bc::iota(bc::make_permutation_iterator(output.begin(), map.begin()),
             bc::make_permutation_iterator(output.end(), map.end()),
             1);
    CHECK_RANGE_EQUAL(int, 5, output, (2, 3, 1, 5, 4));
}

BOOST_AUTO_TEST_SUITE_END()


// Simple development tests

#include "craam/ImMDP.hpp"
#include "craam/Simulation.hpp"
#include "craam/Samples.hpp"
#include "craam/algorithms/values.hpp"

#include "rm/range.hpp"

#include <iostream>
#include <iterator>
#include <random>
#include <cmath>
#include <cassert>

#include <boost/functional/hash.hpp>
#include <iostream>
#include <iterator>

using namespace std;
using namespace craam;
using namespace craam::impl;
using namespace util::lang;

template<class T>
void print_vector(vector<T> vec){
    for(auto&& p : vec){
        cout << p << " ";
    }
}

/**
A simple simulator class. The state represents a position in a chain
and actions move it up and down. The reward is equal to the position.

Representation
~~~~~~~~~~~~~~
- State: position (int)
- Action: change (int)
*/
class Counter{
private:
    default_random_engine gen;
    bernoulli_distribution d;
    const vector<int> actions_list;
    const int initstate;

public:
    using State = int;
    using Action = int;

    /**
    Define the success of each action
    \param success The probability that the action is actually applied
    */
    Counter(double success, int initstate, random_device::result_type seed = random_device{}())
        : gen(seed), d(success), actions_list({1,-1}), initstate(initstate) {};

    int init_state() const {
        return initstate;
    }

    pair<double,int> transition(int pos, int action) {
        int nextpos = d(gen) ? pos + action : pos;
        return make_pair((double) pos, nextpos);
    }

    bool end_condition(const int state){
        return false;
    }

    int action(State state, long index) const{
        return actions_list[index];
    }

    const vector<int>& get_valid_actions(int state) const{
        return actions_list;
    }

    size_t action_count(State) const{return actions_list.size();};
};

/** A counter that terminates at either end as defined by the end state */
class CounterTerminal : public Counter {
public:
    int endstate;

    CounterTerminal(double success, int initstate, int endstate, random_device::result_type seed = random_device{}())
        : Counter(success, initstate, seed), endstate(endstate) {};

    bool end_condition(const int state){
        return (abs(state) >= endstate);
    }
};
// Hash function for the Counter / CounterTerminal EState above
namespace std{
    template<> struct hash<pair<int,int>>{
        size_t operator()(pair<int,int> const& s) const{
            boost::hash<pair<int,int>> h;
            return h(s);
        };
    };
}

using namespace craam::msen;

int main(void){

    const int terminal_state = 8;
    const prec_t discount = 0.9;

    CounterTerminal sim(0.9,0,terminal_state,1);
    RandomPolicy<CounterTerminal> random_pol(sim);

    auto samples = make_samples<CounterTerminal>();
    simulate(sim,samples,random_pol,100,100);
    simulate(sim,samples,[](int){return 1;},10,20);
    simulate(sim,samples,[](int){return -1;},10,20);

    SampleDiscretizerSI<typename CounterTerminal::State, typename CounterTerminal::Action> sd;
    // initialize action values
    sd.add_action(-1); sd.add_action(+1);
    //initialize state values
    for(auto i : util::lang::range(-terminal_state,terminal_state)) sd.add_state(i);

    sd.add_samples(samples);

    SampledMDP smdp;
    smdp.add_samples(*sd.get_discrete());
    auto mdp = smdp.get_mdp();
    auto&& initial = smdp.get_initial();

    auto&& sol = algorithms::solve_mpi(*mdp,discount);

    cout << "Optimal policy: "; print_vector(sol.policy); cout << "Return " <<  sol.total_return(initial) << endl;

    // define observations
    indvec observations(mdp->state_count(), -1);
    size_t last_obs(0), inobs(0);
    cout << "Observations: " << mdp->state_count() << " states  ";
    for(auto i : util::lang::range(size_t(0), mdp->state_count())){
        // check if this is a terminal state
        if(mdp->get_state(i).action_count() == 0 || inobs >= 2){
            if(inobs > 0 && mdp->get_state(i).action_count() == 0){
                last_obs++;
            }
            observations[i] = last_obs++;
            inobs = 0;
        }else {
            observations[i] = last_obs;
            inobs++;
        }
        cout << observations[i] << " ";
    }
    cout << endl;

    MDPI_R mdpi(mdp, observations, initial);
    auto&& randompolicy = mdpi.random_policy(25);

    auto isol = mdpi.solve_reweighted(0, discount, randompolicy);

    isol = mdpi.solve_reweighted(10, discount, randompolicy);

    auto sol_impl = solve_mpi(*mdp, discount, numvec(0), mdpi.obspol2statepol(isol));

    cout << "Implementable pol: "; print_vector(isol);
    cout << "  Return: " << sol_impl.total_return(initial) << endl;

    cout << "Generating implementable policies (randomly) ..." << endl;

    auto max_return = 0.0;
    indvec max_pol(mdpi.obs_count(),-1);

    for(auto i : util::lang::range(0,200)){
        (void)(i);
        auto rand_pol = mdpi.random_policy();

        auto ret = solve_mpi(*mdp, discount, numvec(0), mdpi.obspol2statepol(rand_pol)).total_return(initial);

        if(ret > max_return){
            max_pol = rand_pol;
            max_return = ret;
        }

    }

    cout << "Maximal return " << max_return << endl;
    cout << "Best policy: ";
    print_vector(max_pol);
    cout << endl;

    return 0;

}


//==================================================================================================
/*!
  @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_CSC_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_CSC_HPP_INCLUDED

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

 /*!

    @ingroup group-trigonometric
    Function object implementing csc capabilities

    cosecante of the input in radian.

    @par Semantic:

    For every parameter of floating type T

    @code
    T r = csc(x);
    @endcode

    is similar to:

    @code
    T r = rec(sin(x));
    @endcode

    As most other trigonometric function csc can be called with a second optional parameter
    which is a tag on speed and accuracy (see @ref cos for further details)

    @see cscd, cscpi,

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

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

#endif


// TestVectorExcel.cpp
//
// Test output of a vector in Excel. Here we
// use the Excel Driver object directly.
//
// The output is in cell/numeric format.
//
// (C) Datasim Education BV 2006-2017
//

#include "ExcelDriverlite.hpp"
#include "Utilities.hpp"

#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/ublas/matrix_proxy.hpp>

#include <string>
#include <vector>
#include <list>

double rho = 0.5;
double NormalPdf1d(double x)
{ // Univariate normal density function

	double fac = 1.0 / std::sqrt(2.0 * 3.14159265359);
	return fac * std::exp(-x*x/2);
}

template <typename Vector>
	Vector DiscreteNormalPdf1d(const std::vector<double>& x)
{
	return CreateDiscreteFunction< std::vector<double>>(x, NormalPdf1d);
}

int main()
{

	// C++11 syntax
	//using NumericMatrix = boost::numeric::ublas::matrix<double>;
	// using Vector = std::vector<double>
	typedef boost::numeric::ublas::matrix<double> NumericMatrix;
	typedef std::vector<double> Vector;
	//	auto xarr = CreateMesh(N + 1, 0.0, 10.0);

	std::string sheetName("Vector Case");

	ExcelDriver& excel = ExcelDriver::Instance();
	excel.MakeVisible(true);		// Default is VISIBLE!

	// Labels for columns of the Excel vector.
	// Only labelled values are printed!!

	std::string rowLabel("row1");
	std::list<std::string> colLabels; // C++11: {"C1", "C2", "C3", "C4","C5"};
	colLabels.push_back("C1");
	colLabels.push_back("C2");
	colLabels.push_back("C3");
	colLabels.push_back("C4");
	colLabels.push_back("C5");

	Vector myVector(colLabels.size());
	for (std::size_t i = 0; i < myVector.size(); ++i)
	{

		myVector[i] = static_cast<double>(i);
	}

	try
	{
		long row = 4; long col = 3;
		excel.AddVector<NumericMatrix>(myVector, sheetName, rowLabel, colLabels, row, col);
	}
	catch(std::out_of_range& e)
	{
		std::cout << e.what() << '\n';
	}
	catch (...)
	{
		// Catches everything else
		std::cout << "oop";
	}

	{
		// Using mapping continuous space to discrete space
		std::size_t N = 10;
		auto x = CreateMesh(N, -4.0, 4.0);

		std::vector<double> vec = DiscreteNormalPdf1d<std::vector<double>>(x);

		try
		{

			ExcelDriver& excel = ExcelDriver::Instance();
			std::string sheetName("Bivariate Normal pdf");
			long row = 3; long col = 2;
			//excel.AddVector<NumericMatrix>(vec, sheetName, row, col);
		}
		catch (std::out_of_range& e)
		{
			std::cout << e.what() << '\n';
		}
		catch (...)
		{
			// Catches everything else
			std::cout << "oop";
		}
	}

	return 0;
}


//
// Created by Amir Masoud Abdol on 2021-01-18
//

#define BOOST_TEST_DYN_LINK

#define BOOST_TEST_MODULE ExperimentSetup Tests

#include "nlohmann/json.hpp"
#include <boost/test/unit_test.hpp>
namespace tt = boost::test_tools;
namespace utf = boost::unit_test;

#include "DependentVariable.h"

#include "sample_experiment_setup.h"

using json = nlohmann::ordered_json;

using namespace arma;
using namespace sam;
using namespace std;

BOOST_AUTO_TEST_SUITE(constructor)

  BOOST_AUTO_TEST_CASE( from_an_arma_array ) {

    BOOST_TEST_MESSAGE("Testing the Initialization from arma::Row<float>...");

    arma::Row<float> data(100);
    data.randn();

    DependentVariable dp{data};

    BOOST_TEST(arma::approx_equal(data, dp.measurements(), "absdiff", 0.001));

    // Testing whether stats are initialized, and updated
    BOOST_TEST(dp.mean_ != 0);
    BOOST_TEST(dp.var_ != 0);

  }

  BOOST_AUTO_TEST_CASE( copy_constructor ) {

    BOOST_TEST_MESSAGE("Testing the Copy Constructor...");

    arma::Row<float> data(100);
    data.randn();

    DependentVariable dv{data};

    DependentVariable dv_copy{dv};
    BOOST_TEST(dv.mean_ == dv_copy.mean_);
    BOOST_TEST(arma::approx_equal(dv.measurements(), dv_copy.measurements(),
                                  "absdiff", 0.001));

  }

BOOST_AUTO_TEST_SUITE_END()

BOOST_AUTO_TEST_SUITE( working_with_dependent_variable )

  BOOST_AUTO_TEST_CASE( manipulating_measurements ) {

    arma::Row<float> data(100);
    data.randn();

    DependentVariable dp{data};

    dp.removeMeasurements(arma::uvec{1, 5, 10, 50});
    BOOST_TEST(dp.mean_ != arma::mean(data));
    BOOST_TEST(dp.var_ != arma::var(data));
    BOOST_TEST(dp.nobs_ == 96);
    BOOST_TEST(dp.n_removed_obs == 4);

    dp.addNewMeasurements(arma::Row<float>{50000});
    BOOST_TEST(dp.nobs_ == 97);
    BOOST_TEST(dp.mean_ > arma::mean(data));
    BOOST_TEST(dp.var_ > arma::var(data));
    BOOST_TEST(dp.n_added_obs == 1);

    BOOST_TEST(dp.true_nobs_ == 100);
  }

  BOOST_AUTO_TEST_CASE( indices_operator_test, * utf::expected_failures(1) ) {

    arma::Row<float> data(100);
    data.randn();

    DependentVariable dp{data};

    dp[50] = 50;
    BOOST_TEST(dp[50] == 50);
  }

  BOOST_AUTO_TEST_CASE( manipulating_dps_status) {

    arma::Row<float> data(100);
    data.randn();

    DependentVariable dp{data};

    dp.removeMeasurements(arma::uvec{1, 5, 10, 25, 50});
    BOOST_TEST(dp.isModified());
    BOOST_TEST(not dp.isHacked());

    dp.setHackedStatus(true);
    BOOST_TEST(dp.isHacked());
    BOOST_TEST(dp.isModified());

    dp.setCandidateStatus(true);
    BOOST_TEST(dp.isCandidate());
  }

BOOST_AUTO_TEST_SUITE_END()


/**
 * @file MPCTest.cpp
 * @author Giulio Romualdi
 * @copyright Released under the terms of the LGPLv2.1 or later, see LGPL.TXT
 * @date 2018
 */

// Catch2
#include <catch2/catch.hpp>

// OsqpEigen
#include <OsqpEigen/OsqpEigen.h>

// eigen
#include <Eigen/Dense>

#include <cmath>
#include <fstream>
#include <iostream>

// colors
#define ANSI_TXT_GRN "\033[0;32m"
#define ANSI_TXT_MGT "\033[0;35m" //Magenta
#define ANSI_TXT_DFT "\033[0;0m" //Console default
#define GTEST_BOX "[     cout ] "
#define COUT_GTEST ANSI_TXT_GRN << GTEST_BOX //You could add the Default
#define COUT_GTEST_MGT COUT_GTEST << ANSI_TXT_MGT

void setDynamicsMatrices(Eigen::Matrix<double, 12, 12> &a, Eigen::Matrix<double, 12, 4> &b)
{
    a << 1.,      0.,     0., 0., 0., 0., 0.1,     0.,     0.,  0.,     0.,     0.    ,
        0.,      1.,     0., 0., 0., 0., 0.,      0.1,    0.,  0.,     0.,     0.    ,
        0.,      0.,     1., 0., 0., 0., 0.,      0.,     0.1, 0.,     0.,     0.    ,
        0.0488,  0.,     0., 1., 0., 0., 0.0016,  0.,     0.,  0.0992, 0.,     0.    ,
        0.,     -0.0488, 0., 0., 1., 0., 0.,     -0.0016, 0.,  0.,     0.0992, 0.    ,
        0.,      0.,     0., 0., 0., 1., 0.,      0.,     0.,  0.,     0.,     0.0992,
        0.,      0.,     0., 0., 0., 0., 1.,      0.,     0.,  0.,     0.,     0.    ,
        0.,      0.,     0., 0., 0., 0., 0.,      1.,     0.,  0.,     0.,     0.    ,
        0.,      0.,     0., 0., 0., 0., 0.,      0.,     1.,  0.,     0.,     0.    ,
        0.9734,  0.,     0., 0., 0., 0., 0.0488,  0.,     0.,  0.9846, 0.,     0.    ,
        0.,     -0.9734, 0., 0., 0., 0., 0.,     -0.0488, 0.,  0.,     0.9846, 0.    ,
        0.,      0.,     0., 0., 0., 0., 0.,      0.,     0.,  0.,     0.,     0.9846;

    b << 0.,      -0.0726,  0.,     0.0726,
        -0.0726,  0.,      0.0726, 0.    ,
        -0.0152,  0.0152, -0.0152, 0.0152,
        -0.,     -0.0006, -0.,     0.0006,
        0.0006,   0.,     -0.0006, 0.0000,
        0.0106,   0.0106,  0.0106, 0.0106,
        0,       -1.4512,  0.,     1.4512,
        -1.4512,  0.,      1.4512, 0.    ,
        -0.3049,  0.3049, -0.3049, 0.3049,
        -0.,     -0.0236,  0.,     0.0236,
        0.0236,   0.,     -0.0236, 0.    ,
        0.2107,   0.2107,  0.2107, 0.2107;
}

void setInequalityConstraints(Eigen::Matrix<double, 12, 1> &xMax, Eigen::Matrix<double, 12, 1> &xMin,
                              Eigen::Matrix<double, 4, 1> &uMax, Eigen::Matrix<double, 4, 1> &uMin)
{
    double u0 = 10.5916;

    // input inequality constraints
    uMin << 9.6 - u0,
        9.6 - u0,
        9.6 - u0,
        9.6 - u0;

    uMax << 13 - u0,
        13 - u0,
        13 - u0,
        13 - u0;

    // state inequality constraints
    xMin << -M_PI/6,-M_PI/6,-OsqpEigen::INFTY,-OsqpEigen::INFTY,-OsqpEigen::INFTY,-1.,
        -OsqpEigen::INFTY, -OsqpEigen::INFTY,-OsqpEigen::INFTY,-OsqpEigen::INFTY,
        -OsqpEigen::INFTY,-OsqpEigen::INFTY;

    xMax << M_PI/6,M_PI/6, OsqpEigen::INFTY,OsqpEigen::INFTY,OsqpEigen::INFTY,
        OsqpEigen::INFTY, OsqpEigen::INFTY,OsqpEigen::INFTY,OsqpEigen::INFTY,
        OsqpEigen::INFTY,OsqpEigen::INFTY,OsqpEigen::INFTY;
}

void setWeightMatrices(Eigen::DiagonalMatrix<double, 12> &Q, Eigen::DiagonalMatrix<double, 4> &R)
{
    Q.diagonal() << 0, 0, 10., 10., 10., 10., 0, 0, 0, 5., 5., 5.;
    R.diagonal() << 0.1, 0.1, 0.1, 0.1;
}

void castMPCToQPHessian(const Eigen::DiagonalMatrix<double, 12> &Q, const Eigen::DiagonalMatrix<double, 4> &R, int mpcWindow,
                        Eigen::SparseMatrix<double> &hessianMatrix)
{

    hessianMatrix.resize(12*(mpcWindow+1) + 4 * mpcWindow, 12*(mpcWindow+1) + 4 * mpcWindow);

    //populate hessian matrix
    for(int i = 0; i<12*(mpcWindow+1) + 4 * mpcWindow; i++){
        if(i < 12*(mpcWindow+1)){
            int posQ=i%12;
            float value = Q.diagonal()[posQ];
            if(value != 0)
                hessianMatrix.insert(i,i) = value;
        }
        else{
            int posR=i%4;
            float value = R.diagonal()[posR];
            if(value != 0)
                hessianMatrix.insert(i,i) = value;
        }
    }
}

void castMPCToQPGradient(const Eigen::DiagonalMatrix<double, 12> &Q, const Eigen::Matrix<double, 12, 1> &xRef, int mpcWindow,
                         Eigen::VectorXd &gradient)
{

    Eigen::Matrix<double,12,1> Qx_ref;
    Qx_ref = Q * (-xRef);

    // populate the gradient vector
    gradient = Eigen::VectorXd::Zero(12*(mpcWindow+1) +  4*mpcWindow, 1);
    for(int i = 0; i<12*(mpcWindow+1); i++){
        int posQ=i%12;
        float value = Qx_ref(posQ,0);
        gradient(i,0) = value;
    }
}

void castMPCToQPConstraintMatrix(const Eigen::Matrix<double, 12, 12> &dynamicMatrix, const Eigen::Matrix<double, 12, 4> &controlMatrix,
                                 int mpcWindow, Eigen::SparseMatrix<double> &constraintMatrix)
{
    constraintMatrix.resize(12*(mpcWindow+1)  + 12*(mpcWindow+1) + 4 * mpcWindow, 12*(mpcWindow+1) + 4 * mpcWindow);

    // populate linear constraint matrix
    for(int i = 0; i<12*(mpcWindow+1); i++){
        constraintMatrix.insert(i,i) = -1;
    }

    for(int i = 0; i < mpcWindow; i++)
        for(int j = 0; j<12; j++)
            for(int k = 0; k<12; k++){
                float value = dynamicMatrix(j,k);
                if(value != 0){
                    constraintMatrix.insert(12 * (i+1) + j, 12 * i + k) = value;
                }
            }

    for(int i = 0; i < mpcWindow; i++)
        for(int j = 0; j < 12; j++)
            for(int k = 0; k < 4; k++){
                float value = controlMatrix(j,k);
                if(value != 0){
                    constraintMatrix.insert(12*(i+1)+j, 4*i+k+12*(mpcWindow + 1)) = value;
                }
            }

    for(int i = 0; i<12*(mpcWindow+1) + 4*mpcWindow; i++){
        constraintMatrix.insert(i+(mpcWindow+1)*12,i) = 1;
    }
}

void castMPCToQPConstraintVectors(const Eigen::Matrix<double, 12, 1> &xMax, const Eigen::Matrix<double, 12, 1> &xMin,
                                   const Eigen::Matrix<double, 4, 1> &uMax, const Eigen::Matrix<double, 4, 1> &uMin,
                                   const Eigen::Matrix<double, 12, 1> &x0,
                                   int mpcWindow, Eigen::VectorXd &lowerBound, Eigen::VectorXd &upperBound)
{
    // evaluate the lower and the upper inequality vectors
    Eigen::VectorXd lowerInequality = Eigen::MatrixXd::Zero(12*(mpcWindow+1) +  4 * mpcWindow, 1);
    Eigen::VectorXd upperInequality = Eigen::MatrixXd::Zero(12*(mpcWindow+1) +  4 * mpcWindow, 1);
    for(int i=0; i<mpcWindow+1; i++){
        lowerInequality.block(12*i,0,12,1) = xMin;
        upperInequality.block(12*i,0,12,1) = xMax;
    }
    for(int i=0; i<mpcWindow; i++){
        lowerInequality.block(4 * i + 12 * (mpcWindow + 1), 0, 4, 1) = uMin;
        upperInequality.block(4 * i + 12 * (mpcWindow + 1), 0, 4, 1) = uMax;
    }

    // evaluate the lower and the upper equality vectors
    Eigen::VectorXd lowerEquality = Eigen::MatrixXd::Zero(12*(mpcWindow+1),1 );
    Eigen::VectorXd upperEquality;
    lowerEquality.block(0,0,12,1) = -x0;
    upperEquality = lowerEquality;
    lowerEquality = lowerEquality;

    // merge inequality and equality vectors
    lowerBound = Eigen::MatrixXd::Zero(2*12*(mpcWindow+1) +  4*mpcWindow,1 );
    lowerBound << lowerEquality,
        lowerInequality;

    upperBound = Eigen::MatrixXd::Zero(2*12*(mpcWindow+1) +  4*mpcWindow,1 );
    upperBound << upperEquality,
        upperInequality;
}

void updateConstraintVectors(const Eigen::Matrix<double, 12, 1> &x0,
                             Eigen::VectorXd &lowerBound, Eigen::VectorXd &upperBound)
{
    lowerBound.block(0,0,12,1) = -x0;
    upperBound.block(0,0,12,1) = -x0;
}

double getErrorNorm(const Eigen::Matrix<double, 12, 1> &x,
                    const Eigen::Matrix<double, 12, 1> &xRef)
{
    // evaluate the error
    Eigen::Matrix<double, 12, 1> error = x - xRef;

    // return the norm
    return error.norm();
}

TEST_CASE("MPCTest")
{
    // open the ofstream
    std::ofstream dataStream;
    dataStream.open ("output.txt");

    // set the preview window
    int mpcWindow = 20;

    // allocate the dynamics matrices
    Eigen::Matrix<double, 12, 12> a;
    Eigen::Matrix<double, 12, 4> b;

    // allocate the constraints vector
    Eigen::Matrix<double, 12, 1> xMax;
    Eigen::Matrix<double, 12, 1> xMin;
    Eigen::Matrix<double, 4, 1> uMax;
    Eigen::Matrix<double, 4, 1> uMin;

    // allocate the weight matrices
    Eigen::DiagonalMatrix<double, 12> Q;
    Eigen::DiagonalMatrix<double, 4> R;

    // allocate the initial and the reference state space
    Eigen::Matrix<double, 12, 1> x0;
    Eigen::Matrix<double, 12, 1> xRef;

    // allocate QP problem matrices and vectores
    Eigen::SparseMatrix<double> hessian;
    Eigen::VectorXd gradient;
    Eigen::SparseMatrix<double> linearMatrix;
    Eigen::VectorXd lowerBound;
    Eigen::VectorXd upperBound;

    // set the initial and the desired states
    x0 << 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ;
    xRef <<  0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0;

    // set MPC problem quantities
    setDynamicsMatrices(a, b);
    setInequalityConstraints(xMax, xMin, uMax, uMin);
    setWeightMatrices(Q, R);

    // cast the MPC problem as QP problem
    castMPCToQPHessian(Q, R, mpcWindow, hessian);
    castMPCToQPGradient(Q, xRef, mpcWindow, gradient);
    castMPCToQPConstraintMatrix(a, b, mpcWindow, linearMatrix);
    castMPCToQPConstraintVectors(xMax, xMin, uMax, uMin, x0, mpcWindow, lowerBound, upperBound);

    // instantiate the solver
    OsqpEigen::Solver solver;

    // settings
    solver.settings()->setVerbosity(false);
    solver.settings()->setWarmStart(true);

    // set the initial data of the QP solver
    solver.data()->setNumberOfVariables(12 * (mpcWindow + 1) + 4 * mpcWindow);
    solver.data()->setNumberOfConstraints(2 * 12 * (mpcWindow + 1) +  4 * mpcWindow);
    REQUIRE(solver.data()->setHessianMatrix(hessian));
    REQUIRE(solver.data()->setGradient(gradient));
    REQUIRE(solver.data()->setLinearConstraintsMatrix(linearMatrix));
    REQUIRE(solver.data()->setLowerBound(lowerBound));
    REQUIRE(solver.data()->setUpperBound(upperBound));

    // instantiate the solver
    REQUIRE(solver.initSolver());

    // controller input and QPSolution vector
    Eigen::Vector4d ctr;
    Eigen::VectorXd QPSolution;

    // number of iteration steps
    int numberOfSteps = 50;

    // profiling quantities
    clock_t startTime, endTime;
    double avarageTime = 0;

    for (int i = 0; i < numberOfSteps; i++){
        startTime = clock();

        // solve the QP problem
        REQUIRE(solver.solve());

        // get the controller input
        QPSolution = solver.getSolution();
        ctr = QPSolution.block(12 * (mpcWindow + 1), 0, 4, 1);

        // save data into file
        auto x0Data = x0.data();
        for(int j = 0; j < 12; j++)
            dataStream << x0Data[j] << " ";
        dataStream << std::endl;

        // propagate the model
        x0 = a * x0 + b * ctr;

        // update the constraint bound
        updateConstraintVectors(x0, lowerBound, upperBound);
        REQUIRE(solver.updateBounds(lowerBound, upperBound));

        endTime = clock();

        avarageTime += static_cast<double>(endTime - startTime) / CLOCKS_PER_SEC;
      }

    // close the stream
    dataStream.close();

    std::cout << COUT_GTEST_MGT << "Avarage time = " << avarageTime / numberOfSteps
              << " seconds." << ANSI_TXT_DFT << std::endl;

    REQUIRE(getErrorNorm(x0, xRef) <=  0.001);
}


#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <Eigen/LU>

namespace py = pybind11;
using Eigen::MatrixXd;
using Eigen::VectorXd;
using Eigen::log;
using std::lgamma;

MatrixXd scale_matrix(MatrixXd X, int N, int k, double r, MatrixXd m, MatrixXd S) {
    MatrixXd xsum = X.colwise().sum();
    MatrixXd res = S + X.transpose() * X
                     + r * N / (N + r) * m * m.transpose()
                     - 1/(N+r) * (xsum.transpose() * xsum)
                     - (r / (N + r)) * (m * xsum + xsum.transpose() * m.transpose());
    return res;
}

double signprod(VectorXd dU){
    double res = 1;
    for(int i = 0; i < dU.size(); i++){
        if(dU[i] < 0){
            res *= -1;
        }
    }
    return res;
}

double logdet(MatrixXd S){
    Eigen::PartialPivLU<MatrixXd> lu(S);
    MatrixXd U = lu.matrixLU();
    VectorXd dU = U.diagonal();
    MatrixXd P = lu.permutationP();
    double c = P.determinant() * signprod(dU);
    double d = log(c) + dU.array().abs().log().sum();
    return d;
}

double niw(MatrixXd X, MatrixXd m, MatrixXd S, double r){
    const double PI  =3.141592653589793238463;
    double N = (double)X.rows();
    double k = (double)X.cols();
    // v = k;
    double vprime = k + N;

    MatrixXd Sprime = scale_matrix(X, N, k, r, m, S);

    double num = vprime*k/2*log(2);
    double den = k * k / 2 * log(2);
    for(int i = 0; i < k; i++){
        num += lgamma((double)(vprime - i)/2);
        den += lgamma((double)(k - i)/2);
    }

    double lml = - N*k/2 * (log(2) + log(PI))
                 + k/2*(log(r)-log(N + r))
                 + k/2*logdet(S)
                 - vprime/2*logdet(Sprime)
                 + num - den;
    return lml;
}

PYBIND11_PLUGIN(helper){
    pybind11::module m("helper", "helper functions");
    m.def("scale_matrix", &scale_matrix);
    m.def("niw", &niw);
    m.def("logdet", &logdet);
    return m.ptr();
}


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

#ifndef _ALS_
#define _ALS_
#include "../../utils/options_utils.h"
#include "../../utils/desc_utils.h"
#include "../../utils/boost_log_wrapper.h"
#include "../../core/x-lib.hpp"
#include <cmath>
#include <boost/random.hpp>
#include <boost/thread/mutex.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/bindings/lapack/gesv.hpp>
#include <boost/numeric/bindings/traits/ublas_matrix.hpp>
#include <boost/numeric/bindings/traits/ublas_vector2.hpp>

#define LAMBDA 0.065
#define RANK   5

namespace ublas = boost::numeric::ublas;
namespace lapack = boost::numeric::bindings::lapack;

namespace algorithm {
  namespace sg_simple {
    class als_per_processor_data:public per_processor_data {
    public:
      static double sse;
      double sse_local;
      als_per_processor_data()
	:sse_local(0.0)
      {}
      bool reduce(per_processor_data **per_cpu_array,
		  unsigned long processors)
      {
	sse = 0;
	for(unsigned long i=0;i<processors;i++) {
	  als_per_processor_data * data =
	    static_cast<als_per_processor_data *>(per_cpu_array[i]);
	  sse += data->sse_local;
	  data->sse_local = 0;
	}
	return false;
      }
    } __attribute__((__aligned__(64))) ;

    template <typename F>
    class als_factorization {

    private:

      struct vertex {
	vertex_t degree;
	vertex_t count;
	double feature_vec[RANK];
	double temp_mat[RANK][RANK];
      } __attribute__((__packed__));

      struct update {
	vertex_t target;
	double feature_vec[RANK];
	double rating;
      } __attribute__((__packed__));

      static unsigned long niters;

      // Helpers
      static void copy_vector(double src_vec[RANK], double dst_vec[RANK])
      {
	for (int i = 0; i < RANK; i++)
	  dst_vec[i] = src_vec[i];
      }
      static void zero_vector(double vec[RANK])
      {
	for (int i = 0; i < RANK; i++)
	  vec[i] = 0;
      }
      static void zero_matrix(double mat[RANK][RANK])
      {
	for (int i = 0; i < RANK; i++)
	  for (int j = 0; j < RANK; j++)
	    mat[i][j] = 0;
      }

      static void init_vertex(struct vertex& v)
      {
	v.degree = 0;
	v.count = 0;

	boost::mt19937 generator;
	boost::uniform_int<> distribution(0, 1000);
	for (int i = 0; i < RANK; i++)
	  v.feature_vec[i] =  0.001 * distribution(generator);
      }

      // Solve the system Ax=b using LAPACK library and boost bindings for it
      static void solve(double mat[RANK][RANK], double vec[RANK])
      {
	ublas::matrix<double, ublas::column_major> A(RANK, RANK);
	ublas::vector<double> b(RANK);

	for (int i = 0; i < RANK; i++) {
	  b(i) = vec[i];
	  for (int j = 0; j < RANK; j++) {
	    A(i,j) = mat[i][j];
	  }
	}

	lapack::gesv(A,b);

	for (int i = 0; i < RANK; i++)
	  vec[i] = b(i);
      }

    public:
      static unsigned long vertex_state_bytes() {
	return sizeof(struct vertex);
      }
      static unsigned long split_size_bytes() {
	return sizeof(struct update);
      }

      static unsigned long split_key(unsigned char* buffer, unsigned long jump)
      {
	struct update* u = (struct update*)buffer;
	vertex_t key = u->target;
	key = key >> jump;
	return key;
      }

      static bool init(unsigned char* vertex_state,
		       unsigned long vertex_index,
		       unsigned long bsp_phase,
		       per_processor_data *cpu_state)
      {
	struct vertex* vertices = (struct vertex*)vertex_state;
	init_vertex(*vertices);
	return true;
      }

      static bool apply_one_update(unsigned char* vertex_state,
				   unsigned char* update_stream,
				   per_processor_data *per_cpu_data,
				   unsigned long bsp_phase)
      {
	struct update* u = (struct update*)update_stream;
	struct vertex* vertices = (struct vertex*)vertex_state;
	struct vertex* v = &vertices[x_lib::configuration::map_offset(u->target)];

	if (bsp_phase <= niters)
	  {
	    // Track how many updates (edges) have been processed
	    v->count++;

	    // Initialize data structures for this iteration - zero out feature vector
	    // and temp matrix since we'll be adding to them.
	    if (v->count == 1) {
	      zero_vector(v->feature_vec);
	      zero_matrix(v->temp_mat);
	    }

	    // To compute the new feature vector of vertex v, we need to solve the system: A*feature_vec=b
	    // A is a matrix: A = O * O^T + D
	    // O is a submatrix of the other side of the graph where column vectors are feature vectors of
	    // those vertices that are connected to this vertex
	    // O^T is a transpose of O
	    // D is a diagonal matrix: D = lambda * degree * I, where I is an identity matrix
	    // b is a vector: b = O * r
	    // r is a ratings vector formed from the ratings of outgoing edges

	    // Calculating O*O^T and b
	    for (int i = 0; i < RANK; i++) {
	      v->feature_vec[i] += u->feature_vec[i] * u->rating;
	      for (int j = 0; j < RANK; j++) {
		v->temp_mat[i][j] += u->feature_vec[i] * u->feature_vec[j];
	      }
	    }

	    // Additional procesing after all updates have been gathered
	    if (v->count == v->degree)
	      {
		// Adding D to A
		for (int i = 0; i < RANK; i++) {
		  v->temp_mat[i][i] += LAMBDA * v->degree;
		}

		// Solve to get the new feature vector for the vertex
		solve(v->temp_mat, v->feature_vec);

		// Reset count for the next iteration
		v->count = 0;
	      }

	    return true;
	  }

	// After all iterations are finished, we use one more phase to
	// compute the sum of square errors. This is done on the right side.
	else
	  {
	    double sqerror = u->rating;
	    for (int i = 0; i < RANK; i++)
	      sqerror -= v->feature_vec[i] * u->feature_vec[i];
	    sqerror *= sqerror;

	    static_cast<als_per_processor_data*>(per_cpu_data)->sse_local += sqerror;

	    // And we stop processing
	    return false;
	  }
      }

      static bool generate_update(unsigned char* vertex_state,
				  unsigned char* edge_format,
				  unsigned char* update_stream,
				  per_processor_data* per_processor_data,
				  unsigned long bsp_phase)
      {
	vertex_t src, dst;
	weight_t rating;
	F::read_edge(edge_format, src, dst, rating);

	struct vertex* vertices = (struct vertex*)vertex_state;
	struct vertex* v = &vertices[x_lib::configuration::map_offset(src)];

	// Iteration 0 is also used to count the vertex degree
	if (bsp_phase == 0)
	  {
	    v->degree++;
	  }

	// The graph is bipartite, and it is assumed that lower ids
	// form the left side.
	// The alternating starts by solving the right side in iteration 1,
	// after the left side generates first updates in iteration 0.
	// Therefore, in even numbered iterations we solve the left
	// side, and in odd numbered iterations we solve the right side.
	// Likewise, the left side generates updates in even numbered
	// iterations, and the right side in odd numbered iterations.

	bool leftside = (src < dst) ? true : false;

	if ((bsp_phase % 2 == 0 && leftside) || (bsp_phase % 2 == 1 && !leftside))
	  {
	    struct update* u = (struct update*)update_stream;
	    u->target = dst;
	    u->rating = (double)rating;
	    copy_vector(v->feature_vec, u->feature_vec);
	    return true;
	  }
	else
	  return false;
      }

      static void preprocessing()
      {
	// We double the number, since in one phase we process one side of the
	// bipartite graph
	niters = 2 * vm["als::niters"].as<unsigned long>();
      }

      static void postprocessing()
      {
	BOOST_LOG_TRIVIAL(info) << "ALGORITHM::ALS::SSE " << als_per_processor_data::sse;
	unsigned long nedges = pt.get<unsigned long>("graph.edges");
	double rmse = std::sqrt( als_per_processor_data::sse / (1. * nedges) );
	BOOST_LOG_TRIVIAL(info) << "ALGORITHM::ALS::RMSE " << rmse;
      }

      static per_processor_data *
      create_per_processor_data(unsigned long processor_id)
      {
	return new als_per_processor_data();
      }

      static unsigned long min_super_phases()
      {
	return 1;
      }

      static bool need_init(unsigned long bsp_phase)
      {
	return (bsp_phase == 0);
      }

    };

    // These should be in a cpp file, but it's ok since we only include
    // this header once in driver.cpp
    template<typename F>
    unsigned long als_factorization<F>::niters;

    double als_per_processor_data::sse = 0;

  }
}
#endif