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

/**
 * \file mpc_nominal_model.hpp
 * \brief
 * \details
 *
 * \author Autonomous Space Robotics Lab (ASRL)
 */
#pragma once

#define _USE_MATH_DEFINES  // for M_PI
#include <cmath>

#include <algorithm>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>

#include <Eigen/Dense>
#include <Eigen/Sparse>

#include <tf2_geometry_msgs/tf2_geometry_msgs.h>
#include <rclcpp/rclcpp.hpp>

#include <vtr_logging/logging.hpp>
#include <vtr_pose_graph/id/graph_id.hpp>

#include <geometry_msgs/msg/pose.hpp>

namespace vtr {
namespace path_tracker {

using pose_graph::VertexId;

// Indices for gp dependencies
const int DIST_ALONG_PATH = 0;
const int OMEGA_ACT_KM1 = 1;
const int V_ACT_KM1 = 2;
const int OMEGA_CMD_KM1 = 3;
const int V_CMD_KM1 = 4;
const int OMEGA_CMD_K = 5;
const int V_CMD_K = 6;
const int HEAD_ERROR_K = 7;
const int LAT_ERROR_K = 8;

const int CONTROL_SIZE = 1;
const int STATE_SIZE = 3;
const bool MODEL_INCLUDES_VELOCITY = false;
const bool MODEL_DEP_UKM1 = false;
const int ERROR_SIZE = STATE_SIZE - 1;
const int DIST_DEP_SIZE = 9;

const int POSE_SIZE = 3;
const int VELOCITY_SIZE = 2;

const float K_OMEGA = 0.6;

class mtx_triplet {
 public:
  mtx_triplet(int i_in, int j_in, float v_ij_in)
      : i(i_in), j(j_in), v_ij(v_ij_in) {}

  int i, j;
  float v_ij;
};

void adjust_mtx_triplet_row_col(mtx_triplet &mtx_triplet, int &i_in, int &j_in);

void scalar_mult_mtx_triplet(mtx_triplet &mtx_triplet, float &scalar_p);

typedef std::vector<mtx_triplet> mtx_triplet_list_t;
typedef std::vector<mtx_triplet_list_t> vector_valued_hessian_t;

/**
 * \brief MPC nominal model (including first derivatives)
 * \note x_kp1 = f(x_k,u_k) + g(a), where "a" is the disturbance dependency
 */
class MpcNominalModel {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  /** \brief Constructor, do nothing */
  MpcNominalModel() = default;

  /** \brief Destructor, do nothing */
  ~MpcNominalModel() = default;

  /**
   * \brief Struct to hold state x_k
   * This struct is meant to be used by the MPC algorithm
   */
  typedef struct {
    // Main variables, these variables must be present for the mpc algorithm
    Eigen::VectorXf x_k;  /// x_kp1 = f(x_k, u_k) + g(a_k)
    std::vector<Eigen::VectorXf> x_k_wc;
    Eigen::MatrixXf var_x_k;     ///< variance of x_k, used for robust mpc
    Eigen::MatrixXf dvar_x_dv;   ///< derivative of variance w.r.t. v
    Eigen::MatrixXf dvar_x_dth;  ///< derivative of variance w.r.t. th
    float lateral_uncertainty;
    float heading_uncertainty;

    Eigen::VectorXf
        g_a_k;  ///< Computed / Measured disturbance, in the robot frame
    Eigen::MatrixXf var_g_a_k;  ///< Variance of disturbance, matrix enables
                                ///< correlated disturbances, in robot frame

    Eigen::VectorXf g_a_k_des_frame;  /// \todo: add documentation for these
    Eigen::MatrixXf var_g_a_k_des_frame;

    // Additional variables, these variables will depend on your implementation
    // Tracking errors,
    float dist_along_path_k;
    Eigen::VectorXf tracking_error_k;
    Eigen::VectorXf tracking_error_km1;
    Eigen::VectorXf tracking_error_k_interp;

    // Velocities and Commands
    Eigen::VectorXf velocity_km1;
    Eigen::VectorXf command_k;
    Eigen::VectorXf command_km1;

    // Derivatives
    Eigen::MatrixXf df_dxk, df_dxkm1, df_duk, df_dukm1, df_dvk, df_dvkm1;
    Eigen::MatrixXf dg_dxk, dg_dxkm1, dg_duk, dg_dukm1, dg_dvk, dg_dvkm1;
    Eigen::MatrixXf dg_da;

    // Gradients
    Eigen::MatrixXf grad_x, grad_xkm1, grad_u, grad_ukm1, grad_v, grad_vkm1;

    // Hessians (Jacobians of Gradients)
    Eigen::MatrixXf Jv_gx_mtx;
    vector_valued_hessian_t Jx_gx, Ju_gx, Jv_gx, Jx_gu, Ju_gu, Jv_gu, Jx_gv,
        Ju_gv, Jv_gv;

    bool flg_hessians_cleared;

  } model_state_t;

  typedef std::vector<model_state_t> model_trajectory_t;

  /** Define struct to hold GP data points (for fast computation) **/
  typedef struct {
    Eigen::VectorXf x_meas;
    Eigen::VectorXf g_x_meas;
  } gp_data_t;

  /** \brief Struct to hold an "experience" */
  typedef struct {
    // State variables
    VertexId at_vertex_id;
    VertexId to_vertex_id;

    rclcpp::Time transform_time;
    rclcpp::Time store_time;

    model_state_t x_k;

    float tol_posL, tol_posH;

    float path_curvature;
    float dist_from_vertex;

    Eigen::VectorXf velocity_k;
    bool velocity_is_valid;
    bool disturbance_is_valid;

    Eigen::Matrix<double, 6, 1> full_velocity_k;

    // Variables prepped for GP
    gp_data_t gp_data;

    // Recorded variables for post-processing
    tf2::Transform T_0_v;
    float predicted_wc_lateral_error;
    float predicted_wc_heading_error;
    float v_k_odo;
    float w_k_odo;
    int16_t keypoint_match_count;
    int worst_case_scenario;
    bool mpc_valid;

    std::vector<float> J_iter;
    std::vector<float> u_iter;

    double t_mpc;
    double t_robust;

    bool flg_do_not_delete;

  } experience_t;

  /** \brief Motion Model for the MPC algorithm, using linearization to
   * propagate uncertainty */
  void f_x_linearizedUncertainty(const model_state_t &x_k, model_state_t &x_kp1,
                                 float dt);

  /**
   * \brief Motion Model for the MPC algorithm, using the unscented transform to
   * propagate uncertainty This function is written to function regardless of
   * the system definition It assumes x_kp1 = f(x_k,u_k) + g_a_k with uncertain
   * x_k and g_a_k. x_k and g_a_k must be of the same length. It relies on the
   * proper definition of f_x_linearizedUncertainty()
   */
  bool f_x_unscentedUncertainty(const model_state_t &x_k, model_state_t &x_kp1,
                                float dt);

  // Compute nominal model Jacobians

  void get_gdot(model_state_t &x_k, float d_t);

  void get_Jdot_gdot(model_state_t &x_k, float d_t);

  /// Define functions related to predicting pose sequences

  /** \brief Given a nominal sequence of states with uncertainty, generate worst
   * case trajectories */
  bool generateWorstCaseTrajectories(model_trajectory_t &x_sequence,
                                     const double &robust_control_sigma);

  /**
   * \brief Check if robot has passed desired state in sequence
   * \param v_des: the current desired speed
   * \param x_k: The current pose in the frame of the trunk
   * \param x_desired: the desired pose in the frame of the trunk
   *
   * The robot is considered to have passed a vertex if the x-coordinate of the
   * robot in the frame of the vertex has passed a pose. \return
   */
  bool robot_has_passed_desired_poseNew(const float &v_des,
                                        const Eigen::VectorXf &x_k,
                                        const Eigen::MatrixXf &x_desired);

  /** \brief Compute the distance along the path */
  float computeDistAlongPath(const float &nearest_path_length, const float &e_x,
                             const float &linear_speed);

  /// Define functions related to computing tracking errors

  /**
   * \brief Finds the closest point x_des_interp to x_pred along the line
   * connecting x_des_im1 and x_des_i \param x_des_im1: the desired point behind
   * the robot \param x_des_i: the desired point ahead of the robot \param
   * x_pred: The robots predicted pose \param x_des_interp: the resulting
   * interpolated path waypoint
   */
  void computeInterpolatedDesiredPoseNew(
      const Eigen::MatrixXf &x_des_im1, const Eigen::MatrixXf &x_des_i,
      vtr::path_tracker::MpcNominalModel::model_state_t &x_pred,
      Eigen::MatrixXf &x_des_interp);

  // Compute errors for a pose or sequence of poses
  void compute_sequence_errors(model_trajectory_t &x_sequence,
                               const Eigen::MatrixXf &x_desired);
  Eigen::VectorXf compute_pose_errors(model_state_t &x_state,
                                      const Eigen::MatrixXf &x_desired);
  Eigen::VectorXf compute_pose_errors(model_state_t &x_state,
                                      const Eigen::MatrixXf &x_desired,
                                      const tf2::Vector3 &p_0_k_0,
                                      const tf2::Transform &C_0_k);

  /**
   * \brief MpcNominalModel::compute_pose_errorsNew
   * \param x_state: The state of the robot.
   * \param x_desired: The path vertex (can be i (ahead of the robot) or i-1
   * (behind the robot) \return the error (x,y,theta) between the desired pose
   * and the robot in the frame of the desired pose.
   */
  Eigen::VectorXf compute_pose_errorsNew(const model_state_t &x_state,
                                         const Eigen::MatrixXf &x_desired);

  /** \brief Compute pose errors. Called by compute_pose_errorsNew(with two
   * arguments) */
  Eigen::VectorXf compute_pose_errorsNew(const model_state_t &x_state,
                                         const Eigen::MatrixXf &x_desired,
                                         const tf2::Vector3 &p_0_k_0,
                                         const tf2::Transform &C_0_k);

  /** \brief Extract errors for a sequence of poses */
  void extract_pose_errors(Eigen::MatrixXf &traj_errors,
                           const model_trajectory_t &x_sequence);

  /** \brief Convert pose errors to heading and lateral errors */
  Eigen::MatrixXf compute_pose_to_hl_error_conversion_mtx(float &th_des);

  /// Define functions related to computing disturbances

  /**
   * \brief Compute disturbances based on two sequential poses
   *  Computing for km2 because we only know velocity for km1 at time k
   */
  bool computeDisturbancesForExperienceKm2(experience_t &experience_km1,
                                           const experience_t &experience_k);

  /**
   * \brief MpcNominalModel::computeDisturbancesForExperienceKm2SteamVel
   * Compute model errors based on velocity estimates from STEAM. This is more
   * accurate than computeDisturbancesForExperienceKm2. It requires STEAM to be
   * enabled. To do this, set extrapolate_VO = true in a config file.
   * \param experience_km2: Experience from timestep k-2
   * \param experience_km1: Experience from timestep k-1
   * \return true unless the experience could not be computed for some reason.
   */
  bool computeDisturbancesForExperienceKm2SteamVel(
      experience_t &experience_km2, const experience_t &experience_km1);

  bool compute_disturbance_from_state(Eigen::VectorXf &g_a_k_meas,
                                      const model_state_t &state_km1,
                                      const model_state_t &state_k,
                                      const float &d_t);

  /** \brief Compute the velocity by finite difference using the pose estimates
   * from experience_k and experience_km1.
   *
   * If the difference between time-stamps
   * for these pose estimates is too small, keep the last velocity estimate.
   * \param experience_km2: Experience at time k-2
   * \param experience_km1: Experience at time k-1
   * \param experience_k: Experience at the current time, time k
   */
  void computeVelocitiesForExperienceKm1(const experience_t &experience_km2,
                                         experience_t &experience_km1,
                                         experience_t &experience_k);

  void compute_velocities_from_state(Eigen::VectorXf &velocity,
                                     const Eigen::VectorXf &state_km1,
                                     const Eigen::VectorXf &state_k_act,
                                     const float dt);

  void initialize_experience(experience_t &experience_k, rclcpp::Clock &clock);

  /**
   * \brief MpcNominalModel::initialize_state Set all the values in x_k to zero.
   * \param x_k
   */
  void initialize_state(model_state_t &x_k);

  void computeDisturbanceDependancy(model_state_t &x_k,
                                    const model_state_t &x_km1,
                                    const Eigen::MatrixXf &x_des,
                                    const float &nearest_path_length,
                                    float &d_t);

  /** \brief Extract relevant variables from model_state_t */
  void extract_disturbance_dependencies(const model_state_t &x_input,
                                        Eigen::VectorXf &x_test);

  void set_disturbance_model_zero(model_state_t &x_input);
};

}  // namespace path_tracker
}  // namespace vtr


#pragma once
#include "../Force.hpp"
#include <Eigen/Geometry>

namespace mdk {
    /**
     * Solid walls.
     */
    class SolidWall: public Force {
    public:
        Eigen::Hyperplane<double, 3> wall;

        void asyncPart(Dynamics& dynamics) override;
    };
}


/*
 ============================================================================
 Name        : 	INSIGHTv3.cpp
 Author      : 	Jan Mikelson
 Version     :
 Copyright   : 	This free software is available under the Creative Commons Attribution Share Alike License.
 You are permitted to use, redistribute and adapt this software as long as appropriate credit
 is given to the original author, and all derivative works are distributed under the same
 license or a compatible one.
 For more information, visit http://creativecommons.org/licenses/by-sa/3.0/ or send a letter to
 Creative Commons, 171 2nd Street, Suite 300, San Francisco, California, 94105, USA.
 Description : 	Sequential version of the INSIGHT algorithm. Further details can be found in Lillacci & Khammash 2013
 ============================================================================
 */

#include <ctime>
#include <vector>
#include <iostream>
#include <fstream>
#include <ostream>
#include <sstream>
#include <map>

#include <stdio.h>  /* defines FILENAME_MAX */
#ifdef WINDOWS
#include <direct.h>
#define GetCurrentDir _getcwd
#else
#include <unistd.h>
#define GetCurrentDir getcwd
#endif

char cCurrentPath[FILENAME_MAX];

#include <boost/program_options.hpp>
#include <boost/smart_ptr/make_shared.hpp>
#include <boost/ref.hpp>
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/json_parser.hpp>
#include <Eigen/Dense>

#include "INSIGHTv3.h"
#include "BirthDeathModel.h"
#include "InsightAlgorithm.h"
#include "InsightParticleSamplerFactory.h"
#include "IllegalArgumentException.h"
#include "ParticleEvaluatorInsight.h"
#include "UniformSamplerFactory.h"
#include "ScaledUniformSamplerFactory.h"
#include "ProblemFileReader.h"
#include "ModelDescription.h"
#include "UniformSampler.h"
#include "InsightIteration.h"
#include "PrevPopReader.h"

using namespace INSIGHTv3;
using boost::property_tree::ptree;
using boost::property_tree::read_json;

//Default values
const static std::string PROBLEM_FILE = "";
const static int S = 3;
const static double FINAL_TOLERANCE = 0.03;
const static int NUM_PARTICLES = 1000;
const static std::string OUTPUT_FILE = "";
const static std::string PREVIOUS_POPULATION_FILE = "";
const static std::string PARTICLE_INPUT = "/input";
const static std::string PARTICLE_OUTPUT = "/output";
const static bool IGNORE_WEIGHTS = false;
const static bool VERBOSE = false;
const static int PRINT = 20;
const static double BETA = 0.01;
const static double TOL_KOLMOGOROV = 0.0000001;
const static double KAPPA_KOLMOGOROV = 0.005224; // kappa value for M = 80000 and beta = 0.01
const static double TOL_KAPPA_KOLMOGOROV = 0.01;
const static double PARTICLE_TOLERANCE = 1.0;
const static int NUM_TRAJECTORIES = 10;

std::string problem_file_name;
int s;
double final_tolerance;
int num_particles;
std::string output_file_name;
std::string prev_pop_file;
bool ignore_weights;
bool verbose; // surporess alloutput
int print;
double beta;
double tolerance_kolmogorov;
double kappa_kolmogorov;
double tol_kappa_kolmovorog;
double particle_tolerance;
std::string particle_in;
std::string particle_out;
int num_trajectories;

namespace po = boost::program_options;

void handleOptions(int argc, char * argv[]);
void printInfo(ModelDescription& desc);

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

	handleOptions(argc, argv);

	clock_t tic = clock();
	int seed = time(NULL);
#ifdef DEBUG
	seed = 2;
#endif
	RngPtr rng = boost::make_shared<RandomNumberGenerator>(seed);

	ProblemFileReader reader;
	ModelDescription desc = reader.read(rng, problem_file_name);

	printInfo(desc);
	writer_ptr writer = boost::make_shared<InsightOutputWriter>(
			output_file_name);

	KolmogorovComputerPtr kolmogorov_computer = boost::make_shared<
			KolmogorovComputer>(*desc.original_data, beta, tolerance_kolmogorov,
			kappa_kolmogorov);
	sampler_factory_ptr factory =
			boost::make_shared<ScaledUniformSamplerFactory>(boost::ref(rng),
					desc.model->num_params, *desc.lower_bounds_parameter,
					*desc.upper_bounds_parameter, ignore_weights);

	sampler_ptr prior_sampler = boost::make_shared<UniformSampler>(
			boost::ref(rng), *desc.lower_bounds_parameter,
			*desc.upper_bounds_parameter);
	evaluator_ptr evaluator = boost::make_shared<ParticleEvaluatorInsight>(
			boost::ref(*desc.model), *desc.times, boost::ref(rng),
			kolmogorov_computer); // need to use boost::ref, see boost

	InsightIteration iteration(prior_sampler);
	ToleranceProviderPtr tolerance_provider = boost::make_shared<
			FixedSequenceToleranceProvider>(s, desc.model->model,
			kolmogorov_computer);

	sampler_ptr first_sampler;
	std::vector<double> prev_acceptance_rates;
	std::vector<double> prev_tolerances;
	std::vector<int> prev_num_simulations;

	//prev_num_simulations.push_back(20); // set to number of iterations

	if (prev_pop_file.length() > 0) {
		PrevPopReader prev_pop_reader;
		ParticleSet particle_set;
		prev_pop_reader.readPreviousPop(prev_pop_file, desc.model->num_params,
				&particle_set, &prev_acceptance_rates, &prev_tolerances, &prev_num_simulations);
		first_sampler = factory->createSampler(particle_set,
				kolmogorov_computer->getThresholdForS(s), s);
	} else {
		first_sampler = prior_sampler;
	}
	InsightAlgorithm algorithm(iteration, factory, evaluator, writer,
			tolerance_provider, final_tolerance, num_particles,
			desc.model->num_params, prev_acceptance_rates, prev_tolerances, prev_num_simulations);
	IterationLogger logger(print);

	//algorithm.run(first_sampler, logger);

	algorithm._updateTolerance(&num_trajectories, &particle_tolerance, &logger);

    bool acceptable = 0;
    parameters particle(desc.model->num_params);
    //first_sampler->sampleParticle(&particle);

//open particle input file
std::ifstream infile (particle_in.c_str());

//open particle output file
std::ofstream outfile (particle_out.c_str());
if (!outfile.is_open())
    return -1; // error!

std::string line;
int particle_count=0;

while (std::getline(infile, line)) // parse particles one by one
{

    //using text file
    unsigned first = line.find("[");
    unsigned last = line.find("]");
    std::string strNew = line.substr (first+1,last-first);

    std::istringstream iss(strNew);
    char ch;

    double number=2.0;
    int pcount=0;
    while (pcount<desc.model->num_params) {
        iss >> number;
        iss>>ch; //flush the ,
        particle(pcount)=number;
        pcount++;
    }

    if(verbose)
        std::cout << "particle " << particle_count << ": " << particle << std::endl;

    acceptable = evaluator->isParticleAcceptable(particle);
    if(acceptable)
        outfile << "1\n";
    else
        outfile << "0\n";
}
outfile.close();

	clock_t toc = clock();
	if(verbose) {
        std::cout << "\n\nSMC inference complete." << std::endl;
        std::cout << "\nElapsed time: " << ((double) (toc - tic)) / CLOCKS_PER_SEC
                << "sec." << std::endl;
	}
	return 0;
}

void handleOptions(int argc, char * argv[]) {

	po::options_description desc("Allowed options");
    desc.add_options()("help", "produce help message")("problem_file",
            po::value<std::string>(&problem_file_name)->default_value(
                    PROBLEM_FILE),
            "problem file. A problem file must always be provided!")(
            "number_simulations,s", po::value<int>(&s)->default_value(S),
            "starting number of simulations")("final tolerance,e",
            po::value<double>(&final_tolerance)->default_value(FINAL_TOLERANCE),
            "the final tolerance for the algorithm")("pop_size,n",
            po::value<int>(&num_particles)->default_value(NUM_PARTICLES),
            "define the population size")("output_file,O",
            po::value<std::string>(&output_file_name)->default_value(
                    OUTPUT_FILE), "output file")("previous_pop_file,P",
            po::value<std::string>(&prev_pop_file)->default_value(
                    PREVIOUS_POPULATION_FILE), "previous population file")(
            "ignore_weights,i",
            po::value<bool>(&ignore_weights)->default_value(IGNORE_WEIGHTS),
            "ignore weights")("print,p",
            po::value<int>(&print)->default_value(PRINT),
            "define the frequency of print outs")("beta,b",
            po::value<double>(&beta)->default_value(BETA),
            "the confidence of the kolmogorov test for the rejection rule")(
            "tol_kol,k",
            po::value<double>(&tolerance_kolmogorov)->default_value(
                    TOL_KOLMOGOROV),
            "the tolerance for the algorithm to compute the inverse of the kolmogorov distribution")(
            "kappa,K",
            po::value<double>(&kappa_kolmogorov)->default_value(
                    KAPPA_KOLMOGOROV),
            "Kappa for the computation of the Kolmogorov distance")(
            "tol_kol_kap,T",
            po::value<double>(&tol_kappa_kolmovorog)->default_value(
                    TOL_KAPPA_KOLMOGOROV),
            "the tolerance for the algorithm to compute kappa for the computation of the Kolmogorov distance")("X,X",
                    po::value<std::string>(&particle_in)->default_value(PARTICLE_INPUT),"Particle input file")("o,o",
                    po::value<std::string>(&particle_out)->default_value(PARTICLE_OUTPUT),"Particle output file")("t,t",
		po::value<double>(&particle_tolerance)->default_value(
             PARTICLE_TOLERANCE),"Particle tolerance")("N,N",po::value<int>(&num_trajectories)->default_value(
                     NUM_TRAJECTORIES),"Number of trajectories")("v,v",po::value<bool>(&verbose)->default_value(
                             VERBOSE),"Verbose");

	po::positional_options_description p;
	p.add("problem_file", -1);

	po::variables_map vm;
	po::store(po::parse_command_line(argc, argv, desc), vm);
	po::store(
			po::command_line_parser(argc, argv).options(desc).positional(p).run(),
			vm);
	po::notify(vm);

	if (vm.count("help")) {
		std::cout << desc << "\n";
		abort();
	}

	if (problem_file_name.size() == 0) {
		std::cerr << "Problem file is required" << std::endl;
		abort();
	}

	if (output_file_name.size() == 0) {
		GetCurrentDir(cCurrentPath, sizeof(cCurrentPath));
		output_file_name = std::string(cCurrentPath) + "/resutls.txt";
	}

}

void printInfo(ModelDescription& desc) {

    if(verbose) {
        std::cout << "\nThis is INSIGHT v3.0\n" << std::endl;
        std::cout << "Problem file path:" << std::endl;
        std::cout << problem_file_name << std::endl;
        std::cout << "\nThe new particle population will be saved in:" << std::endl;
        std::cout << output_file_name << std::endl;

        std::cout << "\nWill estimate " << desc.model->model_name << " using "
                << num_particles << " SMC particles." << std::endl;
        std::cout << "\nThe data is assumed to have been created using "
                << desc.model->num_outputs << " outputs," << std::endl;
        std::cout << "measured at " << desc.times->size() << " timepoints."
                << std::endl;

        if (prev_pop_file.length() > 0) {
            std::cout << "A previous population is loaded from file: " << std::endl;
            std::cout << prev_pop_file << std::endl;
        }
        std::cout << "To estimate the density of each population "
                << "Kernel Density Estimation with uniform Kernel" << std::endl;
        if (ignore_weights) {
            std::cout << std::endl << "Particle weights will be ignored by sampler!"
                    << std::endl;
        }
        std::cout << "is used." << std::endl;
        std::cout << "\nThe sequential version of the algorithm is used."
                << std::endl << std::endl;
    }
}


#include "sample_models.hpp"

#include "model_distribution.hpp"
#include "types.hpp"
#include "utils.hpp"

#include <algorithm>
#include <boost/log/trivial.hpp>
#include <random>
#include <unordered_map>
#include <utility>
#include <vector>

namespace FilterModel {

ModelSampler::ModelSampler(const std::vector<category_counts_t> &data,
                           std::default_random_engine &generator, const Options &options)
    : model_distribution(data, generator, options),
      generator(generator),
      alphas(ModelSampler::generate_alphas(data.at(0).size(), options)),
      n_objects(data.size()),
      n_categories(data.at(0).size()){};

std::vector<alpha_t> ModelSampler::sample(double epsilon, const delta_t &delta) {
    std::vector<double> priors(alphas.size(), 1.0 / alphas.size());

    std::vector<std::vector<double>> log_alpha_likelyhoods =
        model_distribution.distribution(alphas, epsilon, delta);

    std::vector<alpha_t> models;
    for (int object_index = 0; object_index < n_objects; ++object_index) {
        std::vector<double> alpha_weights(alphas.size(), 0.0);
        double total_weight = 0.0;
        for (int i = 0; i < alphas.size(); ++i) {
            alpha_weights.at(i) = log_alpha_likelyhoods.at(object_index).at(i);
            // check that log_alpha_likelyhoods is sane.
            // shouldn't be uniform.
            total_weight += std::exp(alpha_weights.at(i));
        }

        for (int i = 0; i < alphas.size(); ++i) {
            alpha_weights.at(i) -= std::log(total_weight);
        }

        std::vector<double> probabilities = exp(alpha_weights);

        assert_probability(probabilities);
        // look at probabilities

        // Select an alpha at random for the current object.
        std::uniform_real_distribution<double> distribution(0.0, 1.0);
        double random = distribution(generator);
        double probability_sum = 0.0;
        int i = 0;
        while (probability_sum < random) {
            probability_sum += probabilities.at(i);
            ++i;
        }
        models.push_back(alphas.at(i - 1));
    }
    return models;
};

std::vector<alpha_t> ModelSampler::generate_alphas(int n_categories, const Options &options) {
    if (options.use_smaller_alphas) {
        return generate_real_alphas();
    } else {
        return generate_all_alphas(n_categories);
    }
}

std::vector<alpha_t> ModelSampler::generate_all_alphas(int n_categories) {
    alpha_t current = std::vector<category_t>(n_categories, 0);
    std::vector<alpha_t> alphas;
    alpha_t end = std::vector<category_t>(n_categories, 1);
    while (current != end) {
        current = alpha_t(current);
        int i = 0;
        while (current.at(i) == 1) {
            current.at(i) = 0;
            ++i;
        }
        current.at(i) = 1;
        alphas.push_back(current);
    }
    return alphas;
}

std::vector<alpha_t> ModelSampler::generate_real_alphas() {
    return {{0, 0, 1}, {0, 1, 0}, {1, 1, 1}, {1, 0, 1}};
}
}  // namespace FilterModel


#include "text_exposition_parser.hpp"
#include <boost/spirit/include/qi_operator.hpp>
#include <boost/spirit/include/qi_real.hpp>
#include <boost/spirit/include/qi_rule.hpp>
#include <boost/spirit/include/qi_char.hpp>
#include <boost/spirit/include/qi_uint.hpp>
#include <boost/spirit/include/qi_skip.hpp>
#include <boost/spirit/include/qi_omit.hpp>
#include <boost/spirit/include/qi_eol.hpp>
#include <boost/spirit/include/qi_list.hpp>
#include <boost/spirit/include/qi_char_class.hpp>
#include <boost/spirit/include/qi_parse_auto.hpp>

#include <boost/fusion/include/std_pair.hpp>
#include <boost/fusion/include/adapt_struct.hpp>
#include <boost/fusion/include/vector.hpp>

#include <utility>
#include <string>
#include <map>

BOOST_FUSION_ADAPT_STRUCT(Metric, name, tags, value, timestamp);

Metrics MetricsParser::parseMetrics(const std::string lines)
{
    using namespace boost::spirit;
    using namespace boost::spirit::qi;

    auto quotedString = '"' >> *(char_ - char_("\"")) >> '"';
    auto metricName = +(alnum | char_("_:"));
    auto metricLabel = +(alnum | char_("_"));

    auto labelPair = metricLabel >> '=' >> quotedString;
    auto labelPairs = -( '{' >> (labelPair % ',') >> '}' );

    auto metricLine = metricName >> labelPairs >> ' ' >> double_ >> omit[*space] >> -ulong_long;

    auto commentLine = char_('#') >> *(char_ - eol);

    auto line = *(omit[commentLine] | omit[eol]) >> metricLine;
    auto parser = *line;

    std::vector<Metric> ms;

    parse(lines.begin(), lines.end(), parser, ms);
    return Metrics(ms);
}


// Copyright (c) 2017 Graphcore Ltd. All rights reserved.
#define BOOST_TEST_MODULE DynamicSliceTest
#include <boost/multi_array.hpp>
#include <boost/test/framework.hpp>
#include <boost/test/unit_test.hpp>
#include <cassert>
#include <iostream>
#include <numeric>
#include <poplar/Engine.hpp>
#include <poplar/Interval.hpp>
#include <poplar/Program.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplibs_support/print.hpp>
#include <popops/DynamicSlice.hpp>
#include <popops/codelets.hpp>
#include <poputil/TileMapping.hpp>
#include <poputil/exceptions.hpp>
#include <sstream>
#include <vector>

using namespace poplar;
using namespace poplar::program;
using namespace poputil;
using namespace popops;
using namespace poplibs_support;

using poplibs_support::toString;

constexpr bool useDSMapper = true;

const OptionFlags options;

#define NUM_DIMS 3

struct TestData {
  std::vector<size_t> tDims, sDims;
  boost::multi_array<float, 3> hInit; // fullsized initialiser
  boost::multi_array<float, 3> hSub;  // subTensor, either in or out
  boost::multi_array<float, 3> hUpdateOut;
  TestData(std::vector<size_t> t, std::vector<size_t> s,
           const std::vector<std::vector<std::vector<float>>> &initialData)
      : tDims(t), sDims(s) {
    assert(t.size() == 3);
    assert(s.size() == 3);
    hInit.resize(boost::extents[t[0]][t[1]][t[2]]);
    hSub.resize(boost::extents[s[0]][s[1]][s[2]]);
    hUpdateOut.resize(boost::extents[t[0]][t[1]][t[2]]);
    for (unsigned i = 0; i != initialData.size(); ++i) {
      for (unsigned j = 0; j != initialData[i].size(); ++j) {
        for (unsigned k = 0; k != initialData[i][j].size(); ++k) {
          hInit[i][j][k] = initialData[i][j][k];
        }
      }
    }
  }
};

// Small 3 test data
static const unsigned dimA = 3, dimB = 4, dimC = 2;
static std::vector<size_t> smallTestShape = {dimA, dimB, dimC};
std::vector<std::vector<std::vector<float>>> smallTestData = {
    {{111, 112}, {121, 122}, {131, 132}, {141, 142}},
    {{211, 212}, {221, 222}, {231, 232}, {241, 242}},
    {{311, 312}, {321, 322}, {331, 332}, {341, 342}}};

// long delay data
#define LONG_OUTER 32
#define MAX_DELAY 200
#define ELEM_PER_TAP 4
static std::vector<size_t> delayTestShape = {LONG_OUTER, MAX_DELAY,
                                             ELEM_PER_TAP};
std::vector<std::vector<std::vector<float>>> GenDelayData() {
  std::vector<std::vector<std::vector<float>>> result;
  result.reserve(LONG_OUTER);
  for (unsigned i = 0; i != LONG_OUTER; ++i) {
    result.emplace_back();
    for (unsigned j = 0; j != MAX_DELAY; ++j) {
      result[i].emplace_back();
      for (unsigned k = 0; k != ELEM_PER_TAP; ++k)
        result[i][j].emplace_back((1 + 3 * i + j) + k * (1.0 / ELEM_PER_TAP));
    }
  }
  return result;
}

// map t's specified dimension across tiles
static void MapAcrossTiles(Graph &graph, size_t tilesPerIPU, const Tensor &t) {
  auto nTilesForT = std::min(t.dim(0), tilesPerIPU);
  auto elemPerSlice = t.numElements() / t.dim(0);
  Graph::TileToTensorMapping map;
  for (unsigned a = 0; a != nTilesForT; ++a) {
    std::vector<Interval> submap;
    auto iBegin = a * elemPerSlice;
    {
      auto iEnd =
          (a == nTilesForT - 1) ? t.numElements() : iBegin + elemPerSlice;
      auto interval = Interval(iBegin, iEnd);
      submap.emplace_back(interval);
      map.emplace_back(submap);
    }
  }
  graph.setTileMapping(t, map);
}

static boost::multi_array<float, 3>
refSlice(const std::vector<size_t> &sShape,
         const boost::multi_array<float, 3> &t,
         const std::vector<size_t> &offsets) {
  auto tShape = t.shape();
  boost::multi_array<float, 3> result(
      boost::extents[sShape[0]][sShape[1]][sShape[2]]);
  for (unsigned a = 0; a != sShape[0]; ++a) {
    for (unsigned b = 0; b != sShape[1]; ++b) {
      for (unsigned c = 0; c != sShape[2]; ++c) {
        auto refA = (offsets[0] + a) % tShape[0];
        auto refB = (offsets[1] + b) % tShape[1];
        auto refC = (offsets[2] + c) % tShape[2];
        auto value = t[refA][refB][refC];
        result[a][b][c] = value;
      }
    }
  }
  return result;
}

static boost::multi_array<float, 3>
refUpdate(const boost::multi_array<float, 3> &t,
          const boost::multi_array<float, 3> &s,
          const std::vector<size_t> &offsets) {
  auto tShape = t.shape();
  auto sShape = s.shape();
  boost::multi_array<float, 3> result(
      boost::extents[tShape[0]][tShape[1]][tShape[2]]);
  result = t;
  for (unsigned a = 0; a != sShape[0]; ++a) {
    for (unsigned b = 0; b != sShape[1]; ++b) {
      for (unsigned c = 0; c != sShape[2]; ++c) {
        auto refA = (offsets[0] + a) % tShape[0];
        auto refB = (offsets[1] + b) % tShape[1];
        auto refC = (offsets[2] + c) % tShape[2];
        auto value = s[a][b][c];
        result[refA][refB][refC] = value;
      }
    }
  }
  return result;
}

static void checkResult(const boost::multi_array<float, 3> &m,
                        const boost::multi_array<float, 3> &ref) {
  auto shape = m.shape();

  std::stringstream ss;
  for (unsigned a = 0; a != shape[0]; ++a) {
    ss << "[" << a << "] {";
    for (unsigned b = 0; b != shape[1]; ++b) {
      std::string sep = "";
      ss << "{";
      for (unsigned c = 0; c != shape[2]; ++c) {
        auto result = m[a][b][c];
        auto refResult = ref[a][b][c];
        ss << sep << result << " == " << refResult;
        sep = ", ";

        BOOST_CHECK_EQUAL(result, refResult);
      }
      ss << "}";
    }
    ss << "}";
  }
  BOOST_TEST_MESSAGE(ss.str());
}

// Check dynamicSliceND() extracts \a sliceSizes elements from the \a sliceDims
// dimensions for all possible offsets.
void sliceTestND(unsigned tilesPerIPU, const std::vector<size_t> &testShape,
                 const std::vector<std::vector<std::vector<float>>> &testBase,
                 const std::vector<std::size_t> &sliceDims,
                 const std::vector<std::size_t> &sliceSizes) {
  BOOST_TEST_MESSAGE(
      "Test " << boost::unit_test::framework::current_test_case().p_name);

  auto device = createTestDevice(TEST_TARGET, 1, tilesPerIPU);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);
  std::vector<size_t> t1Shape = testShape;
  auto tWantedOffsets =
      graph.addVariable(UNSIGNED_INT, {sliceDims.size()}, "wantedOffsets");
  graph.setTileMapping(tWantedOffsets, 0);

  Tensor t1;
  if (useDSMapper) {
    // Test with the old test-specific allocator
    t1 = graph.addVariable(FLOAT, t1Shape, "t1");
    BOOST_TEST_MESSAGE("Created tensor t1: " << t1);
    MapAcrossTiles(graph, tilesPerIPU, t1);
  } else {
    // Test with the new reference tensor allocator.
    t1 = createSliceableTensor(graph, FLOAT, testShape, sliceDims, sliceSizes,
                               2, "t1");
  }
  BOOST_TEST_MESSAGE("t1 is " << t1 << " mapping "
                              << toString(graph.getTileMapping(t1)));

  auto prog = Sequence();
  auto tOut = dynamicSlice(graph, t1, tWantedOffsets, sliceDims, sliceSizes,
                           prog, "DSND");

  const auto tOutShape = tOut.shape();
  BOOST_TEST_MESSAGE("output tensor is "
                     << tOut << " mapping "
                     << toString(graph.getTileMapping(tOut)));

  // Check output Tensor shape is correct
  std::vector<size_t> wantedShape = t1.shape();
  for (unsigned i = 0; i != sliceDims.size(); ++i) {
    wantedShape[sliceDims[i]] = sliceSizes[i];
  }
  for (unsigned d = 0; d != t1.rank(); ++d) {
    auto expectedSize = wantedShape[d] ? wantedShape[d] : t1.dim(d);
    BOOST_CHECK_EQUAL(tOutShape[d], expectedSize);
  }

  graph.createHostWrite("in", t1);
  graph.createHostWrite("selector", tWantedOffsets);
  graph.createHostRead("out", tOut);

  BOOST_TEST_MESSAGE("Creating engine");
  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);

    TestData testData(t1Shape, wantedShape, testBase);

    eng.writeTensor("in", testData.hInit.data(),
                    testData.hInit.data() + t1.numElements());

    std::vector<unsigned> nOffsets(t1.rank(), 1);
    for (auto dim : sliceDims) {
      nOffsets[dim] = t1.dim(dim);
    }
    assert(t1.rank() == NUM_DIMS);
    for (unsigned sliceA = 0; sliceA != nOffsets[0]; ++sliceA) {
      for (unsigned sliceB = 0; sliceB != nOffsets[1]; ++sliceB) {
        for (unsigned sliceC = 0; sliceC != nOffsets[2]; ++sliceC) {
          unsigned offsets[NUM_DIMS] = {sliceA, sliceB, sliceC};
          unsigned hOffsets[NUM_DIMS];
          for (unsigned i = 0; i != sliceDims.size(); ++i) {
            hOffsets[i] = offsets[sliceDims[i]];
          }
          std::vector<size_t> checkOffsets = {{sliceA, sliceB, sliceC}};
          eng.writeTensor("selector", hOffsets, &hOffsets[sliceDims.size()]);
          for (unsigned i = 0; i != testData.hUpdateOut.num_elements(); ++i)
            testData.hUpdateOut.data()[i] = 0.0;
          BOOST_TEST_MESSAGE("Engine run " << toString(checkOffsets));
          eng.run();
          eng.readTensor("out", testData.hSub.data(),
                         testData.hSub.data() + tOut.numElements());
          boost::multi_array<float, 3> refResult =
              refSlice(wantedShape, testData.hInit, checkOffsets);
          checkResult(testData.hSub, refResult);
        }
      }
    }
  });
}

static void subTestSmallSlice(unsigned tilesPerIPU,
                              const std::vector<std::size_t> &sliceDims,
                              const std::vector<std::size_t> &sliceSizes) {
  sliceTestND(tilesPerIPU, smallTestShape, smallTestData, sliceDims,
              sliceSizes);
}

BOOST_AUTO_TEST_SUITE(SingleDim)

// Test empty slice
BOOST_AUTO_TEST_CASE(Slice_Empty) { subTestSmallSlice(5, {}, {}); }

// Test slicing of a single dimension
BOOST_AUTO_TEST_CASE(Slice_5_0_1) { subTestSmallSlice(5, {0}, {1}); }
BOOST_AUTO_TEST_CASE(Slice_5_0_2) { subTestSmallSlice(5, {0}, {2}); }
BOOST_AUTO_TEST_CASE(Slice_5_1_1) { subTestSmallSlice(5, {1}, {1}); }
BOOST_AUTO_TEST_CASE(Slice_5_1_2) { subTestSmallSlice(5, {1}, {2}); }
BOOST_AUTO_TEST_CASE(Slice_5_2_1) { subTestSmallSlice(5, {2}, {1}); }
BOOST_AUTO_TEST_CASE(Slice_5_2_2) { subTestSmallSlice(5, {2}, {2}); }

BOOST_AUTO_TEST_SUITE_END()

// Multidimensional slicing
BOOST_AUTO_TEST_SUITE(MultiDim)

// dimensions 1 & 2
BOOST_AUTO_TEST_CASE(ND_1_1_0) { subTestSmallSlice(5, {0, 1}, {1, 1}); }
// all 3 dimensions
BOOST_AUTO_TEST_CASE(ND_1_1_1) { subTestSmallSlice(5, {0, 1, 2}, {1, 1, 1}); }
// dimensions 0 and 2, producing 2xdimBx2 output
BOOST_AUTO_TEST_CASE(ND_2_0_2) { subTestSmallSlice(5, {0, 2}, {2, 2}); }
// 2x2x2 outputs
BOOST_AUTO_TEST_CASE(ND_2_4_2) {
  // The same result has as for 2_0_2 but with an extra compute set and
  // additional testing of dim1 at all 4 offsets
  subTestSmallSlice(5, {0, 1, 2}, {2, 4, 2});
}

BOOST_AUTO_TEST_SUITE_END()

// large-buffer update
BOOST_AUTO_TEST_SUITE(LargeBuffer)

BOOST_AUTO_TEST_CASE(circTest) {
  auto delayTestData = GenDelayData();
  sliceTestND(24, delayTestShape, delayTestData, {1}, {1});
}

BOOST_AUTO_TEST_SUITE_END()

// Dynamic update
// Check dynamicSliceND() extracts \a sliceSizes elements from the \a sliceDims
// dimensions for all possible offsets.
void updateTestND(unsigned tilesPerIPU, const std::vector<size_t> &testShape,
                  const std::vector<std::vector<std::vector<float>>> &testBase,
                  const std::vector<std::size_t> &sliceDims,
                  const std::vector<std::size_t> &sliceSizes) {
  BOOST_TEST_MESSAGE(
      "Test " << boost::unit_test::framework::current_test_case().p_name);

  auto device = createTestDevice(TEST_TARGET, 1, tilesPerIPU);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);
  std::vector<size_t> t1Shape = testShape;

  std::vector<size_t> subShape = t1Shape;
  for (unsigned i = 0; i != sliceDims.size(); ++i) {
    subShape[sliceDims[i]] = sliceSizes[i];
  }
  Tensor t1, s1;

  if (useDSMapper) {
    // Test with the old test-specific allocator
    t1 = graph.addVariable(FLOAT, t1Shape, "t1");
    BOOST_TEST_MESSAGE("Created tensor t1: " << t1);
    s1 = graph.addVariable(FLOAT, subShape, "s1");
    BOOST_TEST_MESSAGE("Created tensor s1: " << s1);
    MapAcrossTiles(graph, tilesPerIPU, t1);
    MapAcrossTiles(graph, tilesPerIPU, s1);
  } else {
    // Test with the new reference tensor allocator.
    t1 = createSliceableTensor(graph, FLOAT, t1Shape, sliceDims, sliceSizes, 2,
                               "t1");
    s1 = createSliceTensor(graph, t1, sliceDims, sliceSizes, 1, "s1")
             .squeeze({0});
  }
  auto tWantedOffsets =
      graph.addVariable(UNSIGNED_INT, {sliceDims.size()}, "wantedOffsets");
  graph.setTileMapping(tWantedOffsets, 0);
  BOOST_TEST_MESSAGE("t1 is " << t1 << " mapping "
                              << toString(graph.getTileMapping(t1)));
  BOOST_TEST_MESSAGE("s1 is " << s1 << " mapping "
                              << toString(graph.getTileMapping(t1)));

  auto prog = Sequence();

  dynamicUpdate(graph, t1, s1, tWantedOffsets, sliceDims, sliceSizes, prog,
                "DSUpdate");

  graph.createHostWrite("in", t1);
  graph.createHostWrite("update", s1);
  graph.createHostWrite("selector", tWantedOffsets);
  graph.createHostRead("out", t1);

  BOOST_TEST_MESSAGE("Creating engine");
  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);

    TestData testData(t1Shape, subShape, testBase);

    for (unsigned a = 0; a != subShape[0]; ++a) {
      for (unsigned b = 0; b != subShape[1]; ++b) {
        for (unsigned c = 0; c != subShape[2]; ++c) {
          testData.hSub[a][b][c] = testData.hInit[a][b][c] * 0.001;
        }
      }
    }
    eng.writeTensor("update", testData.hSub.data(),
                    testData.hSub.data() + s1.numElements());

    std::vector<unsigned> nOffsets(t1.rank(), 1);
    for (auto dim : sliceDims) {
      nOffsets[dim] = t1.dim(dim);
    }
    assert(t1.rank() == NUM_DIMS);
    for (unsigned sliceA = 0; sliceA != nOffsets[0]; ++sliceA) {
      for (unsigned sliceB = 0; sliceB != nOffsets[1]; ++sliceB) {
        for (unsigned sliceC = 0; sliceC != nOffsets[2]; ++sliceC) {
          unsigned offsets[NUM_DIMS] = {sliceA, sliceB, sliceC};
          unsigned hOffsets[NUM_DIMS];
          for (unsigned i = 0; i != sliceDims.size(); ++i) {
            hOffsets[i] = offsets[sliceDims[i]];
          }
          std::vector<size_t> checkOffsets = {{sliceA, sliceB, sliceC}};
          eng.writeTensor("in", testData.hInit.data(),
                          testData.hInit.data() + t1.numElements());
          eng.writeTensor("selector", hOffsets, &hOffsets[sliceDims.size()]);
          for (unsigned i = 0; i != testData.hUpdateOut.num_elements(); ++i)
            testData.hUpdateOut.data()[i] = 0.0;
          BOOST_TEST_MESSAGE("Engine run " << toString(checkOffsets));
          eng.run();
          eng.readTensor("out", testData.hUpdateOut.data(),
                         testData.hUpdateOut.data() + t1.numElements());

          boost::multi_array<float, 3> refResult =
              refUpdate(testData.hInit, testData.hSub, checkOffsets);
          checkResult(testData.hUpdateOut, refResult);
        }
      }
    }
  });
}

static void testSmallUpdate(unsigned tilesPerIPU,
                            const std::vector<std::size_t> &sliceDims,
                            const std::vector<std::size_t> &sliceSizes) {
  updateTestND(tilesPerIPU, smallTestShape, smallTestData, sliceDims,
               sliceSizes);
}

BOOST_AUTO_TEST_SUITE(Update)

// Test full update
BOOST_AUTO_TEST_CASE(Update_Full) {
  subTestSmallSlice(5, {0, 1, 2}, smallTestShape);
}

BOOST_AUTO_TEST_CASE(Update_Empty) {
  // result should be the same as the full update with offsets 0,0,0
  subTestSmallSlice(5, {}, {});
}

// Test insertion of a single dimension
BOOST_AUTO_TEST_CASE(Update_5_0) { testSmallUpdate(5, {0}, {1}); }
BOOST_AUTO_TEST_CASE(Update_5_1) { testSmallUpdate(5, {1}, {1}); }
BOOST_AUTO_TEST_CASE(Update_5_2) { testSmallUpdate(5, {2}, {1}); }
// Test insertion of a single element
BOOST_AUTO_TEST_CASE(Update_5_element) {
  testSmallUpdate(5, {0, 1, 2}, {1, 1, 1});
}
// Test insertion of a 2x2 element
BOOST_AUTO_TEST_CASE(Update_5_2x2) { testSmallUpdate(5, {0, 1}, {2, 2}); }

BOOST_AUTO_TEST_SUITE_END()

BOOST_AUTO_TEST_SUITE(Misc)

// Two large contiguous regions which exceed the maximum size allowed by the
// 2D vertex fields. Only creates additional vertices for IPU2.
void largeRegions() {
  constexpr static unsigned tilesPerIPU = 1;
  const auto dataType = poplar::HALF;

  auto device =
      createTestDevice(TEST_TARGET, 1, tilesPerIPU,
                       /* compileIPUCode */ true /* for exchange code size */);
  const auto &target = device.getTarget();
  Graph graph(target);
  popops::addCodelets(graph);
  popops::addCodelets(graph);
  Sequence uploadProg, prog, downloadProg;

  const auto numColumns = 32000U;

  const auto sliceOffset = 1;
  auto offset = graph.addConstant(UNSIGNED_INT, {1}, sliceOffset);
  graph.setTileMapping(offset, 0);
  auto t1 = graph.addVariable(dataType, {1, numColumns}, "Base tensor1");
  auto t2 = graph.addVariable(dataType, {1, numColumns}, "Base tensor2");
  auto baseTensor = concat({t1, t2});
  graph.setTileMapping(baseTensor, 0);
  auto subTensor = popops::dynamicSlice(graph, baseTensor, offset, {0}, {1},
                                        prog, "dynamicSlice");

  const auto numElements = baseTensor.numElements();
  const auto firstSliceData = 1.0;
  const auto secondSliceData = -1.0;
  graph.createHostWrite("baseTWr", baseTensor, true);
  graph.createHostWrite("subTWr", subTensor, true);
  graph.createHostRead("subTRd", subTensor, true);

  std::vector<float> hIn(numElements);
  std::fill(hIn.begin(), hIn.begin() + numElements / 2, firstSliceData);
  std::fill(hIn.begin() + numElements / 2, hIn.end(), secondSliceData);
  std::vector<char> rawIn(target.getTypeSize(HALF) * hIn.size());
  poplar::copyFloatToDeviceHalf(target, hIn.data(), rawIn.data(), hIn.size());

  std::vector<float> hOut(subTensor.numElements());
  std::vector<char> rawOut(target.getTypeSize(HALF) * hOut.size());

  Engine e(graph, std::move(prog), options);

  device.bind([&](const Device &d) {
    e.load(d);
    e.writeTensor("baseTWr", rawIn.data(), rawIn.data() + rawIn.size());
    e.writeTensor("subTWr", rawOut.data(), rawOut.data() + rawOut.size());
    e.run();
    e.readTensor("subTRd", rawOut.data(), rawOut.data() + rawOut.size());
  });

  poplar::copyDeviceHalfToFloat(target, rawOut.data(), hOut.data(),
                                hOut.size());

  // Now check tensor
  const auto value = sliceOffset == 1 ? secondSliceData : firstSliceData;
  for (unsigned i = 0; i != hOut.size(); ++i) {
    BOOST_CHECK_EQUAL(hOut[i], value);
  }
}

BOOST_AUTO_TEST_CASE(LargeRegionsToSplitTest) { largeRegions(); }

// Check that slices happen in the best order possible. Note that this currently
// abuses Graph::outputComputeGraph(). If this test breaks because of changes
// there I suggest you just delete this test.
BOOST_AUTO_TEST_CASE(SliceOrder) {
  // Input Tensor size: 100 x 50 x 10
  // Output Tensor size: 90 x 15 x 8

  // dims: [2, 0, 1]
  // sizes: [9, 90, 15]

  // It should be smart enough to realise that it should reorder the slicing
  // so that it slices the dimensions in the order [1, 2, 0] (and
  // idxOrder should be [2, 0, 1]).

  auto device = createTestDevice(TEST_TARGET, 1, 4);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  std::vector<size_t> t1Shape = {100, 50, 10};
  std::vector<std::size_t> sliceDims = {2, 0, 1};
  std::vector<std::size_t> sliceSizes = {9, 90, 15};

  auto offset = graph.addVariable(UNSIGNED_INT, {t1Shape.size()}, "offset");
  graph.setTileMapping(offset, 0);

  Tensor input;
  if (useDSMapper) {
    input = graph.addVariable(FLOAT, t1Shape, "input");
    MapAcrossTiles(graph, graph.getTarget().getTilesPerIPU(), input);
  } else {
    input = createSliceableTensor(graph, FLOAT, t1Shape, sliceDims, sliceSizes,
                                  2, "input");
  }

  auto prog = Sequence();

  auto out = dynamicSlice(graph, input, offset, sliceDims, sliceSizes, prog);

  // Check that the graph is correct... in the ugliest way possible.
  std::stringstream computeGraph;
  graph.outputComputeGraph(computeGraph, {prog});

  // "/dynamicSlice_d1/slice" should be before "/dynamicSlice_d2/slice" and
  // so on.

  std::string cg = computeGraph.str();
  auto d0_idx = cg.find("dynamicSlice_d0/slice");
  auto d1_idx = cg.find("dynamicSlice_d1/slice");
  auto d2_idx = cg.find("dynamicSlice_d2/slice");
  const auto npos = std::string::npos;
  BOOST_CHECK(d0_idx != npos && d1_idx != npos && d2_idx != npos);
  BOOST_CHECK(d1_idx < d2_idx);
  BOOST_CHECK(d2_idx < d0_idx);
}

#if 0
BOOST_AUTO_TEST_CASE(ImbalanceTest) {
  auto T = 4u;
  auto device = createTestDevice(TEST_TARGET, 1, T);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);
  auto N = 1024ul;
  auto M = 1024ul;

  auto t = graph.addVariable(FLOAT, {N, M}, "t");

  mapTensorLinearly(graph, t);
  // Check that no matter which way the tensor is sliced, the result is
  // balanced across tiles.
  Sequence prog;
  auto offset = graph.addVariable(UNSIGNED_INT, {1}, "offset");
  graph.setTileMapping(offset, 0);

  auto s1 = dynamicSlice(graph, t, offset, {0}, {1}, prog, "").flatten();
  BOOST_CHECK_EQUAL(s1.dim(0), M);
  BOOST_CHECK_EQUAL(getTileImbalance(graph, s1), 0);
  auto s2 =
      dynamicSlice(graph, t.transpose(), offset, {0}, {1}, prog, "").flatten();
  BOOST_CHECK_EQUAL(s2.dim(0), N);
  BOOST_CHECK_EQUAL(getTileImbalance(graph, s2), 0);
}
#endif
BOOST_AUTO_TEST_CASE(LargeTensorSlice) {
  // This test should pass with large T - but graph construction becomes
  // slow (a couple of minutes for T=1024)
  auto T = 64u;
  auto device = createTestDevice(TEST_TARGET, 1, T);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);
  auto N = 32 * 1024ul;
  auto M = 2 * T; // multiple elements can be stored on each Tile
  // Map the tensor carefully to ensure balance and minimise edge pointers
  auto tmpOrder = graph.addVariable(HALF, {T, N, M / T}, "t");
  mapTensorLinearly(graph, tmpOrder);
  auto t = tmpOrder.dimShuffle({1, 0, 2}).reshape({N, M});

  Sequence prog;
  auto offset = graph.addVariable(UNSIGNED_INT, {1}, "offset");
  graph.setTileMapping(offset, 0);
  auto s1 = dynamicSlice(graph, t, offset, {0}, {1}, prog, "").flatten();
  BOOST_CHECK_EQUAL(s1.dim(0), M);
  BOOST_CHECK_EQUAL(getTileImbalance(graph, s1), 0);

  OptionFlags engineOptions{{"showExecutionSteps", "true"},
                            {"showVarStorage", "true"}};
  // Actually build the graph to check that it fits onto the target
  // This will fail if many edge pointers or significant exchange is required
  Engine eng(graph, prog, options);

  std::stringstream ss;
  eng.printProfileSummary(ss, engineOptions);
  BOOST_TEST_MESSAGE(ss.str());
}

BOOST_AUTO_TEST_CASE(SliceableTensorFromSlice) {
  // This test checks the expected properties of the returned tensor from
  // createSliceableTensorFromSlice which are that the tensor has the
  // same contiguous regions in each slice on each tile, and that
  // the mapping of each individual slice to tiles is the same as that
  // of the given reference slice.
  constexpr std::size_t numTiles = 64;
  auto device = createTestDevice(TEST_TARGET, 1, numTiles);
  Graph graph(device.getTarget());
  auto slice = graph.addVariable(HALF, {20, 10, 10}, "s");

  // Shuffle contiguous regions.
  slice = slice.dimShuffle({2, 0, 1});
  // Map shuffled stuff.
  mapTensorLinearly(graph, slice, 1, 1);

  const std::vector<std::size_t> dims = {0, 1};
  const std::vector<std::size_t> numSlices = {2, 3};
  std::vector<std::size_t> expectedShape = slice.shape();
  for (std::size_t i = 0; i < dims.size(); ++i) {
    expectedShape[dims[i]] *= numSlices[i];
  }
  // Expect the most sliced dimension to be sliced first as
  // this reduces the work for later slices.
  const std::vector<std::size_t> expectedOrder = {1, 0};
  auto totalNumSlices =
      std::accumulate(numSlices.begin(), numSlices.end(), std::size_t(1),
                      std::multiplies<std::size_t>());
  auto t = createSliceableTensorFromSlice(graph, slice, dims, numSlices, "t");
  const auto tShape = t.shape();

  BOOST_CHECK_EQUAL(t.numElements(), slice.numElements() * totalNumSlices);
  BOOST_CHECK_EQUAL_COLLECTIONS(tShape.begin(), tShape.end(),
                                expectedShape.begin(), expectedShape.end());
  std::vector<Tensor> slices(totalNumSlices);
  std::vector<Tensor *> slicePtrs(totalNumSlices);
  for (std::size_t i = 0; i < 3; ++i) {
    for (std::size_t j = 0; j < 2; ++j) {
      slices[i * 2 + j] = t.slice(i * 20, (i + 1) * 20, 1)
                              .slice(j * 10, (j + 1) * 10, 0)
                              .flatten();
      slicePtrs[i * 2 + j] = &slices[i * 2 + j];
    }
  }
  auto sliceFlat = slice.flatten();
  graph.reorderToSimplify(&sliceFlat, slicePtrs);
  const auto referenceMapping = graph.getTileMapping(sliceFlat);
  for (std::size_t i = 0; i < 3; ++i) {
    for (std::size_t j = 0; j < 2; ++j) {
      // Expect each slice to be contiguous on each tile when reordered
      // to be contiguous with respect to the reference slice on each tile.
      const auto tSlice = slices[i * 2 + j];
      BOOST_CHECK_EQUAL(tSlice.numElements(), slice.numElements());

      const auto mapping = graph.getTileMapping(tSlice);
      BOOST_CHECK(mapping == referenceMapping);
      for (unsigned tile = 0; tile < mapping.size(); ++tile) {
        if (!mapping[tile].empty()) {
          auto contiguousRegions =
              graph.getSortedContiguousRegions(tSlice, mapping[tile]);
          BOOST_CHECK_EQUAL(contiguousRegions.size(), 1);
        }
      }
    }
  }
}

BOOST_AUTO_TEST_CASE(GetSliceMapping) {
  // Check the mapping can be retrieved using a very simple test case
  // just to see that the function runs
  auto device = createTestDevice(TEST_TARGET, 1, 4);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  std::vector<size_t> tShape = {100, 64};
  auto input =
      createSliceableTensor(graph, FLOAT, tShape, {0}, {1}, 2, "input");

  auto mapping = getSliceMapping(graph, input, {0}, {1});
  std::vector<std::vector<Interval>> expectedMapping = {
      {{0, 16}}, {{16, 32}}, {{32, 48}}, {{48, 64}}};

  auto checkResult = [&]() {
    if (mapping.size() != expectedMapping.size()) {
      return false;
    }
    for (unsigned i = 0; i < mapping.size(); i++) {
      if (mapping[i].size() != expectedMapping[i].size()) {
        return false;
      }
      for (unsigned j = 0; j < mapping[i].size(); j++) {
        if (mapping[i][j] != expectedMapping[i][j]) {
          return false;
        }
      }
    }
    return true;
  };
  BOOST_TEST(checkResult());
}

BOOST_AUTO_TEST_SUITE_END()

void multislice(const std::vector<uint32_t> &indicies,
                const std::vector<std::size_t> &indiciesShape,
                bool planAsEmbedding) {
  // This test should pass with large T - but graph construction becomes
  // slow (a couple of minutes for T=1024)
  assert(indiciesShape.size() == 2); // max 2 dims supported by this test
  assert(indicies.size() == indiciesShape[0] * indiciesShape[1]);
  const auto T = 16;        // tiles
  const auto D = 1501 * 10; // dictionary size
  const auto E = 8;         // embedding size
  auto device = createTestDevice(TEST_TARGET, 1, T);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);
  std::vector<std::size_t> sliceDims{0};
  std::vector<std::size_t> sliceSizes{1};

  const auto options = OptionFlags();
  auto plan = SlicePlan();
  if (planAsEmbedding) {
    plan = embedding::plan(graph, FLOAT, D, E, {indicies.size()}, options);
  }
  // Map the tensor carefully to ensure balance and minimise edge pointers
  auto t = createSliceableTensor(graph, FLOAT, {D, E}, sliceDims, sliceSizes,
                                 plan, options, "t");
  Sequence prog;

  auto offsetInit =
      graph.addConstant(UNSIGNED_INT, indiciesShape, indicies.data(), "offset");
  graph.setTileMapping(offsetInit, 0);
  auto offset = createIndicesTensor(graph, sliceDims, indicies.size(), plan,
                                    options, "offset");
  prog.add(Copy(offsetInit, offset));
  auto s = multiSlice(graph, t, offset, sliceDims, sliceSizes, prog, plan,
                      options, "MultisliceTest");

  BOOST_CHECK_EQUAL(s.rank(), t.rank() + 1);
  BOOST_CHECK_EQUAL(s.dim(0), indiciesShape[0]);
  BOOST_CHECK_EQUAL(s.dim(1), indiciesShape[1]);
  BOOST_CHECK_EQUAL(s.dim(2), E);

  graph.createHostWrite("in", t, true);
  graph.createHostRead("out", s, true);
  std::vector<uint32_t> hIn(t.numElements());
  std::vector<uint32_t> hOut(s.numElements());
  std::iota(hIn.begin(), hIn.end(), 0u);
  OptionFlags engineOptions{{"showExecutionSteps", "true"},
                            {"showVarStorage", "true"}};
  // Engine creation will fail for non-cpu targets if many edge pointers or
  // significant exchange is required; this should not happen if
  // createSliceableTensor() has given a good layout
  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in", hIn.data(), hIn.data() + hIn.size());
    eng.run();
    eng.readTensor("out", hOut.data(), hOut.data() + hOut.size());
  });
  unsigned outIdx = 0;
  for (const auto &e : hOut)
    BOOST_TEST_MESSAGE("MSlice Output[" << outIdx++ << "] = " << e);
  for (unsigned i = 0; i != indicies.size(); ++i) {
    auto d = indicies[i];
    for (unsigned elem = 0; elem != E; ++elem) {
      unsigned expected = hIn[d * E + elem];
      BOOST_CHECK_EQUAL(hOut[i * E + elem], expected);
    }
  }
  std::stringstream ss;
  eng.printProfileSummary(ss, engineOptions);
  BOOST_TEST_MESSAGE(ss.str());
}

BOOST_AUTO_TEST_SUITE(MultiSlice)

// test the looping multislice
BOOST_AUTO_TEST_CASE(MultiSlice5) {
  multislice({100, 0, 50, 48, 49}, {5, 1}, false);
}

// test the inlined multislice
BOOST_AUTO_TEST_CASE(MultiSlice2) { multislice({100, 0}, {2, 1}, false); }

// test the fast vertex
BOOST_AUTO_TEST_CASE(MultiSlice10) {
  multislice({2, 1, 2, 1, 80, 70, 60, 50, 40, 30}, {10, 1}, false);
}

// test the looping multislice
BOOST_AUTO_TEST_CASE(MultiSlice5_AsEmbedding) {
  multislice({100, 0, 50, 48, 49}, {5, 1}, true);
}

// test the inlined multislice
BOOST_AUTO_TEST_CASE(MultiSlice2_AsEmbedding) {
  multislice({100, 0}, {2, 1}, true);
}

// test the fast vertex
BOOST_AUTO_TEST_CASE(MultiSlice10_AsEmbedding) {
  multislice({2, 1, 2, 1, 80, 70, 60, 50, 40, 30}, {10, 1}, true);
}

// test heuristic which checks for mapping of a slice.
// if this doesn't kick in we will run out of memory on some
// tiles hence we check for an error constructing the engine.
BOOST_AUTO_TEST_CASE(MultiSlicePoorlyMapped) {
  const unsigned numTiles = 16u;
  const unsigned D = 1501u * 10u;
  const unsigned E = 8u;

  auto device = createTestDevice(TEST_TARGET, 1, numTiles);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  const std::vector<std::size_t> sliceDims{0};
  const std::vector<std::size_t> sliceSizes{1};
  const std::vector<std::uint32_t> indices = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};

  const auto options = OptionFlags();
  const auto plan = SlicePlan();
  // We'll create this with the sliced dimension as the innermost
  // so that our rearrangement is hopefully niceish anyway mostly
  // for test speed purposes.
  auto t = graph.addVariable(FLOAT, {E, D}, "t").transpose();
  // We'll map this with the unsliced dimension as innermost such
  // that the first slice will most certainly reside on a single tile.
  // This is the case where there is plenty of data that will be
  // transferred to a single tile and so we should spread it to
  // slice properly.
  mapTensorLinearly(graph, t, 0, 1);

  Sequence prog;

  auto offsetInit = graph.addConstant(UNSIGNED_INT, {indices.size()},
                                      indices.data(), "offset");
  graph.setTileMapping(offsetInit, 0);
  auto offset = createIndicesTensor(graph, sliceDims, indices.size(), plan,
                                    options, "offset");

  prog.add(Copy(offsetInit, offset));
  auto s = multiSlice(graph, t, offset, sliceDims, sliceSizes, prog, plan,
                      options, "multiSliceTest");

  BOOST_CHECK_EQUAL(s.rank(), t.rank() + 1);
  BOOST_CHECK_EQUAL(s.dim(0), indices.size());
  BOOST_CHECK_EQUAL(s.dim(1), 1);
  BOOST_CHECK_EQUAL(s.dim(2), E);

  graph.createHostWrite("in", t, true);
  graph.createHostRead("out", s, true);
  std::vector<std::uint32_t> hIn(t.numElements());
  std::vector<std::uint32_t> hOut(t.numElements());
  std::iota(hIn.begin(), hIn.end(), 0u);

  Engine e(graph, prog, options);
  device.bind([&](const Device &d) {
    e.load(d);
    e.writeTensor("in", hIn.data(), hIn.data() + hIn.size());
    e.run();
    e.readTensor("out", hOut.data(), hOut.data() + hOut.size());
  });

  unsigned outIdx = 0;
  for (const auto &e : hOut)
    BOOST_TEST_MESSAGE("MSlice Output[" << outIdx++ << "] = " << e);
  for (unsigned i = 0; i != indices.size(); ++i) {
    auto d = indices[i];
    for (unsigned elem = 0; elem != E; ++elem) {
      unsigned expected = hIn[d * E + elem];
      BOOST_CHECK_EQUAL(hOut[i * E + elem], expected);
    }
  }
  std::stringstream ss;
  e.printProfileSummary(ss, {
                                {"showExecutionSteps", "true"},
                                {"showVarStorage", "true"},
                            });
  BOOST_TEST_MESSAGE(ss.str());
}

BOOST_AUTO_TEST_SUITE_END()

void multiupdate(const std::vector<uint32_t> &indicies,
                 const std::vector<std::size_t> &indiciesShape,
                 bool planAsEmbedding, bool accumulate = false,
                 float updateScaling = 1.0,
                 const unsigned E = 8) // embedding size
{
  // This test should pass with large T - but graph construction becomes
  // slow (a couple of minutes for T=1024)
  assert(indiciesShape.size() == 2); // max 2 dims supported by this test
  assert(indicies.size() == indiciesShape[0] * indiciesShape[1]);
  const auto T = 16;        // tiles
  const auto D = 1501 * 10; // dictionary size
  auto device = createTestDevice(TEST_TARGET, 1, T);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);
  std::vector<std::size_t> sliceDims{0};
  std::vector<std::size_t> sliceSizes{1};
  Tensor scale;

  const auto options = OptionFlags();
  auto plan = SlicePlan();
  if (planAsEmbedding) {
    plan = embedding::plan(graph, HALF, D, E, {indicies.size()}, options);
  }

  // Map the tensor carefully to ensure balance and minimise edge pointers
  auto t = createSliceableTensor(graph, HALF, {D, E}, sliceDims, sliceSizes,
                                 plan, options, "t");
  auto s = createSliceTensor(graph, HALF, {D, E}, sliceDims, sliceSizes,
                             indicies.size(), plan, options, "s");
  Sequence prog;
  auto offsetInit =
      graph.addConstant(UNSIGNED_INT, indiciesShape, indicies.data(), "offset");
  graph.setTileMapping(offsetInit, 0);
  auto offset = createIndicesTensor(graph, sliceDims, indicies.size(), plan,
                                    options, "offset");
  prog.add(Copy(offsetInit, offset));
  if (!accumulate) {
    multiUpdate(graph, t, s, offset, sliceDims, sliceSizes, prog, plan, options,
                "MultisliceTest");
  } else {
    scale = graph.addVariable(HALF, {}, "scale");
    graph.setTileMapping(scale, 0);
    multiUpdateAdd(graph, t, s, offset, scale, sliceDims, sliceSizes, prog,
                   plan, options, "MultisliceTest");
  }

  BOOST_CHECK_EQUAL(s.rank(), t.rank() + 1);
  BOOST_CHECK_EQUAL(s.dim(0), indiciesShape[0]);
  BOOST_CHECK_EQUAL(s.dim(1), indiciesShape[1]);
  BOOST_CHECK_EQUAL(s.dim(2), E);

  graph.createHostWrite("inS", s, true);
  graph.createHostWrite("inT", t, true);
  graph.createHostRead("outT", t, true);
  if (accumulate)
    graph.createHostWrite("scale", scale, true);
  std::vector<float> hIn(s.numElements());
  const float outBaseValue = 100.0f;
  std::vector<float> hOut(t.numElements(), outBaseValue);
  // This test checks halves - some of these entries will be >maxHalf so the
  // test may fail if large offsets are indexed
  for (unsigned i = 0; i != hIn.size(); ++i)
    hIn[i] = i + 1.0f;

  auto target = device.getTarget();
  std::vector<char> rawIn(target.getTypeSize(HALF) * hIn.size());
  std::vector<char> rawOut(target.getTypeSize(HALF) * hOut.size());
  std::vector<char> rawScaleIn(target.getTypeSize(HALF) * 1);
  poplar::copyFloatToDeviceHalf(target, &updateScaling, rawScaleIn.data(), 1);
  poplar::copyFloatToDeviceHalf(target, hIn.data(), rawIn.data(), hIn.size());
  poplar::copyFloatToDeviceHalf(target, hOut.data(), rawOut.data(),
                                hOut.size());

  OptionFlags engineOptions{{"showExecutionSteps", "true"},
                            {"showVarStorage", "true"}};
  // Engine creation will fail for non-cpu targets if many edge pointers or
  // significant exchange is required; this should not happen if
  // createSliceableTensor() has given a good layout
  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    if (accumulate) {
      eng.writeTensor("scale", rawScaleIn.data(),
                      rawScaleIn.data() + rawScaleIn.size());
    }
    eng.writeTensor("inT", rawOut.data(), rawOut.data() + rawOut.size());
    eng.writeTensor("inS", rawIn.data(), rawIn.data() + rawIn.size());
    eng.run();
    eng.readTensor("outT", rawOut.data(), rawOut.data() + rawOut.size());
  });
  poplar::copyDeviceHalfToFloat(target, rawOut.data(), hOut.data(),
                                hOut.size());
  unsigned outIdx = 0;
  for (const auto &e : hOut) {
    if (e != outBaseValue)
      BOOST_TEST_MESSAGE("MUpdate Output[" << outIdx << "] = " << e);
    outIdx++;
  }
  std::vector<float> expected(t.numElements(), outBaseValue);
  for (unsigned i = 0; i != indicies.size(); ++i) {
    auto d = indicies[i];
    for (unsigned elem = 0; elem != E; ++elem) {
      if (!accumulate) {
        expected[d * E + elem] = hIn[i * E + elem];
      } else {
        expected[d * E + elem] += updateScaling * hIn[i * E + elem];
      }
    }
  }
  for (unsigned i = 0; i != expected.size(); ++i)
    BOOST_CHECK_EQUAL(hOut[i], expected[i]);

  std::stringstream ss;
  eng.printProfileSummary(ss, engineOptions);
  BOOST_TEST_MESSAGE(ss.str());
}

BOOST_AUTO_TEST_SUITE(MultiUpdate)

// test the looping multiupdate
BOOST_AUTO_TEST_CASE(MultiUpdate5) {
  multiupdate({100, 0, 50, 48, 49}, {5, 1}, false);
}

// test the inlined multiupdate
BOOST_AUTO_TEST_CASE(MultiUpdate2) { multiupdate({100, 0}, {2, 1}, false); }

// test the fast vertex
BOOST_AUTO_TEST_CASE(MultiUpdate10) {
  multiupdate({2, 1, 2, 1, 80, 70, 60, 50, 40, 30}, {10, 1}, false);
}

// test the looping multiupdate
BOOST_AUTO_TEST_CASE(MultiUpdateAdd5) {
  multiupdate({100, 0, 50, 48, 49}, {5, 1}, false, true, 0.5);
}

// test the inlined multiupdate
BOOST_AUTO_TEST_CASE(MultiUpdateAdd2) {
  multiupdate({100, 0}, {2, 1}, false, true, 0.5);
}

BOOST_AUTO_TEST_SUITE_END()

BOOST_AUTO_TEST_SUITE(MultiUpdateSingles)

// test the fast vertex
BOOST_AUTO_TEST_CASE(MultiUpdateAdd10Singles) {
  multiupdate({2, 1, 2, 1, 80, 70, 60, 50, 40, 30}, {10, 1}, false, true, 0.5);
}

BOOST_AUTO_TEST_SUITE_END()

BOOST_AUTO_TEST_SUITE(MultiUpdateMultiples)

// test the fast vertex with multiple updates per tile
BOOST_AUTO_TEST_CASE(MultiUpdateAdd10Multiples) {
  multiupdate({2, 1, 2, 1, 80, 70, 60, 50, 40, 30}, {10, 1}, false, true, 0.5,
              64);
}

// test the looping multiupdate
BOOST_AUTO_TEST_CASE(MultiUpdate5_AsEmbedding) {
  // Currently unhandled
  BOOST_CHECK_THROW(multiupdate({100, 0, 50, 48, 49}, {5, 1}, true),
                    poputil::poplibs_error);
}

// test the inlined multiupdate
BOOST_AUTO_TEST_CASE(MultiUpdate2_AsEmbedding) {
  // Currently unhandled
  BOOST_CHECK_THROW(multiupdate({100, 0}, {2, 1}, true),
                    poputil::poplibs_error);
}

// test the fast vertex
BOOST_AUTO_TEST_CASE(MultiUpdate10_AsEmbedding) {
  // Currently unhandled
  BOOST_CHECK_THROW(
      multiupdate({2, 1, 2, 1, 80, 70, 60, 50, 40, 30}, {10, 1}, true),
      poputil::poplibs_error);
}

// test the looping multiupdate
BOOST_AUTO_TEST_CASE(MultiUpdateAdd5_AsEmbedding) {
  multiupdate({100, 0, 50, 48, 49}, {5, 1}, true, true, 0.5);
}

// test the inlined multiupdate
BOOST_AUTO_TEST_CASE(MultiUpdateAdd2_AsEmbedding) {
  multiupdate({100, 0}, {2, 1}, true, true, 0.5);
}

// test the fast vertex
BOOST_AUTO_TEST_CASE(MultiUpdateAdd10Singles_AsEmbedding) {
  multiupdate({2, 1, 2, 1, 80, 70, 60, 50, 40, 30}, {10, 1}, true, true, 0.5);
}

// test the fast vertex with multiple updates per tile
BOOST_AUTO_TEST_CASE(MultiUpdateAdd10Multiples_AsEmbedding) {
  multiupdate({2, 1, 2, 1, 80, 70, 60, 50, 40, 30}, {10, 1}, true, true, 0.5,
              64);
}

// test heuristic which checks for mapping of a slice.
// if this doesn't kick in we will run out of memory on some
// tiles hence we check for an error constructing the engine.
static void multiUpdatePoorlyMapped(bool accumulate) {
  const unsigned numTiles = 16u;
  const unsigned D = 1501u * 10u;
  const unsigned E = 8u;

  auto device = createTestDevice(TEST_TARGET, 1, numTiles);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  const std::vector<std::size_t> sliceDims{0};
  const std::vector<std::size_t> sliceSizes{1};
  const std::vector<std::uint32_t> indices = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};

  const auto options = OptionFlags();
  const auto plan = SlicePlan();
  // We'll create this with the sliced dimension as the innermost
  // so that our rearrangement is hopefully niceish anyway mostly
  // for test speed purposes.
  auto t = graph.addVariable(FLOAT, {E, D}, "t").transpose();
  // We'll map this with the unsliced dimension as innermost such
  // that the first slice will most certainly reside on a single tile.
  // This is the case where there is plenty of data that will be
  // transferred to a single tile and so we should spread it to
  // slice properly.
  mapTensorLinearly(graph, t, 0, 1);

  auto s = createSliceTensor(graph, FLOAT, {D, E}, sliceDims, sliceSizes,
                             indices.size(), plan, options, "s");

  Sequence prog;

  auto offsetInit = graph.addConstant(UNSIGNED_INT, {indices.size()},
                                      indices.data(), "offset");
  graph.setTileMapping(offsetInit, 0);
  auto offset = createIndicesTensor(graph, sliceDims, indices.size(), plan,
                                    options, "offset");

  prog.add(Copy(offsetInit, offset));
  auto scale = graph.addVariable(FLOAT, {}, "scale");
  graph.setTileMapping(scale, 0);
  if (accumulate) {
    multiUpdateAdd(graph, t, s, offset, scale, sliceDims, sliceSizes, prog,
                   plan, options, "multiUpdateAddTest");
  } else {
    multiUpdate(graph, t, s, offset, sliceDims, sliceSizes, prog, plan, options,
                "multiUpdateTest");
  }

  BOOST_CHECK_EQUAL(s.rank(), t.rank() + 1);
  BOOST_CHECK_EQUAL(s.dim(0), indices.size());
  BOOST_CHECK_EQUAL(s.dim(1), 1);
  BOOST_CHECK_EQUAL(s.dim(2), E);

  graph.createHostWrite("inS", s, true);
  graph.createHostWrite("inT", t, true);
  graph.createHostRead("outT", t, true);
  if (accumulate) {
    graph.createHostWrite("scale", scale, true);
  }
  std::vector<float> hIn(s.numElements());
  const float outBaseValue = 100.0f;
  std::vector<float> hOut(t.numElements(), outBaseValue);
  std::iota(hIn.begin(), hIn.end(), 0.0f);
  const float updateScaling = 0.5f;
  std::vector<float> hScale(1, updateScaling);

  Engine e(graph, prog, options);
  device.bind([&](const Device &d) {
    e.load(d);
    if (accumulate) {
      e.writeTensor("scale", hScale.data(), hScale.data() + hScale.size());
    }
    e.writeTensor("inS", hIn.data(), hIn.data() + hIn.size());
    e.writeTensor("inT", hOut.data(), hOut.data() + hOut.size());
    e.run();
    e.readTensor("outT", hOut.data(), hOut.data() + hOut.size());
  });
  unsigned outIdx = 0;
  for (const auto &e : hOut) {
    if (e != outBaseValue) {
      BOOST_TEST_MESSAGE("MUpdate Output[" << outIdx << "] = " << e);
    }
    outIdx++;
  }
  std::vector<float> expected(t.numElements(), outBaseValue);
  for (unsigned i = 0; i != indices.size(); ++i) {
    auto d = indices[i];
    for (unsigned elem = 0; elem != E; ++elem) {
      if (!accumulate) {
        expected[d * E + elem] = hIn[i * E + elem];
      } else {
        expected[d * E + elem] += updateScaling * hIn[i * E + elem];
      }
    }
  }
  for (unsigned i = 0; i != expected.size(); ++i) {
    BOOST_CHECK_EQUAL(hOut[i], expected[i]);
  }

  std::stringstream ss;
  e.printProfileSummary(ss, {
                                {"showExecutionSteps", "true"},
                                {"showVarStorage", "true"},
                            });
  BOOST_TEST_MESSAGE(ss.str());
}

BOOST_AUTO_TEST_CASE(MultiUpdatePoorlyMapped) {
  multiUpdatePoorlyMapped(false);
}

BOOST_AUTO_TEST_CASE(MultiUpdateAddPoorlyMapped) {
  multiUpdatePoorlyMapped(true);
}

void multiUpdatePoorlyMappedSlices() {

  constexpr static std::size_t E = 24;
  constexpr static std::size_t V = 400;
  constexpr static std::size_t I = 50;
  constexpr static unsigned tilesPerIPU = 16;

  auto device =
      createTestDevice(TEST_TARGET, 1, tilesPerIPU,
                       /* compileIPUCode */ true /* for exchange code size */);
  const auto &target = device.getTarget();
  Graph graph(target);
  popops::addCodelets(graph);

  const std::vector<std::size_t> sliceDims = {0};
  const std::vector<std::size_t> sliceSizes = {1};

  std::vector<std::uint32_t> indices(I);
  std::iota(indices.begin(), indices.end(), 0);

  // Currently if the factor by which we broadcast the slices is greater than
  // or equal to 4 we should rearrange before broadcasting to reduce the code.
  // We will test this by finding a rough bound on exchange code size with this
  // optimisation and testing that it doesn't increase by roughly
  // (unslicedDimSplit * slicedDimSplit) times in future
  // (i.e. that the optimisation is still in place).
  const OptionFlags &sliceOptions = {
      {"planConstraints",
       "{\"unslicedDimSplit\": 4, \"slicedDimSplit\": 4, \"lookupSplit\": 1}"}};

  const SlicePlan &plan = embedding::plan(graph, HALF, V, E, {I}, sliceOptions);

  auto t = createSliceableTensor(graph, HALF, {V, E}, sliceDims, sliceSizes,
                                 plan, sliceOptions, "t");
  auto s = graph.addVariable(HALF, {I, 1, E}, "s");
  mapTensorLinearly(graph, s, 0, 1);

  Sequence prog;

  auto offsetInit = graph.addConstant(UNSIGNED_INT, {indices.size()},
                                      indices.data(), "offsetInit");
  graph.setTileMapping(offsetInit, 0);

  auto offset = createIndicesTensor(graph, sliceDims, indices.size(), plan,
                                    sliceOptions, "offset");

  prog.add(Copy(offsetInit, offset));

  auto scale = graph.addVariable(HALF, {}, "scale");
  graph.setTileMapping(scale, 0);

  multiUpdateAdd(graph, t, s, offset, scale, sliceDims, sliceSizes, prog, plan,
                 sliceOptions, "multiUpdateAddTest");

  BOOST_CHECK_EQUAL(s.rank(), t.rank() + 1);
  BOOST_CHECK_EQUAL(s.dim(0), indices.size());
  BOOST_CHECK_EQUAL(s.dim(1), 1);
  BOOST_CHECK_EQUAL(s.dim(2), E);

  graph.createHostWrite("inS", s, true);
  graph.createHostWrite("inT", t, true);
  graph.createHostRead("outT", t, true);
  graph.createHostWrite("scale", scale, true);
  std::vector<float> hIn(s.numElements());
  const float outBaseValue = 100.0f;
  std::vector<float> hOut(t.numElements(), outBaseValue);
  std::iota(hIn.begin(), hIn.end(), 0.0f);
  const float updateScaling = 0.5f;
  std::vector<float> hScale(1, updateScaling);

  std::vector<char> rawIn(target.getTypeSize(HALF) * hIn.size());
  std::vector<char> rawOut(target.getTypeSize(HALF) * hOut.size());
  std::vector<char> rawScaleIn(target.getTypeSize(HALF) * hScale.size());
  poplar::copyFloatToDeviceHalf(target, hScale.data(), rawScaleIn.data(),
                                hScale.size());
  poplar::copyFloatToDeviceHalf(target, hIn.data(), rawIn.data(), hIn.size());
  poplar::copyFloatToDeviceHalf(target, hOut.data(), rawOut.data(),
                                hOut.size());

  Engine e(graph, prog, options);
  device.bind([&](const Device &d) {
    e.load(d);
    e.writeTensor("scale", rawScaleIn.data(),
                  rawScaleIn.data() + rawScaleIn.size());
    e.writeTensor("inS", rawIn.data(), rawIn.data() + rawIn.size());
    e.writeTensor("inT", rawOut.data(), rawOut.data() + rawOut.size());
    e.run();
    e.readTensor("outT", rawOut.data(), rawOut.data() + rawOut.size());
  });
  poplar::copyDeviceHalfToFloat(target, rawOut.data(), hOut.data(),
                                hOut.size());
  unsigned outIdx = 0;
  for (const auto &e : hOut) {
    if (e != outBaseValue) {
      BOOST_TEST_MESSAGE("MUpdate Output[" << outIdx << "] = " << e);
    }
    outIdx++;
  }
  std::vector<float> expected(t.numElements(), outBaseValue);
  for (unsigned i = 0; i != indices.size(); ++i) {
    auto d = indices[i];
    for (unsigned elem = 0; elem != E; ++elem) {
      expected[d * E + elem] += updateScaling * hIn[i * E + elem];
    }
  }
  for (unsigned i = 0; i != expected.size(); ++i) {
    BOOST_CHECK_EQUAL(hOut[i], expected[i]);
  }

  const auto &graphProfile = e.getGraphProfile();
  const auto exchangeCodeBytes =
      graphProfile["memory"]["byCategory"]["internalExchangeCode"]
                  ["nonInterleaved"]["nonOverlapped"]
                      .sumUint() +
      graphProfile["memory"]["byCategory"]["internalExchangeCode"]
                  ["interleaved"]["nonOverlapped"]
                      .sumUint() +
      graphProfile["memory"]["byCategory"]["internalExchangeCode"]["overflowed"]
                  ["nonOverlapped"]
                      .sumUint();
  const auto vertexStateBytes =
      graphProfile["memory"]["byCategory"]["vertexInstanceState"]
                  ["nonInterleaved"]["nonOverlapped"]
                      .sumUint() +
      graphProfile["memory"]["byCategory"]["copyDescriptor"]["nonInterleaved"]
                  ["nonOverlapped"]
                      .sumUint() +
      graphProfile["memory"]["byCategory"]["vectorListDescriptor"]
                  ["nonInterleaved"]["nonOverlapped"]
                      .sumUint() +
      graphProfile["memory"]["byCategory"]["vertexFieldData"]["nonInterleaved"]
                  ["nonOverlapped"]
                      .sumUint();
  BOOST_TEST_MESSAGE("Total exchange code is " << exchangeCodeBytes);
  BOOST_TEST_MESSAGE("Total vertex state bytes is " << vertexStateBytes);
  BOOST_TEST_MESSAGE("Total exchange code and vertex state bytes is "
                     << exchangeCodeBytes + vertexStateBytes);

  std::stringstream ss;
  e.printProfileSummary(ss, {
                                {"showExecutionSteps", "true"},
                                {"showVarStorage", "true"},
                            });
  BOOST_TEST_MESSAGE(ss.str());
}

BOOST_AUTO_TEST_CASE(MultiUpdateAddPoorlyMappedSlices) {
  multiUpdatePoorlyMappedSlices();
}

BOOST_AUTO_TEST_SUITE_END()

// Build and run a small model to check for cpu-specific target problems
void smallAndSimple() {
  const int num_ipus = 1;

  poplar::Device device = poplar::Device::createCPUDevice();
  bool have_device = false;
  if (device.attach()) {
    have_device = true;
  }
  BOOST_CHECK(have_device);

  Target target = device.getTarget();

  std::cout << "Number of IPUs: " << num_ipus << "\n";

  std::cout << "Creating graph\n";
  Graph graph(device, poplar::replication_factor(1));
  popops::addCodelets(graph);
  popops::SlicePlan plan;

  auto ids = popops::createIndicesTensor(graph, {0}, 2, plan, {}, "ids");
  auto input = popops::createSliceableTensor(graph, INT, {3, 3}, {0}, {1}, plan,
                                             {}, "input");
  auto updates = popops::createSliceTensor(graph, INT, {2, 3}, {0}, {1}, 2,
                                           plan, {}, "input");
  auto scale =
      graph.addVariable(INT, {}, VariableMappingMethod::LINEAR, "scale");

  Sequence sequence;
  popops::multiUpdateAdd(graph, input, updates, ids, scale, {0}, {1}, sequence,
                         plan, {}, "update");

  Engine engine(graph, sequence);
  engine.load(device);
  engine.run();
  BOOST_CHECK(true);
}

BOOST_AUTO_TEST_SUITE(CpuChecks)
BOOST_AUTO_TEST_CASE(SmallAndSimple) { smallAndSimple(); }
BOOST_AUTO_TEST_SUITE_END()


#include <ros/ros.h>
#include <interactive_markers/interactive_marker_server.h>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2/convert.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <sensor_msgs/JointState.h>
#include <Eigen/Dense>
#include <memory>
#include <tf2_geometry_msgs/tf2_geometry_msgs.h>

//std::shared_ptr<ros::Publisher> joint_pub;
ros::Publisher *joint_pub;
std::shared_ptr<ros::NodeHandle> node_handle;

void sendJoint(std::vector<double> angles) {
    sensor_msgs::JointState joint_state;
    //joint_state.name = {"distal_wrist_joint", "proximal_wrist_joint", "tool_rotate_joint"};
    joint_state.name = {"j2", "j3", "j1"};
    joint_state.position = angles;
    joint_state.header.stamp = ros::Time::now();
    static int count = 1;
    joint_state.header.seq = count++;
    joint_pub->publish(joint_state);
}

void processFeedback(const visualization_msgs::InteractiveMarkerFeedbackConstPtr &feedback) {
    Eigen::MatrixXd wrist_joint_transform(3,3);
    wrist_joint_transform.setIdentity();
    wrist_joint_transform << -1, 0, 0,
                            0, 1, 0,
                            0, 0, 1;
    tf2::Quaternion orientation;
    geometry_msgs::Quaternion quat_msg = feedback->pose.orientation;
    tf2::convert(quat_msg, orientation);

    auto orientation_matrix = tf2::Matrix3x3(orientation);
    double roll, pitch, yaw;
    orientation_matrix.getRPY(roll, pitch, yaw);
    std::cout << roll << " " << pitch << " " << yaw << std::endl;
    Eigen::Vector3d orientation_rpy(roll,pitch,yaw);
    Eigen::VectorXd joint_values = wrist_joint_transform * orientation_rpy;
    std::vector<double> data(joint_values.data(), joint_values.data() + joint_values.size());
    sendJoint(data);
}

int main(int argc, char** argv) {
    ros::init(argc, argv, "wrist_marker");
    interactive_markers::InteractiveMarkerServer server("wrist_marker_server");

    node_handle.reset(new ros::NodeHandle("~"));
    ros::NodeHandle &n = *node_handle; //ros::NodeHandle("~");
    //node_handle = make_shared<ros::NodeHandle>{);
    std::string s = "wrist_frame";
    n.getParam("wrist_frame", s);
    double scale = .1;
    n.getParam("scale", scale);

    visualization_msgs::InteractiveMarker int_marker;
    int_marker.header.frame_id = s;
    int_marker.header.stamp = ros::Time::now();
    int_marker.name = "wrist_marker";
    int_marker.description = "Orientation control";
    int_marker.scale = scale;

    visualization_msgs::InteractiveMarkerControl rotate_control;
    rotate_control.orientation.w = 1;
    rotate_control.orientation.x = 1;
    rotate_control.orientation.y = 0;
    rotate_control.orientation.z = 0;
    rotate_control.name = "rotate_x";
    rotate_control.interaction_mode = visualization_msgs::InteractiveMarkerControl::ROTATE_AXIS;
    rotate_control.orientation_mode = visualization_msgs::InteractiveMarkerControl::INHERIT;
    int_marker.controls.push_back(rotate_control);

    rotate_control.orientation.w = 1;
    rotate_control.orientation.x = 0;
    rotate_control.orientation.y = 1;
    rotate_control.orientation.z = 0;
    rotate_control.name = "rotate_y";
    rotate_control.interaction_mode = visualization_msgs::InteractiveMarkerControl::ROTATE_AXIS;
    int_marker.controls.push_back(rotate_control);

    rotate_control.orientation.w = 1;
    rotate_control.orientation.x = 0;
    rotate_control.orientation.y = 0;
    rotate_control.orientation.z = 1;
    rotate_control.name = "rotate_z";
    rotate_control.interaction_mode = visualization_msgs::InteractiveMarkerControl::ROTATE_AXIS;
    int_marker.controls.push_back(rotate_control);

    server.insert(int_marker, &processFeedback);
    server.applyChanges();

    ros::NodeHandle n_pub = ros::NodeHandle("~");
    auto pub = n_pub.advertise<sensor_msgs::JointState>("joint_states", 1, true);
    joint_pub = &pub;

    sendJoint(std::vector<double>(3,0));

    ros::spin();
}


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

#include <iostream>

#include <boost/safe_numerics/safe_integer_range.hpp>
#include <boost/safe_numerics/safe_integer.hpp>

bool test1(){
    std::cout << "test1" << std::endl;
    boost::safe_numerics::safe_signed_range<-64, 63> x, y, z;
    x = 1;
    y = 2;
    z = 3;
    z = x + y;
    z = x - y;

    try{
        short int yi, zi;
        yi = y;
        zi = x + yi;
    }
    catch(std::exception e){
        // none of the above should trap. Mark failure if they do
        std::cout << e.what() << std::endl;
        return false;
    }
    return true;
}

bool test2(){
    std::cout << "test2" << std::endl;
    boost::safe_numerics::safe_unsigned_range<0, 64> x, y, z;
    x = 1;
    y = 2;
    z = 3;

    bool success = false;
    try{
        z = x - y; // should trap here
    }
    catch(std::exception e){
        success = true;
    }
    if(success == false)
        return false;

    try{
        int yi = y;
        z = x + yi; // should trap here
    }
    catch(std::exception e){
        // none of the above should trap. Mark failure if they do
        std::cout << e.what() << std::endl;
        return false;
    }
    return true;
}

bool test3(){
    using namespace boost::safe_numerics;
    std::cout << "test3" << std::endl;
    safe<int> x, y, z;
    x = 1;
    y = 2;
    z = 3;
    try{
        z = x + y;
        z = x - y;
        int yi, zi;
        zi = x + yi;
        z = x + yi;
    }
    catch(std::exception e){
        // none of the above should trap. Mark failure if they do
        std::cout << e.what() << std::endl;
        return false;
    }
    return true;
}

bool test4(){
    std::cout << "test4" << std::endl;
    boost::safe_numerics::safe<unsigned int> x, y, z;
    x = 1;
    y = 2;
    z = 3;
    z = x + y;
    bool success = false;
    try{
        z = x - y; // should trap here
    }
    catch(std::exception e){
        success = true;
    }
    if(success == false)
        return false;
    unsigned int yi, zi;
    zi = x;
    zi = x + yi;
    z = x + yi;
    zi = x + y;
    return true;
}

#include <cstdint>

bool test5(){
    std::cout << "test5" << std::endl;
    boost::safe_numerics::safe<boost::uint64_t> x, y, z;
    x = 1;
    y = 2;
    z = 3;
    z = x + y;
    bool success = false;
    try{
        z = x - y; // should trap here
    }
    catch(std::exception e){
        success = true;
    }
    if(success == false)
        return false;
    boost::uint64_t yi, zi;
    zi = x;
    zi = x + yi;
    z = x + yi;
    zi = x + y;
    return true;
}

int main(int, char *[]){
    bool rval = (
        test1() &&
        test2() &&
        test3() &&
        test4() &&
        test5()
    );
    std::cout << (rval ? "success!" : "failure") << std::endl;
    return rval ? EXIT_SUCCESS : EXIT_FAILURE;
}


// Boost.Geometry (aka GGL, Generic Geometry Library)
// Unit Test

// Copyright (c) 2010 Alfredo Correa
// Copyright (c) 2010-2012 Barend Gehrels, Amsterdam, the Netherlands.

// Use, modification and distribution is subject to the Boost Software License,
// Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#include <boost/config.hpp>
#include <geometry_test_common.hpp>

#include<boost/geometry/geometry.hpp>
#include<boost/geometry/geometries/adapted/boost_array.hpp>
#include<boost/geometry/geometries/adapted/c_array.hpp>
#include<boost/geometry/geometries/adapted/boost_tuple.hpp>
#include<iostream>

BOOST_GEOMETRY_REGISTER_C_ARRAY_CS(cs::cartesian)
BOOST_GEOMETRY_REGISTER_BOOST_ARRAY_CS(cs::cartesian)
BOOST_GEOMETRY_REGISTER_BOOST_TUPLE_CS(cs::cartesian)

#ifndef BOOST_NO_CXX11_HDR_ARRAY
#include<boost/geometry/geometries/adapted/std_array.hpp>
BOOST_GEOMETRY_REGISTER_STD_ARRAY_CS(cs::cartesian)
#endif //BOOST_NO_CXX11_HDR_ARRAY

int test_main(int, char* [])
{
    bg::model::point<double, 3, bg::cs::cartesian> p1(1,2,3);
    double p2[3] = {4,5,6};
    boost::tuple<double, double, double> p3(7,8,9);
    boost::array<double, 3> p4 = {{10,11,12}};
    std::clog << bg::distance(p1, p2) << std::endl;
    std::clog << bg::distance(p2, p3) << std::endl;
    std::clog << bg::distance(p3, p4) << std::endl;

#ifndef BOOST_NO_CXX11_HDR_ARRAY
#ifndef BOOST_NO_CXX11_HDR_INITIALIZER_LIST
    std::array<double, 3> p5 = {13,14,15};
#else
    std::array<double, 3> p5; p5[0] = 13; p5[1] = 14; p5[2] = 15;
#endif // BOOST_NO_CXX11_HDR_INITIALIZER_LIST
    std::clog << bg::distance(p4, p5) << std::endl;
#endif //BOOST_NO_CXX11_HDR_ARRAY

    return 0;
}


/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

/*
 * @file PCGSolver.cpp
 * @brief Preconditioned Conjugate Gradient Solver for linear systems
 * @date Feb 14, 2012
 * @author Yong-Dian Jian
 * @author Sungtae An
 */

#include <gtsam/linear/PCGSolver.h>
#include <gtsam/linear/GaussianFactorGraph.h>
#include <gtsam/linear/Preconditioner.h>
#include <gtsam/linear/VectorValues.h>

#include <boost/algorithm/string.hpp>

#include <algorithm>
#include <iostream>
#include <stdexcept>

using namespace std;

namespace gtsam {

/*****************************************************************************/
void PCGSolverParameters::print(ostream &os) const {
  Base::print(os);
  os << "PCGSolverParameters:" << endl;
  preconditioner_->print(os);
}

/*****************************************************************************/
PCGSolver::PCGSolver(const PCGSolverParameters &p) {
  parameters_ = p;
  preconditioner_ = createPreconditioner(p.preconditioner_);
}

void PCGSolverParameters::setPreconditionerParams(const boost::shared_ptr<PreconditionerParameters> preconditioner) {
  preconditioner_ = preconditioner;
}

void PCGSolverParameters::print(const std::string &s) const {
  std::cout << s << std::endl;;
  std::ostringstream os;
  print(os);
  std::cout << os.str() << std::endl;
}

/*****************************************************************************/
VectorValues PCGSolver::optimize(const GaussianFactorGraph &gfg,
    const KeyInfo &keyInfo, const std::map<Key, Vector> &lambda,
    const VectorValues &initial) {
  /* build preconditioner */
  preconditioner_->build(gfg, keyInfo, lambda);

  /* apply pcg */
  GaussianFactorGraphSystem system(gfg, *preconditioner_, keyInfo, lambda);
  Vector x0 = initial.vector(keyInfo.ordering());
  const Vector sol = preconditionedConjugateGradient(system, x0, parameters_);

  return buildVectorValues(sol, keyInfo);
}

/*****************************************************************************/
GaussianFactorGraphSystem::GaussianFactorGraphSystem(
    const GaussianFactorGraph &gfg, const Preconditioner &preconditioner,
    const KeyInfo &keyInfo, const std::map<Key, Vector> &lambda) :
    gfg_(gfg), preconditioner_(preconditioner), keyInfo_(keyInfo), lambda_(
        lambda) {
}

/*****************************************************************************/
void GaussianFactorGraphSystem::residual(const Vector &x, Vector &r) const {
  /* implement b-Ax, assume x and r are pre-allocated */

  /* reset r to b */
  getb(r);

  /* substract A*x */
  Vector Ax = Vector::Zero(r.rows(), 1);
  multiply(x, Ax);
  r -= Ax;
}

/*****************************************************************************/
void GaussianFactorGraphSystem::multiply(const Vector &x, Vector& AtAx) const {
  /* implement A^T*(A*x), assume x and AtAx are pre-allocated */

  // Build a VectorValues for Vector x
  VectorValues vvX = buildVectorValues(x, keyInfo_);

  // VectorValues form of A'Ax for multiplyHessianAdd
  VectorValues vvAtAx = keyInfo_.x0(); // crucial for performance

  // vvAtAx += 1.0 * A'Ax for each factor
  gfg_.multiplyHessianAdd(1.0, vvX, vvAtAx);

  // Make the result as Vector form
  AtAx = vvAtAx.vector(keyInfo_.ordering());
}

/*****************************************************************************/
void GaussianFactorGraphSystem::getb(Vector &b) const {
  /* compute rhs, assume b pre-allocated */

  // Get whitened r.h.s (A^T * b) from each factor in the form of VectorValues
  VectorValues vvb = gfg_.gradientAtZero();

  // Make the result as Vector form
  b = -vvb.vector(keyInfo_.ordering());
}

/**********************************************************************************/
void GaussianFactorGraphSystem::leftPrecondition(const Vector &x,
    Vector &y) const {
  // For a preconditioner M = L*L^T
  // Calculate y = L^{-1} x
  preconditioner_.solve(x, y);
}

/**********************************************************************************/
void GaussianFactorGraphSystem::rightPrecondition(const Vector &x,
    Vector &y) const {
  // For a preconditioner M = L*L^T
  // Calculate y = L^{-T} x
  preconditioner_.transposeSolve(x, y);
}

/**********************************************************************************/
VectorValues buildVectorValues(const Vector &v, const Ordering &ordering,
    const map<Key, size_t> & dimensions) {
  VectorValues result;

  DenseIndex offset = 0;
  for (size_t i = 0; i < ordering.size(); ++i) {
    const Key key = ordering[i];
    map<Key, size_t>::const_iterator it = dimensions.find(key);
    if (it == dimensions.end()) {
      throw invalid_argument(
          "buildVectorValues: inconsistent ordering and dimensions");
    }
    const size_t dim = it->second;
    result.emplace(key, v.segment(offset, dim));
    offset += dim;
  }

  return result;
}

/**********************************************************************************/
VectorValues buildVectorValues(const Vector &v, const KeyInfo &keyInfo) {
  VectorValues result;
  for ( const KeyInfo::value_type &item: keyInfo ) {
    result.emplace(item.first, v.segment(item.second.start, item.second.dim));
  }
  return result;
}

}