#ifdef MPI_ADAPTED_WARMUP

#include <gtest/gtest.h>
#include <stan/mcmc/cross_chain/mpi_cross_chain_adapter.hpp>
#include <stan/analyze/mcmc/compute_potential_scale_reduction.hpp>
#include <stan/math/torsten/mpi.hpp>
#include <stan/callbacks/stream_logger.hpp>
#include <stan/mcmc/hmc/nuts/adapt_unit_e_nuts.hpp>
#include <boost/random/additive_combine.hpp>
#include <stan/callbacks/writer.hpp>
#include <stan/services/util/mcmc_writer.hpp>
#include <stan/services/util/generate_transitions.hpp>
#include <stan/io/dump.hpp>
#include <fstream>

#include <stan/services/sample/hmc_nuts_diag_e_adapt.hpp>
#include <stan/services/sample/hmc_nuts_dense_e_adapt.hpp>
#include <stan/io/empty_var_context.hpp>
#include <test/test-models/good/mcmc/hmc/common/gauss3D.hpp>
#include <test/unit/services/instrumented_callbacks.hpp>
#include <test/unit/services/check_adaptation.hpp>
#include <iostream>
#include <gtest/gtest.h>

TORSTEN_MPI_SESSION_INIT;

using Eigen::MatrixXd;
using Eigen::Matrix;
using std::vector;
using stan::math::mpi::Session;
using stan::math::mpi::Communicator;
using Eigen::MatrixXd;
using Eigen::VectorXd;
using boost::accumulators::accumulator_set;
using boost::accumulators::stats;
using boost::accumulators::tag::mean;
using boost::accumulators::tag::variance;

class CrossChainAdapterTest : public testing::Test {
 public:
  CrossChainAdapterTest() :
    num_chains(3),
    max_num_windows(5),
    cross_chain_window_size(3),
    comm(Session::inter_chain_comm(num_chains)),
    num_iterations(cross_chain_window_size * max_num_windows),
    cross_chain_rhat(1.1),
    cross_chain_ess(100),
    num_warmup(100)
  {}

  std::stringstream model_log;
  stan::callbacks::logger logger;
  stan::callbacks::writer init, parameter, diagnostic;
  int num_chains;     // must run with mpiexec -n with n>=4
  int max_num_windows;
  int cross_chain_window_size;
  const Communicator& comm;
  int num_iterations;
  double cross_chain_rhat;
  double cross_chain_ess;
  int num_warmup;
};

TEST_F(CrossChainAdapterTest, stepsize) {
  const int n_par = 4;
  Eigen::VectorXd q = Eigen::VectorXd::Zero(n_par);

  stan::mcmc::mpi_cross_chain_adapter adapter;
  adapter.set_cross_chain_adaptation_params(75, 50, num_warmup,
                                            cross_chain_window_size, num_chains,
                                            cross_chain_rhat, cross_chain_ess);
  stan::mcmc::mpi_var_adaptation var_adapt(n_par, num_warmup, cross_chain_window_size);
  adapter.set_cross_chain_metric_adaptation(&var_adapt);

  double chain_stepsize_base = 0.13;
  double chain_stepsize = chain_stepsize_base + comm.rank();
  double harmonic_mean_stepsize = 0.0;
  for (int rank = 0; rank < num_chains; ++rank) {
    harmonic_mean_stepsize += 1.0 / (chain_stepsize_base + rank);
  }
  harmonic_mean_stepsize = num_chains / harmonic_mean_stepsize;

  // double harmonic_mean_stepsize = 0.427066;
  adapter.add_cross_chain_sample(1.1, q);
  EXPECT_FLOAT_EQ(adapter.cross_chain_stepsize(chain_stepsize), chain_stepsize);
  adapter.add_cross_chain_sample(1.2, q);
  EXPECT_FLOAT_EQ(adapter.cross_chain_stepsize(chain_stepsize), chain_stepsize);
  adapter.add_cross_chain_sample(1.2, q);
  adapter.set_cross_chain_adapted(true);
  EXPECT_FLOAT_EQ(adapter.cross_chain_stepsize(chain_stepsize), harmonic_mean_stepsize);
}

TEST_F(CrossChainAdapterTest, gather) {
  const int n_par = 4;
  Eigen::VectorXd q = Eigen::VectorXd::Zero(n_par);
  if (comm.rank() >= 0) {       // only inter-chains
    q(comm.rank()) = comm.rank();
  }

  stan::mcmc::mpi_cross_chain_adapter adapter;
  adapter.set_cross_chain_adaptation_params(75, 50, num_warmup,
                                            cross_chain_window_size, num_chains,
                                            cross_chain_rhat, cross_chain_ess);
  stan::mcmc::mpi_var_adaptation var_adapt(n_par, num_warmup, cross_chain_window_size);
  adapter.set_cross_chain_metric_adaptation(&var_adapt);

  double lp;
  // window 1
  lp = 1.1 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 3.8 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 0.1 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  Eigen::MatrixXd all_lp_draws;
  std::vector<double> all_chain_gather;
  adapter.cross_chain_gather(all_chain_gather, all_lp_draws);
  if (comm.rank() == 0) {
    EXPECT_EQ(all_chain_gather.size(), num_chains * 5);
    EXPECT_FLOAT_EQ(all_chain_gather[0], 5.0/3.0);
    EXPECT_FLOAT_EQ(all_chain_gather[1], 3.6633333333);
    EXPECT_FLOAT_EQ(all_chain_gather[2], 1.1);
    EXPECT_FLOAT_EQ(all_chain_gather[3], 3.8);
    EXPECT_FLOAT_EQ(all_chain_gather[4], 0.1);
    EXPECT_FLOAT_EQ(all_chain_gather[num_chains * 5 - 5], 5.0/3.0 + num_chains - 1);
    EXPECT_FLOAT_EQ(all_chain_gather[num_chains * 5 - 4], 3.6633333333);
    EXPECT_FLOAT_EQ(all_chain_gather[num_chains * 5 - 3], 1.1 + num_chains - 1);
    EXPECT_FLOAT_EQ(all_chain_gather[num_chains * 5 - 2], 3.8 + num_chains - 1);
    EXPECT_FLOAT_EQ(all_chain_gather[num_chains * 5 - 1], 0.1 + num_chains - 1);
  } else {
    EXPECT_EQ(all_chain_gather.size(), 0);
  }
  EXPECT_EQ(adapter.num_active_cross_chain_windows(), 1);

  // window 2
  lp = 2.1 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 4.8 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);

  // only gather at window end
  std::vector<double> empty_chain_gather;
  int n_gather = adapter.cross_chain_gather(empty_chain_gather, all_lp_draws);
  EXPECT_EQ(n_gather, 0);
  EXPECT_EQ(empty_chain_gather.size(), 0);

  lp = 1.1 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  adapter.cross_chain_gather(all_chain_gather, all_lp_draws);
  if (comm.rank() == 0) {
    EXPECT_EQ(all_chain_gather.size(), num_chains * 7);
    EXPECT_FLOAT_EQ(all_chain_gather[0], 13.0/6.0);
    EXPECT_FLOAT_EQ(all_chain_gather[1], 3.23066666);
    EXPECT_FLOAT_EQ(all_chain_gather[2], 8.0/3.0);
    EXPECT_FLOAT_EQ(all_chain_gather[3], 3.6633333333);
    EXPECT_FLOAT_EQ(all_chain_gather[4], 2.1);
    EXPECT_FLOAT_EQ(all_chain_gather[5], 4.8);
    EXPECT_FLOAT_EQ(all_chain_gather[6], 1.1);
    EXPECT_FLOAT_EQ(all_chain_gather[num_chains * 7 - 7], 13.0/6.0 + num_chains - 1);
    EXPECT_FLOAT_EQ(all_chain_gather[num_chains * 7 - 6], 3.23066666);
    EXPECT_FLOAT_EQ(all_chain_gather[num_chains * 7 - 5], 8.0/3.0 + num_chains - 1);
    EXPECT_FLOAT_EQ(all_chain_gather[num_chains * 7 - 4], 3.6633333333);
    EXPECT_FLOAT_EQ(all_chain_gather[num_chains * 7 - 3], 2.1 + num_chains - 1);
    EXPECT_FLOAT_EQ(all_chain_gather[num_chains * 7 - 2], 4.8 + num_chains - 1);
    EXPECT_FLOAT_EQ(all_chain_gather[num_chains * 7 - 1], 1.1 + num_chains - 1);
  } else {
    EXPECT_EQ(all_chain_gather.size(), 0);
  }
  EXPECT_EQ(adapter.num_active_cross_chain_windows(), 2);
}

TEST_F(CrossChainAdapterTest, negative_adapted_window) {
  const int n_par = 4;
  Eigen::VectorXd q = Eigen::VectorXd::Zero(n_par);
  if (comm.rank() >= 0) {       // only inter-chains
    q(comm.rank()) = comm.rank();
  }

  stan::mcmc::mpi_cross_chain_adapter adapter;
  adapter.set_cross_chain_adaptation_params(75, 50, num_warmup,
                                            cross_chain_window_size, num_chains,
                                            cross_chain_rhat, cross_chain_ess);
  stan::mcmc::mpi_var_adaptation var_adapt(n_par, num_warmup, cross_chain_window_size);
  adapter.set_cross_chain_metric_adaptation(&var_adapt);

  std::vector<double> all_chain_gather;
  Eigen::MatrixXd all_lp_draws;

  double lp;
  // window 1
  lp = 1.1 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 3.8 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 0.1 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  int n_gather = adapter.cross_chain_gather(all_chain_gather, all_lp_draws);

  // window 2
  lp = 2.1 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 4.8 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 1.1 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  n_gather = adapter.cross_chain_gather(all_chain_gather, all_lp_draws);

  // standard rhat & ess
  std::vector<std::vector<double>> chain_lp(num_chains);
  for (int chain = 0; chain < num_chains; ++chain) {
    chain_lp[chain].push_back(1.1 + chain);
    chain_lp[chain].push_back(3.8 + chain);
    chain_lp[chain].push_back(0.1 + chain);
    chain_lp[chain].push_back(2.1 + chain);
    chain_lp[chain].push_back(4.8 + chain);
    chain_lp[chain].push_back(1.1 + chain);
  }
  std::vector<const double*> draws(num_chains);
  for (int chain = 0; chain < num_chains; ++chain) {
    draws[chain] = chain_lp[chain].data();
  }

  // compare window 1 of total 2 windows against standard results
  {
    double ess =
      stan::analyze::compute_effective_sample_size(draws, 2 * cross_chain_window_size);
    double rhat =
      stan::analyze::compute_potential_scale_reduction(draws, 2 * cross_chain_window_size);

    // compare window 2
    stan::callbacks::logger logger;
    if (stan::math::mpi::Session::is_in_inter_chain_comm(num_chains)) {
      int adapted_win = adapter.cross_chain_adapted_window(n_gather,
                                                           all_chain_gather,
                                                           all_lp_draws,
                                                           logger);
      if (comm.rank() == 0) {
        EXPECT_FLOAT_EQ(adapter.cross_chain_adapt_rhat()[0], rhat);
        EXPECT_FLOAT_EQ(adapter.cross_chain_adapt_ess()[0], ess);
      }
      EXPECT_EQ(adapted_win, -1);
    }
  }

  // window 3
  lp = 2.5 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 4.5 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 1.5 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  n_gather = adapter.cross_chain_gather(all_chain_gather, all_lp_draws);

  for (int chain = 0; chain < num_chains; ++chain) {
    chain_lp[chain].push_back(2.5 + chain);
    chain_lp[chain].push_back(4.5 + chain);
    chain_lp[chain].push_back(1.5 + chain);
  }
  for (int chain = 0; chain < num_chains; ++chain) {
    draws[chain] = chain_lp[chain].data();
  }

  // compare window 1 of total 3 windows against standard results
  {
    double ess =
      stan::analyze::compute_effective_sample_size(draws, 3 * cross_chain_window_size);
    double rhat =
      stan::analyze::compute_potential_scale_reduction(draws, 3 * cross_chain_window_size);

    if (stan::math::mpi::Session::is_in_inter_chain_comm(num_chains)) {
      int adapted_win = adapter.cross_chain_adapted_window(n_gather,
                                                           all_chain_gather,
                                                           all_lp_draws,
                                                           logger);
      if (comm.rank() == 0) {
        EXPECT_FLOAT_EQ(adapter.cross_chain_adapt_rhat()[0], rhat);
        EXPECT_FLOAT_EQ(adapter.cross_chain_adapt_ess()[0], ess);
      }
      EXPECT_EQ(adapted_win, -2);
    }
  }

  // compare window 2 of total 3 windows against standard results
  // for standard computation we need recollect draws
  {
    std::vector<std::vector<double>> chain_lp(num_chains);
    for (int chain = 0; chain < num_chains; ++chain) {
      chain_lp[chain].push_back(2.1 + chain);
      chain_lp[chain].push_back(4.8 + chain);
      chain_lp[chain].push_back(1.1 + chain);
      chain_lp[chain].push_back(2.5 + chain);
      chain_lp[chain].push_back(4.5 + chain);
      chain_lp[chain].push_back(1.5 + chain);
    }
    std::vector<const double*> draws(num_chains);
    for (int chain = 0; chain < num_chains; ++chain) {
      draws[chain] = chain_lp[chain].data();
    }

    double ess =
      stan::analyze::compute_effective_sample_size(draws, 2 * cross_chain_window_size);
    double rhat =
      stan::analyze::compute_potential_scale_reduction(draws, 2 * cross_chain_window_size);

    if (stan::math::mpi::Session::is_in_inter_chain_comm(num_chains)) {
      int adapted_win = adapter.cross_chain_adapted_window(n_gather,
                                                           all_chain_gather,
                                                           all_lp_draws,
                                                           logger);
      if (comm.rank() == 0) {
        EXPECT_FLOAT_EQ(adapter.cross_chain_adapt_rhat()[1], rhat);
        EXPECT_FLOAT_EQ(adapter.cross_chain_adapt_ess()[1], ess);
      }
      EXPECT_EQ(adapted_win, -2);
    }
  }
}

TEST_F(CrossChainAdapterTest, positive_adapted_window) {
  const int n_par = 4;
  Eigen::VectorXd q = Eigen::VectorXd::Zero(n_par);
  if (comm.rank() >= 0) {       // only inter-chains
    q(comm.rank()) = comm.rank();
  }

  cross_chain_rhat = 1.12;
  cross_chain_ess = 15.0;
  stan::mcmc::mpi_cross_chain_adapter adapter;
  adapter.set_cross_chain_adaptation_params(75, 50, num_warmup,
                                            cross_chain_window_size, num_chains,
                                            cross_chain_rhat, cross_chain_ess);
  stan::mcmc::mpi_var_adaptation var_adapt(n_par, num_warmup, cross_chain_window_size);
  adapter.set_cross_chain_metric_adaptation(&var_adapt);

  std::vector<double> all_chain_gather;
  Eigen::MatrixXd all_lp_draws;
  stan::callbacks::logger logger;

  double lp;
  int n_gather = 0;
  // window 1
  lp = 1.1 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 3.8 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 0.1 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  n_gather = adapter.cross_chain_gather(all_chain_gather, all_lp_draws);

  // window 2
  lp = 2.1 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 4.8 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 1.1 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  n_gather = adapter.cross_chain_gather(all_chain_gather, all_lp_draws);

  if (stan::math::mpi::Session::is_in_inter_chain_comm(num_chains)) {
    int adapted_win = adapter.cross_chain_adapted_window(n_gather,
                                                         all_chain_gather,
                                                         all_lp_draws,
                                                         logger);
    EXPECT_EQ(adapted_win, 0);
  }

  // window 3
  lp = 2.5 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 4.5 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  lp = 1.5 + comm.rank();
  adapter.add_cross_chain_sample(lp, q);
  n_gather = adapter.cross_chain_gather(all_chain_gather, all_lp_draws);

  if (stan::math::mpi::Session::is_in_inter_chain_comm(num_chains)) {
    int adapted_win = adapter.cross_chain_adapted_window(n_gather,
                                                         all_chain_gather,
                                                         all_lp_draws,
                                                         logger);
    EXPECT_EQ(adapted_win, 1);
  }
}

#endif


#include "game_data.hpp"
#include "kernel.hpp"

#include <boost/range/adaptor/map.hpp>
#include <boost/range/algorithm/copy.hpp>

namespace wesnoth {

bool sides::are_allied(int a, int b) {
	auto it = ally_cache_.find(std::make_pair(a,b));
	if (it != ally_cache_.end()) {
		return it->second;
	}
	return ally_cache_[std::make_pair(a,b)] = ally_calculator_(a, b);
}

typedef std::pair<map_location, pathing_node> heap_entry;

struct comp {
	bool operator()(const heap_entry & a, const heap_entry & b) {
		return a.second.turns_left > b.second.turns_left ||
			(a.second.turns_left == b.second.turns_left && a.second.moves_left > b.second.moves_left);
	}
};

static path get_path(const shortest_path_tree & tree, map_location loc) {
	auto it = tree.find(loc);
	assert(it != tree.end());
	std::pair<map_location, pathing_node> pos = *it;
	path ret(1, pos.first);
	while (!(pos.second.pred == pos.first)) {
		ret.push_back(pos.second.pred);
		auto it = tree.find(pos.second.pred);
		assert(it != tree.end());
		pos = *it;
	}
	return ret;
}

// Tries to find a match for a visible enemy at some position with respect to a pathing query
static const unit_rec * get_visible_enemy(map_location neighbor, const pathfind_context::pathing_query & query, bool must_exert_zoc) {
	unit_map::index<by_loc>::type & u_map = query.units->get<by_loc>();
	sides & sides = *query.sides_;
	auto u_it = u_map.find(neighbor);

	if (u_it == u_map.end()) {
		return NULL;
	}

	const unit_rec & u = *u_it;
	if (u.dirty_) {
		u.update();
	}

	if (u.loc_ == neighbor && // if it did not move during update
		(!must_exert_zoc || u.emits_zoc_) && // check the must_exert_zoc flag
		!sides.are_allied(u.side_, *query.moving_side)) // only enemy units block this unit
	{
		if (!query.viewing_side) { // the blocker is definitely visible since we see all
			return & u;
		}
		// we don't see all so check if the blocker is actually visible
		if ((!u.hidden_ || sides.are_allied(u.side_, *query.viewing_side))
				&& !sides.ally_adjusted_fog(neighbor, *query.viewing_side)) {
			return & u;
		}
	}
	return NULL;
}

shortest_path_tree pathfind_context::compute_tree(const pathfind_context::pathing_query & query, boost::optional<map_location> destination) {
	shortest_path_tree result;

	std::deque<heap_entry> priority_queue(1, heap_entry(query.start, pathing_node(query.moves, query.turns, query.start)));

	comp heap_comparator;

	assert(query.sides_);
	sides & sides = *query.sides_;

	while (!priority_queue.empty()) {
		heap_entry he = priority_queue.front();
		const map_location & loc = he.first;

		if (destination && *destination == loc) {
			// If we actually have a destination and we find it in the shortest path tree, we can skip the rest of the computation.
			shortest_path_tree new_result;
			std::pair<map_location, pathing_node> pos = he;

			while(!(pos.second.pred == pos.first)) {
				new_result.insert(pos);
				auto it = result.find(pos.second.pred);
				assert(it != result.end());
				pos = *it;
			}

			new_result.insert(pos); // we need to insert the final self loop node as well to preserve the invariant of the data structure
			return new_result;
		}

		result.insert(he);

		std::pop_heap(priority_queue.begin(), priority_queue.end(), heap_comparator); priority_queue.pop_back();

		BOOST_FOREACH(map_location neighbor, geom_.neighbors(loc)) {
			if (result.find(neighbor) != result.end()) {
				continue; //skip nodes that we already processed
			}

			assert(query.tmap_);
			if (query.tmap_->find(neighbor) == query.tmap_->end()) {
				continue; // skip nodes that are off map
			}

			if(query.viewing_side) {
				if (sides.ally_adjusted_shroud(neighbor, *query.viewing_side)) {
					continue; // skip shrouded nodes
				}
			}

			size_t cost_of_move;

			bool used_first_turn_override = false;

			if (query.first_turn_override_cost_map && (he.second.turns_left == query.turns)) {// check our node to see if it is still the first turn
				cost_of_move = (*query.first_turn_override_cost_map)(neighbor);
				used_first_turn_override = true;
			} else {
				cost_of_move = (query.cost_map ? (*query.cost_map)(neighbor) : 1);
			}

			size_t turns_left = he.second.turns_left;
			size_t moves_left = he.second.moves_left;

			if (cost_of_move > moves_left && turns_left > 0) {
				turns_left--;
				moves_left = query.max_moves;
				if (used_first_turn_override) { //recalculate cost since we had to end the turn
					cost_of_move = (query.cost_map ? (*query.cost_map)(neighbor) : 1);
				}
			}

			if (cost_of_move > moves_left) {
				continue; // we can't actually afford to make the move at all, even after possibly refreshing the moves for a turn
			}
			moves_left = moves_left - cost_of_move;

			// If we should think of the move as being done by a unit as some side, and check for blockers / zocs
			if (query.moving_side) {
				if (get_visible_enemy(neighbor, query, false)) {
					continue;
				}
				// If our unit can get zocced, check if it will be by this move
				if (!query.ignore_zoc && moves_left > 0) {
					BOOST_FOREACH(map_location neighbor2, geom_.neighbors(neighbor)) {
						if (get_visible_enemy(neighbor2, query, true)) {
							moves_left = 0;
							break;
						}
					}
				}
			}

			pathing_node node(moves_left, turns_left, loc);
			priority_queue.push_back(std::make_pair(neighbor, node));
			std::push_heap(priority_queue.begin(), priority_queue.end(), heap_comparator);
		}
	}
	return result;
}

loc_set pathfind_context::reachable_hexes(const pathfind_context::pathing_query & query) {
	loc_set result;
	boost::copy(compute_tree(query) | boost::adaptors::map_keys, std::inserter(result, result.end()));
	return result;
}

std::vector<path> pathfind_context::reachable_hexes_with_paths(const pathfind_context::pathing_query & query) {
	std::vector<path> result;
	shortest_path_tree tree = compute_tree(query);
	BOOST_FOREACH(const shortest_path_tree::value_type & v, tree) {
		result.push_back(get_path(tree, v.first));
	}
	return result;
}

path pathfind_context::shortest_path(map_location end, const pathing_query & query) {
	return get_path(compute_tree(query, end), end);
}

size_t pathfind_context::shortest_path_distance(map_location end, const pathing_query & query) {
	if (end == query.start) {
		return 0;
	}
	shortest_path_tree tree = compute_tree(query, end);
	auto it = tree.find(end);
	assert(it != tree.end());
	return it->second.turns_left - query.turns + 1;
}

} // end namespace wesnoth


//
// Copyright Antoine Leblanc 2010 - 2015
// Distributed under the MIT license.
//
// http://nauths.fr
// http://github.com/nicuveo
// mailto://antoine.jp.leblanc@gmail.com
//

//HHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHH
// Includes

#include <cstdlib>
#include <iostream>
#include <boost/test/unit_test.hpp>
#include "nauths/mcl/mcl.hh"

//HHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHH
// Local code

namespace
{

#define CREATE(R, D, S)                         \
  case mcl::spaces::S:                          \
  {                                             \
    mcl::S color;                               \
    fill(color, alpha);                         \
    return color;                               \
  }                                             \

  void init()
  {
    srand(42);
  }

  template <typename T>
  T random()
  {
    return T((rand() % 10000) / 10000.);
  }

  template <>
  mcl::byte random()
  {
    return (rand() >> 8) % 256;
  }

  template <typename T>
  T normalize(double x)
  {
    return x;
  }

  template <>
  mcl::byte normalize(double x)
  {
    return x * 255;
  }

  template <typename T, std::size_t N>
  void fill(std::array<T, N>& color, double alpha, double factor = 1.0)
  {
    for (std::size_t i = 0; i < N - 1; ++i)
      color[i] = random<T>() * factor;
    color[N - 1] = normalize<T>(alpha);
  }

  mcl::Color make(mcl::SpaceType s, double alpha)
  {
    switch (s)
    {
      BOOST_PP_SEQ_FOR_EACH(CREATE, _, MCL_SPACES)
    }

    return mcl::Color();
  }

}

//HHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHH
// Implementation

BOOST_AUTO_TEST_SUITE(color)

BOOST_AUTO_TEST_CASE(color_00)
{
  init();

  for (mcl::SpaceType s1 : mcl::COLOR_SPACES)
  {
    const double a = random<double>();
    const mcl::Color c1 = make(s1, a);

    switch (s1)
    {
      case MCL_RGBub: BOOST_CHECK_EQUAL(c1.alpha(), (mcl::byte(a * 255)) / 255.0); break;
      case MCL_RGBf:  BOOST_CHECK_EQUAL(c1.alpha(), (float) a);                    break;
      default:        BOOST_CHECK_EQUAL(c1.alpha(), a);                            break;
    }

    BOOST_CHECK_EQUAL(c1.type(), s1);
  }
}

BOOST_AUTO_TEST_CASE(clamp_cmy)
{
  init();

  for (int i = 0; i < 100; ++i)
  {
    mcl::CMY c;

    fill(c, random<double>(), 10.0);
    c.clamp();

    BOOST_REQUIRE(c.c() >=    0.0 and c.c() <=   1.0);
    BOOST_REQUIRE(c.m() >=    0.0 and c.m() <=   1.0);
    BOOST_REQUIRE(c.y() >=    0.0 and c.y() <=   1.0);
  }
}

BOOST_AUTO_TEST_CASE(clamp_cmyk)
{
  init();

  for (int i = 0; i < 100; ++i)
  {
    mcl::CMYK c;

    fill(c, random<double>(), 10.0);
    c.clamp();

    BOOST_REQUIRE(c.c() >=    0.0 and c.c() <=   1.0);
    BOOST_REQUIRE(c.m() >=    0.0 and c.m() <=   1.0);
    BOOST_REQUIRE(c.y() >=    0.0 and c.y() <=   1.0);
    BOOST_REQUIRE(c.k() >=    0.0 and c.k() <=   1.0);
  }
}

BOOST_AUTO_TEST_CASE(clamp_rgb)
{
  init();

  for (int i = 0; i < 100; ++i)
  {
    mcl::RGB c;

    fill(c, random<double>(), 10.0);
    c.clamp();

    BOOST_REQUIRE(c.r() >=    0.0 and c.r() <=   1.0);
    BOOST_REQUIRE(c.g() >=    0.0 and c.g() <=   1.0);
    BOOST_REQUIRE(c.b() >=    0.0 and c.b() <=   1.0);
  }
}

BOOST_AUTO_TEST_CASE(clamp_rgbf)
{
  init();

  for (int i = 0; i < 100; ++i)
  {
    mcl::RGBf c;

    fill(c, random<float>(), 10.f);
    c.clamp();

    BOOST_REQUIRE(c.r() >=    0.f and c.r() <=   1.f);
    BOOST_REQUIRE(c.g() >=    0.f and c.g() <=   1.f);
    BOOST_REQUIRE(c.b() >=    0.f and c.b() <=   1.f);
  }
}

BOOST_AUTO_TEST_CASE(clamp_hsl)
{
  init();

  for (int i = 0; i < 100; ++i)
  {
    mcl::HSL c;

    fill(c, random<double>(), 10.0);
    c.clamp();

    BOOST_REQUIRE(c.h() >=    0.0 and c.h() <= 360.0);
    BOOST_REQUIRE(c.s() >=    0.0 and c.s() <=   1.0);
    BOOST_REQUIRE(c.l() >=    0.0 and c.l() <=   1.0);
  }
}

BOOST_AUTO_TEST_CASE(clamp_hsv)
{
  init();

  for (int i = 0; i < 100; ++i)
  {
    mcl::HSV c;

    fill(c, random<double>(), 10.0);
    c.clamp();

    BOOST_REQUIRE(c.h() >=    0.0 and c.h() <= 360.0);
    BOOST_REQUIRE(c.s() >=    0.0 and c.s() <=   1.0);
    BOOST_REQUIRE(c.v() >=    0.0 and c.v() <=   1.0);
  }
}

BOOST_AUTO_TEST_CASE(clamp_lab)
{
  init();

  for (int i = 0; i < 100; ++i)
  {
    mcl::LAB c;

    fill(c, random<double>(), 10.0);
    c.clamp();

    BOOST_REQUIRE(c.l() >=    0.0 and c.l() <= 100.0);
    BOOST_REQUIRE(c.a() >= -120.0 and c.a() <= 120.0);
    BOOST_REQUIRE(c.b() >= -120.0 and c.b() <= 120.0);
  }
}

BOOST_AUTO_TEST_CASE(clamp_lch)
{
  init();

  for (int i = 0; i < 100; ++i)
  {
    mcl::LCH c;

    fill(c, random<double>(), 10.0);
    c.clamp();

    BOOST_REQUIRE(c.l() >=    0.0 and c.l() <= 100.0);
    BOOST_REQUIRE(c.c() >= -120.0 and c.c() <= 120.0);
    BOOST_REQUIRE(c.h() >=    0.0 and c.h() <= 360.0);
  }
}

BOOST_AUTO_TEST_CASE(clamp_xyz)
{
  init();

  for (int i = 0; i < 100; ++i)
  {
    mcl::XYZ c;

    fill(c, random<double>(), 10.0);
    c.clamp();

    BOOST_REQUIRE(c.x() >=    0.0 and c.x() <= 120.0);
    BOOST_REQUIRE(c.y() >=    0.0 and c.y() <= 120.0);
    BOOST_REQUIRE(c.z() >=    0.0 and c.z() <= 120.0);
  }
}

BOOST_AUTO_TEST_SUITE_END()


// Copyright 2005 The Trustees of Indiana University.

// 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)

//  Authors: Douglas Gregor
//           Andrew Lumsdaine
#include "config.hpp"
#include "graph_types.hpp"
#include <boost/parallel/property_map.hpp>
#include <boost/graph/python/distributed/dehne_gotz_min_spanning_tree.hpp>

namespace boost { namespace graph { namespace distributed { namespace python {

void export_dehne_gotz_min_spanning_tree()
{
  using boost::python::arg;
  using boost::python::def;

  using boost::graph::python::distributed::dense_boruvka_minimum_spanning_tree;
  using boost::graph::python::distributed::merge_local_minimum_spanning_trees;
  using boost::graph::python::distributed::boruvka_then_merge;
  using boost::graph::python::distributed::boruvka_mixed_merge;

#define UNDIRECTED_GRAPH(Name,Type)                     \
  def("dense_boruvka_minimum_spanning_tree",            \
      &dense_boruvka_minimum_spanning_tree<Type>,       \
      (arg("graph"), arg("weight_map")));               \
                                                        \
  def("merge_local_minimum_spanning_trees",             \
      &merge_local_minimum_spanning_trees<Type>,        \
      (arg("graph"), arg("weight_map")));               \
                                                        \
  def("boruvka_then_merge",                             \
      &boruvka_then_merge<Type>,                        \
      (arg("graph"), arg("weight_map")));               \
                                                        \
  def("boruvka_mixed_merge",                            \
      &boruvka_mixed_merge<Type>,                       \
      (arg("graph"), arg("weight_map")));
#define DIRECTED_GRAPH(Name,Type)
#include "graphs.hpp"
}

} } } } // end namespace boost::graph::distributed::python


#include <sstream>
#include <iostream>
#include <fstream>
#include <vector>

#include <boost/program_options.hpp>
#include <boost/program_options/variables_map.hpp>

#include "sparse_vector.h"
#include "error_surface.h"
#include "line_optimizer.h"
#include "b64tools.h"

using namespace std;
namespace po = boost::program_options;

void InitCommandLine(int argc, char** argv, po::variables_map* conf) {
  po::options_description opts("Configuration options");
  opts.add_options()
        ("loss_function,l",po::value<string>(), "Loss function being optimized")
        ("help,h", "Help");
  po::options_description dcmdline_options;
  dcmdline_options.add(opts);
  po::store(parse_command_line(argc, argv, dcmdline_options), *conf);
  bool flag = conf->count("loss_function") == 0;
  if (flag || conf->count("help")) {
    cerr << dcmdline_options << endl;
    exit(1);
  }
}

int main(int argc, char** argv) {
  po::variables_map conf;
  InitCommandLine(argc, argv, &conf);
  const string loss_function = conf["loss_function"].as<string>();
  ScoreType type = ScoreTypeFromString(loss_function);
  LineOptimizer::ScoreType opt_type = LineOptimizer::MAXIMIZE_SCORE;
  if (type == TER || type == AER) {
    opt_type = LineOptimizer::MINIMIZE_SCORE;
  }
  string last_key;
  vector<ErrorSurface> esv;
  while(cin) {
    string line;
    getline(cin, line);
    if (line.empty()) continue;
    size_t ks = line.find("\t");
    assert(string::npos != ks);
    assert(ks > 2);
    string key = line.substr(2, ks - 2);
    string val = line.substr(ks + 1);
    if (key != last_key) {
      if (!last_key.empty()) {
	float score;
        double x = LineOptimizer::LineOptimize(esv, opt_type, &score);
	cout << last_key << "|" << x << "|" << score << endl;
      }
      last_key = key;
      esv.clear();
    }
    if (val.size() % 4 != 0) {
      cerr << "B64 encoding error 1! Skipping.\n";
      continue;
    }
    string encoded(val.size() / 4 * 3, '\0');
    if (!B64::b64decode(reinterpret_cast<const unsigned char*>(&val[0]), val.size(), &encoded[0], encoded.size())) {
      cerr << "B64 encoding error 2! Skipping.\n";
      continue;
    }
    esv.push_back(ErrorSurface());
    esv.back().Deserialize(type, encoded);
  }
  if (!esv.empty()) {
    // cerr << "ESV=" << esv.size() << endl;
    // for (int i = 0; i < esv.size(); ++i) { cerr << esv[i].size() << endl; }
    float score;
    double x = LineOptimizer::LineOptimize(esv, opt_type, &score);
    cout << last_key << "|" << x << "|" << score << endl;
  }
  return 0;
}


/*
 * http://pointclouds.org/documentation/tutorials/ground_based_rgbd_people_detection.php
 * http://ua-ros-pkg.googlecode.com/svn/trunk/arrg/ua_vision/background_filters/src/ground_filter.cpp
 * http://pointclouds.org/documentation/tutorials/passthrough.php
 */
#include <signal.h>
#include <time.h>

#include <fstream>
#include <sstream>

#include <pcl/console/parse.h>
#include <pcl/point_types.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/filters/extract_indices.h>
#include <pcl/sample_consensus/sac_model_plane.h>
#include <pcl/sample_consensus/model_types.h>
#include <pcl/sample_consensus/method_types.h>
#include <pcl/people/ground_based_people_detection_app.h>
#include <pcl/segmentation/sac_segmentation.h>
#include <pcl/visualization/point_cloud_handlers.h>
#include <pcl/filters/passthrough.h>

#include <boost/thread.hpp>
#include <ros/ros.h>
#include <pcl_conversions/pcl_conversions.h>
#include <sensor_msgs/PointCloud2.h>
#include <ground_based_detector/human.h>
#include <ground_based_detector/humanArray.h>

typedef pcl::PointXYZRGBA PointT;
typedef pcl::PointCloud<PointT> PointCloudT;

boost::mutex cloud_mutex;
bool new_cloud_available_flag = false;
PointCloudT::Ptr cloud(new PointCloudT);

bool interrupted = false;
void signal_handler(int signal) {
    interrupted = true;
}

void cloud_cb_ (const sensor_msgs::PointCloud2ConstPtr& callback_cloud) {
  cloud_mutex.lock ();

  pcl::fromROSMsg(*callback_cloud, *cloud);
  new_cloud_available_flag = true;

  cloud_mutex.unlock ();
}

int main (int argc, char** argv) {
  // Initialize ROS
  ros::init (argc, argv, "ground_based_detector");
  ros::NodeHandle nodeHandle;
  ros::Publisher humanPublisher = nodeHandle.advertise<ground_based_detector::humanArray>("ground_based_detector/detected_human_clusters", 20);
  ros::Publisher groundPublisher = nodeHandle.advertise<sensor_msgs::PointCloud2>("ground_based_detector/ground_cloud", 1);

  // Algorithm parameters:
  std::string svm_filename;
  std::string trained_upper_body_classifier_xml_path;
  double min_confidence;
  double min_height;
  double max_height;
  double voxel_size;
  double groud_filter_limit_up;
  double groud_filter_limit_down;
  std::string results_dir;
  Eigen::Matrix3f rgb_intrinsics_matrix;
  rgb_intrinsics_matrix << 525, 0.0, 319.5, 0.0, 525, 239.5, 0.0, 0.0, 1.0; // Kinect RGB camera intrinsics

  // Read parameters
  nodeHandle.getParam("ground_based_detector/svm", svm_filename);
  nodeHandle.getParam("ground_based_detector/min_confidence", min_confidence);
  nodeHandle.getParam("ground_based_detector/min_h", min_height);
  nodeHandle.getParam("ground_based_detector/max_h", max_height);
  nodeHandle.getParam("ground_based_detector/voxel_size", voxel_size);
  nodeHandle.getParam("ground_based_detector/grf_up", groud_filter_limit_up);
  nodeHandle.getParam("ground_based_detector/grf_down", groud_filter_limit_down);
  nodeHandle.getParam("ground_based_detector/results_dir", results_dir);

  // Read Kinect live stream:
  // Create a ROS subscriber for the input point cloud
  ros::Subscriber subscriberCloud = nodeHandle.subscribe<sensor_msgs::PointCloud2>("/camera/depth_registered/points", 1, cloud_cb_);

  // Wait for the first frame:
  while(!new_cloud_available_flag)
    ros::spinOnce();

  new_cloud_available_flag = false;

  // (GPE) Ground plane estimation:

  // (GPE) A. Get lower heigth cloud
  PointCloudT::ConstPtr const_cloud(cloud);
  PointCloudT::Ptr lowerHeigthCloud(new PointCloudT);
  pcl::PassThrough<PointT> pass;
  pass.setInputCloud (const_cloud);
  pass.setFilterFieldName ("y");
  pass.setFilterLimits (groud_filter_limit_up, groud_filter_limit_down); // lower heigth axis(Y) (Y points to ground i.e. center is 0, ceil is negative, and floor is positive)
  pass.filter (*lowerHeigthCloud);

  // (GPE) B. Setup RANSAC on the lower height cloud
  pcl::PointIndices::Ptr ground_indices(new pcl::PointIndices);
  pcl::SACSegmentation<PointT> ground_finder;
  Eigen::Vector3f depth_axis(0.0,0.0,1.0); // depth axis Z
  ground_finder.setOptimizeCoefficients(true);
  ground_finder.setModelType(pcl::SACMODEL_PARALLEL_PLANE);
  ground_finder.setMethodType(pcl::SAC_RANSAC);
  ground_finder.setDistanceThreshold(0.015);
  ground_finder.setAxis(depth_axis);
  ground_finder.setInputCloud(lowerHeigthCloud);

  // (GPE) C. Segment and find ground plane coefficients
  pcl::ModelCoefficients ground_coefficients;
  ground_finder.segment(*ground_indices, ground_coefficients);

  // (GPE) D. Convert pcl::ModelCoefficients to Eigen::VectorXf (for people_detector compatibility)
  Eigen::VectorXf plane_parameters;
  plane_parameters.resize(4);
  plane_parameters << ground_coefficients.values[0], ground_coefficients.values[1], ground_coefficients.values[2], ground_coefficients.values[3];

  // Extract the ground cloud
  PointCloudT::ConstPtr const_lowerHeigthCloud(lowerHeigthCloud);
  PointCloudT::Ptr extractedGroundCloud(new PointCloudT);
  pcl::ExtractIndices<PointT> groundExtractor;
  groundExtractor.setInputCloud (const_lowerHeigthCloud);
  groundExtractor.setIndices(ground_indices);
  groundExtractor.setNegative(false);
  groundExtractor.filter(*extractedGroundCloud);

  sensor_msgs::PointCloud2 ground_cloud_msg;
  pcl::toROSMsg(*extractedGroundCloud, ground_cloud_msg);
  ground_cloud_msg.header.frame_id = "/camera_depth_frame";

  // Create classifier for people detection:
  pcl::people::PersonClassifier<pcl::RGB> person_classifier;
  person_classifier.loadSVMFromFile(svm_filename);                 // load trained SVM

  // People detection app initialization:
  pcl::people::GroundBasedPeopleDetectionApp<PointT> people_detector;    // people detection object
  people_detector.setVoxelSize(voxel_size);                        // set the voxel size
  people_detector.setIntrinsics(rgb_intrinsics_matrix);            // set RGB camera intrinsic parameters
  people_detector.setClassifier(person_classifier);                // set person classifier
  people_detector.setHeightLimits(min_height, max_height);         // set person classifier
  //people_detector.setSensorPortraitOrientation(true);            // set sensor orientation to vertical

  people_detector.setInputCloud(cloud);

  // Main loop:
  signal(SIGINT, signal_handler);
  while (ros::ok() && !interrupted) {
    groundPublisher.publish(ground_cloud_msg);

    std::vector<pcl::people::PersonCluster<PointT> > clusters; // vector containing persons clusters

    people_detector.setGround(plane_parameters); // set floor coefficients

    if (new_cloud_available_flag && cloud_mutex.try_lock()) // if a new cloud is available
    {
      new_cloud_available_flag = false;

      // Perform people detection on the new cloud:
      people_detector.compute(clusters); // perform people detection

      cloud_mutex.unlock();
      ground_based_detector::humanArray published_clusters;

      // draw boxes
      unsigned int k = 0;
      for(std::vector<pcl::people::PersonCluster<PointT> >::iterator it = clusters.begin(); it != clusters.end(); ++it)
      {
        if(it->getPersonConfidence() > min_confidence) // draw only people with confidence above a threshold an ok from opencv upper body detector
        {
          ground_based_detector::human published_human;
          published_human.height = it->getHeight();
          published_human.tcenter_x = it->getTCenter()[0];
          published_human.tcenter_y = it->getTCenter()[1];
          published_human.tcenter_z = it->getTCenter()[2];
          published_human.ttop_x = it->getTTop()[0];
          published_human.ttop_y = it->getTTop()[1];
          published_human.ttop_z = it->getTTop()[2];

          published_clusters.humans.push_back(published_human);
        }

      }
      humanPublisher.publish(published_clusters);
    }

    plane_parameters = people_detector.getGround();
    ros::spinOnce();
  }

  if(interrupted) std::cout << std::endl;
  return 0;
}