//
// Created by Arnie on 2016-11-15.
//

#include "dtm.h"
#include <Eigen/Core>
#include <Eigen/Dense>
#include <sstream>
#include <fstream>
#include <utility>
#include <functional>
#include <algorithm>
#include <cmath>
using namespace std;

VectorXf get_mvn_samples(VectorXf mean, MatrixXf cov) {
  Eigen::SelfAdjointEigenSolver<Eigen::MatrixXf> eigenSolver(cov);
  normal_distribution<float> dist(0, 1);
  default_random_engine gen;
  auto std_norm = [&] (float) {return dist(gen);};
  return mean + eigenSolver.eigenvectors() * eigenSolver
      .eigenvalues().cwiseSqrt().asDiagonal() *
      VectorXf::NullaryExpr(mean.size(), std_norm);
}

static float gaussianVar(float mean, float std_dev) {
  default_random_engine generator;
  normal_distribution<float> distribution(mean, std_dev);
  return distribution(generator);
}

static VectorXf softmax(VectorXf weights) {
  size_t K = weights.size();
  VectorXf soft(K);
  float MAX = weights[0];
  float norm = 0.0;
  for (size_t i = 0; i < K; i++) {
    if (weights(i) > MAX)
      MAX = weights(i);
  }

  for (size_t i = 0; i < K; i++)   {
    norm += exp(weights(i) - MAX);
  }

  for (size_t i = 0; i < K; i++) {
    soft(i) = exp(weights(i) - MAX) / (norm + 0.00001);
    if (soft(i) < 1e-30) {
      soft(i) = 0.0;
    }
  }

  return soft;
}

void DTM::build_alias_table(size_t t, size_t w) {
  AliasSamples term(phi[t].row(w));
  term_alias_samples[t][w] = term.get_samples(K);
}

DTM::DTM(const vector<vector<vector<size_t>>> &data, const vector<string> &
dictionary,
         size_t num_topics, float sgld_a, float sgld_b, float sgld_c,
         float dtm_phi_var, float dtm_eta_var, float dtm_alpha_var) : W(data),
         vocabulary(dictionary), K(num_topics), sgld_a(sgld_a), sgld_b(sgld_b),
         sgld_c(sgld_c), dtm_phi_var(dtm_phi_var), dtm_eta_var(dtm_eta_var),
         dtm_alpha_var(dtm_alpha_var) {
  V = vocabulary.size();
  T = W.size();
  D = vector<size_t>(T);
  Z = vector<vector<vector<size_t>>>(T);
  term_alias_samples =
      vector<vector<vector<size_t>>>
          (T, vector<vector<size_t>> (V, vector<size_t>(K)));
  sample_indices = vector<vector<size_t>>(T, vector<size_t>(V));

  CDK = vector<MatrixXf> (T);
  CWK = vector<MatrixXf> (T, MatrixXf::Zero(V, K));
  CK = MatrixXf::Zero(T, K);
  phi = vector<MatrixXf>(T, MatrixXf::Zero(V, K));
  eta = vector<MatrixXf>(T);
  alpha = MatrixXf::Zero(T, K);

  for (size_t t = 0; t < T; ++t) {
    D[t] = W[t].size();
    Z[t] = vector<vector<size_t>>(D[t]);
    eta[t] = MatrixXf::Zero(D[t], K);
    CDK[t] = MatrixXf::Zero(D[t], K);
    for (size_t d = 0; d < D[t]; ++d) {
      Z[t][d] = vector<size_t>(W[t][d].size());
    }
  }
}

void DTM::initialize(bool init_with_lda) {
  default_random_engine generator;
  uniform_int_distribution<size_t> uniform_topic(0, K - 1);
  u01 = uniform_real_distribution<float>(0, 1);
  float init_alpha = 50.0 / K;
  float init_beta = 0.01;
  for (size_t t = 0; t < T; ++t) {
    for (size_t d = 0; d < D[t]; ++d) {
      size_t N = W[t][d].size();
      for (size_t n = 0; n < N; ++n) {
        size_t w = W[t][d][n];
        size_t k = uniform_topic(generator);
        Z[t][d][n] = k;
        CDK[t](d, k)++;
        CWK[t](w, k)++;
        CK(t, k)++;
        eta[t](d, k) += (1 + init_alpha) / (N + K * init_alpha);
      }
    }
  }

  if (init_with_lda) {
    size_t t = 0;
    for (size_t iter = 0; iter < 50; iter++) {
      cout << "LDA Iter: " << iter << endl;
      for (size_t d = 0; d < D[t]; ++d) {
        for (size_t n = 0; n < W[t][d].size(); n++) {
          size_t k = Z[t][d][n];
          size_t w = W[t][d][n];
          CDK[t](d, k)--;
          CWK[t](w, k)--;
          CK(t, k)--;

          vector<float> prob(K);
          for (k = 0; k < K; k++) {
            prob[k] = (CDK[t](d, k) + init_alpha) * ((CWK[t](w, k) + init_beta)
                / (CK(t, k) + V * init_beta));
          }
          discrete_distribution<size_t> mult(prob.begin(), prob.end());
          k = mult(generator);
          Z[t][d][n] = k;
          CDK[t](d, k)++;
          CWK[t](w, k)++;
          CK(t, k)++;
        }
      }
    }
  }
  for (size_t t = 0; t < T; t++) {
    for (size_t w = 0; w < V; w++) {
      for (size_t k = 0; k < K; k++) {
        phi[t](w, k) = (CWK[0](w, k) + init_beta) / (CK(0, k) + V * init_beta);
      }
      build_alias_table(t, w);
    }
  }
}

void DTM::estimate(size_t num_iters) {
  default_random_engine generator;
  for (size_t iter = 0; iter < num_iters; iter++) {
    cout << "Iteration " << iter << endl;
    float eps = sgld_a * (pow(sgld_b + iter, -sgld_c));
    float xi = gaussianVar(0.0, pow(eps, 2));
    VectorXf xi_vec;
    VectorXf mean(K);
    for (size_t t = 0; t < T; t++) {
      xi_vec = VectorXf::Constant(K, xi);
      for (size_t d = 0; d < D[t]; d++) {
        size_t N = W[t][d].size();
        uniform_int_distribution<size_t> doc_dist(0, N - 1);

        // estimate eta
        VectorXf soft_eta = softmax(eta[t].row(d));
        VectorXf prior_eta = (alpha.row(t) - eta[t].row(d)) / dtm_eta_var;
        VectorXf denom_eta = N * soft_eta;
        VectorXf grad_eta = CDK[t].row(d).transpose() - denom_eta;
        eta[t].row(d) += ((eps / 2) * (grad_eta + prior_eta)) + xi_vec;

        for (size_t n = 0; n < N; n++) {
          for (size_t mh = 0; mh < 4; mh++) {
            size_t k = Z[t][d][n];
            size_t w = W[t][d][n];
            CDK[t](d, k)--;
            CWK[t](w, k)--;
            CK(t, k)--;

            size_t proposal;
            float acceptance_prob = 0.0;
            if (mh % 2 == 0) {
              // Z-proposal
              size_t index = doc_dist(generator);
              proposal = Z[t][d][index];

              acceptance_prob =
                  exp(phi[t](w, proposal)) / exp(phi[t](w, k));
            } else {
              if (sample_indices[t][w] >= K) {
                build_alias_table(t, w);
                sample_indices[t][w] = 0;
              }
              proposal = term_alias_samples[t][w][sample_indices[t][w]];
              sample_indices[t][w]++;
              acceptance_prob = exp(eta[t](d, proposal)) / exp(eta[t](d, k));
            }
            acceptance_prob = acceptance_prob > 1.0 ? 1.0 : acceptance_prob;
            if (u01(generator) >= acceptance_prob) {
              // reject proposal
              proposal = k;
            }
            Z[t][d][n] = proposal;
            CDK[t](d, proposal)++;
            CWK[t](w, proposal)++;
            CK(t, proposal)++;
          }
        }
      }

      xi_vec = VectorXf::Constant(V, xi);
      for (unsigned k = 0; k < K; ++k) {
        // sample phi
        VectorXf soft_phi = softmax(phi[t].col(k));
        VectorXf prior_phi(V);
        if (t == 0) {
          float phi_sigma = 1.0 / ((1.0 / 100) + (1 / dtm_phi_var));
          prior_phi = phi[t + 1].col(k) * (phi_sigma / dtm_phi_var);
          prior_phi = ((2 * prior_phi) - 2 * phi[t].col(k)) / dtm_phi_var;
        } else if (t == T - 1) {
          prior_phi = (phi[t - 1].col(k) - phi[t].col(k)) / dtm_phi_var;
        } else {
          prior_phi = (phi[t + 1].col(k) + phi[t - 1].col(k) - 2 * phi[t].col
              (k)) / dtm_phi_var;
        }

        VectorXf denom_phi = CK(t, k) * soft_phi;
        VectorXf grad_phi = CWK[t].col(k) - denom_phi;

        phi[t].col(k) += ((eps / 2) * (grad_phi + prior_phi)) + xi_vec;
      }

      // sample alpha
      VectorXf alpha_bar(K);
      float alpha_precision = 0.0;  // designed to be a diagonal matrix
      MatrixXf cov = MatrixXf::Identity(K, K);
      if (t == 0) {
        alpha_precision = (1.0 / 100) + (1 / dtm_alpha_var);
        float alpha_sigma = 1.0 / alpha_precision;
        alpha_bar = alpha.row(t+1) * (alpha_sigma / dtm_alpha_var);
      } else if (t == T-1) {
        alpha_bar = (alpha.row(t-1) - alpha.row(t)) / dtm_alpha_var;
        alpha_precision = 1.0 / dtm_alpha_var;
      } else {
        alpha_precision = (2 / dtm_alpha_var);
        alpha_bar = (alpha.row(t+1) - alpha.row(t-1)) / 2;
      }
      VectorXf eta_bar = eta[t].colwise().sum();
      float sigma = 1.0 / (1.0 / alpha_precision + (D[t] / dtm_eta_var));
      cov *= sigma;
      mean = (alpha_bar / alpha_precision + (eta_bar / dtm_eta_var)) * sigma;
      alpha.row(t) = get_mvn_samples(mean, cov);
      if (iter % 5 == 0) {
        diagnosis(t);
      }
    }
  }
}

void DTM::diagnosis(size_t t) {
  float perp = 0.0;
  unsigned N = 0;
  float total_log_likelihood = 0.0;
  vector<VectorXf> softmax_phi(K);
  vector<VectorXf> softmax_eta(D[t]);
  for (size_t k = 0; k < K; ++k) {
    softmax_phi[k] = softmax(phi[t].col(k));
  }
  for (size_t d = 0; d < D[t]; d++) {
    N += W[t][d].size();
    softmax_eta[d] = softmax(eta[t].row(d));
    for (size_t n = 0; n < W[t][d].size(); n++) {
      float likelihood = 0.0;
      size_t w = W[t][d][n];
      for (size_t k = 0; k < K; k++) {
        likelihood += ((softmax_eta[d](k) * (softmax_phi[k](w))));
        if (likelihood < 0)
          std::cout << "Likelihood less than 0, error" << std::endl;
      }
      total_log_likelihood += log(likelihood);
    }
  }
  cout << "Perplexity: " <<  t << "  "
       << exp(-total_log_likelihood / N) << endl;
}

void DTM::save_data(string dir) {
  for (size_t t = 0; t < T; ++t) {
    stringstream sstm;
    sstm << dir << "/time_slice_" << t << ".txt";
    string fname = sstm.str();
    ofstream myfile;

    myfile.open(fname.c_str());
    for (size_t k = 0; k < K; ++k) {
      vector<pair<float, size_t>> ranking;
      for (size_t v = 0; v < V; v++) {
        ranking.push_back(make_pair(phi[t](v, k), v));
      }
      sort(ranking.begin(), ranking.end(),
           std::greater<pair<float, size_t>>());
      myfile << "Topic " << k << "\n";
      for (size_t v = 0; v < 10; v++) {
        size_t w = ranking[v].second;
        myfile << "(" << vocabulary[w] << ", " << phi[t](w, k) << ")" << endl;
      }
      myfile << endl;
    }
    myfile.close();
  }
}


/*
 *  Copyright 2011-2015 Maxim Milakov
 *
 *  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.
 */

#include "roc_result.h"

#include <boost/format.hpp>
#include <algorithm>
#include <numeric>

namespace nnforge
{
	roc_result::roc_result(
		const output_neuron_value_set& predicted_value_set,
		const output_neuron_value_set& actual_value_set,
		float threshold,
		float beta,
		unsigned int segment_count,
		float min_val,
		float max_val)
		: segment_count(segment_count)
		, min_val(min_val)
		, max_val(max_val)
		, threshold(threshold)
		, beta(beta)
		, actual_positive_elem_count(0)
		, actual_negative_elem_count(0)
		, values_for_positive_elems(segment_count)
		, values_for_negative_elems(segment_count)
	{
		float mult = 1.0F / (max_val - min_val);
		float segment_count_f = static_cast<float>(segment_count);
		std::vector<std::vector<float> >::const_iterator predicted_it = predicted_value_set.neuron_value_list.begin();
		for(std::vector<std::vector<float> >::const_iterator actual_it = actual_value_set.neuron_value_list.begin();
			actual_it != actual_value_set.neuron_value_list.end();
			actual_it++, predicted_it++)
		{
			const std::vector<float>& actual_value_list = *actual_it;
			const std::vector<float>& predicted_value_list = *predicted_it;

			std::vector<float>::const_iterator predicted_value_it = predicted_value_list.begin();
			for(std::vector<float>::const_iterator actual_value_it = actual_value_list.begin();
				actual_value_it != actual_value_list.end();
				actual_value_it++, predicted_value_it++)
			{
				float actual_value = *actual_value_it;
				float predicted_value = *predicted_value_it;

				unsigned int bucket_id = std::min<unsigned int>(static_cast<unsigned int>(std::max<float>(std::min<float>((predicted_value - min_val) * mult, 1.0F), 0.0F) * segment_count_f), (segment_count - 1));

				if (actual_value > 0.0F)
				{
					values_for_positive_elems[bucket_id]++;
					actual_positive_elem_count++;
				}
				else
				{
					values_for_negative_elems[bucket_id]++;
					actual_negative_elem_count++;
				}
			}
		}
	}

	float roc_result::get_accuracy() const
	{
		unsigned int starting_segment_id = static_cast<unsigned int>(std::max(std::min((threshold - min_val) / (max_val - min_val), 1.0F), 0.0F) * static_cast<float>(segment_count));

		unsigned int true_positive = std::accumulate(values_for_positive_elems.begin() + starting_segment_id, values_for_positive_elems.end(), 0);
		unsigned int true_negative = std::accumulate(values_for_negative_elems.begin(), values_for_negative_elems.begin() + starting_segment_id, 0);

		return static_cast<float>(true_positive + true_negative) / static_cast<float>(actual_positive_elem_count + actual_negative_elem_count);
	}

	float roc_result::get_f_score() const
	{
		unsigned int starting_segment_id = static_cast<unsigned int>(std::max(std::min((threshold - min_val) / (max_val - min_val), 1.0F), 0.0F) * static_cast<float>(segment_count));

		unsigned int true_positive = std::accumulate(values_for_positive_elems.begin() + starting_segment_id, values_for_positive_elems.end(), 0);
		unsigned int false_positive = std::accumulate(values_for_negative_elems.begin() + starting_segment_id, values_for_negative_elems.end(), 0);
		// unsigned int true_negative = std::accumulate(values_for_negative_elems.begin(), values_for_negative_elems.begin() + starting_segment_id, 0);
		unsigned int false_negative = std::accumulate(values_for_positive_elems.begin(), values_for_positive_elems.begin() + starting_segment_id, 0);

		return (1.0F + beta * beta) * static_cast<float>(true_positive) /
			((1.0F + beta * beta) * static_cast<float>(true_positive) + (beta * beta) * static_cast<float>(false_negative) + static_cast<float>(false_positive));
	}

	float roc_result::get_precision() const
	{
		unsigned int starting_segment_id = static_cast<unsigned int>(std::max(std::min((threshold - min_val) / (max_val - min_val), 1.0F), 0.0F) * static_cast<float>(segment_count));

		unsigned int true_positive = std::accumulate(values_for_positive_elems.begin() + starting_segment_id, values_for_positive_elems.end(), 0);
		unsigned int false_positive = std::accumulate(values_for_negative_elems.begin() + starting_segment_id, values_for_negative_elems.end(), 0);

		return static_cast<float>(true_positive) / (static_cast<float>(true_positive) + static_cast<float>(false_positive));
	}

	float roc_result::get_recall() const
	{
		unsigned int starting_segment_id = static_cast<unsigned int>(std::max(std::min((threshold - min_val) / (max_val - min_val), 1.0F), 0.0F) * static_cast<float>(segment_count));

		unsigned int true_positive = std::accumulate(values_for_positive_elems.begin() + starting_segment_id, values_for_positive_elems.end(), 0);
		unsigned int false_negative = std::accumulate(values_for_positive_elems.begin(), values_for_positive_elems.begin() + starting_segment_id, 0);

		return static_cast<float>(true_positive) / (static_cast<float>(true_positive) + static_cast<float>(false_negative));
	}

	float roc_result::get_auc() const
	{
		std::vector<float> true_positive_rates;
		{
			unsigned int current_positive_elems_count = 0;
			float mult = 1.0F / static_cast<float>(actual_positive_elem_count);
			for(std::vector<unsigned int>::const_reverse_iterator it = values_for_positive_elems.rbegin(); it != values_for_positive_elems.rend(); ++it)
			{
				current_positive_elems_count += *it;
				true_positive_rates.push_back(mult * static_cast<float>(current_positive_elems_count));
			}
		}
		true_positive_rates.push_back(1.0F);

		std::vector<float> false_positive_rates;
		{
			unsigned int current_negative_elems_count = 0;
			float mult = 1.0F / static_cast<float>(actual_negative_elem_count);
			for(std::vector<unsigned int>::const_reverse_iterator it = values_for_negative_elems.rbegin(); it != values_for_negative_elems.rend(); ++it)
			{
				current_negative_elems_count += *it;
				false_positive_rates.push_back(mult * static_cast<float>(current_negative_elems_count));
			}
		}
		false_positive_rates.push_back(1.0F);

		float sum = 0.0F;
		float previous_fpr = 0.0F;
		float previous_tpr = 0.0F;
		std::vector<float>::const_iterator tpr_it = true_positive_rates.begin();
		for(std::vector<float>::const_iterator fpr_it = false_positive_rates.begin(); fpr_it != false_positive_rates.end(); ++fpr_it, ++tpr_it)
		{
			float current_fpr = *fpr_it;
			float current_tpr = *tpr_it;

			if (current_fpr != previous_fpr)
				sum += (current_fpr - previous_fpr) * (previous_tpr + current_tpr) * 0.5F;

			previous_fpr = current_fpr;
			previous_tpr = current_tpr;
		}

		return sum;
	}

	std::ostream& operator<< (std::ostream& out, const roc_result& val)
	{
		out << (boost::format("AUC %|1$.5f|, (using threshold %|2$.3f|) Accuracy %|3$.5f|, Precision %|4$.5f|, Recall %|5$.5f|, F-score %|6$.5f| (beta %|7$.3f|)") % val.get_auc() % val.threshold % val.get_accuracy() % val.get_precision() % val.get_recall() % val.get_f_score() % val.beta).str();

		return out;
	}
}


#include "hoCuNDArray_utils.h"
#include "cuNDArray_fileio.h"
#include "cuNDArray_math.h"
#include "imageOperator.h"
#include "hoPartialDerivativeOperator.h"
#include "hoCuConebeamProjectionOperator.h"
#include "cgSolver.h"
#include "hoCuGPBBSolver.h"
#include "hoCuNCGSolver.h"
#include "hoCuPartialDerivativeOperator.h"
#include "CBSubsetOperator.h"
#include "osSPSSolver.h"
#include "osMOMSolverF.h"
#include <boost/program_options.hpp>
#include "osPDsolver.h"
#include "hdf5_utils.h"
#include "cuEdgeATrousOperator.h"
#include "cuDCTOperator.h"
#include "cuNCGSolver.h"
#include "subselectionOperator.h"
#include "solver_utils.h"
#include "cuSolverUtils.h"
#include "osMOMSolverDual.h"
using namespace std;
using namespace Gadgetron;

namespace po = boost::program_options;

boost::shared_ptr<cuNDArray<float> > calculate_prior(boost::shared_ptr<CBCT_binning>  binning,boost::shared_ptr<CBCT_acquisition> ps, hoCuNDArray<float>& projections, std::vector<size_t> is_dims, floatd3 imageDimensions){
	std::cout << "Calculating FDK prior" << std::endl;
	boost::shared_ptr<CBCT_binning> binning_pics=binning->get_3d_binning();
	std::vector<size_t> is_dims3d = is_dims;
	is_dims3d.pop_back();
	boost::shared_ptr< hoCuConebeamProjectionOperator >
	Ep( new hoCuConebeamProjectionOperator() );
	Ep->setup(ps,binning_pics,imageDimensions);
	Ep->set_codomain_dimensions(ps->get_projections()->get_dimensions().get());
	Ep->set_domain_dimensions(&is_dims3d);
	Ep->set_use_filtered_backprojection(true);
	boost::shared_ptr<hoCuNDArray<float> > prior3d(new hoCuNDArray<float>(&is_dims3d));
	Ep->mult_MH(&projections,prior3d.get());

	hoCuNDArray<float> tmp_proj(*ps->get_projections());
	Ep->mult_M(prior3d.get(),&tmp_proj);
	float s = dot(ps->get_projections().get(),&tmp_proj)/dot(&tmp_proj,&tmp_proj);
	*prior3d *= s;
	boost::shared_ptr<cuNDArray<float> > prior(new cuNDArray<float>(*expand( prior3d.get(), is_dims.back() )));
	std::cout << "Prior complete" << std::endl;
	return prior;
}

boost::shared_ptr<cuNDArray<float>> calculate_weightImage(boost::shared_ptr<CBCT_binning>  binning,boost::shared_ptr<CBCT_acquisition> ps, hoCuNDArray<float>& ho_projections, std::vector<size_t> is_dims, floatd3 imageDimensions){

	cuNDArray<float> projections(ho_projections);
	boost::shared_ptr<CBCT_binning> binning_pics=binning->get_3d_binning();
	std::vector<size_t> is_dims3d = is_dims;
	is_dims3d.pop_back();
	auto Ep = boost::make_shared<cuConebeamProjectionOperator>();
	auto ps2 = boost::make_shared<CBCT_acquisition>(boost::shared_ptr<hoCuNDArray<float>>(),ps->get_geometry());
	Ep->setup(ps2,binning_pics,imageDimensions);
	Ep->set_codomain_dimensions(ps->get_projections()->get_dimensions().get());
	Ep->set_domain_dimensions(&is_dims3d);
	//Ep->set_use_filtered_backprojection(true);
	Ep->offset_correct(&projections);
	//Ep->mult_MH(&projections,prior3d.get());
	//cgSolver<hoCuNDArray<float>> solv;
	cuNCGSolver<float> solv;
	solv.set_non_negativity_constraint(true);
	solv.set_encoding_operator(Ep);
	solv.set_max_iterations(10);
	auto prior3d = solv.solve(&projections);
	//auto prior3d = boost::make_shared<cuNDArray<float>>(is_dims3d);
	write_nd_array(prior3d.get(),"fdk.real");
	cuNDArray<float> tmp_proj(projections);
	clear(&tmp_proj);
	Ep->mult_M(prior3d.get(),&tmp_proj);
	//float s = dot(ps->get_projections().get(),&tmp_proj)/dot(&tmp_proj,&tmp_proj);
	//std::cout << "Scaling " << s << std::endl;
	/*
	 //Ep->offset_correct(&tmp_proj);
	//tmp_proj *= s;
	write_nd_array(&tmp_proj,"projtmp.real");
	tmp_proj -= projections;
	tmp_proj *= float(-1);
	abs_inplace(&tmp_proj);

	write_nd_array(ps->get_projections().get(),"proj.real");
	write_nd_array(&tmp_proj,"projdiff.real");
	std::cout << "Proj size ";
	auto pdims = *tmp_proj.get_dimensions();
	for (auto p : pdims ) std::cout << p << " ";
	std::cout << std::endl;
	//Ep->set_use_filtered_backprojection(false);
	//Ep->mult_MH(&tmp_proj,prior3d.get());
	//solv.set_non_negativity_constraint(false);
	prior3d = solv.solve(&tmp_proj);
	//abs_inplace(prior.get());
	std::cout << "Prior complete" << std::endl;
	 */
	return prior3d;
}

int main(int argc, char** argv)
{
	string acquisition_filename;
	string outputFile;
	uintd3 imageSize;
	floatd3 voxelSize;
	int device;
	unsigned int iterations;
	floatd2 scale_factor;
	unsigned int subsets;
	float rho;
	float tv_weight,pics_weight, wavelet_weight,huber,sigma,dct_weight;

	po::options_description desc("Allowed options");

	desc.add_options()
    						("help", "produce help message")
    						("acquisition,a", po::value<string>(&acquisition_filename)->default_value("acquisition.hdf5"), "Acquisition data")
    						("samples,n",po::value<unsigned int>(),"Number of samples per ray")
    						("output,f", po::value<string>(&outputFile)->default_value("reconstruction.hdf5"), "Output filename")
    						("size,s",po::value<uintd3>(&imageSize)->default_value(uintd3(512,512,1)),"Image size in pixels")
    						("binning,b",po::value<string>(),"Binning file for 4d reconstruction")
    						("SAG","Use exact SAG correction if present")
    						("voxelSize,v",po::value<floatd3>(&voxelSize)->default_value(floatd3(0.488f,0.488f,1.0f)),"Voxel size in mm")
    						("dimensions,d",po::value<floatd3>(),"Image dimensions in mm. Overwrites voxelSize.")
    						("iterations,i",po::value<unsigned int>(&iterations)->default_value(10),"Number of iterations")
    						("device",po::value<int>(&device)->default_value(0),"Number of the device to use (0 indexed)")
	     					 ("downsample,D",po::value<floatd2>(&scale_factor)->default_value(floatd2(1,1)),"Downsample projections this factor")
    						("subsets,u",po::value<unsigned int>(&subsets)->default_value(10),"Number of subsets to use")
    						("TV",po::value<float>(&tv_weight)->default_value(0),"Total variation weight")
    						("PICS",po::value<float>(&pics_weight)->default_value(0),"PICS weight")
    						("Wavelet,W",po::value<float>(&wavelet_weight)->default_value(0),"Weight of the wavelet operator")
    						("Huber",po::value<float>(&huber)->default_value(0),"Huber weight")
    						("use_prior","Use an FDK prior")
    						("sigma",po::value<float>(&sigma)->default_value(0.001),"Sigma for billateral filter")
    						("DCT",po::value<float>(&dct_weight)->default_value(0),"DCT regularization")
    						("3D","Only use binning for selecting valid projections")
    						;

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

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

	std::stringstream command_line_string;
	std::cout << "Command line options:" << std::endl;
	for (po::variables_map::iterator it = vm.begin(); it != vm.end(); ++it){
		boost::any a = it->second.value();
		command_line_string << it->first << ": ";
		if (a.type() == typeid(std::string)) command_line_string << it->second.as<std::string>();
		else if (a.type() == typeid(int)) command_line_string << it->second.as<int>();
		else if (a.type() == typeid(unsigned int)) command_line_string << it->second.as<unsigned int>();
		else if (a.type() == typeid(float)) command_line_string << it->second.as<float>();
		else if (a.type() == typeid(vector_td<float,3>)) command_line_string << it->second.as<vector_td<float,3> >();
		else if (a.type() == typeid(vector_td<int,3>)) command_line_string << it->second.as<vector_td<int,3> >();
		else if (a.type() == typeid(vector_td<unsigned int,3>)) command_line_string << it->second.as<vector_td<unsigned int,3> >();
		else command_line_string << "Unknown type" << std::endl;
		command_line_string << std::endl;
	}
	std::cout << command_line_string.str();

	cudaSetDevice(device);

	//Really weird stuff. Needed to initialize the device?? Should find real bug.
	cudaDeviceManager::Instance()->lockHandle();
	cudaDeviceManager::Instance()->unlockHandle();

	boost::shared_ptr<CBCT_acquisition> ps(new CBCT_acquisition());
	ps->load(acquisition_filename);
	ps->get_geometry()->print(std::cout);
    if (scale_factor[0] != 1 || scale_factor[1] != 1)
        ps->downsample(scale_factor[0],scale_factor[1]);

	float SDD = ps->get_geometry()->get_SDD();
	float SAD = ps->get_geometry()->get_SAD();

	boost::shared_ptr<CBCT_binning> binning(new CBCT_binning());
	if (vm.count("binning")){
		std::cout << "Loading binning data" << std::endl;
		binning->load(vm["binning"].as<string>());
		if (vm.count("3D"))
			binning = binning->get_3d_binning();
	} else binning->set_as_default_3d_bin(ps->get_projections()->get_size(2));
	binning->print(std::cout);

	floatd3 imageDimensions;
	if (vm.count("dimensions")){
		imageDimensions = vm["dimensions"].as<floatd3>();
		voxelSize = imageDimensions/imageSize;
	}
	else imageDimensions = voxelSize*imageSize;

	float lengthOfRay_in_mm = norm(imageDimensions);
	unsigned int numSamplesPerPixel = 3;
	float minSpacing = min(voxelSize)/numSamplesPerPixel;

	unsigned int numSamplesPerRay;
	if (vm.count("samples")) numSamplesPerRay = vm["samples"].as<unsigned int>();
	else numSamplesPerRay = ceil( lengthOfRay_in_mm / minSpacing );

	float step_size_in_mm = lengthOfRay_in_mm / numSamplesPerRay;
	size_t numProjs = ps->get_projections()->get_size(2);
	size_t needed_bytes = 2 * prod(imageSize) * sizeof(float);
	std::vector<size_t> is_dims = to_std_vector((uint64d3)imageSize);

	std::cout << "IS dimensions " << is_dims[0] << " " << is_dims[1] << " " << is_dims[2] << std::endl;
	std::cout << "Image size " << imageDimensions << std::endl;

	is_dims.push_back(binning->get_number_of_bins());

	std::vector<size_t> double_dims = is_dims;
	double_dims.push_back(2);

	//osLALMSolver<cuNDArray<float>> solver;
	osMOMSolverDual<cuNDArray<float>> solver;
	//osAHZCSolver<cuNDArray<float>> solver;
	//osMOMSolverF<cuNDArray<float>> solver;
	//ADMMSolver<cuNDArray<float>> solver;
	solver.set_dump(false);

	boost::shared_ptr<cuNDArray<float>> prior;
	if (vm.count("use_prior") || pics_weight > 0) {
		auto projections = *ps->get_projections();
		prior = calculate_prior(binning,ps,projections,is_dims,imageDimensions);
		//prior = calculate_weightImage(binning,ps,projections,is_dims,imageDimensions);
		solver.set_x0(prior);
	}
	//osPDSolver<cuNDArray<float>> solver;
	/*
  {
  hoCuConebeamProjectionOperator op;
  	auto bin3D = boost::make_shared<CBCT_binning>(binning->get_3d_binning());
  	op.setup(ps,bin3D,imageDimensions);
  	hoCuNDArray<float> proj(*ps->get_projections());
  	op.offset_correct(&proj);
  	op.set_use_filtered_backprojection(true);

  	std::vector<size_t> is_dims3D = to_std_vector((uint64d3)imageSize);

  	hoCuNDArray<float> image(is_dims3D);

  	op.mult_MH(&proj,&image,false);

  	auto cuimage = boost::make_shared<cuNDArray<float>>(image);
  	solver.set_x0(expand(cuimage.get(),is_dims.back()));
  }
	 */
	//solver.set_regularization_iterations(1);
	//osSPSSolver<cuNDArray<float>> solver;
	/*
  if (pics_weight > 0){
  	std::cout << "Calculating PICS prior" << std::endl;
  	hoCuConebeamProjectionOperator op;
  	auto bin3D = boost::make_shared<CBCT_binning>(binning->get_3d_binning());
  	op.setup(ps,bin3D,imageDimensions);
  	hoCuNDArray<float> proj(*ps->get_projections());
  	op.offset_correct(&proj);
  	op.set_use_filtered_backprojection(true);

  	std::vector<size_t> is_dims3D = to_std_vector((uint64d3)imageSize);

  	hoCuNDArray<float> image(is_dims3D);

  	op.mult_MH(&proj,&image,false);
  	auto prior = boost::make_shared<cuNDArray<float>>(image);
  	auto PICS = boost::make_shared<cuTvPicsOperator<float,3>>();
  	PICS->set_prior(prior);
  	PICS->set_weight(pics_weight);
  	solver.add_nonlinear_operator(PICS);

  	write_nd_array(prior.get(),"fdk_prior.real");
  }

	 */

	// Define encoding matrix

	/*auto weight_array = boost::make_shared<cuNDArray<float>>(is_dims);
  {
  	boost::shared_ptr<linearOperator<cuNDArray<float>>> E2(E);
  	cuNDArray<float> tmp_proj(*ps->get_projections()->get_dimensions());
  	fill(&tmp_proj,1.0f);
  	E2->mult_MH(&tmp_proj,weight_array.get(),false);
  	clamp_min(weight_array.get(),1.0f);
  	reciprocal_inplace(weight_array.get());
  	//*weight_array *= *weight_array;
  }
  write_nd_array(weight_array.get(),"weights.real");
	 */

	//hoCuCgDescentSolver<float> solver;

	//osSPSSolver<hoNDArray<float>> solver;
	//hoCuNCGSolver<float> solver;
	//solver.set_domain_dimensions(&is_dims);
	solver.set_max_iterations(iterations);
	solver.set_output_mode(osSPSSolver<cuNDArray<float>>::OUTPUT_VERBOSE);
	solver.set_tau(5e-5);
	solver.set_non_negativity_constraint(true);
	solver.set_huber(huber);

	solver.set_reg_steps(4);
	//solver.set_rho(rho);

	if (tv_weight > 0){

		auto Dx = boost::make_shared<cuPartialDerivativeOperator<float,4>>(0);
		Dx->set_weight(tv_weight);

		Dx->set_domain_dimensions(&is_dims);
		Dx->set_codomain_dimensions(&is_dims);
		/*
  	Dx->set_domain_dimensions(&double_dims);
  	Dx->set_codomain_dimensions(&double_dims);
		 */
		auto Dy = boost::make_shared<cuPartialDerivativeOperator<float,4>>(1);
		Dy->set_weight(tv_weight);
		Dy->set_domain_dimensions(&is_dims);
		Dy->set_codomain_dimensions(&is_dims);
		/*
  	Dy->set_domain_dimensions(&double_dims);
  	Dy->set_codomain_dimensions(&double_dims);
		 */

		auto Dz = boost::make_shared<cuPartialDerivativeOperator<float,4>>(2);
		Dz->set_weight(tv_weight);
		Dz->set_domain_dimensions(&is_dims);
		Dz->set_codomain_dimensions(&is_dims);
		/*
  	Dz->set_domain_dimensions(&double_dims);
  	Dz->set_codomain_dimensions(&double_dims);
		 */

		auto Dt = boost::make_shared<cuPartialDerivativeOperator<float,4>>(3);
		Dt->set_weight(tv_weight);
		Dt->set_domain_dimensions(&is_dims);
		Dt->set_codomain_dimensions(&is_dims);

		auto Dx1 = boost::make_shared<subselectionOperator<cuNDArray<float>>>(Dx,0);
		Dx1->set_domain_dimensions(&double_dims);
		Dx1->set_codomain_dimensions(&is_dims);
		auto Dy1 = boost::make_shared<subselectionOperator<cuNDArray<float>>>(Dy,0);
		Dy1->set_domain_dimensions(&double_dims);
		Dy1->set_codomain_dimensions(&is_dims);
		auto Dz1 = boost::make_shared<subselectionOperator<cuNDArray<float>>>(Dz,0);
		Dz1->set_domain_dimensions(&double_dims);
		Dz1->set_codomain_dimensions(&is_dims);

		auto Dt1 = boost::make_shared<subselectionOperator<cuNDArray<float>>>(Dt,0);
		Dt1->set_domain_dimensions(&double_dims);
		Dt1->set_codomain_dimensions(&is_dims);

		Dx1->set_weight(tv_weight);
		Dy1->set_weight(tv_weight);
		Dz1->set_weight(tv_weight);
		Dt1->set_weight(tv_weight*2);

		//solver.add_regularization_group({Dx,Dy,Dz});
		solver.add_regularization_group({Dx1,Dy1,Dz1});
//        solver.add_regularization_operator(Dt1);
		/*
auto Dt = boost::make_shared<cuPartialDerivativeOperator<float,4>>(3);
	Dt->set_weight(tv_weight);
	Dt->set_domain_dimensions(&is_dims);
	Dt->set_codomain_dimensions(&is_dims);
	solver.add_regularization_operator(Dt);
		 */
		/*
	auto projections = *ps->get_projections();
  	auto prior_weight = calculate_weightImage(binning,ps,projections,is_dims,imageDimensions);
  	//sqrt_inplace(prior_weight.get());
  	std::cout << "Prior min " << min(prior_weight.get()) << std::endl;
  	//*prior_weight -= min(prior_weight.get());
		 *prior_weight /= asum(prior_weight.get())/prior_weight->get_number_of_elements();
		 *prior_weight -= max(prior_weight.get());
		 *prior_weight *= float(-1);
  	//clamp_min(prior_weight.get(),float(1e-2));
  	//reciprocal_inplace(prior_weight.get());

  	write_nd_array(prior_weight.get(),"prior.real");
  	//cudaDeviceReset();
  	auto Wt = boost::make_shared<weightingOperator<cuNDArray<float>>>(prior_weight,Dt);
	Wt->set_weight(tv_weight);
  	Wt->set_domain_dimensions(&is_dims);
  	Wt->set_codomain_dimensions(&is_dims);
		 */
		//lver.add_regularization_operator(Dt);

	}
	/*
  if (pics_weight > 0){

  	auto Dx = boost::make_shared<cuPartialDerivativeOperator<float,4>>(0);
  	Dx->set_weight(pics_weight);
  	Dx->set_domain_dimensions(&is_dims);
  	Dx->set_codomain_dimensions(&is_dims);

  	auto Dy = boost::make_shared<cuPartialDerivativeOperator<float,4>>(1);
  	Dy->set_weight(pics_weight);
  	Dy->set_domain_dimensions(&is_dims);
  	Dy->set_codomain_dimensions(&is_dims);

  	auto Dz = boost::make_shared<cuPartialDerivativeOperator<float,4>>(2);
  	Dz->set_weight(pics_weight);
  	Dz->set_domain_dimensions(&is_dims);
  	Dz->set_codomain_dimensions(&is_dims);

  	solver.add_regularization_group({Dx,Dy,Dz},prior);

  }*/

	/*
	if (tv_weight > 0){

		auto Dx = boost::make_shared<cuDCTDerivativeOperator<float>>(0);
		Dx->set_weight(tv_weight);
		Dx->set_domain_dimensions(&is_dims);
		Dx->set_codomain_dimensions(&is_dims);

		auto Dy = boost::make_shared<cuDCTDerivativeOperator<float>>(1);
		Dy->set_weight(tv_weight);
		Dy->set_domain_dimensions(&is_dims);
		Dy->set_codomain_dimensions(&is_dims);

		auto Dz = boost::make_shared<cuDCTDerivativeOperator<float>>(2);
		Dz->set_weight(tv_weight);
		Dz->set_domain_dimensions(&is_dims);
		Dz->set_codomain_dimensions(&is_dims);

		solver.add_regularization_group({Dx,Dy,Dz});

	}

	 */

	if (dct_weight > 0){
		//auto dctOp = boost::make_shared<identityOperator<cuNDArray<float>>>();
		auto dctOp = boost::make_shared<cuDCTOperator<float>>();
		dctOp->set_domain_dimensions(&is_dims);
		//dctOp->set_codomain_dimensions(&is_dims);
		dctOp->set_weight(dct_weight);

		auto dctOp1 = boost::make_shared<subselectionOperator<cuNDArray<float>>>(dctOp,1);
		dctOp1->set_domain_dimensions(&double_dims);
		dctOp1->set_codomain_dimensions(dctOp->get_codomain_dimensions().get());
		dctOp1->set_weight(dct_weight);
		solver.add_regularization_operator(dctOp1);
		/*
		auto Dt = boost::make_shared<cuPartialDerivativeOperator<float,4>>(3);
		Dt->set_weight(dct_weight);
		Dt->set_domain_dimensions(&is_dims);
		Dt->set_codomain_dimensions(&is_dims);
		auto Dt1 = boost::make_shared<subselectionOperator<cuNDArray<float>>>(Dt,1);
		Dt1->set_domain_dimensions(&double_dims);
		Dt1->set_codomain_dimensions(&is_dims);

		solver.add_regularization_operator(Dt1);
		 */
/*
		auto Dt = boost::make_shared<cuPartialDerivativeOperator<float,4>>(3);
		Dt->set_weight(tv_weight);
		Dt->set_domain_dimensions(&double_dims);
		Dt->set_codomain_dimensions(&double_dims);
		solver.add_regularization_operator(Dt);
		auto Dx = boost::make_shared<cuPartialDerivativeOperator<float,4>>(0);
		Dx->set_weight(tv_weight);

		Dx->set_domain_dimensions(&is_dims);
		Dx->set_codomain_dimensions(&is_dims);
		/*
  	Dx->set_domain_dimensions(&double_dims);
  	Dx->set_codomain_dimensions(&double_dims);
		 */
		/*
		auto Dy = boost::make_shared<cuPartialDerivativeOperator<float,4>>(1);
		Dy->set_weight(tv_weight);
		Dy->set_domain_dimensions(&is_dims);
		Dy->set_codomain_dimensions(&is_dims);
		*/
		/*
  	Dy->set_domain_dimensions(&double_dims);
  	Dy->set_codomain_dimensions(&double_dims);
		 */
/*

		auto Dz = boost::make_shared<cuPartialDerivativeOperator<float,4>>(2);
		Dz->set_weight(tv_weight);
		Dz->set_domain_dimensions(&is_dims);
		Dz->set_codomain_dimensions(&is_dims);
		*/
		/*
  	Dz->set_domain_dimensions(&double_dims);
  	Dz->set_codomain_dimensions(&double_dims);
		 */

/*
		auto Dx1 = boost::make_shared<subselectionOperator<cuNDArray<float>>>(Dx,1);
		Dx1->set_domain_dimensions(&double_dims);
		Dx1->set_codomain_dimensions(&is_dims);
		auto Dy1 = boost::make_shared<subselectionOperator<cuNDArray<float>>>(Dy,1);
		Dy1->set_domain_dimensions(&double_dims);
		Dy1->set_codomain_dimensions(&is_dims);
		auto Dz1 = boost::make_shared<subselectionOperator<cuNDArray<float>>>(Dz,1);
		Dz1->set_domain_dimensions(&double_dims);
		Dz1->set_codomain_dimensions(&is_dims);

		Dx1->set_weight(tv_weight*0.1);
		Dy1->set_weight(tv_weight*0.1);
		Dz1->set_weight(tv_weight*0.1);

*/
	}

	auto E = boost::make_shared<CBSubsetOperator<cuNDArray> >(subsets);

	//E->setup(ps,binning,imageDimensions);
	E->setup(ps,binning,imageDimensions);
	E->set_domain_dimensions(&is_dims);
	E->set_codomain_dimensions(ps->get_projections()->get_dimensions().get());

	solver.set_encoding_operator(E);

	auto projections = boost::make_shared<cuNDArray<float>>(*ps->get_projections());
	std::cout << "Projection norm:" << nrm2(projections.get()) << std::endl;
	//E->set_use_offset_correction(false);

	//boost::shared_ptr<cuNDArray<bool>> mask;
	//mask = E->calculate_mask(projections,0.03f);
	//ps->set_projections(boost::shared_ptr<hoCuNDArray<float>>()); //Clear projections from host memory.

//	E->offset_correct(projections.get());
	//E->set_mask(mask);
	std::cout << "Projection norm:" << nrm2(projections.get()) << std::endl;

	//solver.set_damping(1e-6);

	/*
    boost::shared_ptr<hoCuNDArray<float> > prior;

  if (vm.count("use_prior")) {
  	prior = calculate_prior(binning,ps,projections,is_dims,imageDimensions);
  	solver.set_x0(prior);
  }
	 */

	auto result = solver.solve(projections.get());
	std::cout << "Penguin" << nrm2(result.get()) << std::endl;

	std::cout << "Result sum " << asum(result.get()) << std::endl;

	//apply_mask(result.get(),mask.get());

	std::cout << "Result sum " << asum(result.get()) << std::endl;
	//saveNDArray2HDF5(result.get(),outputFile,imageDimensions,vector_td<float,3>(0),command_line_string.str(),iterations);

	if (wavelet_weight > 0){
		osMOMSolverF<cuNDArray<float>> solverF;
		solverF.set_max_iterations(iterations);
		solverF.set_x0(result);
		solverF.set_encoding_operator(E);

		auto wave = boost::make_shared<cuEdgeATrousOperator<float>>();

		wave->set_sigma(sigma);
		wave->set_domain_dimensions(&is_dims);
		if (binning->get_number_of_bins() == 1)
			wave->set_levels({2,2,2});
		else
			wave->set_levels({2,2,2,2});
		wave->set_weight(wavelet_weight);
		solverF.add_regularization_operator(wave);

		result = solverF.solve(projections.get());

	}

	auto result2 = sum(result.get(),4);
	saveNDArray2HDF5(result.get(),"seperate.hdf5",imageDimensions,floatd3(0,0,0),command_line_string.str(),iterations);
	saveNDArray2HDF5(result2.get(),outputFile,imageDimensions,floatd3(0,0,0),command_line_string.str(),iterations);
	//write_dicom(result.get(),command_line_string.str(),imageDimensions);
	/*
  cuNDArray<float> tmp(W->get_codomain_dimensions());

  W->mult_M(result.get(),&tmp);

  write_nd_array(&tmp,"test.real");
	 */
	//E->set_use_offset_correction(false);
	/*
  linearOperator<cuNDArray<float>> * E_all = E.get();
  auto tmp_proj = projections;
  clear(&tmp_proj);

  E_all->mult_M(result.get(),&tmp_proj);
  tmp_proj -= projections;
  write_nd_array(&tmp_proj,"projection_diffs.real");

  std::cout <<"Projection dimensions ";
  auto pdims = *projections.get_dimensions();
  for (auto d : pdims)
  	std::cout << d << " ";
  std::cout << std::endl;

	 */

}


/********************************************************************************
 * Copyright 2009 The Robotics Group, The Maersk Mc-Kinney Moller Institute,
 * Faculty of Engineering, University of Southern Denmark
 *
 * 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.
 ********************************************************************************/

#include "CubicSplineFactory.hpp"

#include <rw/core/Ptr.hpp>
#include <rw/math/EAA.hpp>
#include <rw/math/Q.hpp>
#include <rw/math/Quaternion.hpp>
#include <rw/math/Transform3D.hpp>
#include <rw/math/Vector2D.hpp>
#include <rw/math/Vector3D.hpp>
#include <rw/math/VectorND.hpp>
#include <rw/trajectory/CubicSplineInterpolator.hpp>
#include <rw/trajectory/SQUADInterpolator.hpp>

#include <Eigen/Sparse>
#include <math.h>

#if EIGEN_VERSION_AT_LEAST(3, 1, 0)
#include <Eigen/SparseCholesky>
#endif

using namespace rw::trajectory;

using namespace rw::math;
using namespace rw::core;

// ###########################################################################
// #                                Q Functions                              #
// ###########################################################################

InterpolatorTrajectory< Q >::Ptr CubicSplineFactory::makeNaturalSpline (QPath::Ptr qpath,
                                                                        double timeStep)
{
    std::vector< double > times;
    for (size_t i = 0; i < qpath->size (); i++) {
        times.push_back (i * timeStep);
    }
    return makeNaturalSpline (*qpath, times);
}

InterpolatorTrajectory< rw::math::Q >::Ptr
CubicSplineFactory::makeNaturalSpline (TimedQPath::Ptr tqpath)
{
    QPath path;
    std::vector< double > times;
    for (const TimedQ& tq : *tqpath) {
        path.push_back (tq.getValue ());
        times.push_back (tq.getTime ());
    }

    return makeNaturalSpline (path, times);
}

InterpolatorTrajectory< rw::math::Q >::Ptr
CubicSplineFactory::makeNaturalSpline (const QPath& qpath, const std::vector< double >& times)
{
    typedef float T;

    if (qpath.size () < 2)
        RW_THROW ("Path must be longer than 1!");
    if (qpath.size () != times.size ())
        RW_THROW ("Length of path and times need to be equal");

    size_t dim = (qpath)[0].size ();    // the number of dimensions of the points
    size_t N   = qpath.size () - 1;     // we have N+1 points, which yields N splines

    typedef Eigen::Matrix< T, Eigen::Dynamic, 1 > Vector;
    // typedef Eigen::Matrix<T, Eigen::Dynamic, 1, 1> Matrix;

    Vector B (N + 1);    // make room for boundary conditions

    Vector Y (N + 1);    // the points that the spline should intersect

    Vector a (dim * (N + 1)), b (dim * N), c (dim * N), d (dim * N);

    Vector H (N);    // duration from point i to i+1
    for (size_t i = 0; i < N; i++) {
        // T timeI0 = (T)((*tqpath)[i]).getTime();
        // T timeI1 = (T)((*tqpath)[i+1]).getTime();
        // H[i] = timeI1-timeI0;
        H[i] = (float) (times[i + 1] - times[i]);
    }

#if EIGEN_VERSION_AT_LEAST(3, 1, 0)
    Eigen::SparseMatrix< T > A ((int) N + 1, (int) N + 1);
    // D[0] = 2*H[0];
    A.insert (0, 0) = 2 * H[0];
    A.insert (0, 1) = H[0];
    A.insert (1, 0) = H[0];
    for (size_t i = 1; i < N; i++) {
        // D[i] = 2*(H[i-1]+H[i]);
        int ei                = (int) i;
        A.insert (ei, ei)     = 2 * (H[i - 1] + H[i]);
        A.insert (ei, ei + 1) = H[i];
        A.insert (ei + 1, ei) = H[i];
    }
    // D[N] = 2*H[N-1];
    A.insert ((int) N, (int) N) = 2 * H[N - 1];

    Eigen::SimplicialLLT< Eigen::SparseMatrix< T > > solver;
    solver.compute (A);

#else
    Eigen::Matrix< T, Eigen::Dynamic, Eigen::Dynamic > A =
        Eigen::Matrix< T, Eigen::Dynamic, Eigen::Dynamic >::Zero ((int) N + 1, (int) N + 1);
    A (0, 0) = 2 * H[0];
    A (0, 1) = H[0];
    A (1, 0) = H[0];
    for (size_t i = 1; i < N; i++) {
        // D[i] = 2*(H[i-1]+H[i]);
        int ei         = (int) i;
        A (ei, ei)     = 2 * (H[i - 1] + H[i]);
        A (ei, ei + 1) = H[i];
        A (ei + 1, ei) = H[i];
    }
    // D[N] = 2*H[N-1];
    A ((int) N, (int) N) = 2 * H[N - 1];
    Eigen::LLT< Eigen::Matrix< T, Eigen::Dynamic, Eigen::Dynamic > > solver;
    solver.compute (A);

#endif
    for (size_t j = 0; j < (size_t) dim; j++) {
        for (size_t i = 0; i < (size_t) N + 1; i++) {
            Y[i] = (T) (qpath[i])[j];
        }

        B[0] = (T) (3.0 * (Y[1] - Y[0]) / H[0]);
        for (size_t i = 1; i < N; i++) {
            B[i] = (T) (3.0 * ((Y[i + 1] - Y[i]) / H[i] - (Y[i] - Y[i - 1]) / H[i - 1]));
        }
        B[N] = (T) (-3.0 * (Y[N] - Y[N - 1]) / H[N - 1]);

        B = solver.solve (B);

        for (size_t i = 0; i < (size_t) N + 1; i++) {
            a[i * dim + j] = Y[i];
        }

        for (size_t i = 0; i < (size_t) N; i++) {
            c[j + i * dim] = B[i];
            b[j + i * dim] = (Y[i + 1] - Y[i]) / H[i] - H[i] * (B[i + 1] + 2 * B[i]) / (T) 3.0;
            d[j + i * dim] = (B[i + 1] - B[i]) / ((T) 3.0 * H[i]);    //   +B[i]+B[i+1];
        }
    }

    // ************** now create the actual trajectory from the calcualted parameters
    InterpolatorTrajectory< Q >::Ptr traj = ownedPtr (new InterpolatorTrajectory< Q > (times[0]));

    Q ba (dim), bb (dim), bc (dim), bd (dim);
    for (size_t i = 0; i < N; i++) {
        for (size_t j = 0; j < dim; j++) {
            ba[j] = a[j + i * dim];
            bb[j] = b[j + i * dim];
            bc[j] = c[j + i * dim];
            bd[j] = d[j + i * dim];
        }
        Interpolator< Q >* iptr = new CubicSplineInterpolator< Q > (ba, bb, bc, bd, H[i]);
        traj->add (ownedPtr (iptr));
    }

    return traj;
}

InterpolatorTrajectory< Q >::Ptr CubicSplineFactory::makeClampedSpline (QPath::Ptr qpath,
                                                                        const rw::math::Q& dqStart,
                                                                        const rw::math::Q& dqEnd,
                                                                        double timeStep)
{
    std::vector< double > times;
    for (size_t i = 0; i < qpath->size (); i++) {
        times.push_back (i * timeStep);
    }
    return makeClampedSpline (*qpath, times, dqStart, dqEnd);
}

InterpolatorTrajectory< rw::math::Q >::Ptr
CubicSplineFactory::makeClampedSpline (TimedQPath::Ptr tqpath, const rw::math::Q& dqStart,
                                       const rw::math::Q& dqEnd)

{
    QPath path;
    std::vector< double > times;
    for (const TimedQ& tq : *tqpath) {
        path.push_back (tq.getValue ());
        times.push_back (tq.getTime ());
    }

    return makeClampedSpline (path, times, dqStart, dqEnd);
}

// ###########################################################################
// #                             Template Functions                          #
// ###########################################################################

template< typename T >
typename InterpolatorTrajectory< T >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< T >& path, double timeStep)
{
    std::vector< double > times;
    for (size_t i = 0; i < path.size (); i++) {
        times.push_back (i * timeStep);
    }
    return makeNaturalSpline (path, times);
}

// ### explicit template specifications
template InterpolatorTrajectory< Vector3D<> >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Vector3D<> >& path, double timeStep);

template InterpolatorTrajectory< Transform3DVector<> >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Transform3DVector<> >& path, double timeStep);

template InterpolatorTrajectory< Q >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Q >& path, double timeStep);

template InterpolatorTrajectory< Quaternion< double > >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Quaternion< double > >& path, double timeStep);

// ########### NEXT FUNCTION ###############

template< typename T >
typename InterpolatorTrajectory< T >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Timed< T > >& path)
{
    Path< T > pathN;
    std::vector< double > times;
    for (const Timed< T >& tq : path) {
        pathN.push_back (tq.getValue ());
        times.push_back (tq.getTime ());
    }

    return makeNaturalSpline (pathN, times);
}

// ### explicit template specifications
template InterpolatorTrajectory< Vector3D<> >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Timed< Vector3D<> > >& path);

template InterpolatorTrajectory< Transform3DVector<> >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Timed< Transform3DVector<> > >& path);

template InterpolatorTrajectory< Q >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Timed< Q > >& path);

template InterpolatorTrajectory< Quaternion< double > >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Timed< Quaternion< double > > >& path);

// ########### NEXT FUNCTION ###############

template< typename T >
typename InterpolatorTrajectory< T >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< T >& path, const std::vector< double >& times)
{
    if (path.size () < 2) {
        RW_THROW ("Path must be longer than 1!");
    }
    if (path.size () != times.size ()) {
        RW_THROW ("Length of path and times need to be equal");
    }

    size_t dim = path[0].size ();     // the number of dimensions of the points
    size_t N   = path.size () - 1;    // we have N+1 points, which yields N splines
    Eigen::VectorXd B (N + 1);        // make room for boundary conditions

    Eigen::VectorXd Y (N + 1);    // the points that the spline should intersect

    Eigen::VectorXd a (dim * (N + 1));
    Eigen::VectorXd b (dim * N);
    Eigen::VectorXd c (dim * N);
    Eigen::VectorXd d (dim * N);

    Eigen::VectorXd H (N);    // duration from point i to i+1
    for (size_t i = 0; i < N; i++) {
        H[i] = (double) (times[i + 1] - times[i]);
    }
#if EIGEN_VERSION_AT_LEAST(3, 1, 0)
    Eigen::SparseMatrix< double > A ((int) N + 1, (int) N + 1);
    A.insert (0, 0) = 2 * H[0];
    A.insert (0, 1) = H[0];
    A.insert (1, 0) = H[0];
    for (size_t i = 1; i < N; i++) {
        int ei                = (int) i;
        A.insert (ei, ei)     = 2 * (H[i - 1] + H[i]);
        A.insert (ei, ei + 1) = H[i];
        A.insert (ei + 1, ei) = H[i];
    }
    // D[N] = 2*H[N-1];
    A.insert ((int) N, (int) N) = 2 * H[N - 1];

    Eigen::SimplicialLLT< Eigen::SparseMatrix< double > > solver;
    solver.compute (A);

#else
    Eigen::Matrix< double, Eigen::Dynamic, Eigen::Dynamic > A =
        Eigen::Matrix< double, Eigen::Dynamic, Eigen::Dynamic >::Zero ((int) N + 1, (int) N + 1);
    A (0, 0) = 2 * H[0];
    A (0, 1) = H[0];
    A (1, 0) = H[0];
    for (size_t i = 1; i < N; i++) {
        int ei         = (int) i;
        A (ei, ei)     = 2 * (H[i - 1] + H[i]);
        A (ei, ei + 1) = H[i];
        A (ei + 1, ei) = H[i];
    }
    A ((int) N, (int) N) = 2 * H[N - 1];
    Eigen::LLT< Eigen::Matrix< double, Eigen::Dynamic, Eigen::Dynamic > > solver;
    solver.compute (A);
#endif
    for (size_t j = 0; j < (size_t) dim; j++) {
        for (size_t i = 0; i < (size_t) N + 1; i++) {
            Y[i] = (double) (path[i])[j];
        }

        B[0] = (double) (3.0 * (Y[1] - Y[0]) / H[0]);
        for (size_t i = 1; i < N; i++) {
            B[i] = (double) (3.0 * ((Y[i + 1] - Y[i]) / H[i] - (Y[i] - Y[i - 1]) / H[i - 1]));
        }
        B[N] = (double) (-3.0 * (Y[N] - Y[N - 1]) / H[N - 1]);

        B = solver.solve (B);

        for (size_t i = 0; i < (size_t) N + 1; i++) {
            a[i * dim + j] = Y[i];
        }

        for (size_t i = 0; i < (size_t) N; i++) {
            c[j + i * dim] = B[i];
            b[j + i * dim] = (Y[i + 1] - Y[i]) / H[i] - H[i] * (B[i + 1] + 2 * B[i]) / (double) 3.0;
            d[j + i * dim] = (B[i + 1] - B[i]) / ((double) 3.0 * H[i]);    //   +B[i]+B[i+1];
        }
    }
    // ************** now create the actual trajectory from the calcualted parameters
    typename InterpolatorTrajectory< T >::Ptr traj =
        ownedPtr (new InterpolatorTrajectory< T > (times[0]));

    T ba, bb, bc, bd;
    for (size_t i = 0; i < N; i++) {
        for (size_t j = 0; j < dim; j++) {
            ba[j] = a[j + i * dim];
            bb[j] = b[j + i * dim];
            bc[j] = c[j + i * dim];
            bd[j] = d[j + i * dim];
        }
        Interpolator< T >* iptr = new CubicSplineInterpolator< T > (ba, bb, bc, bd, H[i]);
        traj->add (ownedPtr (iptr));
    }
    return traj;
}

// ### explicit template specifications
template InterpolatorTrajectory< Vector3D<> >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Vector3D<> >& path,
                                       const std::vector< double >& times);

template InterpolatorTrajectory< Transform3DVector<> >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Transform3DVector<> >& path,
                                       const std::vector< double >& times);

template InterpolatorTrajectory< Quaternion<> >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Quaternion<> >& path,
                                       const std::vector< double >& times);

// ########### NEXT FUNCTION ###############
template< typename T >
typename InterpolatorTrajectory< T >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< T >& path, const T& dStart, const T& dEnd,
                                       double timeStep)
{
    std::vector< double > times;
    for (size_t i = 0; i < path.size (); i++) {
        times.push_back (i * timeStep);
    }
    return makeClampedSpline (path, times, dStart, dEnd);
}
namespace rw { namespace trajectory {
template<>
typename InterpolatorTrajectory< Transform3D<> >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Transform3D<> >& path,
                                       const Transform3D<>& dStart, const Transform3D<>& dEnd,
                                       double timeStep)
{
    RW_THROW ("Clamped Cubic Spline not yet implemented for Transform3D");
    return NULL;
}
}}
template InterpolatorTrajectory< Vector3D<> >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Vector3D<> >& path, const Vector3D<>& dStart,
                                       const Vector3D<>& dEnd, double timeStep);

template InterpolatorTrajectory< Q >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Q >& path, const Q& dStart, const Q& dEnd,
                                       double timeStep);

template InterpolatorTrajectory< Transform3DVector<> >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Transform3DVector<> >& path,
                                       const Transform3DVector<>& dStart,
                                       const Transform3DVector<>& dEnd, double timeStep);

template InterpolatorTrajectory< Quaternion<> >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Quaternion<> >& path, const Quaternion<>& dStart,
                                       const Quaternion<>& dEnd, double timeStep);

// ########### NEXT FUNCTION ###############

template< typename T >
typename InterpolatorTrajectory< T >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Timed< T > >& tpath, const T& dStart,
                                       const T& dEnd)
{
    Path< T > path;
    std::vector< double > times;
    for (const Timed< T >& tq : tpath) {
        path.push_back (tq.getValue ());
        times.push_back (tq.getTime ());
    }

    return makeClampedSpline (path, times, dStart, dEnd);
}

template InterpolatorTrajectory< Vector3D<> >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Timed< Vector3D<> > >& tpath,
                                       const Vector3D<>& dStart, const Vector3D<>& dEnd);

template InterpolatorTrajectory< Q >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Timed< Q > >& path, const Q& dStart,
                                       const Q& dEnd);

template InterpolatorTrajectory< Transform3DVector<> >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Timed< Transform3DVector<> > >& tpath,
                                       const Transform3DVector<>& dStart,
                                       const Transform3DVector<>& dEnd);

template InterpolatorTrajectory< Quaternion<> >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Timed< Quaternion<> > >& tpath,
                                       const Quaternion<>& dStart, const Quaternion<>& dEnd);

// ########### NEXT FUNCTION ###############
template< typename T >
typename InterpolatorTrajectory< T >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< T >& path, const std::vector< double >& times,
                                       const T& dStart, const T& dEnd)
{
    if (path.size () < 2)
        RW_THROW ("Path must be longer than 1!");

    if (path.size () != times.size ())
        RW_THROW ("Length of path and times need to match");

    size_t dim = path[0].size ();     // the number of dimensions of the points
    size_t N   = path.size () - 1;    // we have N+1 points, which yields N splines
    Eigen::VectorXd B (N + 1);        // make room for boundary conditions

    Eigen::VectorXd Y (N + 1);    // the points that the spline should intersect

    Eigen::VectorXd a (dim * (N + 1));
    Eigen::VectorXd b (dim * N);
    Eigen::VectorXd c (dim * N);
    Eigen::VectorXd d (dim * N);
    Eigen::VectorXd H (N);    // duration from point i to i+1

    for (size_t i = 0; i < N; i++) {
        H[i] = (float) (times[i + 1] - times[i]);
    }
#if EIGEN_VERSION_AT_LEAST(3, 1, 0)
    Eigen::SparseMatrix< double > A ((int) N + 1, (int) N + 1);

    A.insert (0, 0) = 2 * H[0];
    A.insert (0, 1) = H[0];
    A.insert (1, 0) = H[0];
    for (size_t i = 1; i < N; i++) {
        int ei                = (int) i;
        A.insert (ei, ei)     = 2 * (H[i - 1] + H[i]);
        A.insert (ei + 1, ei) = H[i];
        A.insert (ei, ei + 1) = H[i];
    }
    A.insert ((int) N, (int) N) = 2 * H[N - 1];

    Eigen::SimplicialLLT< Eigen::SparseMatrix< double > > solver;
    solver.compute (A);

#else
    Eigen::Matrix< double, Eigen::Dynamic, Eigen::Dynamic > A =
        Eigen::Matrix< double, Eigen::Dynamic, Eigen::Dynamic >::Zero ((int) N + 1, (int) N + 1);

    A (0, 0) = 2 * H[0];
    A (0, 1) = H[0];
    A (1, 0) = H[0];
    for (size_t i = 1; i < N; i++) {
        int ei         = (int) i;
        A (ei, ei)     = 2 * (H[i - 1] + H[i]);
        A (ei + 1, ei) = H[i];
        A (ei, ei + 1) = H[i];
    }
    A ((int) N, (int) N) = 2 * H[N - 1];

    Eigen::LLT< Eigen::Matrix< double, Eigen::Dynamic, Eigen::Dynamic > > solver;
    solver.compute (A);

#endif

    for (size_t j = 0; j < (size_t) dim; j++) {
        for (size_t i = 0; i < (size_t) N + 1; i++) {
            Y[i] = (double) (path[i])[j];
        }

        B[0] = (double) (3.0 * (Y[1] - Y[0]) / H[0] - 3 * dStart[j]);
        for (size_t i = 1; i < (std::size_t) B.size () - 1; i++) {
            B[i] = (double) (3.0 * ((Y[i + 1] - Y[i]) / H[i] - (Y[i] - Y[i - 1]) / H[i - 1]));
        }
        B[N] = (double) (3.0 * dEnd[j] - 3.0 * (Y[N] - Y[N - 1]) / H[N - 1]);

        B = solver.solve (B);

        for (size_t i = 0; i < N + 1; i++) {
            a[i * dim + j] = Y[i];
        }

        for (size_t i = 0; i < (size_t) N; i++) {
            c[j + i * dim] = B[i];
            b[j + i * dim] =
                (double) ((Y[i + 1] - Y[i]) / H[i] - H[i] * (B[i + 1] + 2 * B[i]) / 3.0);
            d[j + i * dim] = (double) ((B[i + 1] - B[i]) / (3.0 * H[i]));    //   +B[i]+B[i+1];
        }
    }

    // ************** now create the actual trajectory from the calcualted parameters
    typename InterpolatorTrajectory< T >::Ptr traj =
        ownedPtr (new InterpolatorTrajectory< T > (times[0]));

    T ba=path[0], bb=path[0], bc=path[0], bd=path[0];
    for (size_t i = 0; i < N; i++) {

        for (size_t j = 0; j < dim; j++) {
            ba[j] = a[j + i * dim];
            bb[j] = b[j + i * dim];
            bc[j] = c[j + i * dim];
            bd[j] = d[j + i * dim];
        }
        Interpolator< T >* iptr = new CubicSplineInterpolator< T > (ba, bb, bc, bd, H[i]);
        traj->add (ownedPtr (iptr));
    }

    return traj;
}

template InterpolatorTrajectory< Vector3D<> >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Vector3D<> >& path,
                                       const std::vector< double >& times, const Vector3D<>& dStart,
                                       const Vector3D<>& dEnd);

template InterpolatorTrajectory< Q >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Q >& path,
                                       const std::vector< double >& times, const Q& dStart,
                                       const Q& dEnd);

template InterpolatorTrajectory< Transform3DVector<> >::Ptr CubicSplineFactory::makeClampedSpline (
    const Path< Transform3DVector<> >& path, const std::vector< double >& times,
    const Transform3DVector<>& dStart, const Transform3DVector<>& dEnd);

template InterpolatorTrajectory< Quaternion<> >::Ptr
CubicSplineFactory::makeClampedSpline (const Path< Quaternion<> >& path,
                                       const std::vector< double >& times,
                                       const Quaternion<>& dStart, const Quaternion<>& dEnd);

// ###########################################################################
// #                          Transform3D Functions                          #
// ###########################################################################

InterpolatorTrajectory< rw::math::Transform3DVector<> >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Transform3D<> >& path, double timeStep)
{
    std::vector< double > times;
    for (size_t i = 0; i < path.size (); i++) {
        times.push_back (i * timeStep);
    }
    return makeNaturalSpline (path, times);
}

InterpolatorTrajectory< rw::math::Transform3DVector<> >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< Timed< Transform3D<> > >& path)
{
    Path< Transform3D<> > pathN;
    std::vector< double > times;
    for (const Timed< Transform3D<> >& tq : path) {
        pathN.push_back (tq.getValue ());
        times.push_back (tq.getTime ());
    }

    return makeNaturalSpline (pathN, times);
}

InterpolatorTrajectory< rw::math::Transform3DVector<> >::Ptr
CubicSplineFactory::makeNaturalSpline (const Path< rw::math::Transform3D<> >& path,
                                       const std::vector< double >& times)
{
    Path< Transform3DVector<> > NPath;
    for (const Transform3D<>& trans : path) {
        NPath.push_back (Transform3DVector<> (trans));
    }
    return makeNaturalSpline (NPath, times);
}

// ###########################################################################
// #                             SQUAD Functions                             #
// ###########################################################################

InterpolatorTrajectory< rw::math::Quaternion<> >::Ptr
CubicSplineFactory::makeSQUAD (const Path< rw::math::Quaternion<> >& path, double timeStep)
{
    std::vector< double > times;
    for (size_t i = 0; i < path.size (); i++) {
        times.push_back (i * timeStep);
    }
    return makeSQUAD (path, times);
}
InterpolatorTrajectory< rw::math::Quaternion<> >::Ptr
CubicSplineFactory::makeSQUAD (const Path< Timed< rw::math::Quaternion<> > >& tpath)
{
    Path< rw::math::Quaternion<> > pathN;
    std::vector< double > times;
    for (const Timed< rw::math::Quaternion<> >& tq : tpath) {
        pathN.push_back (tq.getValue ());
        times.push_back (tq.getTime ());
    }
    return makeSQUAD (pathN, times);
}

InterpolatorTrajectory< rw::math::Quaternion<> >::Ptr
CubicSplineFactory::makeSQUAD (const Path< rw::math::Quaternion<> >& path,
                               const std::vector< double >& times)
{
    InterpolatorTrajectory< Quaternion<> >::Ptr traj =
        ownedPtr (new InterpolatorTrajectory< Quaternion<> > (times[0]));
    size_t N = path.size () - 1;
    std::vector< Quaternion<> > s (path.size ());
    s[0] = path[0];
    for (size_t i = 1; i < s.size () - 1; i++) {
        const Quaternion<>& qn1 = path[i - 1];
        const Quaternion<>& q0  = path[i];
        const Quaternion<>& q1  = path[i + 1];

        s[i] = q0 * exp (-1 * (ln (q0.inverse () * q1) + ln (q0.inverse () * qn1)) * (1.0 / 4.0));
    }
    s.back () = path.back ();

    for (size_t i = 0; i < N - 1; i++) {
        double duration = times[i + 1] - times[i];
        Interpolator< Quaternion<> >* iptr =
            new SQUADInterpolator< double > (path[i], path[i + 1], s[i], s[i + 1], duration);
        traj->add (ownedPtr (iptr));
    }
    return traj;
}


#ifndef QST_RBM_HPP
#define QST_RBM_HPP
#include <iostream>
#include <Eigen/Dense>
#include <random>
#include <fstream>

namespace qst{

// RBM class
class Rbm{

    int nv_;                        // Number of visible units
    int nh_;                        // Number of hidden units
    int npar_;                      // Number of parameters
    int nchains_;                   // Number of sampling chains

    Eigen::MatrixXd v_;             // Visible states
    Eigen::MatrixXd h_;             // Hidden states
    Eigen::MatrixXd probv_given_h_; // Visible probabilities
    Eigen::MatrixXd probh_given_v_; // Hidden probabilities

    Eigen::MatrixXd W_;             // Weights
    Eigen::VectorXd b_;             // Visible fields
    Eigen::VectorXd c_;             // Hidden fields

    Eigen::VectorXd gamma_;         // Container for hidden contribution to quantities

    std::mt19937 rgen_;             // Random number generator

public:
    // Contructor
    Rbm(Parameters &par):nv_(par.nv_),nh_(par.nh_){
        npar_=nv_+nh_+nv_*nh_;
        nchains_ = par.nc_;
        v_.setZero(nchains_,nv_);
        h_.setZero(nchains_,nh_);
        probv_given_h_.resize(nchains_,nv_);
        probh_given_v_.resize(nchains_,nh_);
        W_.resize(nh_,nv_);
        b_.resize(nv_);
        c_.resize(nh_);
        gamma_.resize(nh_);
        rgen_.seed(13579);
        //std::random_device rd;
        //rgen_.seed(rd());
    }

    // Private members access functions
    inline int Nvisible()const{
        return nv_;
    }
    inline int Nhidden()const{
        return nh_;
    }
    inline int Npar()const{
        return npar_;
    }
    inline int Nchains(){
        return nchains_;
    }
    inline Eigen::VectorXd VisibleStateRow(int s){
        return v_.row(s);
    }

    // Set the visible layer state
    inline void SetVisibleLayer(Eigen::MatrixXd v){
        v_=v;
    }

    // Initialize the network parameters
    void InitRandomPars(int seed,double sigma){
        std::default_random_engine generator(seed);
        std::normal_distribution<double> distribution(0,sigma);
        for(int i=0;i<nh_;i++){
            for(int j=0;j<nv_;j++){
                W_(i,j)=distribution(generator);
            }
        }
        //b_.setZero();
        //c_.setZero();
        for(int j=0;j<nv_;j++){
            b_(j)=distribution(generator);
        }
        for(int i=0;i<nh_;i++){
            c_(i)=distribution(generator);
        }
    }

    // Compute derivative of the effective visible energy
    Eigen::VectorXd VisEnergyGrad(const Eigen::VectorXd & v){
        Eigen::VectorXd der(npar_);
        int p=0;
        logistic(W_*v+c_,gamma_);
        for(int i=0;i<nh_;i++){
            for(int j=0;j<nv_;j++){
                der(p)=gamma_(i)*v(j);
                p++;
            }
        }
        for(int j=0;j<nv_;j++){
            der(p)=v(j);
            p++;
        }
        for(int i=0;i<nh_;i++){
            der(p)=gamma_(i);
            p++;
        }
        return -der;
    }

    // Return the probability for state v
    inline double prob(const Eigen::VectorXd & v){
        ln1pexp(W_*v+c_,gamma_);
        return std::exp(v.dot(b_)+gamma_.sum());
    }

    // Conditional Probabilities
    void ProbHiddenGivenVisible(const Eigen::MatrixXd &v,Eigen::MatrixXd &probs){
        logistic((v*W_.transpose()).rowwise() + c_.transpose(),probs);
    }
    void ProbVisibleGivenHidden(const Eigen::MatrixXd &h,Eigen::MatrixXd &probs){
        logistic((h*W_).rowwise() + b_.transpose(),probs);
    }

    // Sample one layer
    void SampleLayer(Eigen::MatrixXd & hv,const Eigen::MatrixXd & probs){
        std::uniform_real_distribution<double> distribution(0,1);
        for(int s=0;s<hv.rows();s++){
            for(int i=0;i<hv.cols();i++){
                hv(s,i)=distribution(rgen_)<probs(s,i);
            }
        }
    }

    // Perform k steps of Gibbs sampling
    void Sample(int steps){
        for(int k=0;k<steps;k++){
            ProbHiddenGivenVisible(v_,probh_given_v_);
            SampleLayer(h_,probh_given_v_);
            ProbVisibleGivenHidden(h_,probv_given_h_);
            SampleLayer(v_,probv_given_h_);
        }
    }

    // Get RBM parameters
    Eigen::VectorXd GetParameters(){
        Eigen::VectorXd pars(npar_);
        int p=0;
        for(int i=0;i<nh_;i++){
            for(int j=0;j<nv_;j++){
                pars(p)=W_(i,j);
                p++;
            }
        }
        for(int j=0;j<nv_;j++){
            pars(p)=b_(j);
            p++;
        }
        for(int i=0;i<nh_;i++){
            pars(p)=c_(i);
            p++;
        }
        return pars;
    }

    // Set RBM parameters
    void SetParameters(const Eigen::VectorXd & pars){
        int p=0;
        for(int i=0;i<nh_;i++){
            for(int j=0;j<nv_;j++){
                W_(i,j)=pars(p);
                p++;
            }
        }
        for(int j=0;j<nv_;j++){
            b_(j)=pars(p);
            p++;
        }
        for(int i=0;i<nh_;i++){
            c_(i)=pars(p);
            p++;
        }
    }

    // Read weights from file
    void LoadWeights(std::ifstream &fin){
        for(int i=0;i<nh_;i++){
            for(int j=0;j<nv_;j++){
                fin >> W_(i,j);
            }
        }
        for(int j=0;j<nv_;j++){
            fin >> b_(j);
        }
        for(int i=0;i<nh_;i++){
            fin >> c_(i);
        }
    }

    // Read weights from file
    void SaveWeights(std::ofstream &fout){
        for(int i=0;i<nh_;i++){
            for(int j=0;j<nv_;j++){
                fout << W_(i,j) << " ";
            }
            fout << std::endl;
        }
        fout << std::endl;
        for(int j=0;j<nv_;j++){
            fout <<  b_(j) << " ";
        }
        fout << std::endl <<std::endl;
        for(int i=0;i<nh_;i++){
            fout << c_(i) << " ";
        }
        fout << std::endl <<std::endl;
    }

    // Functions
    inline void logistic(const Eigen::VectorXd & x,Eigen::VectorXd & y){
        for(int i=0;i<x.size();i++){
            y(i)=1./(1.+std::exp(-x(i)));;
        }
    }
    inline void logistic(const Eigen::MatrixXd & x,Eigen::MatrixXd & y){
        for(int i=0;i<x.rows();i++){
            for(int j=0;j<x.cols();j++){
                //y(i,j)=logistic(x(i,j));
                y(i,j)=1./(1.+std::exp(-x(i,j)));
            }
        }
    }
    inline double ln1pexp(double x){
        if(x>30){
            return x;
        }
        return std::log1p(std::exp(x));
    }
    void ln1pexp(const Eigen::VectorXd & x,Eigen::VectorXd & y){
        for(int i=0;i<x.size();i++){
            y(i)=ln1pexp(x(i));
        }
    }

 };

}

#endif


//[ TArray
///////////////////////////////////////////////////////////////////////////////
//  Copyright 2008 Eric Niebler. 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)
//
// This example constructs a mini-library for linear algebra, using
// expression templates to eliminate the need for temporaries when
// adding arrays of numbers. It duplicates the TArray example from
// PETE (http://www.codesourcery.com/pooma/download.html)

#include <iostream>
#include <boost/mpl/int.hpp>
#include <boost/proto/core.hpp>
#include <boost/proto/context.hpp>
namespace mpl = boost::mpl;
namespace proto = boost::proto;
using proto::_;

// This grammar describes which TArray expressions
// are allowed; namely, int and array terminals
// plus, minus, multiplies and divides of TArray expressions.
struct TArrayGrammar
  : proto::or_<
        proto::terminal< int >
      , proto::terminal< int[3] >
      , proto::plus< TArrayGrammar, TArrayGrammar >
      , proto::minus< TArrayGrammar, TArrayGrammar >
      , proto::multiplies< TArrayGrammar, TArrayGrammar >
      , proto::divides< TArrayGrammar, TArrayGrammar >
    >
{};

template<typename Expr>
struct TArrayExpr;

// Tell proto that in the TArrayDomain, all
// expressions should be wrapped in TArrayExpr<> and
// must conform to the TArrayGrammar
struct TArrayDomain
  : proto::domain<proto::generator<TArrayExpr>, TArrayGrammar>
{};

// Here is an evaluation context that indexes into a TArray
// expression, and combines the result.
struct TArraySubscriptCtx
  : proto::callable_context< TArraySubscriptCtx const >
{
    typedef int result_type;

    TArraySubscriptCtx(std::ptrdiff_t i)
      : i_(i)
    {}

    // Index array terminals with our subscript. Everything
    // else will be handled by the default evaluation context.
    int operator ()(proto::tag::terminal, int const (&data)[3]) const
    {
        return data[this->i_];
    }

    std::ptrdiff_t i_;
};

// Here is an evaluation context that prints a TArray expression.
struct TArrayPrintCtx
  : proto::callable_context< TArrayPrintCtx const >
{
    typedef std::ostream &result_type;

    TArrayPrintCtx() {}

    std::ostream &operator ()(proto::tag::terminal, int i) const
    {
        return std::cout << i;
    }

    std::ostream &operator ()(proto::tag::terminal, int const (&arr)[3]) const
    {
        return std::cout << '{' << arr[0] << ", " << arr[1] << ", " << arr[2] << '}';
    }

    template<typename L, typename R>
    std::ostream &operator ()(proto::tag::plus, L const &l, R const &r) const
    {
        return std::cout << '(' << l << " + " << r << ')';
    }

    template<typename L, typename R>
    std::ostream &operator ()(proto::tag::minus, L const &l, R const &r) const
    {
        return std::cout << '(' << l << " - " << r << ')';
    }

    template<typename L, typename R>
    std::ostream &operator ()(proto::tag::multiplies, L const &l, R const &r) const
    {
        return std::cout << l << " * " << r;
    }

    template<typename L, typename R>
    std::ostream &operator ()(proto::tag::divides, L const &l, R const &r) const
    {
        return std::cout << l << " / " << r;
    }
};

// Here is the domain-specific expression wrapper, which overrides
// operator [] to evaluate the expression using the TArraySubscriptCtx.
template<typename Expr>
struct TArrayExpr
  : proto::extends<Expr, TArrayExpr<Expr>, TArrayDomain>
{
    typedef proto::extends<Expr, TArrayExpr<Expr>, TArrayDomain> base_type;

    TArrayExpr( Expr const & expr = Expr() )
      : base_type( expr )
    {}

    // Use the TArraySubscriptCtx to implement subscripting
    // of a TArray expression tree.
    int operator []( std::ptrdiff_t i ) const
    {
        TArraySubscriptCtx const ctx(i);
        return proto::eval(*this, ctx);
    }

    // Use the TArrayPrintCtx to display a TArray expression tree.
    friend std::ostream &operator <<(std::ostream &sout, TArrayExpr<Expr> const &expr)
    {
        TArrayPrintCtx const ctx;
        return proto::eval(expr, ctx);
    }
};

// Here is our TArray terminal, implemented in terms of TArrayExpr
// It is basically just an array of 3 integers.
struct TArray
  : TArrayExpr< proto::terminal< int[3] >::type >
{
    explicit TArray( int i = 0, int j = 0, int k = 0 )
    {
        (*this)[0] = i;
        (*this)[1] = j;
        (*this)[2] = k;
    }

    // Here we override operator [] to give read/write access to
    // the elements of the array. (We could use the TArrayExpr
    // operator [] if we made the subscript context smarter about
    // returning non-const reference when appropriate.)
    int &operator [](std::ptrdiff_t i)
    {
        return proto::value(*this)[i];
    }

    int const &operator [](std::ptrdiff_t i) const
    {
        return proto::value(*this)[i];
    }

    // Here we define a operator = for TArray terminals that
    // takes a TArray expression.
    template< typename Expr >
    TArray &operator =(Expr const & expr)
    {
        // proto::as_expr<TArrayDomain>(expr) is the same as
        // expr unless expr is an integer, in which case it
        // is made into a TArrayExpr terminal first.
        return this->assign(proto::as_expr<TArrayDomain>(expr));
    }

    template< typename Expr >
    TArray &printAssign(Expr const & expr)
    {
        *this = expr;
        std::cout << *this << " = " << expr << std::endl;
        return *this;
    }

private:
    template< typename Expr >
    TArray &assign(Expr const & expr)
    {
        // expr[i] here uses TArraySubscriptCtx under the covers.
        (*this)[0] = expr[0];
        (*this)[1] = expr[1];
        (*this)[2] = expr[2];
        return *this;
    }
};

int main()
{
    TArray a(3,1,2);

    TArray b;

    std::cout << a << std::endl;
    std::cout << b << std::endl;

    b[0] = 7; b[1] = 33; b[2] = -99;

    TArray c(a);

    std::cout << c << std::endl;

    a = 0;

    std::cout << a << std::endl;
    std::cout << b << std::endl;
    std::cout << c << std::endl;

    a = b + c;

    std::cout << a << std::endl;

    a.printAssign(b+c*(b + 3*c));

    return 0;
}
//]


/*!
 * @file demo.cpp
 * @author  (Aleksander Rybin)
 * @brief Threadpool example demo program for parallel matrix x vector multiply
 * @date 2021-09-15
 */

// ---------------------
// includes: STL
#include <algorithm>
#include <functional>
#include <iomanip>
#include <iostream>
#include <random>
#include <vector>

// ---------------------
// includes: boost
#include <boost/program_options.hpp>

// ---------------------
// includes: local
#include "threadpool_example/threadpool.hpp"

std::vector<int> generate(std::size_t size, std::function<int()> gen);

void printMatrix(const std::vector<int>& matrix, std::size_t matrixRows, std::size_t matrixCols);
void printVector(const std::vector<int>& vector);

int main(int argc, char* argv[]) {
    namespace po = boost::program_options;

    po::options_description options("Allowed programm options");

    std::size_t threadsNum;
    bool        verboseOutput;

    std::size_t matrixRows;
    std::size_t matrixCols;

    // clang-format off
    options.add_options()
        ("help,h", "produce help message")
        ("threads_num,t", po::value<std::size_t>(&threadsNum)->default_value(std::thread::hardware_concurrency() - 1), "number of threads to run")
        ("verbose,v", po::value<bool>(&verboseOutput)->default_value(false), "verbose output")
        ("matrix_rows", po::value<std::size_t>(&matrixRows)->default_value(10000), "number of rows in multiplied matrix")
        ("matrix_cols", po::value<std::size_t>(&matrixCols)->default_value(10000), "number of cols in multiplied matrix");
    // clang-format on

    po::variables_map vm;

    try {
        po::store(po::parse_command_line(argc, argv, options), vm);
        po::notify(vm);
    } catch (const po::error& e) {
        std::cerr << e.what() << '\n';
        return EXIT_FAILURE;
    }

    if (vm.count("help")) {
        std::cout << options << std::endl;
        return EXIT_SUCCESS;
    }

    try {
        std::random_device                 randomDevice;
        std::mt19937                       mersenneEngine{randomDevice()};
        std::uniform_int_distribution<int> distribution{-9, 9};
        auto gen = [&distribution, &mersenneEngine] { return distribution(mersenneEngine); };

        std::vector<int> matrix = generate(matrixRows * matrixCols, gen);
        std::vector<int> vector = generate(matrixCols, gen);

        std::vector<int> result(matrixRows);

        if (verboseOutput) {
            printMatrix(matrix, matrixRows, matrixCols);
            printVector(vector);
        }

        try {
            threadpool_example::ThreadPool thp(threadsNum);

            std::size_t rowsPerThread = matrixRows / threadsNum;
            for (std::size_t i = 0; i < threadsNum; ++i) {
                thp.enqueue([=, &matrix, &vector, &result] {
                    for (std::size_t j = i * rowsPerThread; j < (i + 1) * rowsPerThread; ++j) {
                        int currentResult = 0;
                        for (std::size_t k = 0; k < matrixCols; ++k) {
                            currentResult += matrix[j * matrixCols + k] * vector[k];
                        }
                        result[j] = currentResult;
                    }
                });
            }
        } catch (const std::system_error& e) {
            std::cerr << "thread error" << std::endl;
            return EXIT_FAILURE;
        }

        if (verboseOutput) {
            printVector(result);
        }
    } catch (const std::bad_alloc& e) {
        std::cerr << "bad alock" << std::endl;
        return EXIT_FAILURE;
    } catch (const std::exception& e) {
        std::cerr << e.what() << std::endl;
        return EXIT_FAILURE;
    } catch (...) {
        std::cerr << "unknown error" << std::endl;
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}

std::vector<int> generate(std::size_t size, std::function<int()> gen) {
    std::vector<int> result(size);
    std::generate(result.begin(), result.end(), gen);
    return result;
}

void printMatrix(const std::vector<int>& matrix, std::size_t matrixRows, std::size_t matrixCols) {
    std::cout << "\n";
    for (std::size_t i = 0; i < matrixRows; ++i) {
        for (std::size_t j = 0; j < matrixCols; ++j) {
            std::cout << std::setw(2) << matrix[i * matrixCols + j] << ' ';
        }
        std::cout << "\n";
    }
    std::cout << std::endl;
}

void printVector(const std::vector<int>& vector) {
    std::cout << "\n";
    for (std::size_t i = 0; i < vector.size(); ++i) {
        std::cout << std::setw(2) << vector[i] << ' ';
    }
    std::cout << std::endl;
}