/*
Copyright (c) 2016 Bastien Durix

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

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

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

/**
 *  \brief 2D skeletonization
 *  \author Bastien Durix
 */

#include <boost/program_options.hpp>
#include <opencv2/core.hpp>
#include <opencv2/core/matx.hpp>
#include <opencv2/core/types.hpp>
#include <opencv2/imgproc.hpp>
#include <string>
#include <time.h>
#include <iostream>
#include <iomanip>
#include <chrono>
#include <stdlib.h>

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/imgcodecs/imgcodecs.hpp>

#include <shape/DiscreteShape.h>
#include <boundary/DiscreteBoundary.h>
#include <skeleton/Skeletons.h>

#include <algorithm/extractboundary/NaiveBoundary.h>
#include <algorithm/skeletonization/SpherePropagation2D.h>
#include <algorithm/skinning/Filling.h>
#include <algorithm/evaluation/ReprojError.h>

#include <displayopencv/DisplayShapeOCV.h>
#include <displayopencv/DisplayBoundaryOCV.h>
#include <displayopencv/DisplaySkeletonOCV.h>

#include <fileio/SkeletonFile.h>
#include <fileio/BoundaryFile2D.h>

std::tuple<double,double,int,int> EvalSkel(const shape::DiscreteShape<2>::Ptr dissh,
									   const boundary::DiscreteBoundary<2>::Ptr disbnd,
									   const skeleton::GraphSkel2d::Ptr skel)
{
	shape::DiscreteShape<2>::Ptr shp(new shape::DiscreteShape<2>(dissh->getWidth(),dissh->getHeight()));
	algorithm::skinning::Filling(shp,skel);

	double res = algorithm::evaluation::SymDiffArea(dissh,shp);
	double res2 = algorithm::evaluation::HausDist(skel,disbnd,dissh->getFrame());

	std::list<unsigned int> lnod;
	skel->getAllNodes(lnod);
	unsigned int nbbr = 0;
	for(std::list<unsigned int>::iterator it = lnod.begin(); it != lnod.end(); it++)
	{
		unsigned int deg = skel->getNodeDegree(*it);
		if(deg != 2)
			nbbr += deg;
	}
	nbbr /= 2;
	std::tuple<double,double,int,int> result = std::make_tuple(res*100.0,res2,skel->getNbNodes(),nbbr);

	return result;
}

int main(int argc, char** argv)
{
	std::string imgfile, fileimg;
    std::string nodfile, edgfile, delfile, bndfile;
	bool output = false;
	bool resize = false;
	int rows = 100;
	int cols = 100;
	double epsilon;

	boost::program_options::options_description desc("OPTIONS");

	desc.add_options()
		("help", "Help message")
		("imgfile", boost::program_options::value<std::string>(&imgfile)->default_value("mask.png"), "Binary image file (*.png)")
		("output", boost::program_options::value<bool>(&output)->implicit_value(true), "Returns output images")
		("epsilon", boost::program_options::value<double>(&epsilon)->default_value(1.0), "Skeleton precision")
		("fileimg", boost::program_options::value<std::string>(&fileimg)->default_value("skeleton.png"), "Skeleton img file")
		("resize", boost::program_options::value<bool>(&resize)->implicit_value(false), "Resize and center the image")
		("rows", boost::program_options::value<int>(&rows)->default_value(100), "Resized image rows")
		("cols", boost::program_options::value<int>(&cols)->default_value(100), "Resized image cols")
		("nodfile", boost::program_options::value<std::string>(&nodfile)->default_value("nod.txt"), "Skeleton out node file")
		("edgfile", boost::program_options::value<std::string>(&edgfile)->default_value("edg.txt"), "Skeleton out edge file")
		("delfile", boost::program_options::value<std::string>(&delfile)->default_value("del.txt"), "Delaunay triangles associated to each node")
		("bndfile", boost::program_options::value<std::string>(&bndfile)->default_value("bnd.txt"), "Out boundary file")
		;

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

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

	time_t start,end;
	double diff;

	cv::Mat shpimggray = cv::imread(imgfile,cv::ImreadModes::IMREAD_GRAYSCALE);
	cv::Mat shpimg;
	cv::threshold(shpimggray,shpimg,125,255,cv::THRESH_BINARY);

	if(resize)
	{
		srand(time(NULL));
		if(rows < shpimg.rows || cols < shpimg.cols)
		{
			std::cout << shpimg.rows << " " << shpimg.cols << std::endl;
			std::cout << "New size should be higher than current image" << std::endl;
			return -1;
		}

		cv::Mat resh_shpimg(rows, cols, shpimg.type(), cv::Scalar(0));

		int diffx = cols - shpimg.cols;
		int diffy = rows - shpimg.rows;

		cv::Rect rect(diffx/2, diffy/2, shpimg.cols, shpimg.rows);
		cv::Mat roi = resh_shpimg(rect);
		shpimg.copyTo(roi);

		resh_shpimg.copyTo(shpimg);
		cv::imwrite("resized.png", resh_shpimg);
	}

	// topological closure of the binary shape
	shape::DiscreteShape<2>::Ptr dissh = shape::DiscreteShape<2>::Ptr(new shape::DiscreteShape<2>(shpimg.cols,shpimg.rows));
	cv::Mat cpymat(shpimg.rows,shpimg.cols,CV_8U,&dissh->getContainer()[0]);
	shpimg.copyTo(cpymat);
	cv::Mat image(shpimg.rows,shpimg.cols,CV_8UC3,cv::Scalar(255,255,255));

	boundary::DiscreteBoundary<2>::Ptr disbnd = algorithm::extractboundary::NaiveBoundary(dissh);
    std::map<unsigned int, std::vector<unsigned int> > deldat;

	auto start0 = std::chrono::steady_clock::now();
	algorithm::skeletonization::propagation::OptionsSphProp options(2.0*epsilon);
	skeleton::GraphSkel2d::Ptr grskelpropag = algorithm::skeletonization::propagation::SpherePropagation2D(disbnd, deldat, options);
	auto duration0 = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - start0);
	std::tuple<double,double,int,int> respropag = EvalSkel(dissh,disbnd,grskelpropag);
	int t0 = duration0.count();
	double A0 = std::get<0>(respropag); // sym area diff
	double H0 = std::get<1>(respropag); // Hausdorff dist
	int N0 = std::get<2>(respropag); // nb nodes
	int B0 = std::get<3>(respropag); // nb branches

	//std::cout << "Skeleton estmated in " << t0 << "ms." << std::endl;
	//std::cout << "Symetric difference area over reference area: " << A0 << "\%" << std::endl;
	//std::cout << "Hausdorff distance to reference: " << H0 << "px (epsilon=" << epsilon << "px)" << std::endl;
	//std::cout << "Number of branches: " << B0 << std::endl;
	//std::cout << "Number of nodes: " << N0 << std::endl;

	shape::DiscreteShape<2>::Ptr shppropag(new shape::DiscreteShape<2>(dissh->getWidth(),dissh->getHeight()));
	algorithm::skinning::Filling(shppropag,grskelpropag);

	cv::Mat imagepropag(shpimg.rows,shpimg.cols,CV_8UC3,cv::Scalar(0,0,0));
	displayopencv::DisplayDiscreteShape(dissh,imagepropag,shppropag->getFrame(),cv::Scalar(255,0,0));
	// displayopencv::DisplayDiscreteShape(shppropag,imagepropag,shppropag->getFrame(),cv::Scalar(125,125,125));
	// displayopencv::DisplayDiscreteBoundary(disbnd,imagepropag,dissh->getFrame(),cv::Scalar(0,0,0));
	displayopencv::DisplayGraphSkeleton(grskelpropag,imagepropag,dissh->getFrame(),cv::Scalar(255,255,255));

	if(output)
	{
		cv::imwrite(fileimg, imagepropag);
        fileio::writeSkeleton(grskelpropag, deldat, nodfile, edgfile, delfile);
        fileio::WriteBoundary2D(disbnd, bndfile);
	}

	return 0;
}


// Copyright 2014, Max Planck Society.
// Distributed under the BSD 3-Clause license.
// (See accompanying file LICENSE.txt or copy at
// http://opensource.org/licenses/BSD-3-Clause)

//!@file
//! Mex wrapper file for Grassmann PCA

// this first include is a small workaround for clang503/xcode5.1: algorithm should be included prior to mex.h
#include <algorithm>

#include "mex.h"

#include <boost/numeric/ublas/storage.hpp>

#include <include/grassmann_pca.hpp>
#include <include/grassmann_pca_with_trimming.hpp>
#include <include/private/boost_ublas_external_storage.hpp>
#include <include/private/boost_ublas_row_iterator.hpp>

template <class output_type, class input_type>
output_type get_matlab_array_value_dispatch(mxArray const* array_, size_t index_row, size_t index_column)
{
  input_type* p_array = static_cast<input_type*>(mxGetData(array_));
  if(p_array == 0)
  {
    throw std::runtime_error("Unable to retrieve a pointer to the typed array");
  }

  // column major
  return p_array[index_row + mxGetN(array_) * index_column];
}

template <class output_type>
output_type get_matlab_array_value(mxArray const* array_, size_t index_row, size_t index_column)
{
  assert(index_row < mxGetM(array_));
  assert(index_column < mxGetN(array_));

  mxClassID classId = mxGetClassID(array_);
  switch(classId)
  {
  case mxDOUBLE_CLASS:
    return get_matlab_array_value_dispatch<output_type, double>(array_, index_row, index_column);
  case mxSINGLE_CLASS:
    return get_matlab_array_value_dispatch<output_type, float>(array_, index_row, index_column);
  default:
    throw std::runtime_error("Unable to dispatch to the correct type");
  }
}

struct s_algorithm_configuration
{
  size_t rows;
  size_t columns;

  size_t max_dimension;
  size_t max_iterations;
  size_t max_chunk_size;
  size_t nb_processors;
  size_t nb_pca_steps;
  double trimming_percentage;

  mxArray *initial_vectors;

  // we should also add the initial value of the basis vectors
};

template <class input_array_type>
bool grassmann_pca_dispatch(
  mxArray const* X,
  s_algorithm_configuration const& algorithm_configuration,
  mxArray *outputMatrix)
{
  namespace ub = boost::numeric::ublas;

  using namespace grassmann_averages_pca;
  using namespace grassmann_averages_pca::details::ublas_helpers;

  typedef external_storage_adaptor<input_array_type> input_storage_t;
  typedef ub::matrix<input_array_type, ub::column_major, input_storage_t> input_matrix_t;

  typedef external_storage_adaptor<input_array_type> output_storage_t;
  typedef ub::matrix<input_array_type, ub::row_major, output_storage_t> output_matrix_t; // this is in fact column_major, it should be in accordance with the
                                                                               // dimension of the matrix output_basis_vectors (we take the transpose of it)

  const size_t &dimension = algorithm_configuration.columns;
  const size_t &nb_elements = algorithm_configuration.rows;
  const size_t &max_dimension = algorithm_configuration.max_dimension;
  const size_t &max_iterations = algorithm_configuration.max_iterations;

  // input data matrix, external storage.
  input_storage_t input_storage(nb_elements * dimension, static_cast<input_array_type *>(mxGetData(X)));
  input_matrix_t input_data(nb_elements, dimension, input_storage);

  // output data matrix, also external storage for uBlas
  output_storage_t storageOutput(dimension * max_dimension, static_cast<input_array_type *>(mxGetData(outputMatrix)));
  output_matrix_t output_basis_vectors(max_dimension, dimension, storageOutput);

  size_t nb_pca_steps = algorithm_configuration.nb_pca_steps;

  // this is the form of the data extracted from the storage
  typedef ub::vector<input_array_type> data_t;
  typedef grassmann_pca< data_t > grassmann_pca_t;

  typedef row_iter<const input_matrix_t> const_input_row_iter_t;
  typedef row_iter<output_matrix_t> output_row_iter_t;

  // main instance
  grassmann_pca_t instance;

  if(algorithm_configuration.nb_processors > 0)
  {
    if(!instance.set_nb_processors(algorithm_configuration.nb_processors))
    {
      mexWarnMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Incorrect number of processors. Please consult the documentation.");
      return false;
    }
  }

  if(algorithm_configuration.max_chunk_size > 0)
  {
    if(!instance.set_max_chunk_size(algorithm_configuration.max_chunk_size))
    {
	    mexWarnMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Incorrect chunk size. Please consult the documentation.");
      return false;
    }
  }

  // initialisation vector if given
  std::vector<data_t> init_vectors;
  if(algorithm_configuration.initial_vectors != 0)
  {
    init_vectors.resize(max_dimension);
    input_storage_t input_init_vector_storage(max_dimension*dimension, static_cast<input_array_type*>(mxGetData(algorithm_configuration.initial_vectors)));
    input_matrix_t input_init_vector_data(dimension, max_dimension, input_init_vector_storage);
    for(size_t index = 0;
        index < max_dimension;
        index++)
    {
      init_vectors[index] = ub::column(input_init_vector_data, index);
    }

    // if the initial vectors are set, we avoid the computation of the regular PCA.
    nb_pca_steps = 0;
  }

  if(!instance.set_nb_steps_pca(nb_pca_steps))
  {
    mexWarnMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Incorrect number of regular PCA steps (%d). Please consult the documentation.", nb_pca_steps);
    return false;
  }

  return instance.batch_process(
    max_iterations,
    max_dimension,
    const_input_row_iter_t(input_data, 0),
    const_input_row_iter_t(input_data, input_data.size1()),
    output_row_iter_t(output_basis_vectors, 0),
    algorithm_configuration.initial_vectors ? &init_vectors: 0);

}

template <class input_array_type>
bool grassmann_pca_trimming_dispatch(
  mxArray const* X,
  s_algorithm_configuration const& algorithm_configuration,
  mxArray *outputMatrix)
{
  namespace ub = boost::numeric::ublas;

  using namespace grassmann_averages_pca;
  using namespace grassmann_averages_pca::details::ublas_helpers;

  typedef external_storage_adaptor<input_array_type> input_storage_t;
  typedef ub::matrix<input_array_type, ub::column_major, input_storage_t> input_matrix_t;

  typedef external_storage_adaptor<input_array_type> output_storage_t;
  typedef ub::matrix<input_array_type, ub::row_major, output_storage_t> output_matrix_t; // this is in fact column_major, it should be in accordance with the
                                                                               // dimension of the matrix output_basis_vectors (we take the transpose of it)

  const size_t &dimension = algorithm_configuration.columns;
  const size_t &nb_elements = algorithm_configuration.rows;
  const size_t &max_dimension = algorithm_configuration.max_dimension;
  const size_t &max_iterations = algorithm_configuration.max_iterations;

  // input data matrix, external storage.
  input_storage_t input_storage(nb_elements*dimension, static_cast<input_array_type*>(mxGetData(X)));
  input_matrix_t input_data(nb_elements, dimension, input_storage);

  // output data matrix, also external storage for uBlas
  output_storage_t storageOutput(dimension * max_dimension, static_cast<input_array_type *>(mxGetData(outputMatrix)));
  output_matrix_t output_basis_vectors(max_dimension, dimension, storageOutput);

  size_t nb_pca_steps = algorithm_configuration.nb_pca_steps;

  // this is the form of the data extracted from the storage
  typedef ub::vector<input_array_type> data_t;
  typedef grassmann_pca_with_trimming< data_t > grassmann_pca_with_trimming_t;

  typedef row_iter<const input_matrix_t> const_input_row_iter_t;
  typedef row_iter<output_matrix_t> output_row_iter_t;

  // main instance
  grassmann_pca_with_trimming_t instance(algorithm_configuration.trimming_percentage / 100);

  if(algorithm_configuration.nb_processors > 0)
  {
    if(!instance.set_nb_processors(algorithm_configuration.nb_processors))
    {
      mexWarnMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Incorrect number of processors. Please consult the documentation.");
      return false;
    }
  }

  if(algorithm_configuration.max_chunk_size > 0)
  {
    if(!instance.set_max_chunk_size(algorithm_configuration.max_chunk_size))
    {
	    mexWarnMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Incorrect chunk size. Please consult the documentation.");
      return false;
    }
  }

  // initialisation vector if given
  std::vector<data_t> init_vectors;
  if(algorithm_configuration.initial_vectors != 0)
  {
    init_vectors.resize(max_dimension);
    input_storage_t input_init_vector_storage(max_dimension*dimension, static_cast<input_array_type*>(mxGetData(algorithm_configuration.initial_vectors)));
    input_matrix_t input_init_vector_data(dimension, max_dimension, input_init_vector_storage);
    for(size_t index = 0;
        index < max_dimension;
        index++)
    {
      init_vectors[index] = ub::column(input_init_vector_data, index);
    }

    // if the initial vectors are set, we avoid the computation of the regular PCA.
    nb_pca_steps = 0;
  }

  if(!instance.set_nb_steps_pca(nb_pca_steps))
  {
    mexWarnMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Incorrect number of regular PCA steps (%d). Please consult the documentation.", nb_pca_steps);
    return false;
  }

  //std::string projectpath = "/Users/irita/Documents/Master/Research/Tracking/Moving Camera - code/Learn_Piecewise_Subspace/TGA-PCA/";
  return instance.batch_process(
    max_iterations,
    max_dimension,
    const_input_row_iter_t(input_data, 0),
    const_input_row_iter_t(input_data, input_data.size1()),
    output_row_iter_t(output_basis_vectors, 0),
    algorithm_configuration.initial_vectors ? &init_vectors: 0);

}

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{

  // arguments checking
  if (nrhs < 1 || nrhs > 4)
  {
	  mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Incorrect number of arguments. Please consult the documentation.");
  }

  const mxArray* const X = prhs[0];
  assert(X);

  // checking the format of the data
  if (!mxIsDouble(X) && !mxIsSingle(X))
  {
	  mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Unsupported input format (floating point required)");
  }

  if(mxIsComplex(X))
  {
    mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Unsupported format (scalar data required)");
  }

  s_algorithm_configuration config;

  // Check the dimensions of inputs
  config.rows = mxGetM(X);
  config.columns = mxGetN(X);

  size_t dimension = config.columns;
  config.initial_vectors  = 0;

  // third argument is the optional trimming percentage
  bool b_trimming = false;
  config.trimming_percentage = -1;
  if(nrhs >= 2)
  {
    const mxArray* const trimmingArray = prhs[1];
    if(!mxIsNumeric(trimmingArray))
    {
      mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Erroneous argument for the trimming percentage (non numeric argument)");
    }

    if(mxIsEmpty(trimmingArray))
    {
      mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Erroneous argument for the trimming percentage (empty value)");
    }

    if(mxGetNumberOfElements(trimmingArray) > 1)
    {
      mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Erroneous argument for the trimming percentage (non scalar)");
    }

    mxClassID classId = mxGetClassID(trimmingArray);
    if(classId == mxDOUBLE_CLASS || classId == mxSINGLE_CLASS)
    {
      //mexErrMsgTxt("Erroneous argument for the maximal dimension specification (floating point type)");
    }

    b_trimming = true;
    config.trimming_percentage = mxGetScalar(trimmingArray);

    if(config.trimming_percentage < 0 || config.trimming_percentage > 100)
    {
      mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Erroneous argument for the trimming percentage (not within the range [0, 100])");
    }

    b_trimming = config.trimming_percentage > 0;
  }

  config.max_iterations = config.rows; // by default the number of points
  config.max_chunk_size = std::numeric_limits<size_t>::max();
  config.nb_processors = 1;
  config.max_dimension = dimension;
  config.nb_pca_steps = 3;

  if(nrhs == 3)
  {
    const mxArray* const algorithmConfiguration = prhs[2];

    if(!mxIsStruct(algorithmConfiguration))
    {
      mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Erroneous argument for the algorithm configuration (not a structure)");
    }

    mxArray *nb_iteration_array = mxGetField(algorithmConfiguration, 0, "nb_iterations_max");
    if(nb_iteration_array != 0)
    {
      config.max_iterations = static_cast<int>(mxGetScalar(nb_iteration_array) + 0.5);
    }

    mxArray *max_chunk_size_array = mxGetField(algorithmConfiguration, 0, "max_chunk_size");
    if(max_chunk_size_array != 0)
    {
      config.max_chunk_size = static_cast<size_t>(mxGetScalar(max_chunk_size_array) + 0.5);
    }

    mxArray *nb_processing_threads_array = mxGetField(algorithmConfiguration, 0, "nb_processing_threads");
    if(nb_processing_threads_array != 0)
    {
      config.nb_processors = static_cast<size_t>(mxGetScalar(nb_processing_threads_array) + 0.5);
    }

    mxArray *nb_max_dimensions = mxGetField(algorithmConfiguration, 0, "max_dimensions");
    if(nb_max_dimensions != 0)
    {
      config.max_dimension = static_cast<size_t>(mxGetScalar(nb_max_dimensions) + 0.5);
    }

    config.initial_vectors = mxGetField(algorithmConfiguration, 0, "initial_vectors");
    if(config.initial_vectors != 0)
    {
      if (!mxIsDouble(config.initial_vectors) && !mxIsSingle(config.initial_vectors))
      {
        mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Unsupported input format for initial directions (floating point required)");
      }
      if(mxIsComplex(config.initial_vectors))
      {
        mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Unsupported format for initial directions (scalar data required)");
      }

      if(mxGetM(config.initial_vectors) != dimension)
      {
        mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Error in the dimension of the initial values");
      }

      if(mxGetN(config.initial_vectors) != config.max_dimension)
      {
        mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "Error in the number of the initial values provided. Should be equal to \"max_dimensions\"");
      }

    }

    mxArray *nb_pca_steps = mxGetField(algorithmConfiguration, 0, "nb_pca_steps");
    if(nb_pca_steps != 0)
    {
      config.nb_pca_steps = static_cast<size_t>(mxGetScalar(nb_pca_steps) + 0.5);
    }

  }

  plhs[0] = mxCreateNumericMatrix(dimension, config.max_dimension, mxGetClassID(X), mxREAL);
  mxArray *outputMatrix = plhs[0];
  assert(outputMatrix);

  bool result = false;
  switch(mxGetClassID(X))
  {
  case mxDOUBLE_CLASS:
  {
    if(!b_trimming)
    {
      result = grassmann_pca_dispatch<double>(X, config, outputMatrix);
    }
    else
    {
      result = grassmann_pca_trimming_dispatch<double>(X, config, outputMatrix);
    }

    break;
  }
  case mxSINGLE_CLASS:
  {
    if(!b_trimming)
    {
      result = grassmann_pca_dispatch<float>(X, config, outputMatrix);
    }
    else
    {
      result = grassmann_pca_trimming_dispatch<float>(X, config, outputMatrix);
    }

    break;
  }
  default:
    break;
  }

  if(!result)
  {
    mexErrMsgIdAndTxt("GrassmannAveragesPCA:configuration", "An error occurred in the call of the function.");
  }

}


// (C) Copyright David James Pearce 2003. Permission to copy, use,
// modify, sell and distribute this software is granted provided this
// copyright notice appears in all copies. This software is provided
// "as is" without express or implied warranty, and with no claim as
// to its suitability for any purpose.
//
// Email: d.pearce@doc.ic.ac.uk
//
// Notes
// =====
// This class implements an ordered linked list, based on the STL
// slist class.  It provides, I believe, an O (log N) complexity result.
//  The ordered asspect means that we can determine in constant
// time the relative ordering of two nodes in the list.  The data
// structure was first proposed by Dietz and Sleator [1] and later
// reformulated [2].
//
// This container has the rather interesting and unusual property
// that it is useful to instantiate it to contain void. This would
// indicate that the user is only interested in the structural
// information which the container provides. I.e. the relative
// ordering of nodes.
//
// There are potentially some portability issues as I rely on
// the use of uint32_t and uint64_t and I haven't tested the code
// on different architectures.
//
// Bibliography
// ============
// [1] Two algorithms for Maintaining order in a List, Paul F. Dietz
//     and Daniel D. Sleator, in Proc. 19th Annual ACM Symp. on Theory
//     of Computing (STOC'87), ACM, New York, 1987, pp. 365-372.
//
// [2] Two Simplified Algorithms for Maintaining Order in a List,
//     M. A. Bender, R. Cole, E. D. Demaine, M. F.-Colton and J. Zito,
//     In Proc. of the 10th Annual European Symposium on Algorithms,
//     Lecture Notes in Computer Science, 2002, pages 152-164.

#ifndef ORDERED_SLIST_HPP
#define ORDERED_SLIST_HPP

#include <cassert>
#include <functional>
#include <boost/iterator_adaptors.hpp>
#include <boost/iterator/transform_iterator.hpp>
#include <boost/cstdint.hpp>

// now a slight hack to get unsigned int limits
// on GCC < 3

#if __GNUC__ >= 3
#include <ext/slist>
#include <ext/functional>
// #include <limits>
// #define std::numeric_limits<unsigned int>::max() OL_UINT_MAX
#include <limits.h>
// for some reason UINT_MAX is 4294967295 ?
#define OL_Mm1_CONSTANT UINT_MAX
using __gnu_cxx::slist;
using __gnu_cxx::select2nd;
#else
#include <slist>
#include <limits.h>
#define OL_Mm1_CONSTANT UINT_MAX
using std::slist;
#endif

#ifdef OL_GENERATE_STATS
extern unsigned int ol_ncreated;
extern unsigned int ol_nrelabels;
#endif

class ol_dummy_class {
  // empty!
};

template<class T = ol_dummy_class>
class ordered_slist {
private:
  typedef slist<std::pair<uint32_t, T> > internal_list_t;

  // ==================
  // the basic iterator
  // ==================

  class _iterator {
  public:
    typedef typename internal_list_t::const_iterator::iterator_category iterator_category;
    typedef typename internal_list_t::const_iterator::difference_type difference_type;
    typedef T value_type;
    typedef T* pointer;
    typedef T& reference;
  private:
    typename internal_list_t::iterator _iter;
  public:
    _iterator(void) {}
    _iterator(typename internal_list_t::iterator const &iter) : _iter(iter) {
    }

    T &operator*() { return _iter->second; }
    T *operator->() { return &_iter->second; }

    _iterator const &operator++(void) {
      ++_iter; return *this;
    }
    _iterator operator++(int) {
      _iterator r(*this);++_iter;return r;
    }

    typename internal_list_t::iterator const &base(void) const {
      return _iter;
    }

    bool operator==(_iterator const &src) const {
      return _iter == src._iter;
    }

    bool operator!=(_iterator const &src) const {
      return _iter != src._iter;
    }
  };

  // ==================
  // the const iterator
  // ==================

  class _const_iterator {
  public:
    typedef typename internal_list_t::const_iterator::iterator_category iterator_category;
    typedef typename internal_list_t::const_iterator::difference_type difference_type;
    typedef T value_type;
    typedef T const * pointer;
    typedef T const & reference;
  private:
    typename internal_list_t::const_iterator _iter;
  public:
    _const_iterator(void) {}
    _const_iterator(typename internal_list_t::const_iterator const &iter) : _iter(iter) {
    }

    _const_iterator(_iterator const &src) : _iter(src.base()) {
    }

    T const &operator*() const { return _iter->second; }
    T const *operator->() const { return &_iter->second; }

    _const_iterator const &operator++(void) {
      ++_iter; return *this;
    }
    _const_iterator &operator++(int) {
      _const_iterator r(*this);++_iter;return r;
    }
    bool operator==(_const_iterator const &src) const {
      return _iter == src._iter;
    }

    bool operator!=(_const_iterator const &src) const {
      return _iter != src._iter;
    }

    typename internal_list_t::const_iterator const &base(void) const {
      return _iter;
    }
  };

public:
  typedef T value_type;
  typedef T *pointer;
  typedef T &reference;
  typedef T const &const_reference;
  typedef typename internal_list_t::size_type size_type;
  typedef typename internal_list_t::difference_type difference_type; // is this correct ?
  typedef _iterator iterator;
  typedef _const_iterator const_iterator;
private:
  struct ol_filler {
    typedef T argument_type;
    typedef std::pair<uint32_t,T> result_type;
    result_type const &operator()(argument_type const &x) const {
      return std::make_pair(0,x);
    }
    result_type &operator()(argument_type &x) {
      return std::make_pair(0,x);
    }
  };
  typedef typename boost::transform_iterator<ol_filler,iterator> fillout_iterator;
private:
  internal_list_t _list;
  typename internal_list_t::iterator _last;
public:

  // ------------
  // constructors
  // ------------

  ordered_slist(void) : _list() {
  }

  ordered_slist(size_t n) : _list(n) {
    // assign initial priorities
    _init_order();
#ifdef OL_GENERATE_STATS
    ol_ncreated += n;
#endif
  }

  ordered_slist(size_t n, T const &t)
  : _list(n,std::pair<uint32_t,T>(0,t)) {
    _init_order();
#ifdef OL_GENERATE_STATS
    ol_ncreated += n;
#endif
  }

  ordered_slist(ordered_slist const &src)
    : _list(src._list) {
    if(!_list.empty()) {
      _last = _list.begin();
      typename internal_list_t::iterator i(_list.begin());
      while(i != _list.end()) { _last=i;++i; }
    }
#ifdef OL_GENERATE_STATS
    ol_ncreated += _list.size();
#endif
  }

  ordered_slist const &operator=(ordered_slist const &src) {
    if(&src != this) {
      _list = src._list;
      if(!_list.empty()) {
	_last = _list.begin();
	typename internal_list_t::iterator i(_list.begin());
	while(i != _list.end()) { _last=i;++i; }
      }
#ifdef OL_GENERATE_STATS
    ol_ncreated += _list.size();
#endif
    }
  }

  template <class InputIterator>
  ordered_slist(InputIterator f, InputIterator l)
    : _list(fillout_iterator(f),fillout_iterator(l)) {
    _init_order();
#ifdef OL_GENERATE_STATS
    ol_ncreated += _list.size();
#endif
  }

  // ----------------------
  // const member functions
  // ----------------------

  size_type size(void) const { return _list.size(); }

  size_type max_size(void) const {
    // this is probably an error
    return _list.max_size();
  }

  bool empty(void) const { return _list.empty(); }

  uint32_t order(const_iterator a) const {
    return vb(a.base());
  }

  bool order_lt(const_iterator a, const_iterator b) const {
    return vb(a.base()) < vb(b.base());
  }

  const_iterator begin(void) const {
    return const_iterator(_list.begin());
  }
  const_iterator end(void) const {
    return const_iterator(_list.end());
  }

  const_reference front(void) const {
    return _list.front().second;
  }

  const_iterator previous(const_iterator pos) const {
    return const_iterator(_list.previous(pos.base()));
  }

  // --------------------------
  // non-const member functions
  // --------------------------

  void swap(ordered_slist &src) {
    _list.swap(src._list);
  }

  void push_front(T const &x) {
    if(!_list.empty()) {
      uint64_t vb_last = vb(_last);
      uint64_t vbs_last = vbs(_last);

      if((vb_last+1) == vbs_last) {
	// need to do some relabeling
	relabel(_last);
	vb_last = vb(_last);
	vbs_last = vbs(_last);
      }

      uint64_t tmp = vb_last;
      tmp = ((uint64_t(vbs_last) + tmp) / 2) + _list.begin()->first;
      _list.push_front(std::make_pair(tmp,x));
    } else {
      _list.push_front(std::make_pair(0,x));
      _last = _list.begin();
    }
#ifdef OL_GENERATE_STATS
    ol_ncreated ++;
#endif
  }

  iterator begin(void) { return iterator(_list.begin()); }
  iterator end(void) { return iterator(_list.end()); }

  reference front(void) {
    return _list.front().second;
  }

  iterator previous(iterator pos) {
    return iterator(_list.previous(pos.base()));
  }

  iterator insert_after(iterator pos) {
    return insert_after(pos,T());
  }

  iterator insert_after(iterator pos, T const &data) {
    assert(pos != end());
    typename internal_list_t::iterator x(pos.base());
    uint64_t vb_x = vb(x);
    uint64_t vbs_x = vbs(x);

    if((vb_x + 1) >= vbs_x) {
      relabel(x);
      vb_x = vb(x);
      vbs_x = vbs(x);
      assert(vb_x < vbs_x);
      assert((vb_x + 1) < vbs_x);
    }

    uint64_t ly = vbs_x;
    ly = (ly + vb_x) / 2 + _list.begin()->first;

    x = _list.insert_after(pos.base(),
				    std::make_pair(ly,data));
    if(pos.base() == _last) { _last = x; }

#ifdef OL_GENERATE_STATS
    ol_ncreated ++;
#endif

    return iterator(x);
  }

  iterator erase_after(iterator pos) {
    typename internal_list_t::iterator t(pos.base());
    if(++t == _last) {
      _last=pos.base();
    }
    return iterator(_list.erase_after(pos.base()));
  }

  iterator erase_after(iterator before_first, iterator last) {
    if(last == _list.end()) {
      _last=before_first.base();
    }

    return iterator(_list.erase_after(before_first.base(), last.base()));
  }

  // quite a lot left to do here.

private:
  // ----------------
  // helper functions
  // ----------------

  uint64_t vb(typename internal_list_t::const_iterator x) const {
    return x->first - (_list.begin()->first);
  }

  uint64_t vbs(typename internal_list_t::const_iterator x) const {
    if(++x == _list.end()) {
      return uint64_t(OL_Mm1_CONSTANT) + 1;
    } else {
      return vb(x);;
    }
  }

  void _init_order(void) {
    uint64_t tmp(OL_Mm1_CONSTANT);
    uint32_t gap = (tmp+1) / size();
    uint32_t val(gap);
    for(typename internal_list_t::iterator i(_list.begin());
	i!=_list.end();++i, val += gap) {
      i->first = val;
      _last=i;
    }
  }

  void relabel(typename internal_list_t::iterator pos) {
    // --------------------------------------------
    // pre: pos points to the first node which must
    //      be relabeled. We can easily infer the
    //      order value for the position before as
    //      it is pos->first - 1 (because otherwise
    //      we would not have needed to call
    //      relabel!)
    // --------------------------------------------

    typename internal_list_t::iterator ip(pos);
    typename internal_list_t::iterator jp(pos);
    _advance(ip,1); // as i = 1
    _advance(jp,2); // as j = 2
    unsigned int v0(pos->first);
    unsigned int i(1),j(2),nj;

    while(w(j,jp,v0) <= (4*w(i,ip,v0))) {
      i=i+1;nj=std::min(2*i,size());
      _advance(ip,1);
      _advance(jp,nj-j);
      j=nj;
    }

    // now perform the reassignment
    ip=pos; _advance(ip,1);
    uint32_t gap = w(j,jp,v0) / j;
    uint32_t val(gap);

    for(;ip!=jp;val+=gap,_advance(ip,1)) {
      ip->first = val + v0;
    }
#ifdef OL_GENERATE_STATS
    ol_nrelabels++;
#endif
  }

  uint64_t w(unsigned int index,
	     typename internal_list_t::const_iterator x,
	     unsigned int v0) const {
    if(index == size()) {
      return uint64_t(OL_Mm1_CONSTANT)+1;
    } else {
      return x->first - v0;
    }
  }

  void _advance(typename internal_list_t::iterator &j, unsigned int n) {
    while(n > 0) {
      if(++j == _list.end()) { j=_list.begin(); }
      n=n-1;
    }
  }
};

// -------------------------------------------
// The following specialisation is for a void
// element type. This may seem pointless, but
// there are situations when you just want a
// "priority space" and don't need a container
// as such.  Personally, I think the idea of
// a container which doesn't contain anything
// as quite interesting ...
// -------------------------------------------

template<>
class ordered_slist<void>
  : public ordered_slist<ol_dummy_class> {
public:
  ordered_slist(void)
    : ordered_slist<ol_dummy_class>() {}

  ordered_slist(size_t n)
    : ordered_slist<ol_dummy_class>(n) {}

  template <class InputIterator>
  ordered_slist(InputIterator f, InputIterator l)
    : ordered_slist<ol_dummy_class>(f,l) {}

  void push_front(void) {
    ordered_slist<ol_dummy_class>::push_front(ol_dummy_class());
  }
};

#endif


#include <iostream>
#include <algorithm>
#include <cmath>
#include <vector>
#include <map>
#include <set>
#include <tuple>
#include <stdbool.h>
#include <bitset>
#include <string>
#include <boost/multiprecision/cpp_int.hpp>

namespace mp = boost::multiprecision;
using namespace std;

int main(void) {
    int n,q;
    cin >> n;
    cin >> q;

    string s;
    cin >> s;

    // 事前計算
    vector<int> ans(s.length() - 1);
    int sum = 0;
    vector<int> ans_sum(s.length()); // 何回出たか
    for(int i = 0 ; i < ans.size() ; ++i) {
        ans_sum[i] = sum;
        if (s[i] == 'A' && s[i + 1] == 'C') {
            ans[i] = 1;
            sum++;
        } else {
            ans[i] = 0;
        }
    }
    ans_sum[ans.size()] = sum; // 最後

    for(int i = 0 ; i < q ; ++i) {
        int l,r;
        cin >> l;
        cin >> r;

        int a = ans_sum[l - 1];
        int b = ans_sum[r - 1];
        auto count = b - a;
        cout << count << endl;
    }
    return 0;
}


#include "Common/Common.h"
#include "Demos/Visualization/MiniGL.h"
#include "Demos/Visualization/Selection.h"
#include "Simulation/TimeManager.h"
#include <Eigen/Dense>
#include "Simulation/SimulationModel.h"
#include "Simulation/TimeStepController.h"
#include <iostream>
#include "Demos/Visualization/Visualization.h"
#include "Utils/Logger.h"
#include "Utils/Timing.h"
#include "Utils/FileSystem.h"
#include "Demos/Common/DemoBase.h"
#include "Demos/Common/TweakBarParameters.h"
#include "Simulation/Simulation.h"

// Enable memory leak detection
#if defined(_DEBUG) && !defined(EIGEN_ALIGN)
	#define new DEBUG_NEW
#endif

using namespace PBD;
using namespace Eigen;
using namespace std;
using namespace Utilities;

void initParameters();
void timeStep ();
void buildModel ();
void createHelix(const Vector3r &position, const Matrix3r &orientation, Real radius, Real height, Real totalAngle, int nPoints);
void render ();
void reset();
void TW_CALL setStretchingStiffness(const void *value, void *clientData);
void TW_CALL getStretchingStiffness(void *value, void *clientData);
void TW_CALL setShearingStiffness1(const void *value, void *clientData);
void TW_CALL getShearingStiffness1(void *value, void *clientData);
void TW_CALL setShearingStiffness2(const void *value, void *clientData);
void TW_CALL getShearingStiffness2(void *value, void *clientData);
void TW_CALL setBendingStiffness1(const void *value, void *clientData);
void TW_CALL getBendingStiffness1(void *value, void *clientData);
void TW_CALL setBendingStiffness2(const void *value, void *clientData);
void TW_CALL getBendingStiffness2(void *value, void *clientData);
void TW_CALL setTwistingStiffness(const void *value, void *clientData);
void TW_CALL getTwistingStiffness(void *value, void *clientData);

DemoBase *base;

const unsigned int nParticles = 50;
const Real helixRadius = 0.5;
const Real helixHeight = -5.0;
const Real helixTotalAngle = static_cast<Real>(10.0*M_PI);
const Matrix3r helixOrientation = AngleAxisr(-static_cast<Real>(0.5 * M_PI), Vector3r(0,1,0)).toRotationMatrix();
bool drawFrames = false;
Real shearingStiffness1 = 1.0;
Real shearingStiffness2 = 1.0;
Real stretchingStiffness = 1.0;
Real bendingStiffness1 = 0.5;
Real bendingStiffness2 = 0.5;
Real twistingStiffness = 0.5;

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

	base = new DemoBase();
	base->init(argc, argv, "Cosserat rods demo");

	SimulationModel *model = new SimulationModel();
	model->init();
	Simulation::getCurrent()->setModel(model);

	buildModel();

	initParameters();

	// OpenGL
	MiniGL::setClientIdleFunc (timeStep);
	MiniGL::addKeyFunc('r', reset);
	MiniGL::setClientSceneFunc(render);
	MiniGL::setViewport (40.0f, 0.1f, 500.0f, Vector3r (5.0, 10.0, 30.0), Vector3r (5.0, 0.0, 0.0));

	TwAddVarRW(MiniGL::getTweakBar(), "Draw frames", TW_TYPE_BOOLCPP, &drawFrames, " label='Draw frames' group=Simulation");
	TwAddVarCB(MiniGL::getTweakBar(), "Stretching stiffness", TW_TYPE_REAL, setStretchingStiffness, getStretchingStiffness, model, " label='Stretching stiffness'  min=0.0 max=1.0 step=0.1 precision=4 group='Stretch shear constraints' ");
	TwAddVarCB(MiniGL::getTweakBar(), "Shearing stiffness 1", TW_TYPE_REAL, setShearingStiffness1, getShearingStiffness1, model, " label='Shearing stiffness 1'  min=0.0 max=1.0 step=0.1 precision=4 group='Stretch shear constraints' ");
	TwAddVarCB(MiniGL::getTweakBar(), "Shearing stiffness 2", TW_TYPE_REAL, setShearingStiffness2, getShearingStiffness2, model, " label='Shearing stiffness 2'  min=0.0 max=1.0 step=0.1 precision=4 group='Stretch shear constraints' ");
	TwAddVarCB(MiniGL::getTweakBar(), "Bending stiffness 1", TW_TYPE_REAL, setBendingStiffness1, getBendingStiffness1, model, " label='Bending stiffness 1'  min=0.0 max=1.0 step=0.1 precision=4 group='Bend twist constraints' ");
	TwAddVarCB(MiniGL::getTweakBar(), "Bending stiffness 2", TW_TYPE_REAL, setBendingStiffness2, getBendingStiffness2, model, " label='Bending stiffness 2'  min=0.0 max=1.0 step=0.1 precision=4 group='Bend twist constraints' ");
	TwAddVarCB(MiniGL::getTweakBar(), "Twisting stiffness", TW_TYPE_REAL, setTwistingStiffness, getTwistingStiffness, model, " label='Twisting stiffness'  min=0.0 max=1.0 step=0.1 precision=4 group='Bend twist constraints' ");

	MiniGL::mainLoop ();

	Utilities::Timing::printAverageTimes();
	Utilities::Timing::printTimeSums();

	delete Simulation::getCurrent();
	delete base;
	delete model;

	return 0;
}

void initParameters()
{
	TwRemoveAllVars(MiniGL::getTweakBar());
	TweakBarParameters::cleanup();

	MiniGL::initTweakBarParameters();

	TweakBarParameters::createParameterGUI();
	TweakBarParameters::createParameterObjectGUI(base);
	TweakBarParameters::createParameterObjectGUI(Simulation::getCurrent());
	TweakBarParameters::createParameterObjectGUI(Simulation::getCurrent()->getModel());
	TweakBarParameters::createParameterObjectGUI(Simulation::getCurrent()->getTimeStep());
}

void reset()
{
	Utilities::Timing::printAverageTimes();
	Utilities::Timing::reset();

	Simulation::getCurrent()->reset();
	base->getSelectedParticles().clear();

	Simulation::getCurrent()->getModel()->cleanup();
	buildModel();
}

void timeStep ()
{
	const Real pauseAt = base->getValue<Real>(DemoBase::PAUSE_AT);
	if ((pauseAt > 0.0) && (pauseAt < TimeManager::getCurrent()->getTime()))
		base->setValue(DemoBase::PAUSE, true);

	if (base->getValue<bool>(DemoBase::PAUSE))
		return;

	// Simulation code
	SimulationModel *model = Simulation::getCurrent()->getModel();
	const unsigned int numSteps = base->getValue<unsigned int>(DemoBase::NUM_STEPS_PER_RENDER);
	for (unsigned int i = 0; i < numSteps; i++)
	{
		START_TIMING("SimStep");
		Simulation::getCurrent()->getTimeStep()->step(*model);
		STOP_TIMING_AVG;
	}
}

void buildModel ()
{
	TimeManager::getCurrent ()->setTimeStepSize (static_cast<Real>(0.005));

	createHelix(Vector3r(0, 0, 0), helixOrientation, helixRadius, helixHeight, helixTotalAngle, nParticles);
}

void renderLineModels()
{
	// Draw simulation model
	SimulationModel *model = Simulation::getCurrent()->getModel();
	const ParticleData &pd = model->getParticles();
	const OrientationData &od = model->getOrientations();
	float red[4] = { 0.8f, 0.0f, 0.0f, 1 };
	float green[4] = { 0.0f, 0.8f, 0.0f, 1 };
	float blue[4] = { 0.0f, 0.0f, 0.8f, 1 };

	for (unsigned int i = 0; i < model->getLineModels().size(); i++)
	{
		LineModel *lineModel = model->getLineModels()[i];

		for(unsigned int e=0; e<lineModel->getEdges().size(); e++)
		{
			const unsigned int indexOffset = lineModel->getIndexOffset();
			const unsigned int indexOffsetQuaternions = lineModel->getIndexOffsetQuaternions();
			const unsigned int i1 = lineModel->getEdges()[e].m_vert[0] + indexOffset;
			const unsigned int i2 = lineModel->getEdges()[e].m_vert[1] + indexOffset;
			const unsigned int iq = lineModel->getEdges()[e].m_quat + indexOffsetQuaternions;
			const Vector3r &v1 = pd.getPosition(i1);
			const Vector3r &v2 = pd.getPosition(i2);
			const Quaternionr &q = od.getQuaternion(iq);

			MiniGL::drawSphere(v1, 0.07f, blue);
			if( e == lineModel->getEdges().size() -1 ) MiniGL::drawSphere(v2, 0.07f, blue);
			if(drawFrames) MiniGL::drawCylinder(v1, v2, blue, 0.01f);
			else MiniGL::drawCylinder(v1, v2, blue, 0.07f);

			//draw coordinate frame at the center of the edges
			if(drawFrames)
			{
				Vector3r vm = 0.5 * (v1 + v2);
				Real scale = static_cast<Real>(0.15);
				Vector3r d1 = q._transformVector(Vector3r(1, 0, 0)) * scale;
				Vector3r d2 = q._transformVector(Vector3r(0, 1, 0)) * scale;
				Vector3r d3 = q._transformVector(Vector3r(0, 0, 1)) * scale;
				MiniGL::drawCylinder(vm, vm + d1, red,   0.01f);
				MiniGL::drawCylinder(vm, vm + d2, green, 0.01f);
				MiniGL::drawCylinder(vm, vm + d3, blue,  0.01f);
			}
		}
	}
}

void render ()
{
	base->render();

	renderLineModels();
}

/** Create a particle model and orientation model
*/
void createHelix(const Vector3r &position, const Matrix3r &orientation, Real radius, Real height, Real totalAngle, int nPoints)
{
	int nQuaternions = nPoints - 1;
	vector<Vector3r> points(nPoints);
	vector<Quaternionr> quaternions(nQuaternions);

	//init particles
	for (int i = 0; i<nPoints; i++)
	{
		points[i].x() = radius * std::cos(totalAngle / ((Real)nPoints) * (Real)i);
		points[i].y() = radius * std::sin(totalAngle / ((Real)nPoints) * (Real)i);
		points[i].z() = height / ((Real)nPoints) * (Real)i;

		points[i] = orientation * points[i] + position;
	}

	//init quaternions
	Vector3r from(0, 0, 1);
	for(int i=0; i<nQuaternions; i++)
	{
		Vector3r to = (points[i + 1] - points[i]).normalized();
		Quaternionr dq = Quaternionr::FromTwoVectors(from, to);
		if(i == 0) quaternions[i] = dq;
		else quaternions[i] = dq * quaternions[i - 1];
		from = to;
	}

	vector<unsigned int> indices(2 * nPoints - 1);
	vector<unsigned int> indicesQuaternions(nQuaternions);

	for(int i=0; i < nPoints -1; i++)
	{
		indices[2 * i] = i;
		indices[2 * i + 1] = i + 1;
	}

	for (int i = 0; i < nQuaternions; i++)
	{
		indicesQuaternions[i] = i;
	}

	SimulationModel *model = Simulation::getCurrent()->getModel();
	model->addLineModel(nPoints, nQuaternions, &points[0], &quaternions[0], &indices[0], &indicesQuaternions[0]);

	ParticleData &pd = model->getParticles();
	const int nPointsTotal = pd.getNumberOfParticles();
	for (int i = nPointsTotal - 1; i > nPointsTotal - nPoints; i--)
	{
		pd.setMass(i, 1.0);
	}

	// Set mass of points to zero => make it static
	pd.setMass(nPointsTotal - nPoints, 0.0);

	OrientationData &od = model->getOrientations();
	const unsigned int nQuaternionsTotal = od.getNumberOfQuaternions();
	for(unsigned int i = nQuaternionsTotal - 1; i > nQuaternionsTotal - nQuaternions; i--)
	{
		od.setMass(i, 1.0);
	}

	// Set mass of quaternions to zero => make it static
	od.setMass(nQuaternionsTotal - nQuaternions, 0.0);

	// init constraints
	const size_t rodNumber = model->getLineModels().size() - 1;
	const unsigned int offset = model->getLineModels()[rodNumber]->getIndexOffset();
	const unsigned int offsetQuaternions = model->getLineModels()[rodNumber]->getIndexOffsetQuaternions();
	const size_t nEdges = model->getLineModels()[rodNumber]->getEdges().size();
	const LineModel::Edges &edges = model->getLineModels()[rodNumber]->getEdges();

	//stretchShear constraints
	for(unsigned int i=0; i < nEdges; i++)
	{
		const unsigned int v1 = edges[i].m_vert[0] + offset;
		const unsigned int v2 = edges[i].m_vert[1] + offset;
		const unsigned int q1 = edges[i].m_quat + offsetQuaternions;
		model->addStretchShearConstraint(v1, v2, q1, stretchingStiffness, shearingStiffness1, shearingStiffness2);
	}

	//bendTwist constraints
	for(unsigned int i=0; i < nEdges - 1; i++)
	{
		const unsigned int q1 = edges[i].m_quat + offsetQuaternions;
		const unsigned int q2 = edges[i + 1].m_quat + offsetQuaternions;
		model->addBendTwistConstraint(q1, q2, twistingStiffness, bendingStiffness1, bendingStiffness2);
	}

// 	LOG_INFO << "Number of particles: " << nPoints;
// 	LOG_INFO << "Number of quaternions: " << nQuaternions;
}

void TW_CALL setStretchingStiffness(const void *value, void *clientData)
{
	stretchingStiffness = *(const Real *)(value);
	((SimulationModel*)clientData)->setConstraintValue<StretchShearConstraint, Real, &StretchShearConstraint::m_stretchingStiffness>(stretchingStiffness);
}

void TW_CALL getStretchingStiffness(void *value, void *clientData)
{
	*(Real *)(value) = stretchingStiffness;
}

void TW_CALL setShearingStiffness1(const void *value, void *clientData)
{
	shearingStiffness1 = *(const Real *)(value);
	((SimulationModel*)clientData)->setConstraintValue<StretchShearConstraint, Real, &StretchShearConstraint::m_shearingStiffness1>(shearingStiffness1);
}

void TW_CALL getShearingStiffness1(void *value, void *clientData)
{
	*(Real *)(value) = shearingStiffness1;
}

void TW_CALL setShearingStiffness2(const void *value, void *clientData)
{
	shearingStiffness2 = *(const Real *)(value);
	((SimulationModel*)clientData)->setConstraintValue<StretchShearConstraint, Real, &StretchShearConstraint::m_shearingStiffness2>(shearingStiffness2);
}

void TW_CALL getShearingStiffness2(void *value, void *clientData)
{
	*(Real *)(value) = shearingStiffness2;
}

void TW_CALL setBendingStiffness1(const void *value, void *clientData)
{
	bendingStiffness1 = *(const Real *)(value);
	((SimulationModel*)clientData)->setConstraintValue<BendTwistConstraint, Real, &BendTwistConstraint::m_bendingStiffness1>(bendingStiffness1);
}

void TW_CALL getBendingStiffness1(void *value, void *clientData)
{
	*(Real *)(value) = bendingStiffness1;
}

void TW_CALL setBendingStiffness2(const void *value, void *clientData)
{
	bendingStiffness2 = *(const Real *)(value);
	((SimulationModel*)clientData)->setConstraintValue<BendTwistConstraint, Real, &BendTwistConstraint::m_bendingStiffness2>(bendingStiffness2);
}

void TW_CALL getBendingStiffness2(void *value, void *clientData)
{
	*(Real *)(value) = bendingStiffness2;
}

void TW_CALL setTwistingStiffness(const void *value, void *clientData)
{
	twistingStiffness = *(const Real *)(value);
	((SimulationModel*)clientData)->setConstraintValue<BendTwistConstraint, Real, &BendTwistConstraint::m_twistingStiffness>(twistingStiffness);
}

void TW_CALL getTwistingStiffness(void *value, void *clientData)
{
	*(Real *)(value) = twistingStiffness;
}


#pragma once
#include <Eigen/Core>
#include <cstddef>
#include <string>

/// Generate random impulses.
/// @param len length of returned buffer
/// @param num_nonzero The number of values to randomise
/// @param seedval Seed for rng.
Eigen::VectorXf generate_random(size_t len, size_t num_nonzero,
                                long int seedval);

/// Simple convolution implementation.
/// @param len output and input lengths
/// @param out[out] output data
/// @param out[in] input data
/// @param ir_len length of impulse response
/// @param ir impulse response data
void convolve(size_t len, float *out, float *in, size_t ir_len, float *ir);

/// Fade up linearly over a block.
void fade_up(float *out, float *in, size_t n);

/// Fade down linearly over a block.
void fade_down(float *out, float *in, size_t n);

/// Write a float array to a file, readable by numpy.loadtext
void write_array(const std::string &base, const std::string &name, size_t n,
                 float *x);


#pragma once

#include <crab/domains/discrete_domains.hpp>
#include <crab/support/debug.hpp>
#include <crab/support/stats.hpp>

#include <boost/range/iterator_range.hpp>

#include <algorithm>
#include <unordered_map>
#include <vector>

/**
 * A standard union-find equipped with lattice operations (inclusion,
 * lub, meet, etc).
 **/

namespace crab {
namespace domains {

template <class Variable, class Domain> class equivalence_class {
public:
  using variable_t = Variable;
  using domain_t = Domain;

private:
  std::size_t m_rank;
  Domain m_val;

public:
  explicit equivalence_class(Domain val) : m_rank(0), m_val(val) {}

  std::size_t &get_rank() { return m_rank; }

  std::size_t get_rank() const { return m_rank; }

  Domain &get_domain() { return m_val; }

  const Domain &get_domain() const { return m_val; }

}; // end class equivalence_class

/*
  Partition a set of variables into equivalence classes. Apart from
  lattice operations, the domain provides standard union-find
  operations:

  - make(v)
  - union(v1, v2)
  - find(v)

  In addition, each equivalence class has associated an abstract state
  of type Domain.

  The method "top()" returns an union-find in which all variables are
  in the same equivalence class. Since the domain does not know the
  universe of variables, the method "is_top()" only succeeds if the
  union_find_domain state was created by calling directly "top()".
  The method "bottom()" represents an unreachable/failure abstract
  state. Do not use "top()" if you intend to generate an empty
  union-find. Instead, use the constructor without parameters to
  produce an empty union-find.
*/
template <class Variable, class Domain> class union_find_domain {
public:
  using variable_t = Variable;
  using domain_t = Domain;
  enum class lattice_val { bottom, top, neither_top_nor_bot };

private:
  using union_find_domain_t = union_find_domain<Variable, Domain>;
  using variable_set_t = ikos::discrete_domain<variable_t>;
  using parents_map_t = std::unordered_map<variable_t, variable_t>;
  using equivalence_class_t = equivalence_class<variable_t, Domain>;
  using equiv_class_vars_t = std::unordered_map<variable_t, variable_set_t>;
  using classes_map_t = std::unordered_map<variable_t, equivalence_class_t>;

  // immediate link to the parent
  parents_map_t m_parents;
  // each root has its equivalence class
  classes_map_t m_classes;
  lattice_val m_val;

  void clear() {
    m_parents.clear();
    m_classes.clear();
  }

  // Build a map from representative to a set with all the variables
  // in the equivalence class.
  equiv_class_vars_t equiv_classes_vars() {
    equiv_class_vars_t res;
    for (auto &kv : m_parents) {
      variable_t rep = find(kv.second);
      auto it = res.find(rep);
      if (it == res.end()) {
        variable_set_t varset = variable_set_t::bottom();
        varset += kv.first;
        res.insert({rep, varset});
      } else {
        variable_set_t &varset = it->second;
        varset += kv.first;
      }
    }
    return res;
  }

  // Return the representative after joining all variables in vars.
  // Pre-condition: if dom != none then forall v \in vars:: contains(v)
  boost::optional<variable_t>
  merge_vars(const variable_set_t &vars,
             boost::optional<domain_t> dom = boost::none) {
    auto it = vars.begin();
    auto et = vars.end();
    if (it == et) {
      return boost::none;
    }
    variable_t v = *it;
    if (dom && !contains(v)) {
      make(v, *dom);
    }
    ++it;

    if (it == et) {
      return (contains(v) ? boost::optional<variable_t>(find(v)) : boost::none);
    }

    boost::optional<variable_t> repr;
    for (; it != et; ++it) {
      if (dom && !contains(*it)) {
        make(*it, *dom);
      }
      repr = join(v, *it);
    }
    return repr;
  }

  void print(crab_os &o) const {
    o << "({";
    for (auto it = m_parents.begin(), et = m_parents.end(); it != et;) {
      variable_t key = it->first, value = it->second;
      o << key << " -> " << value;
      ++it;
      if (it != et) {
        o << ", ";
      }
    }
    o << "}, {";
    for (auto it = m_classes.begin(), et = m_classes.end(); it != et;) {
      variable_t rep = it->first;
      equivalence_class_t ec = it->second;
      o << rep << " -> " << ec.get_domain();
      ++it;
      if (it != et) {
        o << ", ";
      }
    }
    o << "})";
  }

  void get_all_variables(variable_set_t &out) const {
    for (auto &kv: m_parents) {
      out += kv.first;
    }
  }

public:
  union_find_domain(lattice_val val = lattice_val::neither_top_nor_bot)
      : m_val(val) {}
  union_find_domain(const union_find_domain_t &o) = default;
  union_find_domain(union_find_domain_t &&o) = default;
  union_find_domain_t &operator=(const union_find_domain_t &o) = default;
  union_find_domain_t &operator=(union_find_domain_t &&o) = default;

  /* =============================================================
   *    Standard union-find operations: make, find, and union
   * =============================================================
   */

  // Pre-condition: !contains(v)
  void make(const variable_t &v, Domain val) {
    if (is_bottom()) {
      CRAB_ERROR("calling union_find_domain::make on bottom");
    }
    if (is_top()) {
      CRAB_ERROR("calling union_find_domain::make on top");
    }
    if (contains(v)) {
      CRAB_ERROR("variable already exists when called union_find_domain::make");
    }
    m_parents.insert({v, v});
    m_classes.insert({v, equivalence_class_t(val)});
  }

  // Pre-condition: contains(v)
  // NOTE: it is not "const" because it does path-compression
  // TODO: make non-recursive
  variable_t &find(const variable_t &v) {
    if (!contains(v)) {
      assert(false);
      CRAB_ERROR("called union_find_domain::find on a non-existing variable ",
                 v, " in ", *this);
    }
    variable_t &parent = m_parents.at(v);
    if (parent == v) {
      return parent;
    } else {
      return parent = find(parent);
    }
  }

  // Return the representative after merging x and y in the same
  // equivalence class.
  // Pre-condition: contains(x) && contains(y)
  variable_t join(const variable_t &x, const variable_t &y) {
    if (!contains(x)) {
      CRAB_ERROR("called union_find_domain::join on a non-existing variable ",
                 x, " in ", *this);
    }
    if (!contains(y)) {
      CRAB_ERROR("called union_find_domain::join on a non-existing variable ",
                 y, " in ", *this);
    }

    variable_t rep_x = find(x);
    variable_t rep_y = find(y);
    if (rep_x != rep_y) {
      equivalence_class_t &ec_x = m_classes.at(rep_x);
      equivalence_class_t &ec_y = m_classes.at(rep_y);
      domain_t &dom_x = ec_x.get_domain();
      domain_t &dom_y = ec_y.get_domain();
      if (ec_x.get_rank() > ec_y.get_rank()) {
        dom_x = dom_x | dom_y;
        m_parents.at(rep_y) = rep_x;
        m_classes.erase(rep_y);
        return rep_x;
      } else {
        dom_y = dom_y | dom_x;
        m_parents.at(rep_x) = rep_y;
        if (ec_x.get_rank() == ec_y.get_rank()) {
          ec_y.get_rank()++;
        }
        m_classes.erase(rep_x);
        return rep_y;
      }
    } else {
      return rep_x;
    }
  }

  bool contains(const variable_t &v) const {
    return m_parents.find(v) != m_parents.end();
  }

  // Pre-condition: contains(v)
  equivalence_class_t &get_equiv_class(const variable_t &v) {
    return m_classes.at(find(v));
  }

  void remove_equiv_class(const variable_t &v) {
    if (!contains(v)) {
      return;
    }
    variable_t rep_v = find(v);
    m_classes.erase(rep_v);

    equiv_class_vars_t map = equiv_classes_vars();
    auto it = map.find(rep_v);
    if (it != map.end()) {
      for (auto v : it->second) {
        m_parents.erase(v);
      }
    }
  }

  /* =============================================================
   *                     Abstract domain API
   * =============================================================
   */

  static union_find_domain_t top() {
    return union_find_domain_t(lattice_val::top);
  }

  static union_find_domain_t bottom() {
    return union_find_domain_t(lattice_val::bottom);
  }

  bool is_bottom() const { return (m_val == lattice_val::bottom); }

  bool is_top() const { return (m_val == lattice_val::top); }

  void set_to_top() {
    clear();
    m_val = lattice_val::top;
  }

  void set_to_bottom() {
    clear();
    m_val = lattice_val::bottom;
  }

  void set(const variable_t &x, domain_t dom) {
    if (is_bottom() || is_top()) {
      return;
    }

    if (!contains(x)) {
      // Create a singleton equivalence class
      make(x, dom);
    } else {
      // Modify the abstract state of the whole equivalence class
      variable_t rep_x = find(x);
      equivalence_class_t &ec_x = m_classes.at(rep_x);
      ec_x.get_domain() = dom;
    }
  }

  // Pre-condition: contains(x)
  const domain_t &operator[](const variable_t &x) {
    if (is_bottom()) {
      CRAB_ERROR("called union_find_domain::operator[] on bottom");
    }
    if (is_top()) {
      CRAB_ERROR("called union_find_domain::operator[] on top");
    }
    if (!contains(x)) {
      CRAB_ERROR("union_find_domain::operator[] on a non-existing variable ", x,
                 " in ", *this);
    }

    variable_t rep_x = find(x);
    const equivalence_class_t &ec_x = m_classes.at(rep_x);
    return ec_x.get_domain();
  }

  // Return true if *this is a refined partitioning of o
  //    forall x,y \in vars(o) :: o.find(x) != o.find(y) =>
  //                              this.find(x) != this.find(y)
  bool operator<=(const union_find_domain_t &o) const {
    if (is_bottom() || o.is_top()) {
      return true;
    } else if (is_top() || o.is_bottom()) {
      return false;
    } else {
      union_find_domain_t left(*this);
      union_find_domain_t right(o);
      equiv_class_vars_t right_equiv_classes = right.equiv_classes_vars();
      equiv_class_vars_t left_equiv_classes = left.equiv_classes_vars();
      for (auto &kv : right_equiv_classes) {
        variable_t right_repr = kv.first;
        variable_set_t &right_vars = kv.second;
        for (variable_t right_v :
             boost::make_iterator_range(right_vars.begin(), right_vars.end())) {
          if (!left.contains(right_v)) {
            return false;
          }
          variable_t left_repr = left.find(right_v);
          variable_set_t &left_vars = left_equiv_classes[left_repr];
          if (!(left_vars <= right_vars)) {
            return false;
          }
          const domain_t &left_dom =
              left.get_equiv_class(left_repr).get_domain();
          const domain_t &right_dom =
              right.get_equiv_class(right_repr).get_domain();
          if (!(left_dom <= right_dom)) {
            return false;
          }
        }
      }
      return true;
    }
  }

  union_find_domain_t operator|(const union_find_domain_t &o) const {
    if (is_bottom()) {
      return o;
    } else if (o.is_bottom()) {
      return *this;
    } else if (is_top() | o.is_top()) {
      union_find_domain_t res;
      return res;
    } else {
      CRAB_LOG("union-find",
               crab::outs() << "Join \n\t" << *this << "\n\t" << o << "\n";);
      union_find_domain_t left(*this);
      union_find_domain_t right(o);

      // Keep only common variables
      variable_set_t left_vars = variable_set_t::bottom();
      left.get_all_variables(left_vars);
      for (auto v: left_vars) {
	if (!right.contains(v)) {
	  left.forget(v);
	}
      }
      variable_set_t right_vars = variable_set_t::bottom();
      right.get_all_variables(right_vars);
      for (auto v: right_vars) {
	if (!left.contains(v)) {
	  right.forget(v);
	}
      }

      equiv_class_vars_t right_equiv_classes = right.equiv_classes_vars();
      equiv_class_vars_t left_equiv_classes = left.equiv_classes_vars();

      // Merge equivalence classes from the left to the right while
      // joining the domains associated to the equivalence classes.
      //
      // The merging on the right is needed so that right_dom is updated.
      for (auto &kv : left_equiv_classes) {
        boost::optional<variable_t> right_repr = right.merge_vars(kv.second);
        if (!right_repr)
          continue; // this shouldn't happen
        domain_t &left_dom = left.get_equiv_class(kv.first).get_domain();
        const domain_t &right_dom =
            right.get_equiv_class(*right_repr).get_domain();
        left_dom = left_dom | right_dom;
      }

      // Merge equivalence classes from the right to the left while
      // joining the domains associated to the equivalence classes.
      for (auto &kv : right_equiv_classes) {
        boost::optional<variable_t> left_repr = left.merge_vars(kv.second);
        if (!left_repr)
          continue; // this shouldn't happen
        domain_t &left_dom = left.get_equiv_class(*left_repr).get_domain();
        const domain_t &right_dom =
            right.get_equiv_class(kv.first).get_domain();
        left_dom = left_dom | right_dom;
      }
      CRAB_LOG("union-find", crab::outs() << "Res=" << left << "\n";);
      return left;
    }
  }

  union_find_domain_t operator&(const union_find_domain_t &o) const {
    if (is_bottom() || o.is_top()) {
      return *this;
    } else if (o.is_bottom() || is_top()) {
      return o;
    } else {
      /* TOIMPROVE: this operation is imprecise because we still merge
       * the equivalence classes.
       */
      union_find_domain_t left(*this);
      union_find_domain_t right(o);
      equiv_class_vars_t right_equiv_classes = right.equiv_classes_vars();
      equiv_class_vars_t left_equiv_classes = left.equiv_classes_vars();

      // Merge equivalence classes from the right to the left while
      // applying the meet
      //
      // The merging on the right is needed so that right_dom is updated.
      for (auto &kv : left_equiv_classes) {
        domain_t &left_dom = left.get_equiv_class(kv.first).get_domain();
        // add variables on the right if they do not exist
        boost::optional<variable_t> right_repr =
            right.merge_vars(kv.second, left_dom);
        if (!right_repr)
          continue; // this shouldn't happen
        const domain_t &right_dom =
            right.get_equiv_class(*right_repr).get_domain();
        left_dom = left_dom & right_dom;
        if (left_dom.is_bottom()) {
          left.set_to_bottom();
          return left;
        }
      }

      // Merge equivalence classes from the right to the left while
      // applying the meet
      for (auto &kv : right_equiv_classes) {
        const domain_t &right_dom =
            right.get_equiv_class(kv.first).get_domain();
        boost::optional<variable_t> left_repr =
            left.merge_vars(kv.second, right_dom);
        if (!left_repr)
          continue; // this shouldn't happen
        domain_t &left_dom = left.get_equiv_class(*left_repr).get_domain();
        left_dom = left_dom & right_dom;
        if (left_dom.is_bottom()) {
          left.set_to_bottom();
          return left;
        }
      }
      return left;
    }
  }

  union_find_domain_t operator||(const union_find_domain_t &o) const {
    // finite domain
    return *this | o;
  }
  union_find_domain_t operator&&(const union_find_domain_t &o) const {
    // finite domain
    return *this & o;
  }

  // Add y into the equivalence class of x
  void add(const variable_t &x, const variable_t &y) {
    if (is_bottom() || is_top()) {
      return;
    }
    if (!contains(x)) {
      return;
    }
    if (contains(y)) {
      forget(y);
    }
    variable_t rep_x = find(x);
    m_parents.insert({y, rep_x});
  }

  std::vector<variable_t> get_equiv_vars_by_var(const variable_t &x) {
    std::vector<variable_t> out;
    if (!contains(x)) {
      return out;
    }
    variable_t rep_x = find(x);
    out.reserve(m_parents.size());
    for (auto &kv : m_parents) {
      if (kv.second == rep_x) {
        out.push_back(kv.first);
      }
    }
    return out;
  }

  void forget(const variable_t &v) {
    if (is_bottom() || is_top()) {
      return;
    }

    if (contains(v)) {
      if (m_parents.at(v) != v) {
        // v is not a representative
        variable_t rep_v = find(v);
        for (auto &kv : m_parents) {
          if (kv.second == v) {
            m_parents.at(kv.first) = rep_v;
          }
        }
      } else {
        // v is the representative of the equivalence class
        boost::optional<variable_t> new_rep;
        for (auto &kv : m_parents) {
          if (kv.first != v && kv.second == v) {
            if (!new_rep) {
              // choose the first one as representative
              new_rep = kv.first;
              m_classes.insert({*new_rep, m_classes.at(v)});
            }
            m_parents.at(kv.first) = *new_rep;
          }
        }
        m_classes.erase(v);
      }
      m_parents.erase(v);
    }
  }

  void project(const std::vector<variable_t> &variables) {
    if (is_bottom() || is_top()) {
      return;
    }

    // First, collect all variables to be forgotten
    std::vector<variable_t> sorted_variables(variables);
    std::sort(sorted_variables.begin(), sorted_variables.end());
    std::vector<variable_t> variables_to_forget;
    variables_to_forget.reserve(m_parents.size());
    for (auto &kv : m_parents) {
      auto lower = std::lower_bound(sorted_variables.begin(),
                                    sorted_variables.end(), kv.first);
      if (lower == sorted_variables.end() || kv.first < *lower) { // not found
        variables_to_forget.push_back(kv.first);
      }
    }
    // Forget variables
    for (auto &v : variables_to_forget) {
      forget(v);
    }
  }

  void rename(const std::vector<variable_t> &old_variables,
              const std::vector<variable_t> &new_variables) {
    if (is_top() || is_bottom()) {
      return;
    }
    if (old_variables.size() != new_variables.size()) {
      CRAB_ERROR(
          "union_find_domain::rename with input vectors of different sizes");
    }

    auto rename_var = [this](const variable_t &x, const variable_t &y) {
      if (contains(y)) {
        CRAB_ERROR("union_find_domain::rename assumes that ", y,
                   " does not exist");
      }
      if (contains(x)) {
        // tuples where x is the left operand
        std::vector<typename parents_map_t::value_type> left_x;
        for (auto &kv : m_parents) {
          if (kv.first == x) {
            left_x.push_back(kv);
          } else if (kv.second == x) {
            kv.second = y;
          }
        }
        for (auto &kv : left_x) {
          m_parents.erase(kv.first);
          m_parents.insert({y, (kv.second == x ? y : kv.second)});
        }
        auto it = m_classes.find(x);
        if (it != m_classes.end()) {
          equivalence_class_t equiv_class = it->second;
          m_classes.erase(it);
          m_classes.insert({y, std::move(equiv_class)});
        }
      }
    };

    for (unsigned i = 0, sz = old_variables.size(); i < sz; ++i) {
      rename_var(old_variables[i], new_variables[i]);
    }
  }

  void write(crab_os &o) const {
    if (is_top()) {
      o << "{}";
    } else if (is_bottom()) {
      o << "_|_";
    } else {

      // Sort everything to print always the same thing
      auto sorted_equiv_classes = [](const equiv_class_vars_t &unsorted_map) {
	     std::vector<std::pair<variable_t, std::vector<variable_t>>> sorted_eq_classes;
	     for (auto &kv: unsorted_map) {
	       std::vector<variable_t> sorted_eq_class(kv.second.begin(), kv.second.end());
	       std::sort(sorted_eq_class.begin(), sorted_eq_class.end());
	       sorted_eq_classes.emplace_back(std::make_pair(kv.first, sorted_eq_class));
	     }
	     std::sort(sorted_eq_classes.begin(), sorted_eq_classes.end(),
		       [](const std::pair<variable_t, std::vector<variable_t>> &p1,
			  const std::pair<variable_t, std::vector<variable_t>> &p2) {
			 return p1.first < p2.first;
		       });
	     return sorted_eq_classes;
      };

      auto print_vars_vector = [&o](const std::vector<variable_t>& vars) {
				 o << "{";
				 for (auto it = vars.begin(), et = vars.end(); it!=et;) {
				   o << *it;
				   ++it;
				   if (it != et) {
				     o << ",";
				   }
				 }
				 o << "}";
			       };

      union_find_domain_t tmp(*this);
      CRAB_LOG("union-find-print", tmp.print(o); o << "\n";);
      equiv_class_vars_t ec_vars = tmp.equiv_classes_vars();
      auto sorted_ec_vars = sorted_equiv_classes(ec_vars);
      o << "{";
      for (auto it = sorted_ec_vars.begin(), et = sorted_ec_vars.end(); it != et;) {
	auto &p = *it;
        variable_t &rep = tmp.find(p.first);
	print_vars_vector(p.second);
	o << "=>" << tmp.m_classes.at(rep).get_domain();
        ++it;
        if (it != et) {
          o << ",";
        }
      }
      o << "}";
    }
  }

  friend class crab::crab_os &operator<<(crab::crab_os &o,
                                         const union_find_domain_t &dom) {
    dom.write(o);
    return o;
  }
}; // end class union_find_domain

} // end namespace domains
} // end namespace crab


#include <Mahi/Robo.hpp>
#include <Mahi/Util.hpp>
#include <Params.hpp>
#include <Eigen/Dense>

using Eigen::VectorXd;
using Eigen::MatrixXd;

class Model{
public:
    // stores the state of the robot
    struct State{
        VectorXd Q;   // current position of the robot
        VectorXd Qd;  // current velocities of the robot
        VectorXd Qdd; // current accelerations of the robot
        VectorXd Tau; // current torque outputs of the robot
    };

    State m_s{VectorXd::Zero(num_joints), VectorXd::Zero(num_joints), VectorXd::Zero(num_joints), VectorXd::Zero(num_joints)}; // stores robot state

    /// Constructor
    Model();
    /// sets the zero conditions of the robot
    void zero();
    /// runs one iteration of the update rule, integrating forward in time
    void update(double dt);
    /// integrates the state of the robot according to the latest dt
    void integrate_state(VectorXd Qdd, double dt);
private:
    /// returns the vector of damping terms given the robot state
    VectorXd B();
    /// returns the vector of friction terms given the robot state
    VectorXd Fk();
};


#ifndef __EARTH_ENVIRONMENT_HH__
#define __EARTH_ENVIRONMENT_HH__
/********************************* TRICK HEADER *******************************
PURPOSE:
      (Describe the EarthEnvironment Module Variables and Algorithm)
LIBRARY DEPENDENCY:
      ((../src/EarthEnvironment.cpp))
PROGRAMMERS:
      (((Chun-Hsu Lai) () () () ))
*******************************************************************************/
#include <armadillo>
#include <functional>
#include "aux.hh"
#include "cadac_constants.hh"
#include "numerical_constants.hh"

#include "env/atmosphere.hh"
#include "env/atmosphere76.hh"
#include "env/atmosphere_nasa2002.hh"
#include "env/atmosphere_weatherdeck.hh"

#include "env/wind.hh"
#include "env/wind_constant.hh"
#include "env/wind_no.hh"
#include "env/wind_tabular.hh"

#include "dm_delta_ut.hh"
#include "time_management.hh"
#include "vehicle.hh"

class EarthEnvironment : public FH_module
{
    TRICK_INTERFACE(EarthEnvironment);

public:
    EarthEnvironment();

    ~EarthEnvironment();

    void atmosphere_use_public();
    void atmosphere_use_nasa();
    void atmosphere_use_weather_deck(char *filename);

    void set_no_wind();
    void set_constant_wind(double dvae, double dir, double twind,
                           double vertical_wind);
    void set_tabular_wind(char *filename, double twind, double vertical_wind);

    void set_no_wind_turbulunce();
    void set_wind_turbulunce(double turb_length, double turb_sigma, double taux1,
                             double taux1d, double taux2, double taux2d,
                             double tau, double gauss_value);

    virtual void init(LaunchVehicle *VehicleIn);
    virtual void algorithm(LaunchVehicle *VehicleIn);

private:
    time_management *time;

    /* Constants */
    cad::Atmosphere *atmosphere;
    cad::Wind *wind;

    /* Function declaration */
    arma::mat RNP();
    arma::vec AccelHarmonic(arma::vec3 SBII, arma::mat33 TEI,
                            int n_max, int m_max);
};

#endif  // __EARTH_ENVIRONMENT_HH__


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

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

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

#ifndef __JLANN_DS_HPP__
#define __JLANN_DS_HPP__

#include <Eigen/Dense>
#include <ANN/ANN.h>
#include <iostream>
#include <cmath>
#include <random>

#include <sys/resource.h>

namespace jlann{

typedef Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> DMatrix;
typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> FMatrix;

template<typename T>
class JLAnnDS {
public:
	typedef Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> JLMatrix;
	typedef Eigen::Matrix<T, 1, Eigen::Dynamic, Eigen::RowMajor> JLVector;

	JLAnnDS(JLMatrix& dataPoints, int numOfDs=1, int projectedDimension=0, int searchSize=0):
		originalPoints(dataPoints),
		numberOfDs(numOfDs),
		projDimension(projectedDimension),
		projSearchSize(searchSize) {

		numberOfPoints = originalPoints.rows();
		dimension = originalPoints.cols();

		if (projDimension<=0) {
			projDimension = (int)std::ceil(std::log2(numberOfPoints)/(std::log2(std::log2(numberOfPoints))));
		}
		if (numberOfDs<=0) {
			numberOfDs = 1;
		}
		if (projSearchSize<=0) {
			projSearchSize = (int)std::ceil(std::sqrt(numberOfPoints));
		}

		annIdx = new ANNidx[projSearchSize];
		dists = new ANNdist[projSearchSize];
		annQueryPt = annAllocPt(projDimension);

		//createDs();
	}

	void createDs() {
		this->bdTree.reserve(numberOfDs);
		this->projectionMatrices.reserve(numberOfDs);
        std::mt19937 generator;
        std::normal_distribution<double> distribution(0.0,1.0);
		struct rlimit limit;

		for (int currentDs=0; currentDs<numberOfDs; currentDs++) {
			JLMatrix* projMatrix = new JLMatrix(dimension, projDimension);

        	for (int i=0; i<dimension; i++) {
        	    for (int j=0; j<projDimension; j++) {
        	        (*projMatrix)(i, j) = distribution(generator);
        	    }
        	}
			JLMatrix projectedData = originalPoints*(*projMatrix);

        	ANNpointArray point_array = annAllocPts(numberOfPoints, projDimension);
			auto dataPointer = projectedData.data();
       	    for (int i=0; i<numberOfPoints; i++) {
	        	for (int j=0; j<projDimension; j++) {
        	        point_array[i][j] = (*dataPointer);
					++dataPointer;
        	    }
        	}

			this->projectionMatrices.push_back(projMatrix);
			this->pointArray.push_back(point_array);
			ANNbd_tree* bdtree = new ANNbd_tree(
        	    point_array,
        	    numberOfPoints,
        	   	projDimension,
				2
        	);
			this->bdTree.push_back(bdtree);
		}
	}

	std::vector<int> findNearestNeighbors(JLMatrix& queries, double epsilon) {
		std::vector<int> nns;
		std::vector<double> nnDists;
		nns.reserve(queries.rows());
		nnDists.reserve(queries.rows());
		for (int i=0; i<queries.rows(); i++){
			nns.push_back(-1);
			nnDists.push_back(-1);
		}
		for (int currentDs=0; currentDs<numberOfDs; currentDs++) {
			JLMatrix projectedQueries = queries * (*projectionMatrices[currentDs]);
			auto queriesData = projectedQueries.data();
			for (int i=0; i<projectedQueries.rows(); i++) {
				auto queryPoint = queries.row(i);
				for (int j=0; j<projDimension; j++) {
					annQueryPt[j] = *(queriesData);
					queriesData++;
				}

	        	this->bdTree[currentDs]->annkSearch(
	        	    annQueryPt,
	        	    projSearchSize,
	        	    annIdx,
	        	    dists,
	        	    epsilon
	        	);

	        	for (int k=0; k<projSearchSize; k++) {
	        	    if (annIdx[k]==-1) {
	        	        break;
	        	    }
					double distance = (queryPoint-originalPoints.row(annIdx[k])).norm();
	        	    if (distance<nnDists[i] || nnDists[i]<0) {
	        	        nnDists[i] = distance;
	        	        nns[i] = annIdx[k];
	        	    }
	        	}
			}
		}
		return nns;
	}

	int findNearestNeighbor(JLVector& queryPoint, double epsilon) {

		int minDistance=-1;
		int nearestNeighborIdx=-1;
		for (int currentDs=0; currentDs<numberOfDs; currentDs++) {
        	for (int i=0; i<projDimension; i++) {
				annQueryPt[i] = 0;
        	    for (int j=0; j<dimension; j++) {
        	        annQueryPt[i] += queryPoint(j) * (*projectionMatrices[currentDs])(j, i);
        	    }
        	}
        	this->bdTree[currentDs]->annkSearch(
        	    annQueryPt,
        	    projSearchSize,
        	    annIdx,
        	    dists,
        	    epsilon
        	);

        	for (int i=0; i<projSearchSize; i++) {
        	    if (annIdx[i]==-1) {
        	        break;
        	    }
				double distance = (queryPoint-originalPoints.row(annIdx[i])).norm();
        	    if (distance<minDistance || minDistance<0) {
        	        minDistance = distance;
        	        nearestNeighborIdx = annIdx[i];
        	    }
        	}
		}

		return nearestNeighborIdx;
	}

	~JLAnnDS() {
		for (int i=0; i<pointArray.size(); i++) {
			annDeallocPts(pointArray[i]);
		}
		for (int i=0; i<projectionMatrices.size(); i++) {
			delete projectionMatrices[i];
		}
		for (int i=0; i<bdTree.size(); i++) {
			delete bdTree[i];
		}
		annDeallocPt(annQueryPt);
		delete []annIdx;
		delete []dists;
	}

private:
	std::vector<JLMatrix*> projectionMatrices;
	std::vector<ANNbd_tree*> bdTree;
	JLMatrix& originalPoints;
	std::vector<ANNpointArray> pointArray;
	ANNidxArray annIdx;
	ANNdistArray dists;
	ANNpoint annQueryPt;
	int numberOfDs;
	int projDimension;
	int projSearchSize;
	int dimension;
	int numberOfPoints;
}; //class JLAnnDs
} //namespace jlann
#endif //__JLANN_DS_HPP__


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

#ifndef MTL_MATRIX_SOLVE_INCLUDE
#define MTL_MATRIX_SOLVE_INCLUDE

#include <boost/numeric/mtl/utility/tag.hpp>
#include <boost/numeric/mtl/utility/category.hpp>
#include <boost/numeric/mtl/concept/collection.hpp>
#include <boost/numeric/mtl/operation/lu.hpp>
#include <boost/numeric/mtl/interface/umfpack_solve.hpp>
#include <boost/numeric/mtl/interface/vpt.hpp>

namespace mtl { namespace matrix {

    namespace detail {

	template <typename Matrix, typename Vector>
	Vector inline solve(const Matrix& A, const Vector& b, tag::dense)
	{
		vampir_trace<3034> tracer;
	    return lu_solve(A, b);
	}

#     ifdef MTL_HAS_UMFPACK
	template <typename Value, typename Parameters, typename Vector>
	Vector inline solve(const Matrix& A, const Vector& b, tag::compressed2D)
	{
		vampir_trace<3035> tracer;
	    Vector x(num_cols(A));
	    umfpack_solve(A, x, b);
	    return x;
	}
#     endif
    }

template <typename Matrix, typename Vector>
Vector inline solve(const Matrix& A, const Vector& b)
{
    return detail::solve(A, b, typename category<Coll>::type());
}

}} // namespace mtl::matrix

#endif // MTL_MATRIX_SOLVE_INCLUDE


//
// Copyright © 2017 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//

#include "Pooling2d.hpp"
#include "TensorBufferArrayView.hpp"

#include <armnn/Exceptions.hpp>
#include <armnn/Types.hpp>

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

#include <limits>
#include <algorithm>
#include <functional>

namespace
{
    using PoolingAlgorithm = armnn::PoolingAlgorithm;

    float DefaultInitializer(PoolingAlgorithm algorithm)
    {
        switch (algorithm)
        {
            case PoolingAlgorithm::Max:
            {
                return std::numeric_limits<float>::lowest();
            }
            case PoolingAlgorithm::Average:
            case PoolingAlgorithm::L2:
            {
                return 0.0f;
            }
            default:
            {
                throw armnn::InvalidArgumentException("Unsupported pooling algorithm");
            }
        }
    }

    using Accumulator = std::function<void(float & accu, float value)>;

    Accumulator GetAccumulator(PoolingAlgorithm algorithm)
    {
        switch (algorithm)
        {
            case PoolingAlgorithm::Max:
            {
                return [](float & accu, float value) {
                    if (value > accu) {
                        accu = value;
                    }
                };
            }

            case PoolingAlgorithm::Average:
            {
                return [](float & accu, float value) {
                    accu += value;
                };
            }

            case PoolingAlgorithm::L2:
            {
                return [](float & accu, float value) {
                    accu += (value*value);
                };
            }

            default:
            {
                throw armnn::InvalidArgumentException("Unsupported pooling algorithm");
            }
        }
    }

    using Executor = std::function<void(float & accumulated, float kernelSize)>;

    Executor GetExecutor(PoolingAlgorithm algorithm)
    {
        switch (algorithm)
        {
            case PoolingAlgorithm::Max:
            {
                return [](float & accumulated, float kernelSize) {};
            }

            case PoolingAlgorithm::Average:
            {
                return [](float & accumulated, float kernelSize) {
                    accumulated /= kernelSize;
                };
            }

            case PoolingAlgorithm::L2:
            {
                return [](float & accumulated, float kernelSize) {
                    accumulated = sqrtf(accumulated / kernelSize);
                };
            }

            default:
            {
                throw armnn::InvalidArgumentException("Unsupported pooling algorithm");
            }
        }
    }

    bool OnPaddingOnly(int start, int end, int maxRange, int padding)
    {
        if (end <= 0 || start > (maxRange - padding))
        {
            return true;
        }
        else
        {
            return false;
        }
    }

    bool ClampRange(int & start, int & end, int maxRange)
    {
        if (start < 0 || end > maxRange)
        {
            start = std::min(std::max(start, 0), maxRange);
            end   = std::min(std::max(end, 0), maxRange);
            return true;
        }
        else
        {
            return false;
        }
    }
}

using namespace armnnUtils;

namespace armnn
{

void Pooling2d(const float* in,
               float* out,
               const TensorInfo& inputInfo,
               const TensorInfo& outputInfo,
               const Pooling2dDescriptor& params)
{
    const DataLayoutIndexed dataLayout = params.m_DataLayout;
    auto channelsIndex = dataLayout.GetChannelsIndex();
    auto heightIndex = dataLayout.GetHeightIndex();
    auto widthIndex = dataLayout.GetWidthIndex();

    const int batchSize    = boost::numeric_cast<int>(outputInfo.GetShape()[0]);
    const int channels     = boost::numeric_cast<int>(outputInfo.GetShape()[channelsIndex]);
    const int heightOutput = boost::numeric_cast<int>(outputInfo.GetShape()[heightIndex]);
    const int widthOutput  = boost::numeric_cast<int>(outputInfo.GetShape()[widthIndex]);
    const int heightInput  = boost::numeric_cast<int>(inputInfo.GetShape()[heightIndex]);
    const int widthInput   = boost::numeric_cast<int>(inputInfo.GetShape()[widthIndex]);
    const int padLeft      = boost::numeric_cast<int>(params.m_PadLeft);
    const int padRight     = boost::numeric_cast<int>(params.m_PadRight);
    const int padTop       = boost::numeric_cast<int>(params.m_PadTop);
    const int padBottom    = boost::numeric_cast<int>(params.m_PadBottom);
    const int strideX      = boost::numeric_cast<int>(params.m_StrideX);
    const int strideY      = boost::numeric_cast<int>(params.m_StrideY);
    const int poolHeight   = boost::numeric_cast<int>(params.m_PoolHeight);
    const int poolWidth    = boost::numeric_cast<int>(params.m_PoolWidth);

    float defaultInitializer = DefaultInitializer(params.m_PoolType);

    Accumulator accumulate = GetAccumulator(params.m_PoolType);
    Executor execute       = GetExecutor(params.m_PoolType);

    TensorBufferArrayView<const float> input(inputInfo.GetShape(), in, dataLayout);
    TensorBufferArrayView<float> output(outputInfo.GetShape(), out, dataLayout);

    // Check supported padding methods outside the loop to simplify
    // the inner loop.
    if (params.m_PaddingMethod != PaddingMethod::Exclude &&
        params.m_PaddingMethod != PaddingMethod::IgnoreValue)
    {
        throw armnn::InvalidArgumentException("Unsupported padding type");
    }

    for (int n = 0; n < batchSize; n++)
    {
        for (int c = 0; c < channels; c++)
        {
            for (int yOutput = 0; yOutput < heightOutput; yOutput++)
            {
                for (int xOutput = 0; xOutput < widthOutput; xOutput++)
                {
                    int hstart = (yOutput * strideY) - padTop;
                    int wstart = (xOutput * strideX) - padLeft;
                    int hend = hstart + poolHeight;
                    int wend = wstart + poolWidth;

                    // Clamp the pooling region inside the valid input area (which includes the padding).
                    // This is necessary because the final pooling in a row may overlap beyond the padding.
                    hend = std::min(hend, heightInput + padBottom);
                    wend = std::min(wend, widthInput + padRight);

                    float result = defaultInitializer;
                    float poolAreaSize = boost::numeric_cast<float>((hend - hstart) * (wend - wstart));

                    // Special case: when the pooling kernel is over a padding region and the padding
                    //               size is larger or equal to the kernel and the kernel only covers
                    //               padding and no real values, then we initialize the result as zero
                    //               by convention. This is because we need to choose a value here and
                    //               all values we have are padding, which we ignore.
                    if (OnPaddingOnly(hstart, hend, heightInput, padBottom) ||
                        OnPaddingOnly(wstart, wend, widthInput, padRight))
                    {
                        result = 0.0f;
                    }

                    bool clamped = ClampRange(wstart, wend, widthInput);
                    clamped |= ClampRange(hstart, hend, heightInput);

                    if (clamped && params.m_PaddingMethod == PaddingMethod::Exclude)
                    {
                        // When we exclude the padding, it means we calculate with a smaller
                        // kernel size, so I changed the divisor here.
                        poolAreaSize = boost::numeric_cast<float>((hend - hstart) * (wend - wstart));
                    }

                    for (auto yInput = hstart; yInput < hend; yInput++)
                    {
                        for (auto xInput = wstart; xInput < wend; xInput++)
                        {
                            float inval = input.Get(boost::numeric_cast<unsigned int>(n),
                                                    boost::numeric_cast<unsigned int>(c),
                                                    boost::numeric_cast<unsigned int>(yInput),
                                                    boost::numeric_cast<unsigned int>(xInput));

                            accumulate(result, inval);
                        }
                    }

                    execute(result, poolAreaSize);

                    output.Get(boost::numeric_cast<unsigned int>(n),
                               boost::numeric_cast<unsigned int>(c),
                               boost::numeric_cast<unsigned int>(yOutput),
                               boost::numeric_cast<unsigned int>(xOutput)) = result;
                }
            }
        }
    }
}

} //namespace armnn


#include "bomd_electronic.hpp"
#include "electronic.hpp"
#include "constants.hpp"
#include "util.hpp"
#include <armadillo>
#include <complex>
#include <cmath>

ConfigBlockReader ElectronicBomd::setup_reader()
{
    using types = ConfigBlockReader::types;
    ConfigBlockReader reader{"bomd-electronic"};
    reader.add_entry("dtc", types::DOUBLE);
    reader.add_entry("amplitude_file", "cs.hdf5");

    return reader;
}

void ElectronicBomd::get_reader_data(ConfigBlockReader& reader) {
  {
    double in_dtc;
    reader.get_data("dtc", in_dtc);  // in fs
    dtc = in_dtc * (1e-15 / AU2SI_TIME);  // fs -> a.u.
  }

  reader.get_data("amplitude_file", amplitude_file);

  /* added in BOMD::add_qm_keys() */
  reader.get_data("min_state", min_state);

  {
    size_t excited_states;
    /* added in BOMD::add_qm_keys() */
    reader.get_data("excited_states", excited_states);

    nstates = excited_states + 1 - min_state;
  }

  if (nstates < 2){
    throw std::logic_error("Doesn't make sense to propagate the electronic wavefunction on only 1 surface!");
  }

  if (!(min_state <= active_state && active_state <= min_state + nstates)){
    throw std::range_error("Active state not in the range of propagating states!");
  }

  active_state -= min_state;

  energy.set_size(nstates);

  U.set_size(nstates, nstates);
  T.set_size(nstates, nstates);
  V.set_size(nstates, nstates);

  phases.set_size(nstates);
  phases.ones();

  {
    arma::cx_mat initial(nstates, nstates, arma::fill::eye);
    c = Electronic(initial);
  }
}

// This is our primary hook into the Gromacs (or other) MD loop
double ElectronicBomd::update_gradient(void){
  qm->update();

  // get gradients, energies, and overlap
  {
    PropMap props{};
    props.emplace(QMProperty::qmgradient, {min_state + active_state}, &qm_grd);
    props.emplace(QMProperty::mmgradient, {min_state + active_state}, &mm_grd);
    props.emplace(QMProperty::energies, util::range(min_state, min_state + nstates), &energy);
    props.emplace(QMProperty::wfoverlap, &U);
    qm->get_properties(props);
  }

  // write amplitudes first so we include the initial conditions
  c.get().save(arma::hdf5_name(amplitude_file,
                               "/amps/" + std::to_string(call_idx()),
                               arma::hdf5_opts::replace));
  saveh5(U, "overlapraw");
  saveh5(c.get(), "amps");

  Electronic::phase_match(U, phases);
  saveh5(U, "overlap");
  T = util::logmat_unitary(U) / dtc;
  V = arma::diagmat(energy);

  // compute max time step for electronic propagation (Jain 2016 eqs. 20, 21)
  {
    double dtq_ =
      std::min(dtc/20, // make sure 20dtq < dtc
               std::min(dtc,
                        std::min(0.02 / T.max(),
                                 0.02 / arma::max( V.diag() - arma::mean(V.diag()) )
                                 )
                        ));

    dtq = dtc / std::round(dtc / dtq_);
  }

  const size_t n_steps = (size_t) dtc / dtq;
  const std::complex<double> I(0,1);

  // Propagate electronic coefficients
  for (size_t nt = 0; nt < n_steps; nt++){
    c.advance_exact(V - I*T, dtq);
  }

  return energy(active_state);
}