// Copyright András Vukics 2006–2020. Distributed under the Boost Software License, Version 1.0. (See accompanying file LICENSE.txt)

#ifndef   CPPQEDELEMENTS_UTILS_AVERAGINGUTILS_TCC_INCLUDED
#define   CPPQEDELEMENTS_UTILS_AVERAGINGUTILS_TCC_INCLUDED

#include "AveragingUtils.h"

#include "DensityOperator.h"
#include "LazyDensityOperator.h"
#include "NegPT.tcc"

#include "Algorithm.h"
#include "MathExtensions.h"
#include "MultiIndexIterator.h"

#include <boost/bind.hpp>

#include <boost/range/adaptor/transformed.hpp>

#include <algorithm>

template<int RANK>
auto
ReducedDensityOperator<RANK>::helper(const Dimensions& dim, bool offDiagonals, const KeyLabels& subsequent) -> const KeyLabels
{
  typedef cpputils::MultiIndexIterator<RANK> Iterator;
  using std::ostringstream;

  KeyLabels res;
  Iterator i(dim-1,cpputils::mii::begin);

  // Diagonals
  for (; i!=i.getEnd(); ++i) {
    ostringstream ss;
    ss<<"rho_"<<*i<<';'<<*i;
    res.push_back(ss.str());
  }

  if (!offDiagonals) return res;

  // Offdiagonals
  for (i.setToBegin(); i!=i.getEnd(); ++i)
    for (Iterator j(std::next(i)); j!=j.getEnd(); ++j) {
      {
        ostringstream ss;
        ss<<"real["<<"rho_"<<*i<<';'<<*j<<']';
        res.push_back(ss.str());
      }
      {
        ostringstream ss;
        ss<<"imag["<<"rho_"<<*i<<';'<<*j<<']';
        res.push_back(ss.str());
      }
  }

  res.insert(res.end(),subsequent.begin(),subsequent.end());

  return res;

}

template<int RANK>
ReducedDensityOperator<RANK>::ReducedDensityOperator(const std::string& label, const Dimensions& dim, bool offDiagonals, const KeyLabels& subsequent) :
  DimensionsBookkeeper<RANK>(dim),
  Base(label,helper(getDimensions(),offDiagonals,subsequent)),
  offDiagonals_(offDiagonals)
{
}

template<int RANK>
const typename ReducedDensityOperator<RANK>::Averages
ReducedDensityOperator<RANK>::average_v(structure::NoTime, const LazyDensityOperator& matrix) const
{
  return deflate(matrix,offDiagonals_);
}

template<int RANK, typename V>
const typename ReducedDensityOperatorNegativity<RANK,V>::Averages
ReducedDensityOperatorNegativity<RANK,V>::average_v(structure::NoTime t, const LazyDensityOperator& matrix) const
{
  Averages res(Base::average_v(t,matrix));
  res.resize(res.size()+1);
  return res;
}

template<int RANK, typename V>
void ReducedDensityOperatorNegativity<RANK,V>::process_v(Averages& averages) const
{
  quantumdata::DensityOperator<RANK> rho(getDimensions());
  linalg::CMatrix matrix(rho.matrixView()); // references the same data
  const size_t dim=getTotalDimension();

  {
    int idx=0;
    for (int i=0; i<dim; ++i, ++idx) matrix(i,i)=averages(idx);
    for (int i=0; i<dim; ++i) for (int j=i+1; j<dim; ++j, idx+=2) matrix(j,i)=conj(matrix(i,j)=dcomp(averages(idx),averages(idx+1)));
  }

  averages(averages.size()-1)=
#ifndef   DO_NOT_USE_FLENS
    quantumdata::negPT(rho,V())
#else  // DO_NOT_USE_FLENS
    0
#endif // DO_NOT_USE_FLENS
    ;

}

template<int RANK, bool IS_TIME_DEPENDENT>
averagingUtils::Collecting<RANK,IS_TIME_DEPENDENT>::Collecting(const Collection& collection)
  : Base(collection.front().getTitle(),
         cpputils::concatenateGrow<KeyLabels>( collection | boost::adaptors::transformed(bind(&Element::getLabels,_1)) )),
    collection_(collection.clone())
{}

template<int RANK, bool IS_TIME_DEPENDENT>
averagingUtils::Collecting<RANK,IS_TIME_DEPENDENT>::Collecting(const Collecting& collecting)
  : Base(collecting),
    collection_(collecting.collection_.clone())
{}

namespace averagingUtils { namespace details {

double convert(structure::OneTime t) {return t ;}
double convert(structure:: NoTime  ) {return 0.;}

} } // averagingUtils::details

template<int RANK, bool IS_TIME_DEPENDENT>
auto
averagingUtils::Collecting<RANK,IS_TIME_DEPENDENT>::average_v(Time t, const LazyDensityOperator& matrix) const -> const Averages
{
  Averages res(this->nAvr()); res=0;
  return cpputils::concatenate( collection_ | boost::adaptors::transformed(bind(&Element::average,_1,details::convert(t),boost::cref(matrix))) , res);

}

template<int RANK, bool IS_TIME_DEPENDENT>
void
averagingUtils::Collecting<RANK,IS_TIME_DEPENDENT>::process_v(Averages& avr) const
{
  struct Helper
  {
    static void doIt(const Element& eav, Averages& avr, ptrdiff_t& l, ptrdiff_t& u)
    {
      using blitz::Range;
      if ((u=l+eav.nAvr())>l) {
        Averages temp(avr(Range(l+1,u)));
        eav.process(temp);
      }
      std::swap(l,u);
    }
  };

  ptrdiff_t l=-1, u=0;
  for_each(collection_,boost::bind(Helper::doIt,_1,avr,l,u));
}

template<int RANKFROM, int RANKTO, bool IS_TIME_DEPENDENT>
auto
averagingUtils::Transferring<RANKFROM,RANKTO,IS_TIME_DEPENDENT>::average_v(Time t, const LazyDensityOperator&) const -> const Averages
{
  return averaged_->average(details::convert(t),ldo_);
}

/* The earlier solution fails with the C++11 standard due to a "bug" in Boost.Assign, cf. https://svn.boost.org/trac/boost/ticket/7364

template<int RANK>
struct ReducedDensityOperator<RANK>::Helper
{
  typedef cpputils::MultiIndexIterator<RANK> Iterator;

  Helper(const Dimensions& dim) : i_(Dimensions(size_t(0)),dim-1,cpputils::mii::begin), j_(i_), real_(true), offDiagonals_(false) {++i_; ++j_;}

  Helper() : i_(Dimensions(size_t(0)),Dimensions(size_t(0)),cpputils::mii::begin), j_(i_), real_(true), offDiagonals_(false) {}

  const std::string operator()()
  {
    using namespace std;
    stringstream ss(stringstream::out);

    // Diagonals
    if (!offDiagonals_ && i_!=i_.getEnd()) {
      ss<<"rho_"<<*i_<<';'<<*i_;
      ++i_;
    }
    else if (!offDiagonals_ && i_==i_.getEnd()) {
      offDiagonals_=true;
      i_.setToBegin();
    }

    // Offdiagonals
    if (offDiagonals_) {
      ss<<(real_ ? "real[" : "imag[")<<"rho_"<<*i_<<','<<*j_<<']';
      if (real_)
        real_=false;
      else {
        real_=true;
        ++j_;
        if (j_==j_.getEnd())
          ++(j_=++i_);
      }
    }

    return ss.str();
  }

private:
  Iterator i_, j_;
  bool real_, offDiagonals_;

};

template<int RANK>
ReducedDensityOperator<RANK>::ReducedDensityOperator(const std::string& label, const Dimensions& dim, bool offDiagonals, const KeyLabels& subsequent) :
  DimensionsBookkeeper<RANK,true>(dim),
  Base(label,boost::assign::repeat_fun((offDiagonals ? mathutils::sqr(getTotalDimension()) : getTotalDimension())-1,
                                       Helper(getDimensions())).range(subsequent)),
  offDiagonals_(offDiagonals)
{
}

*/

#endif // CPPQEDELEMENTS_UTILS_AVERAGINGUTILS_TCC_INCLUDED


/*
 * Software License Agreement (New BSD License)
 *
 * Copyright (c) 2013, Keith Leung, Felipe Inostroza
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     * Neither the name of the Advanced Mining Technology Center (AMTC), the
 *       Universidad de Chile, nor the names of its contributors may be
 *       used to endorse or promote products derived from this software without
 *       specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE AMTC, UNIVERSIDAD DE CHILE, OR THE COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF
 * THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <boost/lexical_cast.hpp>
#include <libconfig.h++>
#include "RBPHDFilter.hpp"
#include <stdio.h>

/**
 * \class Simulator2d
 * \brief A 2d SLAM Simulator using the RB-PHD Filter
 * \author Keith Leung
 */
class Simulator2d{

public:

  EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

  Simulator2d(){
    pFilter_ = NULL;
  }

  ~Simulator2d(){

    if(pFilter_ != NULL){
      delete pFilter_;
    }

  }

  /** Read the simulator configuration file */
  bool readConfigFile(){

    libconfig::Config cfg;
    const char* fileName= "cfg/simulator2d.cfg";
    try{
      cfg.readFile( fileName );
    }catch( libconfig::FileIOException &ex){
      printf("\nCannot read file: %s\n\n", fileName);
      return false;
    }catch( libconfig::ParseException &ex){
      const char* error = ex.getError();
      int line = ex.getLine();
      printf("\n%s LINE %d\n\n", error, line);
      return false;
    }
    kMax_ = cfg.lookup("timesteps");
    dT_ = cfg.lookup("sec_per_timestep");

    nSegments_ = cfg.lookup("Trajectory.nSegments");
    max_dx_ = cfg.lookup("Trajectory.max_dx_per_sec");
    max_dy_ = cfg.lookup("Trajectory.max_dy_per_sec");
    max_dz_ = cfg.lookup("Trajectory.max_dz_per_sec");
    min_dx_ = cfg.lookup("Trajectory.min_dx_per_sec");
    vardx_ = cfg.lookup("Trajectory.vardx");
    vardy_ = cfg.lookup("Trajectory.vardy");
    vardz_ = cfg.lookup("Trajectory.vardz");

    nLandmarks_ = cfg.lookup("Landmarks.nLandmarks");
    varlmx_ = cfg.lookup("Landmarks.varlmx");
    varlmy_ = cfg.lookup("Landmarks.varlmy");

    rangeLimitMax_ = cfg.lookup("Measurement.rangeLimitMax");
    rangeLimitMin_ = cfg.lookup("Measurement.rangeLimitMin");
    rangeLimitBuffer_ = cfg.lookup("Measurement.rangeLimitBuffer");
    Pd_ = cfg.lookup("Measurement.probDetection");
    c_ = cfg.lookup("Measurement.clutterIntensity");
    varzr_ = cfg.lookup("Measurement.varzr");
    varzb_ = cfg.lookup("Measurement.varzb");

    nParticles_ = cfg.lookup("Filter.nParticles");
    pNoiseInflation_ = cfg.lookup("Filter.processNoiseInflationFactor");
    zNoiseInflation_ = cfg.lookup("Filter.measurementNoiseInflationFactor");
    innovationRangeThreshold_ = cfg.lookup("Filter.innovationRangeThreshold");
    birthGaussianWeight_ = cfg.lookup("Filter.birthGaussianWeight");
    newGaussianCreateInnovMDThreshold_ = cfg.lookup("Filter.newGaussianCreateInnovMDThreshold");
    importanceWeightingMeasurementLikelihoodMDThreshold_ = cfg.lookup("Filter.importanceWeightingMeasurementLikelihoodMDThreshold");
    effNParticleThreshold_ = cfg.lookup("Filter.effectiveNumberOfParticlesThreshold");
    minInterSampleTimesteps_ = cfg.lookup("Filter.minInterSampleTimesteps");
    gaussianMergingThreshold_ = cfg.lookup("Filter.gaussianMergingThreshold");
    gaussianMergingCovarianceInflationFactor_ = cfg.lookup("Filter.gaussianMergingCovarianceInflationFactor");
    gaussianPruningThreshold_ = cfg.lookup("Filter.gaussianPruningThreshold");
    reportTimingInfo_ = cfg.lookup("Filter.reportTimingInfo");
    importanceWeightingEvalPointCount_ = cfg.lookup("Filter.importanceWeightingEvalPointCount");
    useClusterProcess_ = cfg.lookup("Filter.useClusterProcess");

    logToFile_ = cfg.lookup("Computation.logToFile");

    return true;
  }

  /** Generate a random trajectory in 2d space */
  void generateTrajectory(int randSeed = 0){

    printf("Generating trajectory with random seed = %d\n", randSeed);
    srand48( randSeed);

    int seg = 0;
    OdometryMotionModel2d::TState::Mat Q;
    Q << vardx_, 0, 0, 0, vardy_, 0, 0, 0, vardz_;
    Q = dT_ * Q * dT_;
    OdometryMotionModel2d motionModel(Q);
    OdometryMotionModel2d::TInput input_k(0, 0, 0, 0, 0, 0, 0);
    OdometryMotionModel2d::TState pose_k(0, 0, 0, 0, 0, 0, 0);
    OdometryMotionModel2d::TState pose_km(0, 0, 0, 0, 0, 0, 0);
    groundtruth_displacement_.reserve( kMax_ );
    groundtruth_pose_.reserve( kMax_ );
    groundtruth_displacement_.push_back(input_k);
    groundtruth_pose_.push_back(pose_k);

    for( int k = 1; k < kMax_; k++ ){

      if( k <= 50 ){
	double dx = 0;
	double dy = 0;
	double dz = 0;
	input_k = OdometryMotionModel2d::TInput(dx, dy, dz,
						0, 0, 0, k);
      }else if( k >= kMax_ / nSegments_ * seg ){
	seg++;
	double dx = drand48() * max_dx_ * dT_;
	while( dx < min_dx_ * dT_ ){
	  dx = drand48() * max_dx_ * dT_;
	}
	double dy = (drand48() * max_dy_ * 2 - max_dy_) * dT_;
	double dz = (drand48() * max_dz_ * 2 - max_dz_) * dT_;
	input_k = OdometryMotionModel2d::TInput(dx, dy, dz,
						Q(0,0), Q(1,1), Q(2,2), k);
      }

      groundtruth_displacement_.push_back(input_k);
      groundtruth_displacement_.back().setTime(k);

      OdometryMotionModel2d::TState x_k;
      motionModel.step(x_k, groundtruth_pose_[k-1], input_k);
      groundtruth_pose_.push_back( x_k );
      groundtruth_pose_.back().setTime(k);

    }

  }

  /** Generate odometry measurements */
  void generateOdometry(){

    odometry_.reserve( kMax_ );
    OdometryMotionModel2d::TInput zero;
    OdometryMotionModel2d::TInput::Vec u0;
    u0.setZero();
    zero.set(u0, 0);
    odometry_.push_back( zero );

    OdometryMotionModel2d::TState::Mat Q;
    Q << vardx_, 0, 0, 0, vardy_, 0, 0, 0, vardz_;
    OdometryMotionModel2d motionModel(Q);
    deadReckoning_pose_.reserve( kMax_ );
    deadReckoning_pose_.push_back( groundtruth_pose_[0] );

    for( int k = 1; k < kMax_; k++){

      OdometryMotionModel2d::TInput in = groundtruth_displacement_[k];
      in.setCov(Q);
      OdometryMotionModel2d::TInput out;
      RandomVecMathTools<OdometryMotionModel2d::TInput>::sample(in, out);

      odometry_.push_back( out );

      OdometryMotionModel2d::TState p;
      motionModel.step(p, deadReckoning_pose_[k-1], odometry_[k]);
      p.setTime(k);
      deadReckoning_pose_.push_back( p );
    }

  }

  /** Generate landmarks */
  void generateLandmarks(){

    RangeBearingModel measurementModel( varzr_, varzb_);
    RangeBearingModel::TPose pose;

    groundtruth_landmark_.reserve(nLandmarks_);

    int nLandmarksCreated = 0;
    for( int k = 1; k < kMax_; k++ ){

      if( k >= kMax_ / nLandmarks_ * nLandmarksCreated){

	RangeBearingModel::TPose pose;
	RangeBearingModel::TMeasurement measurementToCreateLandmark;
	RangeBearingModel::TMeasurement::Vec z;
	double r = drand48() * rangeLimitMax_;
	double b = drand48() * 2 * PI;
	z << r, b;
	measurementToCreateLandmark.set(z);
	RangeBearingModel::TLandmark lm;

	measurementModel.inverseMeasure( groundtruth_pose_[k],
					 measurementToCreateLandmark,
					 lm);

	groundtruth_landmark_.push_back(lm);

	nLandmarksCreated++;

      }

    }

  }

  /** Generate landmark measurements */
  void generateMeasurements(){

    RangeBearingModel measurementModel( varzr_, varzb_ );
    RangeBearingModel::TMeasurement::Mat R;
    measurementModel.getNoise(R);
    measurementModel.config.rangeLimMax_ = rangeLimitMax_;
    measurementModel.config.rangeLimMin_ = rangeLimitMin_;
    measurementModel.config.probabilityOfDetection_ = Pd_;
    measurementModel.config.uniformClutterIntensity_ = c_;
    double meanClutter = measurementModel.clutterIntensityIntegral();

    double expNegMeanClutter = exp( -meanClutter );
    double poissonPmf[100];
    double poissonCmf[100];
    double mean_pow_i = 1;
    double i_factorial = 1;
    poissonPmf[0] = expNegMeanClutter;
    poissonCmf[0] = poissonPmf[0];
    for( int i = 1; i < 100; i++){
      mean_pow_i *= meanClutter;
      i_factorial *= i;
      poissonPmf[i] = mean_pow_i / i_factorial * expNegMeanClutter;
      poissonCmf[i] = poissonCmf[i-1] + poissonPmf[i];
    }

    for( int k = 1; k < kMax_; k++ ){

      groundtruth_pose_[k];

      // Real detections
      for( int m = 0; m < groundtruth_landmark_.size(); m++){

	bool success;
	RangeBearingModel::TMeasurement z_m_k;
	success = measurementModel.sample( groundtruth_pose_[k],
					   groundtruth_landmark_[m],
					   z_m_k);
	if(success){

	  if(z_m_k.get(0) <= rangeLimitMax_ && z_m_k.get(0) >= rangeLimitMin_ && drand48() <= Pd_){
	    z_m_k.setTime(k);
	    z_m_k.setCov(R);
	    measurements_.push_back( z_m_k );
	  }
	}

      }

      // False alarms
      double randomNum = drand48();
      int nClutterToGen = 0;
      while( randomNum > poissonCmf[ nClutterToGen ] ){
	nClutterToGen++;
      }
      for( int i = 0; i < nClutterToGen; i++ ){

	double r = drand48() * rangeLimitMax_;
	while(r < rangeLimitMin_)
	  r = drand48() * rangeLimitMax_;
	double b = drand48() * 2 * PI - PI;
	RangeBearingModel::TMeasurement z_clutter;
	RangeBearingModel::TMeasurement::Vec z;
	z << r, b;
	z_clutter.set(z, k);
	measurements_.push_back(z_clutter);

      }

    }

  }

  /** Data Logging */
  void exportSimData(){

    double t;

    FILE* pGTPoseFile;
    pGTPoseFile = fopen("data/gtPose.dat", "w");
    OdometryMotionModel2d::TState::Vec x;
    for(int i = 0; i < groundtruth_pose_.size(); i++){
      groundtruth_pose_[i].get(x, t);
      fprintf( pGTPoseFile, "%f   %f   %f   %f\n", t, x(0), x(1), x(2));
    }
    fclose(pGTPoseFile);

    FILE* pGTLandmarkFile;
    pGTLandmarkFile = fopen("data/gtLandmark.dat", "w");
    RangeBearingModel::TLandmark::Vec m;
    for(int i = 0; i < groundtruth_landmark_.size(); i++){
      groundtruth_landmark_[i].get(m);
      fprintf( pGTLandmarkFile, "%f   %f\n", m(0), m(1));
    }
    fclose(pGTLandmarkFile);

    FILE* pOdomFile;
    pOdomFile = fopen("data/odometry.dat","w");
    OdometryMotionModel2d::TInput::Vec u;
    for(int i = 0; i < odometry_.size(); i++){
      odometry_[i].get(u, t);
      fprintf( pOdomFile, "%f   %f   %f   %f\n", t, u(0), u(1), u(2));
    }
    fclose(pOdomFile);

    FILE* pMeasurementFile;
    pMeasurementFile = fopen("data/measurement.dat", "w");
    RangeBearingModel::TMeasurement::Vec z;
    for(int i = 0; i < measurements_.size(); i++){
      measurements_[i].get(z, t);
      fprintf( pMeasurementFile, "%f   %f   %f\n", t, z(0), z(1) );
    }
    fclose(pMeasurementFile);

    FILE* pDeadReckoningFile;
    pDeadReckoningFile = fopen("data/deadReckoning.dat", "w");
    OdometryMotionModel2d::TState::Vec odo;
    for(int i = 0; i < deadReckoning_pose_.size(); i++){
      deadReckoning_pose_[i].get(odo, t);
      fprintf( pDeadReckoningFile, "%f   %f   %f   %f\n", t, odo(0), odo(1), odo(2));
    }
    fclose(pDeadReckoningFile);

  }

  /** RB-PHD Filter Setup */
  void setupRBPHDFilter(){

    pFilter_ = new RBPHDFilter<OdometryMotionModel2d,
			       StaticProcessModel<Landmark2d>,
			       RangeBearingModel,
			       RangeBearingKalmanFilter>( nParticles_ );

    // configure robot motion model
    OdometryMotionModel2d::TState::Mat Q;
    Q << vardx_, 0, 0, 0, vardy_, 0, 0, 0, vardz_;
    Q *= pNoiseInflation_;
    pFilter_->getProcessModel()->setNoise(Q);

    // configure landmark process model
    Landmark2d::Mat Q_lm;
    Q_lm << varlmx_, 0, 0, varlmy_;
    pFilter_->getLmkProcessModel()->setNoise(Q_lm);

    // configure measurement model
    RangeBearingModel::TMeasurement::Mat R;
    R << varzr_, 0, 0, varzb_;
    R *= zNoiseInflation_;
    pFilter_->getMeasurementModel()->setNoise(R);
    pFilter_->getMeasurementModel()->config.probabilityOfDetection_ = Pd_;
    pFilter_->getMeasurementModel()->config.uniformClutterIntensity_ = c_;
    pFilter_->getMeasurementModel()->config.rangeLimMax_ = rangeLimitMax_;
    pFilter_->getMeasurementModel()->config.rangeLimMin_ = rangeLimitMin_;
    pFilter_->getMeasurementModel()->config.rangeLimBuffer_ = rangeLimitBuffer_;

    // configure the Kalman filter for landmark updates
    pFilter_->getKalmanFilter()->config.rangeInnovationThreshold_ = innovationRangeThreshold_;

    // configure the filter
    pFilter_->config.birthGaussianWeight_ = birthGaussianWeight_;
    pFilter_->setEffectiveParticleCountThreshold(effNParticleThreshold_);
    pFilter_->config.minInterSampleTimesteps_ = minInterSampleTimesteps_;
    pFilter_->config.newGaussianCreateInnovMDThreshold_ = newGaussianCreateInnovMDThreshold_;
    pFilter_->config.importanceWeightingMeasurementLikelihoodMDThreshold_ = importanceWeightingMeasurementLikelihoodMDThreshold_;
    pFilter_->config.importanceWeightingEvalPointCount_ = importanceWeightingEvalPointCount_;
    pFilter_->config.gaussianMergingThreshold_ = gaussianMergingThreshold_;
    pFilter_->config.gaussianMergingCovarianceInflationFactor_ = gaussianMergingCovarianceInflationFactor_;
    pFilter_->config.gaussianPruningThreshold_ = gaussianPruningThreshold_;
    pFilter_->config.reportTimingInfo_ = reportTimingInfo_;
    pFilter_->config.useClusterProcess_ = useClusterProcess_;
  }

  /** Run the simulator */
  void run(){

    printf("Running simulation\n\n");

    //////// Initialization at first timestep //////////

    FILE* pParticlePoseFile;
    if(logToFile_){
      pParticlePoseFile = fopen("data/particlePose.dat", "w");
      fprintf( pParticlePoseFile, "Timesteps: %d\n", kMax_);
      fprintf( pParticlePoseFile, "nParticles: %d\n\n", pFilter_->getParticleCount());
    }
    FILE* pLandmarkEstFile;
    if(logToFile_){
      pLandmarkEstFile = fopen("data/landmarkEst.dat", "w");
      fprintf( pLandmarkEstFile, "Timesteps: %d\n", kMax_);
      fprintf( pLandmarkEstFile, "nParticles: %d\n\n", pFilter_->getParticleCount());
    }
    OdometryMotionModel2d::TState x_i;
    int zIdx = 0;

    if(logToFile_){
      fprintf( pParticlePoseFile, "k = 0\n");
      for(int i = 0; i < pFilter_->getParticleCount(); i++){
	pFilter_->getParticleSet()->at(i)->getPose(x_i);
	fprintf( pParticlePoseFile, "%f   %f   %f   1.0\n", x_i.get(0), x_i.get(1), x_i.get(2));
      }
      for(int i = 0; i < pFilter_->getParticleCount(); i++){
	fprintf( pLandmarkEstFile, "Timestep: %d\tParticle: %d\tMap Size: %d\n", 0, i, 0);
      }
    }

    boost::timer::auto_cpu_timer *stepTimer = NULL;
    boost::timer::auto_cpu_timer *processTimer = new boost::timer::auto_cpu_timer(6, "Total run time: %ws\n");

    /////////// Run simulator from k = 1 to kMax_ /////////

    for(int k = 1; k < kMax_; k++){

      if(reportTimingInfo_){
	stepTimer = new boost::timer::auto_cpu_timer(6, "Step time: %ws\n");
      }

      if( k % 1 == 0)
	printf("k = %d\n", k);

      if(logToFile_){
	fprintf( pParticlePoseFile, "k = %d\n", k);
      }

      ////////// Prediction Step //////////
      pFilter_->predict( odometry_[k], k );

      if( k <= 100){
	for( int i = 0; i < nParticles_; i++)
	  pFilter_->setParticlePose(0, groundtruth_pose_[k]);
      }

      // Prepare measurement vector for update
      std::vector<RangeBearingModel::TMeasurement> Z;
      double kz = measurements_[ zIdx ].getTime();
      while( kz == k ){
	Z.push_back( measurements_[zIdx] );
	zIdx++;
	if(zIdx >= measurements_.size())
	  break;
	kz = measurements_[ zIdx ].getTime();
      }

      ////////// Update Step //////////
      pFilter_->update(Z, k);

      // Log particle poses
      if(logToFile_){
	for(int i = 0; i < pFilter_->getParticleCount(); i++){
	  pFilter_->getParticleSet()->at(i)->getPose(x_i);
	  double w = pFilter_->getParticleSet()->at(i)->getWeight();
	  fprintf( pParticlePoseFile, "%f   %f   %f   %f\n", x_i.get(0), x_i.get(1), x_i.get(2), w);
	}
	fprintf( pParticlePoseFile, "\n");
      }

      // Log landmark estimates
      if(logToFile_){
	for(int i = 0; i < pFilter_->getParticleCount(); i++){
	  int mapSize = pFilter_->getGMSize(i);
	  fprintf( pLandmarkEstFile, "Timestep: %d\tParticle: %d\tMap Size: %d\n", k, i, mapSize);
	  for( int m = 0; m < mapSize; m++ ){
	    RangeBearingModel::TLandmark::Vec u;
	    RangeBearingModel::TLandmark::Mat S;
	    double w;
	    pFilter_->getLandmark(i, m, u, S, w);

	    fprintf( pLandmarkEstFile, "%f   %f      ", u(0), u(1));
	    fprintf( pLandmarkEstFile, "%f   %f   ", S(0,0), S(0,1));
	    fprintf( pLandmarkEstFile, "%f   %f   ", S(1,0), S(1,1));
	    fprintf( pLandmarkEstFile, "   %f\n", w );
	  }
	  fprintf( pLandmarkEstFile, "\n");
	}
	fprintf( pLandmarkEstFile, "\n");
      }

      if(reportTimingInfo_){
	delete stepTimer;
	stepTimer = NULL;
      }

    }

    if(reportTimingInfo_){
      delete processTimer;
      processTimer = NULL;
    }

    if(logToFile_){
      fclose(pParticlePoseFile);
      fclose(pLandmarkEstFile);
    }
  }

private:

  int kMax_; /**< number of timesteps */
  double dT_;

  // Trajectory
  int nSegments_;
  double max_dx_;
  double max_dy_;
  double max_dz_;
  double min_dx_;
  double vardx_;
  double vardy_;
  double vardz_;
  std::vector<OdometryMotionModel2d::TInput> groundtruth_displacement_;
  std::vector<OdometryMotionModel2d::TState> groundtruth_pose_;
  std::vector<OdometryMotionModel2d::TInput> odometry_;
  std::vector<OdometryMotionModel2d::TState> deadReckoning_pose_;

  // Landmarks
  int nLandmarks_;
  std::vector<RangeBearingModel::TLandmark> groundtruth_landmark_;
  double varlmx_;
  double varlmy_;

  // Range-Bearing Measurements
  double rangeLimitMax_;
  double rangeLimitMin_;
  double rangeLimitBuffer_;
  double Pd_;
  double c_;
  double varzr_;
  double varzb_;
  std::vector<RangeBearingModel::TMeasurement> measurements_;

  // Filters
  RangeBearingKalmanFilter kf_;
  RBPHDFilter<OdometryMotionModel2d,
	      StaticProcessModel<Landmark2d>,
	      RangeBearingModel,
	      RangeBearingKalmanFilter> *pFilter_;
  int nParticles_;
  double pNoiseInflation_;
  double zNoiseInflation_;
  double innovationRangeThreshold_;
  double birthGaussianWeight_;
  double newGaussianCreateInnovMDThreshold_;
  double importanceWeightingMeasurementLikelihoodMDThreshold_;
  double effNParticleThreshold_;
  int minInterSampleTimesteps_;
  double gaussianMergingThreshold_;
  double gaussianMergingCovarianceInflationFactor_;
  double gaussianPruningThreshold_;
  int importanceWeightingEvalPointCount_;
  bool reportTimingInfo_;
  bool useClusterProcess_;

  bool logToFile_;
};

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

  int initRandSeed = 0;
  if( argc == 2 ){
    initRandSeed = boost::lexical_cast<int>(argv[1]);
  }

  Simulator2d sim;
  if( !sim.readConfigFile() ){
    return -1;
  }
  sim.generateTrajectory( initRandSeed );
  sim.generateOdometry();
  sim.generateLandmarks();
  sim.generateMeasurements();
  sim.exportSimData();

  sim.setupRBPHDFilter();

  sim.run();

  return 0;

}


#include <memory>
#include <chrono>
#include <tuple>
#include <unistd.h>
#include <gflags/gflags.h>
#include <string.h>
#include <assert.h>
#include <Eigen/StdVector>

#include "drake/solvers/snopt_solver.h"
#include "drake/systems/analysis/simulator.h"
#include "drake/systems/framework/diagram.h"
#include "drake/systems/framework/diagram_builder.h"
#include "drake/multibody/parsing/parser.h"
#include <drake/multibody/inverse_kinematics/inverse_kinematics.h>
#include "drake/geometry/geometry_visualization.h"
#include "drake/solvers/solve.h"

#include "common/find_resource.h"
#include "systems/trajectory_optimization/dircon/dircon.h"
#include "multibody/kinematic/world_point_evaluator.h"
#include "solvers/nonlinear_constraint.h"
#include "multibody/kinematic/distance_evaluator.h"
#include "multibody/multibody_utils.h"
#include "multibody/visualization_utils.h"
#include "multibody/kinematic/kinematic_constraints.h"
#include "solvers/nonlinear_cost.h"
#include "examples/Spirit/spirit_utils.h"

using drake::AutoDiffXd;
using drake::multibody::MultibodyPlant;
using drake::geometry::SceneGraph;
using drake::multibody::Parser;
using drake::trajectories::PiecewisePolynomial;

using Eigen::Vector3d;
using Eigen::VectorXd;
using Eigen::Matrix3d;
using Eigen::MatrixXd;

namespace dairlib {

using systems::trajectory_optimization::DirconModeSequence;
using systems::trajectory_optimization::DirconMode;
using systems::trajectory_optimization::Dircon;
using systems::trajectory_optimization::KinematicConstraintType;

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

/// Get a nominal Spirit Stand (i.e. zero hip ad/abduction motor torque, toes below motors) for initializing
template <typename T>
void nominalSpiritStand( MultibodyPlant<T>& plant, Eigen::VectorXd& xState, double height){
  //Get joint name dictionaries
  auto positions_map = dairlib::multibody::makeNameToPositionsMap(plant);
  auto velocities_map = dairlib::multibody::makeNameToVelocitiesMap(plant);

  // Initialize state vector and add orienation and goal height
  xState = Eigen::VectorXd::Zero(plant.num_positions() + plant.num_velocities());
  xState(positions_map.at("base_qw")) = 1;
  xState(positions_map.at("base_z")) = height;

  //Calculate the inverse kinematics to get the joint angles for toe placement
  double upperLegLength = 0.206; // length of the upper leg link
  double hipLength = 0.10098; //absOffset ToeLocation
  double hip2toeZLength = sqrt(height*height - hipLength*hipLength);// length of lower legs (2dof)
  double theta1 = asin(hip2toeZLength/(2*upperLegLength )) ; //upper angle
  double theta2 = 2*theta1 ; //lower angle
  double alpha = asin(hipLength/(height)); //hip angle

  int upperInd, kneeInd, hipInd;
  int mirroredFlag;
  for (int j = 0; j < 4; j++){
    upperInd = j * 2;
    kneeInd = j * 2 + 1;
    hipInd = j + 8;
    xState(positions_map.at("joint_" + std::to_string(upperInd))) = theta1 ;
    xState(positions_map.at("joint_" + std::to_string(kneeInd))) = theta2 ;
    mirroredFlag = -1;
    if ( hipInd > 9 ){
      mirroredFlag = 1;
    }
    xState(positions_map.at("joint_" + std::to_string(hipInd))) = mirroredFlag * alpha;
  }
}

void ikSpiritStand(drake::multibody::MultibodyPlant<double>& plant,
                   Eigen::VectorXd& xState,
                   Eigen::Matrix<bool,1,4> contactSequence,
                   double com_height,
                   double leg_height,
                   double roll,
                   double pitch,
                   double eps){
  drake::multibody::InverseKinematics ik(plant);
  const auto& world_frame = plant.world_frame();
  const auto& body_frame  = plant.GetFrameByName("body");

  // Position of the com
  Vector3d pos = {0, 0, com_height};
  // Orientation of body
  drake::math::RotationMatrix<double> orientation;
  orientation = drake::math::RotationMatrix<double>::MakeXRotation(-roll) *
                drake::math::RotationMatrix<double>::MakeYRotation(-pitch) ;

  const double n_x = plant.num_positions() + plant.num_velocities();
  const double n_q = plant.num_positions();

  // Constrain body location and pose
  ik.AddPositionConstraint(body_frame, Vector3d(0, 0, 0), world_frame,
                           pos,
                           pos);
  ik.AddOrientationConstraint(body_frame, orientation,
                              world_frame, drake::math::RotationMatrix<double>(), 0);

  // Constrain toes
  for(int toe = 0; toe < 4; toe ++){
    constrainToe(plant, ik, toe, contactSequence[toe], orientation, com_height, leg_height, eps);
  }

  // Add initial guess and solve
  Eigen::VectorXd initial_guess;
  dairlib::nominalSpiritStand(plant, initial_guess, com_height);
  ik.get_mutable_prog()->SetInitialGuess(ik.q(), initial_guess.head(n_q));
  const auto result = Solve(ik.prog());
  const auto q_sol = result.GetSolution(ik.q());
  std::cout << (result.is_success() ? "IK Success" : "IK Fail") << std::endl;

  // Return
  VectorXd x_const = VectorXd::Zero(n_x);
  x_const.head(n_q) = q_sol;
  xState = x_const;
}

void constrainToe(drake::multibody::MultibodyPlant<double> & plant,
                  drake::multibody::InverseKinematics &ik,
                  int toe, bool inContact, drake::math::RotationMatrix<double> orientation,
                  double com_height, double leg_height, double eps){
  // Get frames
  const auto& toe_frame = dairlib::getSpiritToeFrame(plant, toe);
  const auto& world_frame = plant.world_frame();
  const auto& body_frame  = plant.GetFrameByName("body");
  Vector3d eps_vec = {eps, eps, eps};

  // Chose distance between toe and com
  const double hip_width = 0.07;
  const double body_length = 0.2263;
  Vector3d toe_offset;
  switch (toe) {
    case 0:
      toe_offset = {body_length, hip_width, -abs(leg_height)};
      break;
    case 1:
      toe_offset = {-body_length, hip_width, -abs(leg_height)};
      break;
    case 2:
      toe_offset = {body_length, -hip_width, -abs(leg_height)};
      break;
    case 3:
      toe_offset = {-body_length, -hip_width, -abs(leg_height)};
      break;
    default:
      break;
  }

  if(inContact){
    // Toe is in contact with ground

    // Constrain toe to be under hip with distance constraint
    toe_offset(2) = 0; // vector between body and hip in body frame
    // Location hip in world frame
    Vector3d hip_in_world = orientation.transpose() * toe_offset + Vector3d({0, 0, abs(com_height)});
    // Vector between hip and toe in world frame
    Vector3d  hip_to_toe_world = {0, 0, -hip_in_world(2)};
    // Vector between body and toe in body frame
    Vector3d body_to_toe_body = toe_offset + orientation * hip_to_toe_world;
    //ik.AddPositionConstraint(body_frame, body_to_toe_body, toe_frame, -eps_vec, eps_vec);

    // Toe is in contact with ground
    ik.AddPositionConstraint(toe_frame, Vector3d(0, 0, 0), world_frame,
                             toe_offset,
                             toe_offset);

  }
  else{
    // Toe is under hip in body frame
    ik.AddPositionConstraint(body_frame, toe_offset, toe_frame,
                             Vector3d(0, 0, 0),
                             Vector3d(0, 0, 0));
  }
}

template <typename T>
void nominalSpiritStandConstraint(
    drake::multibody::MultibodyPlant<T>& plant,
    dairlib::systems::trajectory_optimization::Dircon<T>& trajopt,
    double height,
    std::vector<int> knotPoints,
    double eps){
  int nx = plant.num_positions() + plant.num_velocities();
  VectorXd x_state(nx);
  dairlib::nominalSpiritStand( plant, x_state,  height); //get desired x_state

  auto positions_map = dairlib::multibody::makeNameToPositionsMap(plant);

  for(int knot_index : knotPoints){
    auto xi = trajopt.state(knot_index);
    for(int joint = 0; joint< 12; joint ++){
      std::string joint_string = "joint_"+std::to_string(joint);
      double joint_val = x_state[positions_map.at(joint_string)];
      trajopt.AddBoundingBoxConstraint(joint_val-eps, joint_val+eps, xi(positions_map.at(joint_string)));
    }
  }

}

template <typename T>
const drake::multibody::Frame<T>& getSpiritToeFrame( MultibodyPlant<T>& plant, u_int8_t toeIndex ){
  assert(toeIndex<4);
  return plant.GetFrameByName( "toe" + std::to_string(toeIndex) );
}

template <typename T>
std::unique_ptr<dairlib::multibody::WorldPointEvaluator<T>> getSpiritToeEvaluator(
                      MultibodyPlant<T>& plant,
                      const Eigen::Vector3d toePoint ,
                      u_int8_t toeIndex,
                      const Eigen::Vector3d normal ,
                      const Eigen::Vector3d offset ,
                      bool xy_active,
                      double mu
                      ){
  assert(toeIndex<4); // Check that toeIndex is an actual Spirit leg
  auto toe_eval =  std::make_unique<dairlib::multibody::WorldPointEvaluator<T>>(
        plant,
        toePoint,
        getSpiritToeFrame(plant, toeIndex ) ,
        normal,
        offset,
        xy_active );
  if(!std::isinf(mu)){
    toe_eval->set_frictional(); toe_eval->set_mu(mu);
  }
  return toe_eval;
}

template <typename T>
std::tuple<
                std::vector<std::unique_ptr<DirconMode<T>>>,
                std::vector<std::unique_ptr<dairlib::multibody::WorldPointEvaluator<T>>> ,
                std::vector<std::unique_ptr<dairlib::multibody::KinematicEvaluatorSet<T>>>
          > createSpiritModeSequence(
          MultibodyPlant<T>& plant, // multibodyPlant
          std::vector<Eigen::Matrix<bool,1,4>> modeSeqVect, // bool matrix describing toe contacts as true or false e.g. {{1,1,1,1},{0,0,0,0}} would be a full support mode and flight mode
          std::vector<int> knotpointVect, // Matrix of knot points for each mode
          std::vector<Eigen::Vector3d> normals,
          std::vector<Eigen::Vector3d> offsets,
          std::vector<double> mus,
          std::vector<double> minTs,
          std::vector<double> maxTs ){

  // std::cout<<modeSeqMat<<std::endl;
  // std::cout<<knotpointVect<<std::endl;

  const double toeRadius = 0.02;
  const Vector3d toeOffset(toeRadius,0,0); // vector to "contact point"
  int num_modes = modeSeqVect.size();
  assert( num_modes == knotpointVect.size() );
  assert( num_modes == mus.size() );

  std::vector<std::unique_ptr<dairlib::multibody::WorldPointEvaluator<T>>> toeEvals;
  std::vector<std::unique_ptr<dairlib::multibody::KinematicEvaluatorSet<T>>> toeEvalSets;
  std::vector<std::unique_ptr<DirconMode<T>>> modeVector;
  // DirconModeSequence<T> sequence = DirconModeSequence<T>(plant);

  for (int iMode = 0; iMode<num_modes; iMode++)
  {

    toeEvalSets.push_back( std::move( std::make_unique<dairlib::multibody::KinematicEvaluatorSet<T>>(plant) ));
    for ( int iLeg = 0; iLeg < 4; iLeg++ ){
      if (modeSeqVect.at(iMode)(iLeg)){
        toeEvals.push_back( std::move( getSpiritToeEvaluator(plant, toeOffset, iLeg, normals.at(iMode), offsets.at(iMode), true, mus.at(iMode) )) );
        (toeEvalSets.back())->add_evaluator(  (toeEvals.back()).get()  ); //add evaluator to the set if active //Works ish
      }
    }
    auto dumbToeEvalPtr = (toeEvalSets.back()).get() ;
    int num_knotpoints = knotpointVect.at(iMode);
    modeVector.push_back(std::move( std::make_unique<DirconMode<T>>( *dumbToeEvalPtr , num_knotpoints, minTs.at(iMode), maxTs.at(iMode) )));
    // DirconMode<T> modeDum = DirconMode<T>( *dumbToeEvalPtr , num_knotpoints );
    // sequence.AddMode(  &modeDum  ); // Add the evaluator set to the mode sequence

  }

  return {std::move(modeVector), std::move(toeEvals), std::move(toeEvalSets)};
  // return {std::move(toeEvals), std::move(toeEvalSets)};
}

template <typename T>
std::tuple<
                std::vector<std::unique_ptr<DirconMode<T>>>,
                std::vector<std::unique_ptr<dairlib::multibody::WorldPointEvaluator<T>>> ,
                std::vector<std::unique_ptr<dairlib::multibody::KinematicEvaluatorSet<T>>>
          > createSpiritModeSequence(
          MultibodyPlant<T>& plant, // multibodyPlant
          const dairlib::ModeSequenceHelper& msh ){

return createSpiritModeSequence(plant, msh.modes , msh.knots , msh.normals , msh.offsets, msh.mus, msh.minTs, msh.maxTs);
}

// **********************************************************
//   SETSPIRITJOINTLIMITS

template <typename T>
void setSpiritJointLimits(drake::multibody::MultibodyPlant<T> & plant,
                          dairlib::systems::trajectory_optimization::Dircon<T>& trajopt,
                          int iJoint,
                          double minVal,
                          double maxVal  ){

  auto positions_map = multibody::makeNameToPositionsMap(plant);
  auto velocities_map = multibody::makeNameToVelocitiesMap(plant);
  // std::cout<<"joint_" + std::to_string(iJoint)<<std::endl;
  int N_knotpoints = trajopt.N();
  for(int i = 0; i<N_knotpoints;i++){
    auto xi = trajopt.state(i);
    trajopt.AddBoundingBoxConstraint(minVal,maxVal,xi(positions_map.at("joint_" + std::to_string(iJoint))));
  }
}

template <typename T>
void setSpiritJointLimits(drake::multibody::MultibodyPlant<T> & plant,
                          dairlib::systems::trajectory_optimization::Dircon<T>& trajopt,
                          std::vector<int> iJoints,
                          std::vector<double> minVals,
                          std::vector<double> maxVals  ){
  int numJoints = iJoints.size();
  assert(numJoints==minVals.size());
  assert(numJoints==maxVals.size());
  for (int i = 0; i<numJoints; i++){
    setSpiritJointLimits( plant,
                          trajopt,
                          iJoints[i],
                          minVals[i],
                          maxVals[i] );
  }
}
template <typename T>
void setSpiritJointLimits(drake::multibody::MultibodyPlant<T> & plant,
                          dairlib::systems::trajectory_optimization::Dircon<T>& trajopt,
                          std::vector<int> iJoints,
                          double minVal,
                          double maxVal  ){

  for (int iJoint: iJoints){
    setSpiritJointLimits( plant,
                          trajopt,
                          iJoint,
                          minVal,
                          maxVal );
  }
}
template <typename T>
void setSpiritJointLimits(drake::multibody::MultibodyPlant<T> & plant,
                          dairlib::systems::trajectory_optimization::Dircon<T>& trajopt ){
  // Upper doesn't need a joint limit for now, but we may eventually want to
  // change this to help the optimizations
  double minValUpper = -2 * M_PI ;
  double maxValUpper =  2 * M_PI ;

  // Lower limits are set to 0 and pi in the URDF might be better to include
  // the few degrees that come with the body collision and to remove a few to stay
  // away from singularity
  double minValLower =  0;
  double maxValLower =  M_PI-0.2;

  // The URDF defines symmetric limits if asymmetric constraints we need to
  // add a mirror since the hips are positive in the same direction
  double abKinLimit = 0.707 ; //From the URDF
  double minValAbduc = -abKinLimit ;
  double maxValAbduc =  abKinLimit ;

  // Upper 0,2,4,6
  std::vector<int> indicesUpper = {0,2,4,6};
  // Lower 1,3,5,7
  std::vector<int> indicesLower = {1,3,5,7};
  // Hips  8,9,10,11
  std::vector<int> indicesAbduc = {8,9,10,11};

//See above: Upper doesn't need joint limit since there isnt a meaningful set
setSpiritJointLimits( plant,
                     trajopt,
                     indicesUpper,
                     minValUpper,
                     maxValUpper  );

setSpiritJointLimits( plant,
                     trajopt,
                     indicesLower,
                     minValLower ,
                     maxValLower );

setSpiritJointLimits( plant,
                     trajopt,
                     indicesAbduc,
                     minValAbduc,
                     maxValAbduc  );

}
//   \SETSPIRITJOINTLIMITS

// **********************************************************
//   SETSPIRITACTUATIONLIMITS
template <typename T>
void setSpiritActuationLimits(drake::multibody::MultibodyPlant<T> & plant,
                          dairlib::systems::trajectory_optimization::Dircon<T>& trajopt,
                          double actuatorLimit){
  auto actuators_map = multibody::makeNameToActuatorsMap(plant);

  // Upper 0,2,4,6
  std::vector<int> indicesUpper = {0,2,4,6};
  // Lower 1,3,5,7
  std::vector<int> indicesLower = {1,3,5,7};
  // Hips  8,9,10,11
  std::vector<int> indicesAbduc = {8,9,10,11};

  double mechanicalReductionKnee = 1.5;

  int N_knotpoints = trajopt.N();
  for(int iKnot = 0; iKnot<N_knotpoints;iKnot++){
    auto ui = trajopt.input(iKnot);
    for (int iMotor : indicesUpper){
      trajopt.AddBoundingBoxConstraint(
        -actuatorLimit,
         actuatorLimit,
         ui(actuators_map.at("motor_" + std::to_string(iMotor))));
    }
    for (int iMotor : indicesLower){
      trajopt.AddBoundingBoxConstraint(
        -actuatorLimit*mechanicalReductionKnee,
         actuatorLimit*mechanicalReductionKnee,
         ui(actuators_map.at("motor_" + std::to_string(iMotor))));
    }
    for (int iMotor : indicesAbduc){
      trajopt.AddBoundingBoxConstraint(
        -actuatorLimit,
         actuatorLimit,
         ui(actuators_map.at("motor_" + std::to_string(iMotor))));
    }
  }

}
//   \SETSPIRITACTUATIONLIMITS

// **********************************************************
//   SETSPIRITSYMMETRY

template <typename T>
void makeSaggitalSymmetric(
        drake::multibody::MultibodyPlant<T> & plant,
        dairlib::systems::trajectory_optimization::Dircon<T>& trajopt){

  // Get position and velocity dictionaries
  auto positions_map = multibody::makeNameToPositionsMap(plant);
  auto velocities_map = multibody::makeNameToVelocitiesMap(plant);
  int num_knotpoints = trajopt.N();
  /// Symmetry constraints
  for (int i = 0; i < num_knotpoints; i++){

    auto xi = trajopt.state(i);
    // Symmetry constraints (mirrored hips all else equal across)
    trajopt.AddLinearConstraint( xi( positions_map.at("joint_8") ) == -xi( positions_map.at("joint_10") ) );
    trajopt.AddLinearConstraint( xi( positions_map.at("joint_9") ) == -xi( positions_map.at("joint_11") ) );
    // // Front legs equal and back legs equal constraints
    trajopt.AddLinearConstraint( xi( positions_map.at("joint_0") ) ==  xi( positions_map.at("joint_4") ) );
    trajopt.AddLinearConstraint( xi( positions_map.at("joint_1") ) ==  xi( positions_map.at("joint_5") ) );
    trajopt.AddLinearConstraint( xi( positions_map.at("joint_2") ) ==  xi( positions_map.at("joint_6") ) );
    trajopt.AddLinearConstraint( xi( positions_map.at("joint_3") ) ==  xi( positions_map.at("joint_7") ) );
  }
}
template <typename T>
void setSpiritSymmetry(
        drake::multibody::MultibodyPlant<T> & plant,
        dairlib::systems::trajectory_optimization::Dircon<T>& trajopt,
        std::vector<std::string> symmetries){

  bool hasSaggital = false;
  bool hasForeAft  = false;
  hasSaggital = std::any_of(symmetries.begin(), symmetries.end(), [](const std::string & str) {
                                                        return str == "sagittal";
                                                    });

  hasForeAft = std::any_of(symmetries.begin(), symmetries.end(), [](const std::string & str) {
                                                        return str == "foreaft";
                                                    });
  if (hasSaggital&&hasForeAft){
    // Avoids having overconstrained system of symmetry
    std::cerr << "Warning: Simultaneous Saggital and ForeAft not implemented yet. Only adding saggital constraints." << std::endl;
    makeSaggitalSymmetric(plant,trajopt);
  }else if (hasSaggital){
    makeSaggitalSymmetric(plant,trajopt);
  }else if (hasForeAft){
    std::cerr << "Warning: ForeAft not implemented yet. No constraints added" << std::endl;
  }else{
    std::cerr << "Warning: Failed to add symmetric constraint. Check spelling or implemenation. No constraints added" << std::endl;
  }

}

template <typename T>
void setSpiritSymmetry(drake::multibody::MultibodyPlant<T> & plant,
                       dairlib::systems::trajectory_optimization::Dircon<T>& trajopt,
                       std::string symmetry,
                       bool ignoreWarning){

  static bool hasRun = false;
  if(hasRun && !ignoreWarning){
    std::cerr << "Warning: Running different symmetries seperately could cause system to be overconstrained. Use vector overload to avoid this or set ignoreWarning flag if not a problem."
              << std::endl;
  }else{
    hasRun = true;
  }

  std::vector<std::string> symmetries;
  symmetries.push_back(symmetry);
  setSpiritSymmetry(
                    plant,
                    trajopt,
                    symmetries);
}
//   \SETSPIRITSYMMETRY

template <typename T>
double calcWork(
    drake::multibody::MultibodyPlant<T> & plant,
    drake::trajectories::PiecewisePolynomial<double>& x_traj,
    drake::trajectories::PiecewisePolynomial<double>& u_traj){

  auto velocities_map = multibody::makeNameToVelocitiesMap(plant);
  auto actuator_map = multibody::makeNameToActuatorsMap(plant);
  int n_q = plant.num_positions();

  double work = 0;
  std::vector<double> knot_points = x_traj.get_segment_times();

  for(int knot_index = 0; knot_index < knot_points.size()-1; knot_index++){
      auto u_low = u_traj.value(knot_points[knot_index]);
      auto u_up  = u_traj.value(knot_points[knot_index + 1]);
      auto x_low = x_traj.value(knot_points[knot_index]);
      auto x_up  = x_traj.value(knot_points[knot_index + 1]);

    for (int joint = 0; joint < 12; joint++){

      double actuation_low = u_low(actuator_map.at("motor_" + std::to_string(joint)));
      double actuation_up  =  u_up(actuator_map.at("motor_" + std::to_string(joint)));

      double velocity_low = x_low(n_q + velocities_map.at("joint_" + std::to_string(joint) +"dot"));
      double velocity_up  =  x_up(n_q + velocities_map.at("joint_" + std::to_string(joint) +"dot"));

      double pow_low = abs(actuation_low*velocity_low);
      double pow_up  = abs(actuation_up*velocity_up);

      // trapazoidal integration
      work += (knot_points[knot_index + 1] - knot_points[knot_index])/2.0 * (pow_low + pow_up);
    }
  }
  return work;
}

template <typename T>
double calcMechanicalWork(
    drake::multibody::MultibodyPlant<T> & plant,
    std::vector<drake::trajectories::PiecewisePolynomial<double>>& x_trajs,
    drake::trajectories::PiecewisePolynomial<double>& u_traj){

  auto velocities_map = multibody::makeNameToVelocitiesMap(plant);
  auto actuator_map = multibody::makeNameToActuatorsMap(plant);
  int n_q = plant.num_positions();

  double work = 0;
  for(const auto& x_traj : x_trajs) {
    std::vector<double> knot_points = x_traj.get_segment_times();
    for (int knot_index = 0; knot_index < knot_points.size() - 1; knot_index++) {
      auto u_low = u_traj.value(knot_points[knot_index]);
      auto u_up = u_traj.value(knot_points[knot_index + 1]);
      auto x_low = x_traj.value(knot_points[knot_index]);
      auto x_up = x_traj.value(knot_points[knot_index + 1]);

      for (int joint = 0; joint < 12; joint++) {

        double actuation_low = u_low(actuator_map.at("motor_" + std::to_string(joint)));
        double actuation_up = u_up(actuator_map.at("motor_" + std::to_string(joint)));

        double velocity_low = x_low(n_q + velocities_map.at("joint_" + std::to_string(joint) + "dot"));
        double velocity_up = x_up(n_q + velocities_map.at("joint_" + std::to_string(joint) + "dot"));

        double pow_low = abs(actuation_low * velocity_low);
        double pow_up = abs(actuation_up * velocity_up);

        // trapazoidal integration
        work += (knot_points[knot_index + 1] - knot_points[knot_index]) / 2.0 * (pow_low + pow_up);
      }
    }
  }
  return work;
}

template <typename T>
double calcElectricalWork(
    drake::multibody::MultibodyPlant<T> & plant,
    std::vector<drake::trajectories::PiecewisePolynomial<double>>& x_trajs,
    drake::trajectories::PiecewisePolynomial<double>& u_traj,
    double efficiency){

  auto velocities_map = multibody::makeNameToVelocitiesMap(plant);
  auto actuator_map = multibody::makeNameToActuatorsMap(plant);
  int n_q = plant.num_positions();

  double work = 0;
  double Q;
  for(const auto& x_traj : x_trajs) {
    std::vector<double> knot_points = x_traj.get_segment_times();
    for (int knot_index = 0; knot_index < knot_points.size() - 1; knot_index++) {
      auto u_low = u_traj.value(knot_points[knot_index]);
      auto u_up = u_traj.value(knot_points[knot_index + 1]);
      auto x_low = x_traj.value(knot_points[knot_index]);
      auto x_up = x_traj.value(knot_points[knot_index + 1]);

      for (int joint = 0; joint < 12; joint++) {
        if(joint == 1 or joint == 3 or joint == 5 or joint == 7){
          Q = Q_knee;
        }
        else{
          Q = Q_not_knee;
        }

        double actuation_low = u_low(actuator_map.at("motor_" + std::to_string(joint)));
        double actuation_up = u_up(actuator_map.at("motor_" + std::to_string(joint)));

        double velocity_low = x_low(n_q + velocities_map.at("joint_" + std::to_string(joint) + "dot"));
        double velocity_up = x_up(n_q + velocities_map.at("joint_" + std::to_string(joint) + "dot"));

        double pow_low = actuation_low * velocity_low + Q * actuation_low * actuation_low;
        double pow_up = actuation_up * velocity_up+ Q * actuation_up * actuation_up;

        double battery_pow_low = positivePart(pow_low) + efficiency * negativePart(pow_low);
        double battery_pow_up = positivePart(pow_up) + efficiency * negativePart(pow_up);

        // trapazoidal integration
        work += (knot_points[knot_index + 1] - knot_points[knot_index]) / 2.0 * (battery_pow_low + battery_pow_up);
      }
    }
  }
  return work;
}

template <typename T>
double calcVelocityInt(
    drake::multibody::MultibodyPlant<T> & plant,
    drake::trajectories::PiecewisePolynomial<double>& x_traj){

  auto velocities_map = multibody::makeNameToVelocitiesMap(plant);
  int n_v = plant.num_velocities();

  double vel_int = 0;
  std::vector<double> knot_points = x_traj.get_segment_times();

  for(int knot_index = 0; knot_index < knot_points.size()-1; knot_index++){
    Eigen::VectorXd x_low = x_traj.value(knot_points[knot_index]);
    Eigen::VectorXd x_up  = x_traj.value(knot_points[knot_index + 1]);

    Eigen::VectorXd velocity_low = x_low.tail(n_v);
    Eigen::VectorXd velocity_up  =  x_up.tail(n_v);

    double vel_sq_low = velocity_low.transpose() * velocity_low;
    double vel_sq_up  = velocity_up.transpose()  * velocity_up;

    // trapazoidal integration
    vel_int += (knot_points[knot_index + 1] - knot_points[knot_index])/2.0 * (vel_sq_low  + vel_sq_up);
  }

  return vel_int;
}

template <typename T>
double calcVelocityInt(
    drake::multibody::MultibodyPlant<T> & plant,
    std::vector<drake::trajectories::PiecewisePolynomial<double>>& x_trajs){

  auto velocities_map = multibody::makeNameToVelocitiesMap(plant);
  int n_v = plant.num_velocities();

  double vel_int = 0;

  for(const auto& x_traj : x_trajs){
    std::vector<double> knot_points = x_traj.get_segment_times();

    for(int knot_index = 0; knot_index < knot_points.size()-1; knot_index++){
      Eigen::VectorXd x_low = x_traj.value(knot_points[knot_index]);
      Eigen::VectorXd x_up  = x_traj.value(knot_points[knot_index + 1]);

      Eigen::VectorXd velocity_low = x_low.tail(n_v);
      Eigen::VectorXd velocity_up  =  x_up.tail(n_v);

      double vel_sq_low = velocity_low.transpose() * velocity_low;
      double vel_sq_up  = velocity_up.transpose()  * velocity_up;

      // trapazoidal integration
      vel_int += (knot_points[knot_index + 1] - knot_points[knot_index])/2.0 * (vel_sq_low  + vel_sq_up);
    }

  }
  return vel_int;
}

template <typename T>
double calcTorqueInt(
    drake::multibody::MultibodyPlant<T> & plant,
    drake::trajectories::PiecewisePolynomial<double>& u_traj){

  auto actuator_map = multibody::makeNameToActuatorsMap(plant);

  double act_int = 0;
  std::vector<double> knot_points = u_traj.get_segment_times();

  for(int knot_index = 0; knot_index < knot_points.size()-1; knot_index++){
    Eigen::VectorXd u_low = u_traj.value(knot_points[knot_index]);
    Eigen::VectorXd u_up  = u_traj.value(knot_points[knot_index + 1]);

    double act_sq_low = u_low.transpose() * u_low;
    double act_sq_up  = u_up.transpose() * u_up;

    // trapazoidal integration
    act_int += (knot_points[knot_index + 1] - knot_points[knot_index])/2.0 * (act_sq_low + act_sq_up);
  }
  return act_int;
}

template <typename T>
std::vector<drake::solvers::Binding<drake::solvers::Cost>> AddWorkCost(drake::multibody::MultibodyPlant<T> & plant,
                 dairlib::systems::trajectory_optimization::Dircon<T>& trajopt,
                 double cost_work_gain){
  std::vector<drake::solvers::Binding<drake::solvers::Cost>> cost_joint_work_bindings;
  int n_q = plant.num_positions();
  auto velocities_map = multibody::makeNameToVelocitiesMap(plant);
  auto actuator_map = multibody::makeNameToActuatorsMap(plant);

  double Q = 0;
  // Loop through each joint
  for (int joint = 0; joint < 12; joint++) {
    if(joint == 1 or joint == 3 or joint == 5 or joint == 7)
      Q = Q_knee;
    else
      Q = Q_not_knee;
    auto joint_work_cost = std::make_shared<JointWorkCost>(plant,  Q, cost_work_gain,4);
    int act_int = actuator_map.at("motor_" + std::to_string(joint));
    int vel_int = n_q + velocities_map.at("joint_" + std::to_string(joint) +"dot");
    // Loop through each mode
    for (int mode_index = 0; mode_index < trajopt.num_modes(); mode_index++) {
      for (int knot_index = 0; knot_index < trajopt.mode_length(mode_index)-1; knot_index++) {
        auto u_i = trajopt.input(trajopt.get_mode_start(mode_index) + knot_index)(act_int);
        auto x_i   = trajopt.state_vars(mode_index, knot_index)(vel_int);
        auto u_ip = trajopt.input(trajopt.get_mode_start(mode_index) + knot_index+1)(act_int);
        auto x_ip   = trajopt.state_vars(mode_index, knot_index+1)(vel_int);

        auto hi = trajopt.timestep(trajopt.get_mode_start(mode_index) + knot_index);
        drake::solvers::VectorXDecisionVariable variables(5);
        variables(0) = hi[0];
        variables(1) = u_i;
        variables(2) = u_ip;
        variables(3) = x_i;
        variables(4) = x_ip;
        cost_joint_work_bindings.push_back(trajopt.AddCost(joint_work_cost, variables));
      } // knot point loop
    } // Mode loop
  } // Joint loop
  return cost_joint_work_bindings;
}

JointWorkCost::JointWorkCost(const drake::multibody::MultibodyPlant<double>& plant,
                             const double &Q,
                             const double &cost_work,
                             const double &alpha,
                             const std::string &description)
    :solvers::NonlinearCost<double>(5,description), plant_(plant){
  Q_ = Q;
  cost_work_ = cost_work;
  alpha_ = alpha;
  n_q_ = plant_.num_positions();
  n_v_ = plant_.num_velocities();
  n_u_ = plant_.num_actuators();
  DRAKE_DEMAND(alpha_ > 0);
}

double JointWorkCost::relu_smooth(const double x) const{
  return  x>5 ? x : 1/alpha_ * log(1 + exp(alpha_ * x));
}

void JointWorkCost::EvaluateCost(const Eigen::Ref<const drake::VectorX<double>> &x,
                  drake::VectorX<double> *y) const{

double h_i = x(0);
//auto u_i_all = x.segment(1, n_u_);
//auto u_ip_all = x.segment(1+n_u_, n_u_);
//auto x_i_all = x.segment(1+n_u_+n_u_, n_q_+n_v_);
//auto x_ip_all = x.segment(1+n_u_+n_u_+n_q_+n_v_, n_q_+n_v_);

double u_i = x(1);
double u_ip = x(2);
double v_i = x(3);
double v_ip = x(4);
double work_low = cost_work_ * relu_smooth(u_i * v_i + Q_ * u_i * u_i);
double work_up = cost_work_ * relu_smooth(u_ip * v_ip + Q_ * u_ip * u_ip);

drake::VectorX<double> rv(1);
rv[0] = 1/2.0 * h_i *(work_low + work_up);
(*y) = rv;
};

double positivePart(double x){
  return(std::max(x,0.0));
}

double negativePart(double x){
  return(std::min(x,0.0));
}

template void nominalSpiritStand(
    drake::multibody::MultibodyPlant<double>& plant,
    Eigen::VectorXd& xState,
    double height); // NOLINT

template void nominalSpiritStandConstraint(
    drake::multibody::MultibodyPlant<double>& plant,
    dairlib::systems::trajectory_optimization::Dircon<double>& trajopt,
    double height,
    std::vector<int> knotPoints,
    const double eps); // NOLINT

template const drake::multibody::Frame<double>& getSpiritToeFrame(
    drake::multibody::MultibodyPlant<double>& plant,
    u_int8_t toeIndex ); // NOLINT

template std::unique_ptr<multibody::WorldPointEvaluator<double>> getSpiritToeEvaluator(
                      MultibodyPlant<double>& plant,
                      const Eigen::Vector3d toePoint ,
                      u_int8_t toeIndex,
                      const Eigen::Vector3d normal ,
                      const Eigen::Vector3d offset ,
                      bool xy_active,
                      double mu
                      ); // NOLINT

template std::tuple<  std::vector<std::unique_ptr<dairlib::systems::trajectory_optimization::DirconMode<double>>>,
             std::vector<std::unique_ptr<multibody::WorldPointEvaluator<double>>> ,
             std::vector<std::unique_ptr<multibody::KinematicEvaluatorSet<double>>>
          >
    createSpiritModeSequence(
          drake::multibody::MultibodyPlant<double>& plant, // multibodyPlant
          std::vector<Eigen::Matrix<bool,1,4>> modeSeqVect, // bool matrix describing toe contacts as true or false e.g. {{1,1,1,1},{0,0,0,0}} would be a full support mode and flight mode
          std::vector<int> knotpointVect, // Matrix of knot points for each mode
          std::vector<Eigen::Vector3d> normals,
          std::vector<Eigen::Vector3d> offsets,
          std::vector<double> mus ,
          std::vector<double> minTs,
          std::vector<double> maxTs); // NOLINT

template std::tuple<
                std::vector<std::unique_ptr<DirconMode<double>>>,
                std::vector<std::unique_ptr<dairlib::multibody::WorldPointEvaluator<double>>> ,
                std::vector<std::unique_ptr<dairlib::multibody::KinematicEvaluatorSet<double>>>
          > createSpiritModeSequence(
          MultibodyPlant<double>& plant, // multibodyPlant
          const dairlib::ModeSequenceHelper& msh );//NOLINT

template void setSpiritJointLimits(
          drake::multibody::MultibodyPlant<double> & plant,
          dairlib::systems::trajectory_optimization::Dircon<double>& trajopt,
          int iJoint,
          double minVal,
          double maxVal  );

template void setSpiritJointLimits(
          drake::multibody::MultibodyPlant<double> & plant,
          dairlib::systems::trajectory_optimization::Dircon<double>& trajopt,
          std::vector<int> iJoints,
          std::vector<double> minVals,
          std::vector<double> maxVals  );

template void setSpiritJointLimits(
          drake::multibody::MultibodyPlant<double> & plant,
          dairlib::systems::trajectory_optimization::Dircon<double>& trajopt,
          std::vector<int> iJoints,
          double minVal,
          double maxVal  );

template void setSpiritJointLimits(
          drake::multibody::MultibodyPlant<double> & plant,
          dairlib::systems::trajectory_optimization::Dircon<double>& trajopt );

template void setSpiritActuationLimits(
          drake::multibody::MultibodyPlant<double> & plant,
          dairlib::systems::trajectory_optimization::Dircon<double>& trajopt,
          double actuatorLimit);

template void setSpiritSymmetry(
        drake::multibody::MultibodyPlant<double> & plant,
        dairlib::systems::trajectory_optimization::Dircon<double>& trajopt,
        std::string symmetry,
        bool ignoreWarning);

template void setSpiritSymmetry(
        drake::multibody::MultibodyPlant<double> & plant,
        dairlib::systems::trajectory_optimization::Dircon<double>& trajopt,
        std::vector<std::string> symmetries);

template double calcWork(
    drake::multibody::MultibodyPlant<double> & plant,
    drake::trajectories::PiecewisePolynomial<double>& x_traj,
    drake::trajectories::PiecewisePolynomial<double>& u_traj);

template double calcMechanicalWork(
    drake::multibody::MultibodyPlant<double> & plant,
    std::vector<drake::trajectories::PiecewisePolynomial<double>>& x_traj,
    drake::trajectories::PiecewisePolynomial<double>& u_traj);

template double calcElectricalWork(
    drake::multibody::MultibodyPlant<double> & plant,
    std::vector<drake::trajectories::PiecewisePolynomial<double>>& x_traj,
    drake::trajectories::PiecewisePolynomial<double>& u_traj,
    double efficiency);

template double calcVelocityInt(
    drake::multibody::MultibodyPlant<double> & plant,
    drake::trajectories::PiecewisePolynomial<double>& x_traj);

template double calcVelocityInt(
    drake::multibody::MultibodyPlant<double> & plant,
    std::vector<drake::trajectories::PiecewisePolynomial<double>>& x_trajs);

template double calcTorqueInt(
    drake::multibody::MultibodyPlant<double> & plant,
    drake::trajectories::PiecewisePolynomial<double>& u_traj);

template std::vector<drake::solvers::Binding<drake::solvers::Cost>> AddWorkCost(drake::multibody::MultibodyPlant<double> & plant,
                 dairlib::systems::trajectory_optimization::Dircon<double>& trajopt,
                 double cost_work_gain);

}//namespace dairlib


#include "Setup.hpp"
#include "math/Common.hpp"

#include <cerrno>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <exception>
#include <iostream>
#include <locale>
#include <memory>
#include <random>
#include <regex>
#include <string>
#include <unordered_map>
#include <vector>

#include <mpreal.h>
#include <unistd.h>
#include <boost/algorithm/string.hpp>
#include <boost/date_time/date_facet.hpp>
#include <boost/date_time/time_facet.hpp>
#include <boost/format.hpp>
#include <boost/program_options.hpp>
#include <readline/history.h>
#include <readline/readline.h>

static const std::unordered_map<mpfr_rnd_t, std::string> rmodeNameMap = {
    {mpfr_rnd_t::MPFR_RNDN, "N"},
    {mpfr_rnd_t::MPFR_RNDZ, "Z"},
    {mpfr_rnd_t::MPFR_RNDU, "U"},
    {mpfr_rnd_t::MPFR_RNDD, "D"},
    {mpfr_rnd_t::MPFR_RNDA, "A"},
    {mpfr_rnd_t::MPFR_RNDF, "F"},
    {mpfr_rnd_t::MPFR_RNDNA, "NA"},
};

static const std::unordered_map<std::string, mpfr_rnd_t> rmodeNameMap2 = {
    {"N", mpfr_rnd_t::MPFR_RNDN},
    {"Z", mpfr_rnd_t::MPFR_RNDZ},
    {"U", mpfr_rnd_t::MPFR_RNDU},
    {"D", mpfr_rnd_t::MPFR_RNDD},
    {"A", mpfr_rnd_t::MPFR_RNDA},
    {"F", mpfr_rnd_t::MPFR_RNDF},
    {"NA", mpfr_rnd_t::MPFR_RNDNA},
};

static const std::unordered_map<Associativity, std::string> associativityNameMap = {
    {Associativity::Left, "Left"},
    {Associativity::Right, "Right"},
    {Associativity::Any, "Any"},
};

mpfr_rnd_t strToRmode(const std::string value)
{
  const auto iter = rmodeNameMap2.find(boost::to_upper_copy(value));
  if(iter != rmodeNameMap2.cend())
  {
    return iter->second;
  }
  else
  {
    throw std::domain_error("Invalid rounding mode");
  }
}

static std::istream& operator>>(std::istream& stream, mpfr_rnd_t& result)
{
  std::string value;
  stream >> value;
  result = strToRmode(value);
  return stream;
}

void resolveEnvironmentVariables(std::vector<std::string>& result)
{
  const char* pTmp;

  if((pTmp = std::getenv("KALK_PREC")) != nullptr)
  {
    result.push_back("KALK_PREC");
    result.push_back(pTmp);
  }

  if((pTmp = std::getenv("KALK_RMODE")) != nullptr)
  {
    result.push_back("KALK_RMODE");
    result.push_back(pTmp);
  }

  if((pTmp = std::getenv("KALK_DIGITS")) != nullptr)
  {
    result.push_back("KALK_DIGITS");
    result.push_back(pTmp);
  }

  if((pTmp = std::getenv("KALK_OBASE")) != nullptr)
  {
    result.push_back("KALK_OBASE");
    result.push_back(pTmp);
  }

  if((pTmp = std::getenv("KALK_IBASE")) != nullptr)
  {
    result.push_back("KALK_IBASE");
    result.push_back(pTmp);
  }

  if((pTmp = std::getenv("KALK_BASE")) != nullptr)
  {
    result.push_back("KALK_BASE");
    result.push_back(pTmp);
  }

  if((pTmp = std::getenv("KALK_JUXTA")) != nullptr)
  {
    result.push_back("KALK_JUXTA");
    result.push_back(pTmp);
  }

  if((pTmp = std::getenv("KALK_DATE_OFMT")) != nullptr)
  {
    result.push_back("KALK_DATE_OFMT");
    result.push_back(pTmp);
  }

  if((pTmp = std::getenv("KALK_INTERACTIVE")) != nullptr)
  {
    result.push_back("KALK_INTERACTIVE");
    result.push_back(pTmp);
  }

  if((pTmp = std::getenv("KALK_VERBOSE")) != nullptr)
  {
    result.push_back("KALK_VERBOSE");
    result.push_back(pTmp);
  }
}

static void handleResult(const DefaultValueType* value, bool verbose)
{
  results.push_back(*value);
  if(verbose)
  {
    printValue(*value);
  }

  if(!defaultUninitializedVariableCache.empty())
  {
    std::cout << "*** Warning: Uninitialized variable(s)" << std::endl;
    while(!defaultUninitializedVariableCache.empty())
    {
      auto iter = defaultUninitializedVariableCache.begin();
      defaultVariables.erase(iter->first);
      defaultUninitializedVariableCache.erase(iter);
    }
  }
}

static void evaluate(std::string expression, ExpressionParser& expressionParser, bool verbose = true)
{
  constexpr char kWhitespaceCharacters[] = " \t\v\n\r\f";
  bool end                               = false;
  while(!end && expression.find_first_not_of(kWhitespaceCharacters) != std::string::npos)
  {
    const auto result = expressionParser.Evaluate(expression);
    handleResult(result->As<const DefaultValueType*>(), verbose);
    if(!(end = expressionParser.GetCurrent() != ';'))
    {
      expression = (expressionParser.GetRemaining() + 1u);
    }
  }
}

void list(const std::string& searchPattern)
{
  const std::regex regex(searchPattern);
  bool isPrevEmptyLine = true;

  std::cerr << "Unary operators" << std::endl;
  for(const auto& i : unaryOperatorInfoMap)
  {
    const auto& entry = std::get<0u>(i);
    if(entry == nullptr)
    {
      if(!isPrevEmptyLine)
      {
        std::cerr << std::endl;
        isPrevEmptyLine = true;
      }

      continue;
    }

    const auto identifier    = std::string(1u, entry->GetIdentifier());
    const auto precedence    = entry->GetPrecedence();
    const auto associativity = associativityNameMap.at(entry->GetAssociativity());
    const auto& title        = std::get<1u>(i);
    const auto& description  = std::get<2u>(i);
    if(std::regex_match(identifier.begin(), identifier.end(), regex) || std::regex_match(title.begin(), title.end(), regex) ||
       std::regex_match(description.begin(), description.end(), regex))
    {
      std::cout << (boost::format("  %|1$-5|%|2$-5|%|3$-9|%|4$-20|%|5$|") % identifier % precedence % associativity % title % description) << std::endl;
      isPrevEmptyLine = false;
    }
  }
  std::cerr << std::endl;

  isPrevEmptyLine = true;
  std::cerr << "Binary operators" << std::endl;
  for(const auto& i : binaryOperatorInfoMap)
  {
    const auto& entry = std::get<0u>(i);
    if(entry == nullptr)
    {
      if(!isPrevEmptyLine)
      {
        std::cerr << std::endl;
        isPrevEmptyLine = true;
      }

      continue;
    }

    const auto& identifier   = entry->GetIdentifier();
    const auto precedence    = entry->GetPrecedence();
    const auto associativity = associativityNameMap.at(entry->GetAssociativity());
    const auto& title        = std::get<1u>(i);
    const auto& description  = std::get<2u>(i);
    if(std::regex_match(identifier.begin(), identifier.end(), regex) || std::regex_match(title.begin(), title.end(), regex) ||
       std::regex_match(description.begin(), description.end(), regex))
    {
      std::cout << (boost::format("  %|1$-6|%|2$-5|%|3$-9|%|4$-25|%|5$|") % identifier % precedence % associativity % title % description) << std::endl;
      isPrevEmptyLine = false;
    }
  }
  std::cerr << std::endl;

  isPrevEmptyLine = true;
  std::cerr << "Functions" << std::endl;
  for(const auto& i : functionInfoMap)
  {
    const auto& entry = std::get<0u>(i);
    if(entry == nullptr)
    {
      if(!isPrevEmptyLine)
      {
        std::cerr << std::endl;
        isPrevEmptyLine = true;
      }

      continue;
    }

    const auto& identifier = entry->GetIdentifier();
    const auto argMinCount = ((entry->GetMinArgumentCount() != FunctionToken::GetArgumentCountMaxLimit()) ? std::to_string(entry->GetMinArgumentCount()) : "-");
    const auto argMaxCount = ((entry->GetMaxArgumentCount() != FunctionToken::GetArgumentCountMaxLimit()) ? std::to_string(entry->GetMaxArgumentCount()) : "-");
    const auto& title      = std::get<1u>(i);
    const auto& description = std::get<2u>(i);
    if(std::regex_match(identifier.begin(), identifier.end(), regex) || std::regex_match(title.begin(), title.end(), regex) ||
       std::regex_match(description.begin(), description.end(), regex))
    {
      std::cout << (boost::format("  %|1$-15|%|2$-5|%|3$-5|%|4$-27|%|5$|") % identifier % argMinCount % argMaxCount % title % description) << std::endl;
      isPrevEmptyLine = false;
    }
  }
  std::cerr << std::endl;

  isPrevEmptyLine = true;
  std::cerr << "Variables" << std::endl;
  for(const auto& i : variableInfoMap)
  {
    const auto& entry = std::get<0u>(i);
    if(entry == nullptr)
    {
      if(!isPrevEmptyLine)
      {
        std::cerr << std::endl;
        isPrevEmptyLine = true;
      }

      continue;
    }

    const auto& identifier  = entry->GetIdentifier();
    const auto& title       = std::get<1u>(i);
    const auto& description = std::get<2u>(i);
    if(std::regex_match(identifier.begin(), identifier.end(), regex) || std::regex_match(title.begin(), title.end(), regex) ||
       std::regex_match(description.begin(), description.end(), regex))
    {
      std::cout << (boost::format("  %|1$-16|%|2$-29|%|3$|") % identifier % title % description) << std::endl;
      isPrevEmptyLine = false;
    }
  }
}

static void printOptions()
{
  std::cerr << "Options" << std::endl;
  std::cerr << (boost::format("  %|1$-26|%|2$|") % "Precision" % options.precision) << std::endl;
  std::cerr << (boost::format("  %|1$-26|%|2$|") % "Rounding mode" % rmodeNameMap.at(options.roundingMode)) << std::endl;
  std::cerr << (boost::format("  %|1$-26|%|2$|") % "Output precision" % options.digits) << std::endl;
  std::cerr << (boost::format("  %|1$-26|%|2$|") % "Output base" % options.output_base) << std::endl;
  std::cerr << (boost::format("  %|1$-26|%|2$|") % "Input base" % options.input_base) << std::endl;
  std::cerr << (boost::format("  %|1$-26|%|2$|") % "Juxtaposition precedence" % options.jpo_precedence) << std::endl;
  std::cerr << (boost::format("  %|1$-26|%|2$|") % "Date output format" % options.date_ofmt) << std::endl;
  std::cerr << (boost::format("  %|1$-26|%|2$|") % "Seed" % options.seed) << std::endl;
  std::cerr << std::endl;
}

static void printVersion() { std::cout << (boost::format("%1% v%2%") % PROJECT_NAME % PROJECT_VERSION) << std::endl; }

static void printUsage(const boost::program_options::options_description& desc)
{
  std::cerr << (boost::format("%1% -[xprnbBjdzZilvVh] expr...") % PROJECT_EXECUTABLE) << std::endl;
  std::cerr << desc << std::endl;
}

int main(int argc, char* argv[])
{
  std::vector<std::string> envs;
  resolveEnvironmentVariables(envs);

  boost::program_options::options_description namedEnvDescs;
  namedEnvDescs.add_options()("KALK_PREC", boost::program_options::value<mpfr_prec_t>(&options.precision)->default_value(defaultOptions.precision));
  namedEnvDescs.add_options()("KALK_RMODE", boost::program_options::value<mpfr_rnd_t>(&options.roundingMode)->default_value(defaultOptions.roundingMode));
  namedEnvDescs.add_options()("KALK_DIGITS", boost::program_options::value<int>(&options.digits)->default_value(defaultOptions.digits));
  namedEnvDescs.add_options()("KALK_OBASE", boost::program_options::value<int>(&options.output_base));
  namedEnvDescs.add_options()("KALK_IBASE", boost::program_options::value<int>(&options.input_base));
  namedEnvDescs.add_options()("KALK_BASE", boost::program_options::value<int>()->default_value(defaultOptions.output_base)->notifier([](int value) {
    options.output_base = value;
    options.input_base  = value;
  }));
  namedEnvDescs.add_options()("KALK_JUXTA", boost::program_options::value<int>()->default_value(defaultOptions.jpo_precedence)->notifier([](int value) {
    options.jpo_precedence = Math::Sign(value);
  }));
  namedEnvDescs.add_options()("KALK_DATE_OFMT", boost::program_options::value<std::string>(&options.date_ofmt)->default_value(defaultOptions.date_ofmt));
  namedEnvDescs.add_options()("KALK_INTERACTIVE", boost::program_options::value<bool>(&options.interactive)->default_value(defaultOptions.interactive));
  namedEnvDescs.add_options()("KALK_VERBOSE", boost::program_options::value<std::string>()->default_value(""));

  boost::program_options::variables_map envVariableMap;
  boost::program_options::store(boost::program_options::command_line_parser(envs)
                                    .options(namedEnvDescs)
                                    .extra_parser([](const std::string& value) { return std::make_pair(value, std::string()); })
                                    .run(),
                                envVariableMap);
  boost::program_options::notify(envVariableMap);

  boost::program_options::options_description namedArgDescs("Options");
  namedArgDescs.add_options()("expr,x", boost::program_options::value<std::vector<std::string>>(), "Add an expression");
  namedArgDescs.add_options()("prec,p", boost::program_options::value<mpfr_prec_t>(&options.precision), "Set precision");
  namedArgDescs.add_options()("rmode,r", boost::program_options::value<mpfr_rnd_t>(&options.roundingMode), "Set rounding mode (N, Z, U, D, A, F, NA)");
  namedArgDescs.add_options()("digits,n", boost::program_options::value<int>(&options.digits), "Set output precision (Number of digits)");
  namedArgDescs.add_options()("obase,b", boost::program_options::value<int>(&options.output_base), "Set output base");
  namedArgDescs.add_options()("ibase,B", boost::program_options::value<int>(&options.input_base), "Set input base");
  namedArgDescs.add_options()("base",
                              boost::program_options::value<int>()->notifier([](int value) { options.input_base = options.output_base = value; }),
                              "Set output and input base");
  namedArgDescs.add_options()("juxta,j",
                              boost::program_options::value<int>()->notifier([](int value) { options.jpo_precedence = Math::Sign(value); }),
                              "Set juxtaposition operator precedence (-1, 0, 1)");
  namedArgDescs.add_options()("date_ofmt,d", boost::program_options::value<std::string>(&options.date_ofmt), "Set date output format");
  namedArgDescs.add_options()("seed,z", boost::program_options::value<unsigned int>(&options.seed), "Set random seed (number)");
  namedArgDescs.add_options()("seedstr,Z",
                              boost::program_options::value<std::string>()->notifier([](std::string value) {
                                std::hash<std::string> hasher;
                                options.seed = value.empty() ? defaultOptions.seed : static_cast<unsigned int>(hasher(value));
                              }),
                              "Set random seed (string)");
  namedArgDescs.add_options()("interactive,i", boost::program_options::value<bool>(&options.interactive)->implicit_value(true), "Enable interactive mode");
  namedArgDescs.add_options()("list,l", boost::program_options::value<std::string>()->implicit_value(".*"), "List available operators/functions/variables");
  namedArgDescs.add_options()("verbose,v", boost::program_options::value<std::string>()->default_value("")->implicit_value("op"), "Enable verbose mode");
  namedArgDescs.add_options()("version,V", "Print version");
  namedArgDescs.add_options()("help,h", "Print usage");

  boost::program_options::positional_options_description positionalArgDescs;
  positionalArgDescs.add("expr", -1);

  boost::program_options::variables_map argVariableMap;
  boost::program_options::store(boost::program_options::command_line_parser(argc, argv).options(namedArgDescs).positional(positionalArgDescs).run(),
                                argVariableMap);
  boost::program_options::notify(argVariableMap);

  if(argVariableMap.count("help") > 0u)
  {
    printUsage(namedArgDescs);
    std::exit(EXIT_SUCCESS);
  }

  if(argVariableMap.count("version") > 0u)
  {
    printVersion();
    std::exit(EXIT_SUCCESS);
  }

  mpfr::random((options.seed != 0u) ? options.seed : std::random_device()());

  if((options.precision < MPFR_PREC_MIN) || (options.precision > MPFR_PREC_MAX))
  {
    std::cerr << (boost::format("*** Error: Precision out of range: %1% (%2% - %3%)") % options.precision % MPFR_PREC_MIN % MPFR_PREC_MAX) << std::endl;
    std::exit(EXIT_FAILURE);
  }

  mpfr::mpreal::set_default_prec(options.precision);
  mpfr::mpreal::set_default_rnd(options.roundingMode);

  auto dateFacet = new boost::posix_time::time_facet(options.date_ofmt.c_str());
  std::cout.imbue(std::locale(std::cout.getloc(), dateFacet));

  const bool verboseOptions = envVariableMap["KALK_VERBOSE"].as<const std::string&>().find_first_of("oO") != std::string::npos ||
                              argVariableMap["verbose"].as<const std::string&>().find_first_of("oO") != std::string::npos;
  const bool verbosePipe = envVariableMap["KALK_VERBOSE"].as<const std::string&>().find_first_of("pP") != std::string::npos ||
                           argVariableMap["verbose"].as<const std::string&>().find_first_of("pP") != std::string::npos;

  if(verboseOptions)
  {
    printOptions();
  }

  ExpressionParser expressionParser;
  InitDefaultExpressionParser(expressionParser);

  CommandParser commandParser;
  InitCommandParser(commandParser);

  if(argVariableMap.count("list") > 0u)
  {
    list(argVariableMap["list"].as<const std::string&>());
    std::exit(EXIT_SUCCESS);
  }

  std::unique_ptr<FILE, decltype(&std::fclose)> file_stdin(nullptr, &std::fclose);
  bool hasPipedData = std::cin.rdbuf()->in_avail() != -1 && isatty(fileno(stdin)) == 0;
  if(hasPipedData)
  {
    std::string input;
    while(std::getline(std::cin, input))
    {
      evaluate(input, expressionParser, verbosePipe);
    }

    if(options.interactive)
    {
      const char* ttyFileName = ttyname(fileno(stdout));
      if(ttyFileName == nullptr)
      {
        std::cerr << "*** Error: " << std::strerror(errno) << std::endl;
        std::exit(EXIT_FAILURE);
      }

      FILE* tmpFile;
      if((tmpFile = freopen(ttyFileName, "r", stdin)) == nullptr)
      {
        std::cerr << "*** Error: " << std::strerror(errno) << std::endl;
        std::exit(EXIT_FAILURE);
      }

      file_stdin.reset(tmpFile);
    }
  }

  if(argVariableMap.count("expr") == 0u && !options.interactive && !hasPipedData)
  {
    std::cerr << "*** Error: No expression specified" << std::endl;
    std::exit(EXIT_FAILURE);
  }

  if(argVariableMap.count("expr") > 0u)
  {
    results.clear();

    const auto& exprs = argVariableMap["expr"].as<const std::vector<std::string>&>();
    for(auto& expr : exprs)
    {
      evaluate(expr, expressionParser);
    }
  }

  if(options.interactive)
  {
    results.clear();

    char* tmpInput;
    while(!quit && (tmpInput = readline("> ")) != nullptr)
    {
      auto tmpPtr       = std::unique_ptr<char, decltype(&std::free)>(tmpInput, &std::free);
      std::string input = tmpInput;

      boost::trim_left(input);
      if(!input.empty())
      {
        if(input.front() == '/')
        {
          try
          {
            commandParser.Execute(input.erase(0, 1));
          }
          catch(const SyntaxError& e)
          {
            std::cerr << "*** Command error: " << e.what() << std::endl;
          }
          catch(const std::runtime_error& e)
          {
            std::cerr << "*** Command error: " << e.what() << std::endl;
          }
          catch(const std::domain_error& e)
          {
            std::cerr << "*** Command error: " << e.what() << std::endl;
          }
        }
        else
        {
          try
          {
            evaluate(input, expressionParser);
          }
          catch(const SyntaxError& e)
          {
            std::cerr << "*** Expression error: " << e.what() << std::endl;
          }
          catch(const std::runtime_error& e)
          {
            std::cerr << "*** Expression error: " << e.what() << std::endl;
          }
          catch(const std::domain_error& e)
          {
            std::cerr << "*** Expression error: " << e.what() << std::endl;
          }
        }

        add_history(tmpInput);
      }
    }
  }

  std::exit(EXIT_SUCCESS);
}