/*
 * modules/cmoea/mp_cmoea_nsga2_mpi.hpp
 *
 *  Created on: Mar 10, 2017
 *      Author: Joost Huizinga
 *
 * This is the file for the massively parallel version of C-MOEA.
 * The main idea behind MP C-MOEA, is that there is no longer a master and a set of workers.
 * Instead, each worker will create its own population, and occasionally exchange part of the,
 * population with other workers.
 */

#ifndef MODULES_CMOEA_CMOEA_NSGA2_MPI_HPP_
#define MODULES_CMOEA_CMOEA_NSGA2_MPI_HPP_

// Standard includes
#include <algorithm>
#include <limits>

// Boost includes
#include <boost/foreach.hpp>
#include <boost/multi_array.hpp>
#include <boost/array.hpp>
#include <boost/fusion/algorithm/iteration/for_each.hpp>
#include <boost/fusion/include/for_each.hpp>
#include <boost/spirit/include/karma.hpp>
#include <boost/mpi.hpp>
#include <boost/mpi/environment.hpp>
#include <boost/mpi/detail/point_to_point.hpp>

// Sferes includes
#include <sferes/stc.hpp>
#include <sferes/ea/ea.hpp>
#include <sferes/fit/fitness.hpp>
#include <sferes/dbg/dbg.hpp>
#include <sferes/ea/dom_sort_basic.hpp>
#include <sferes/ea/common.hpp>
#include <sferes/ea/crowd.hpp>

// Module includes
#include <modules/misc/common_compare.hpp>
#include <modules/debugext/dbgext.hpp>

// Local includes
#include "mpi_util.hpp"
#include "nsga_util.hpp"

namespace karma = boost::spirit::karma;

// Debug statements become a bit too long without this statement
using namespace dbg;
using namespace dbgext;
#define DBO dbo("ea") << "["<< _world->rank() << "] "
#define DBOF dbo("eafit") << "["<< _world->rank() << "] "
#define DBOO dbo("eaobj") << "["<< _world->rank() << "] "
#define DBO_MPI dbo("mpi") << "["<< _world->rank() << "] "

namespace sferes
{
namespace ea
{

/**
 * This is the main class for massively parallel C-MOEA.
 *
 * Usually, this part is only executed by the master.
 */
SFERES_EA(MPCmoeaNsga2Mpi, Ea){
public:
    /** PARAMETERS **/

    // The type of Pareto domination sort to use.
    // Currently available types are:
    // - sferes::ea::dom_sort_basic_f
	//   (defined in sferes/ea/dom_sort_basic.hpp)
    //   Sorts according to pareto dominance and will add individuals from the
	//   highest to the lowest layer, with crowding as a tie-breaker in the last
	//   layer to be added.
    // - sferes::ea::dom_sort_no_duplicates_f
	//   (defined in sferes/ea/dom_sort_no_duplicates.hpp)
    //   Same as dom_sort_basic, accept that, for each front, only one
	//   individual per pareto location is added. Other individuals at the same
	//   location will be bumped to the next layer.
    typedef typename Params::ea::dom_sort_f dom_sort_f;

    // The type non dominated comparator to use
    // Currently available types are:
    // - sferes::ea::_dom_sort_basic::non_dominated_f
    //   (defined in sferes/ea/dom_sort_basic.hpp)
    //   Regular comparisons based on dominance
    // - sferes::ea::cmoea_nsga::prob_dom_f<Params>
    //   Comparisons based on probabilistic sorting, where some objectives can
    //   be stronger than others.
    typedef typename Params::cmoea_nsga::non_dom_f non_dom_f;

    // The index used to temporarily store the category
    // This index should hold a dummy value, as it will be overwritten
    // constantly. The default would be 0.
    static const size_t obj_index = Params::cmoea::obj_index;

    // The number of objectives (bins) used in the map
    size_t nr_of_bins;

    // The size of each bin
    static const size_t bin_size = Params::cmoea::bin_size;

    // The number of individuals initially generated to fill the archive
    // If equal to the bin_size, every initially generated individual is added
    // to every bin of every category.
    // The init_size has to be greater than or equal to the bin_size
    static const size_t init_size = Params::pop::init_size;

    // Very large initial populations may cause CMOEA to run out of memory.
    // To avoid this, you can add the initial populations in init_batch batches
    // of init_size.
    //static const size_t init_batch = Params::pop::init_batch;

    // Max batch size to avoid memory issues
    static const size_t max_batch_size = Params::pop::max_batch_size;

    // Individuals generated each epoch per bin (needs to be divisible by 2)
    static const size_t offspring_per_generation = Params::pop::select_size;

    // The number of replicates for each bin
    static const int replicates = Params::cmoea::replicates;

    // Due to memory limitations, you may not always be able to send the entire
    // population at once. max_send_size determines the maximum size of the
    // population to send at once.
    static const size_t max_send_size = Params::cmoea::max_send_size;

    // The total number of bins, including replicates
    int bins_total;

    /** TYPE DEFINITIONS **/

    // The type of the phenotype
    typedef Phen phen_t;

    // A smart pointer to the phenotype
    typedef boost::shared_ptr<Phen> phen_ptr_t;
    typedef phen_ptr_t indiv_t;

    // The population type
    typedef typename std::vector<phen_ptr_t> pop_t;

    // Iterator for the population
    typedef typename std::vector<phen_ptr_t>::iterator pop_it_t;

    // The type of the Pareto front
    typedef typename std::vector<pop_t> front_t;

    // The type of a collections of bins
    typedef typename std::vector<pop_t> bins_t;

    // The type of the mpi environment
    typedef typename boost::mpi::environment mpi_env_t;

    // The type of the mpi communicator
    typedef typename boost::mpi::communicator mpi_com_t;

    // The type of serialization archive for sending
    typedef typename boost::mpi::packed_oarchive mpi_oarchive_t;

    // The type of serialization archive for receiving
    typedef typename boost::mpi::packed_iarchive mpi_iarchive_t;

    // The exact type of non-dominated sort function to use
    typedef fill_dom_sort_f<phen_ptr_t, dom_sort_f, non_dom_f> sort_f;

    // Maps MPI instances to bin indices
    typedef std::map<int, std::vector<int> > inst_bin_map_t;

    // Friend declaration required for Sferes
    SFERES_EA_FRIEND(MPCmoeaNsga2Mpi);

    /**
     * Constructor for MP-CMOEA.
     *
     * Initializes the MPI environment, and calculates some relevant statistics
     * for this run.
     */
    MPCmoeaNsga2Mpi(){
        dbg::trace trace("ea", DBG_HERE);

        // Initialize MPI
		static char* argv[] = {(char*)"sferes2", 0x0};
		char** argv2 = (char**) malloc(sizeof(char*) * 2);
		int argc = 1;
		argv2[0] = argv[0];
		argv2[1] = argv[1];
		dbg::out(dbg::info, "mpi") << "Initializing MPI..." << std::endl;
		_env = boost::shared_ptr<mpi_env_t>(new mpi_env_t(argc, argv2, true));
		dbg::out(dbg::info, "mpi") << "MPI initialized" << std::endl;
		_world = boost::shared_ptr<mpi_com_t>(new mpi_com_t());

		nr_of_bins = Params::cmoea::nb_of_bins;
		bins_total = nr_of_bins * replicates;

		// Lower bound on the number of bins per instance
		_min_bins_inst = bins_total / _world->size();

		// Upper bound on the number of bins per instance
		_max_bins_inst = _min_bins_inst + 1;

		// The number of bins remaining after each instance is assigned its
		// bins_per_instance
		_overfull_inst = bins_total % _world->size();

		// The total number of bins assigned to over-full instances
		_overfull_bins = _max_bins_inst * _overfull_inst;

		// The number of bins this instance has to handle
		int my_nr_bins;
		if(_world->rank() < _overfull_inst){
			my_nr_bins = _max_bins_inst;
			_first_bin_id = _world->rank() * _max_bins_inst;
		} else {
			my_nr_bins = _min_bins_inst;
			int other_instances = (_world->rank()-_overfull_inst);
			int other_bins = other_instances * _min_bins_inst;
			_first_bin_id = _overfull_bins + other_bins;
		}

		// Create the necessary bins
		for( size_t i=0; i<my_nr_bins; ++i){
			_bins.push_back(pop_t());
		}

		// Print parameters
		DBO << "Objectives: " << nr_of_bins << dbe;
		DBO << "Min bins per instance: " << _min_bins_inst << dbe;
		DBO << "Max bins per instance: " << _max_bins_inst << dbe;
		DBO << "Over-full instances: " << _overfull_inst << dbe;
		DBO << "Over-full bins: " << _overfull_bins << dbe;
		DBO << "First bin id: " << _first_bin_id << dbe;
		DBO << "My bins: " << _bins.size() << dbe;
    }

    int get_rank(){
    	return _world->rank();
    }

    /**
     * Creates a random initial population.
     *
     * The number of generated individuals is determined by
     * init_batch * init_size, but the final population size is determined by
     * bin_size.
     */
    void random_pop(){
        dbg::trace trace("ea", DBG_HERE);
        parallel::init();

        // Clear the population
        this->_pop.clear();

        // Define bins
        pop_t send_pop;
        pop_t recv_pop;
        bins_t new_indiv_bins (_bins.size());

        // Calculate the number of batches
        unsigned nr_of_batches = (init_size / max_batch_size);
        unsigned remainder = init_size % max_batch_size;
        if(remainder) ++nr_of_batches;

        // Generate and evaluate new random individuals
        for(size_t bin_i=0; bin_i<_bins.size(); ++bin_i){
        	DBO << "Pop for bin: " << bin_i << dbe;
        	for(unsigned j=0; j<nr_of_batches; ++j){
        		// Calculate batch size
            	unsigned batch_size = max_batch_size;
            	if(j==(nr_of_batches-1) && remainder > 0){
            		batch_size = remainder;
            	}

        		// Generate the new random individuals
        		new_indiv_bins[bin_i].resize(batch_size);
        		int i = 0;
        		BOOST_FOREACH(phen_ptr_t& indiv, new_indiv_bins[bin_i]){
        			DBO << "Random indiv: " << i++ << dbe;
        			indiv = phen_ptr_t(new phen_t());
        			indiv->random();
        		}

            	// Evaluate the new individuals
        		DBO << "Evaluating pop" << std::endl;
            	_evaluate(new_indiv_bins[bin_i]);
        	}
        }

        // Exchange individuals
        for(size_t offset=0; offset<bin_size; offset+=max_send_size){
        	send_pop.clear();
        	recv_pop.clear();
        	for(size_t bin_i=0; bin_i<_bins.size(); ++bin_i){
        		pop_it_t it_slice_begin = new_indiv_bins[bin_i].begin();
        		pop_it_t it_end = new_indiv_bins[bin_i].end();
        		std::advance (it_slice_begin, offset);
        		pop_it_t it_slice_end = it_slice_begin;
        		for(size_t i=0; i<max_send_size && it_slice_end != it_end; ++i){
        			++it_slice_end;
        		}
        		send_pop.insert(send_pop.end(), it_slice_begin, it_slice_end);
        	}

        	g_exchange_v(_world, send_pop, recv_pop);

        	// Add all individuals to the same population and perform selection
        	for(size_t bin_i=0; bin_i<_bins.size(); ++bin_i){
        		this->_pop.clear();

            	// Copy the current individuals to the population for processing
            	compare::merge(this->_pop, _bins[bin_i]);

        		// Copy the new individuals to the population for processing
        		compare::merge(this->_pop, send_pop);

        		// Add the received individuals to the population for processing
        		compare::merge(this->_pop, recv_pop);

        		// Perform selection on the current population
        		_survivor_selection(bin_i);

        		// Copy the current population back to the appropriate bin
        		_copy_to_bin(bin_i, this->_pop);
        	}
        }

        // Copy everything back to the current population for serialization
        _copy_all_to_pop(this->_pop);
    }

    /**
     * Runs one epoch of our evolutionary algorithm.
     */
    void epoch(){
        dbg::trace trace("ea", DBG_HERE);
        DBO << "Starting epoch: " << this->_gen << dbe;

        // Clear the population
        this->_pop.clear();

        // Define bins
        pop_t send_pop;
        pop_t recv_pop;
        bins_t new_indiv_bins (_bins.size());

        // Generate and evaluate new individuals
        for(size_t bin_i=0; bin_i<_bins.size(); ++bin_i){
        	for (size_t i = 0; i < (offspring_per_generation/2); ++i){
        		DBO << "New indiv: " << i << dbe;
        		phen_ptr_t p1 = selection(_bins[bin_i]);
        		phen_ptr_t p2 = selection(_bins[bin_i]);
        		phen_ptr_t i1, i2;
        		p1->cross(p2, i1, i2);
        		i1->mutate();
        		i2->mutate();
        		new_indiv_bins[bin_i].push_back(i1);
        		new_indiv_bins[bin_i].push_back(i2);
        	}

        	// Evaluate the new individuals
        	_evaluate(new_indiv_bins[bin_i]);
        }

#ifdef EA_EVAL_ALL
        for(size_t bin_i=0; bin_i<_bins.size(); ++bin_i){
        	_evaluate(_bins[bin_i]);
        }
#endif

//        // Decide which individuals to send
//        for(size_t bin_i=0; bin_i<_bins.size(); ++bin_i){
//
//        	// Send the current (old) population
//        	//compare::merge(send_bins[bin_i], _bins[bin_i]);
//
//        	// Send the newly created individuals
//        	compare::merge(send_bins[bin_i], new_indiv_bins[bin_i]);
//        }

        /** OLD MP-CMOEA EXCHANGE START **/
//        // Get the set of all partners that we need to interact with
//        std::set<int> partners;
//        std::map<int, std::vector<int> > recv_bin_map;
//        std::map<int, std::vector<int> > send_bin_map;
//        std::map<int, int> recv_map;
//        for(size_t bin_i=0; bin_i<_bins.size(); ++bin_i){
//        	_get_partners(bin_i, partners, recv_bin_map, send_bin_map, recv_map);
//        }
//
//#if defined(DBG_ENABLED)
//        {
//        	DBO_MPI << "Partners of instance: " << _world->rank() << dbe;
//        	std::set<int>::iterator it;
//        	for (it = partners.begin(); it != partners.end(); ++it){
//        		DBO_MPI << "  partner: " << *it << dbe;
//        	}
//        }
//#endif
//        {
//        	std::set<int>::iterator it;
//        	for (it = partners.begin(); it != partners.end(); ++it){
//        		_exchange_send(*it, recv_bin_map, send_bin_map, send_bins, recv_bins, recv_map);
//        	}
//        	for (it = partners.begin(); it != partners.end(); ++it){
//        		_exchange_recv(*it, recv_bin_map, send_bin_map, send_bins, recv_bins, recv_map);
//        	}
//        }
//        _wait_all();
		/** OLD MP-CMOEA EXCHANGE END **/

//        // Exchange individuals with partners in the correct order
//        while(!partners.empty()){
//        	int partner = *partners.begin();
//        	partners.erase(partner);
//        	_exchange(partner, recv_bin_map, send_bin_map, send_bins, recv_bins, recv_map);
//        }

        // Exchange individuals
        for(size_t bin_i=0; bin_i<_bins.size(); ++bin_i){
        	compare::merge(send_pop, new_indiv_bins[bin_i]);
        }

        g_exchange_v(_world, send_pop, recv_pop);

        // Add all individuals to the same population and perform selection
        for(size_t bin_i=0; bin_i<_bins.size(); ++bin_i){
        	this->_pop.clear();

        	// Copy the current individuals to the population for processing
        	compare::merge(this->_pop, _bins[bin_i]);

        	// Copy the new individuals to the population for processing
        	compare::merge(this->_pop, send_pop);

        	// Add the received individuals to the population for processing
        	compare::merge(this->_pop, recv_pop);

        	// Perform selection on the current population
        	_survivor_selection(bin_i);

        	// Copy the current population back to the appropriate bin
        	_copy_to_bin(bin_i, this->_pop);
        }

        // Copy everything back to the current population for serialization
        _copy_all_to_pop(this->_pop);

        DBO << "Epoch " << this->_gen << " complete." << dbe;
    }

    void barrier(){
    	_world->barrier();
    }

protected:
    // Pointer to the MPI environment
    boost::shared_ptr<mpi_env_t> _env;

    // Pointer to the MPI world
    boost::shared_ptr<mpi_com_t> _world;

    std::vector<boost::mpi::packed_oarchive*> _send_buffer;
    std::vector<size_t*> _send_buffer_sizes;
    std::vector<MPI_Request*> _outstanding_requests;

    // Vector of bins stores all bins governed by this instance.
    bins_t _bins;

    // The global index of the first bin this instance is responsible for
    int _first_bin_id;

    // The lower bound on the number of bins handled by any one instance
    int _min_bins_inst;

    // The upper bound on the number of bins handled by any one instance
    // Because bins are distributed as evenly over instances as possible, will
    // always be _min_bins_inst + 1
    int _max_bins_inst;

    // The number of "over-full instances" (MPI instances responsible for
    // _max_bins_inst bins).
    int _overfull_inst;

    // The total number of bins allocated to "over-full instances"
    int _overfull_bins;

    /**
     * Evaluates all individuals in the provided population.
     *
     * @param pop: The population to evaluate.
     */
    void _evaluate(pop_t& pop){
    	dbg::trace trace("ea", DBG_HERE);
    	this->_eval.eval(pop, 0, pop.size(), this->_fit_proto);
#if defined(DBG_ENABLED)
    	for(size_t i=0; i<pop.size(); ++i){
    		DBOF << "Fit indiv: " << i << dbe;
    		for(size_t j=0; j<nr_of_bins; ++j){
    			DBOF << "  bin " << j <<
    					": " << pop[i]->fit().getBinFitness(j) << dbe;
    		}
    	}
#endif
    }

    /**
     * Performs survivor selection on the bin at the provided index.
     *
     * @param bin_i: The local bin index of the bin to evaluate.
     */
    void _survivor_selection(size_t bin_i){
    	dbg::trace trace("ea", DBG_HERE);

    	// Determine the type of bin that we govern based on our rank
    	size_t bin_obj = (_first_bin_id + bin_i) % nr_of_bins;
    	DBO << "Performing selection for bin: " << bin_obj << dbe;
    	_cat_to_obj(this->_pop, bin_obj);

    	// Apply modifiers
    	DBO << "  Applying modifiers" << dbe;
    	this->apply_modifier();

#if defined(DBG_ENABLED)
    	for(size_t i=0; i<this->_pop.size(); ++i){
    		DBOO << "  Pop before selection: " << i << dbe;
    		for(size_t j=0; j<this->_pop[i]->fit().objs().size(); ++j){
    			DBOO << "    obj " << j <<
    					": " << this->_pop[i]->fit().objs()[j] << dbe;
    		}
    	}
#endif

    	// Perform selection
    	DBO << "  Performing selection" << dbe;
    	pop_t ptmp;
    	sort_f()(this->_pop, ptmp, bin_size);
    	this->_pop = ptmp;

#if defined(DBG_ENABLED)
    	for(size_t i=0; i<this->_pop.size(); ++i){
    		DBOO << "  Pop after selection: " << i << dbe;
    		for(size_t j=0; j<this->_pop[i]->fit().objs().size(); ++j){
    			DBOO << "    obj " << j <<
    					": " << this->_pop[i]->fit().objs()[j] << dbe;
    		}
    	}
#endif

    }

    /**
     * Initialize the current bins given a population.
     *
     * @param pop: The population with which to initialize the current bins.
     */
    void _set_pop(const pop_t& pop) {
        dbg::trace trace("ea", DBG_HERE);
        size_t size_expected = _bins.size()*bin_size;
        DBO << "Pop size from archive: " << pop.size() << dbe;
        DBO << "Pop size expected: " << size_expected << dbe;
        dbg::assertion(DBG_ASSERTION(pop.size() == size_expected));
        size_t k=0;
        for(size_t i=0; i<_bins.size(); ++i){
        	for(size_t j=0; j<bin_size; ++j){
        		_bins[i].push_back(pop[k]);
        		++k;
        	}
        }
    }

    /**
     * Copy all individuals from the provided population to the bin at the
     * provided bin index.
     *
     * @param bin_i: The local bin index of the bin to copy to.
     * @param pop: The population of individuals to copy to that bin.
     */
    void _copy_to_bin(size_t bin_i, const pop_t& pop){
    	dbg::trace trace("ea", DBG_HERE);
    	_bins[bin_i].clear();
    	compare::merge(_bins[bin_i], pop);
    }

    /**
     * Copy all individuals of all bins to the provided population.
     *
     * @param pop (out): The population to which to copy the current bins.
     */
    void _copy_all_to_pop(pop_t& pop){
    	dbg::trace trace("ea", DBG_HERE);
    	pop.clear();
    	for(size_t bin_i=0; bin_i<_bins.size(); ++bin_i){
    		compare::merge(pop, _bins[bin_i]);
    	}
    }

    /**
     * Copy the performance on the task associated with the provided bin index
     * to one of the current objectives of the all individuals of the provided
     * population.
     *
     * @param pop: The population of individuals for which to copy task
     *     performance.
     * @param bin_i: The global bin index of the task performance that needs to
     *     be copied.
     */
    void _cat_to_obj(pop_t& pop, size_t bin_i){
        dbg::trace trace("ea", DBG_HERE);
        for(size_t i=0; i<pop.size(); ++i){
        	pop[i]->fit().set_obj(obj_index, pop[i]->fit().getBinFitness(bin_i));
        }
    }

    /**
     * Calculates which other MPI instances we need to communicate with, and
     * what data needs to be exchanged.
     *
     * @param local_bin_i: The local bin index for which we want to exchange
     *     information.
     * @param partners (out): Reference to the set of partners (mpi instances)
     *     with whom we need to communicate.
     * @param recv_bin_map (out): Reference to a map mapping instance rank to
     *     a vector of global bin indices of bins that we will receive from that
     *     instance.
     * @param send_bin_map (out): Reference to a map mapping instance rank to a
     *     vector of global bin indices that we will send to that instance.
     * @param recv_map (out): Map indicating which of the received bins should
     *     be added to which of the bins governed by this instance.
     */
    void _get_partners(int local_bin_i,
    		std::set<int>& partners,
			inst_bin_map_t& recv_bin_map,
			inst_bin_map_t& send_bin_map,
			std::map<int, int>& recv_map){
    	dbg::trace trace("ea", DBG_HERE);

    	// Determine with which bins to exchange individuals
    	int global_bin_i = _get_global_bin_i(local_bin_i);
    	int offset = (this->_gen % (bins_total-1)) + 1;
    	int bin_id1 =  compare::mod(global_bin_i + offset, bins_total);
    	int bin_id2 =  compare::mod(global_bin_i - offset, bins_total);
    	int inst1 = _get_instance_from_bin(bin_id1);
    	int inst2 = _get_instance_from_bin(bin_id2);

    	send_bin_map[inst1].push_back(global_bin_i);
    	send_bin_map[inst2].push_back(global_bin_i);
    	recv_bin_map[inst1].push_back(bin_id1);
    	recv_bin_map[inst2].push_back(bin_id2);
    	recv_map[bin_id1] = global_bin_i;
    	recv_map[bin_id2] = global_bin_i;

    	partners.insert(inst1);
    	partners.insert(inst2);
    }

    /**
     * Exchange all necessary bins with the indicated partner.
     *
     * @param partner: Rank of the MPI instance with whom to exchange data.
     * @param recv_bin_map: The map of MPI instances and bins to receive.
     * @param send_bin_map: The map on MPI instances and bins to send.
     * @param send_bins: Vector containing the actual bins to send.
     * @param recv_bins (out): Vector for storing received bins.
     * @param recv_map: Map indicating which received bins should be added to
     *     which of the local bins (in global bin ids).
     */
    void _exchange(int partner,
    		inst_bin_map_t& recv_bin_map,
			inst_bin_map_t& send_bin_map,
			const bins_t& send_bins,
			bins_t& recv_bins,
			const std::map<int, int>& recv_map){
    	dbg::trace trace("ea", DBG_HERE);
    	DBO_MPI << "Exchanging with partner: " << partner << dbe;

    	// These are the bins I am going to send, and in the order in which I
    	// will send them
    	std::vector<int> bins_to_send = send_bin_map[partner];
    	compare::sort(bins_to_send);

    	// These are the bins I am going to receive, and in the order in which
    	// I will receive them
    	std::vector<int> bins_to_receive = recv_bin_map[partner];
    	compare::sort(bins_to_receive);

    	if(_world->rank() == partner){
    		DBO_MPI << "Performing local exchange" << dbe;
    		for(size_t i=0; i<bins_to_send.size(); ++i){
    			int send_bin_i = _get_local_bin_i(bins_to_send[i]);
    			int recv_bin_i = _get_local_bin_i(bins_to_receive[i]);
    			recv_bins[recv_bin_i] = send_bins[send_bin_i];
    		}
    	} else if(_world->rank() < partner){
    		DBO_MPI << "My rank is lower, sending first" << dbe;
    		for(size_t i=0; i<bins_to_send.size(); ++i){
    			DBO_MPI << "Sending bin: " << bins_to_send[i] << dbe;
    			int local_bin_i = _get_local_bin_i(bins_to_send[i]);
    			DBO_MPI << "Local bin index: " << local_bin_i << dbe;
    			_send(partner, send_bins[local_bin_i]);
    		}
    		for(size_t i=0; i<bins_to_receive.size(); ++i){
    			DBO_MPI << "Receiving bin: " << bins_to_receive[i] << dbe;
    			int global_bin_i = recv_map.at(bins_to_receive[i]);
    			DBO_MPI << "Adding it to bin: " << global_bin_i << dbe;
    			int local_bin_i = _get_local_bin_i(global_bin_i);
    			DBO_MPI << "Local bin index: " << local_bin_i << dbe;
    			_recv(partner, recv_bins[local_bin_i]);
    		}
    	} else {
    		DBO_MPI << "My rank is higher, receiving first" << dbe;
    		for(size_t i=0; i<bins_to_receive.size(); ++i){
    			DBO_MPI << "Receiving bin: " << bins_to_receive[i] << dbe;
    			int global_bin_i = recv_map.at(bins_to_receive[i]);
    			DBO_MPI << "Adding it to bin: " << global_bin_i << dbe;
    			int local_bin_i = _get_local_bin_i(global_bin_i);
    			DBO_MPI << "Local bin index: " << local_bin_i << dbe;
    			_recv(partner, recv_bins[local_bin_i]);
    		}
    		for(size_t i=0; i<bins_to_send.size(); ++i){
    			DBO_MPI << "Sending bin: " << bins_to_send[i] << dbe;
    			int local_bin_i = _get_local_bin_i(bins_to_send[i]);
    			DBO_MPI << "Local bin index: " << local_bin_i << dbe;
    			_send(partner, send_bins[local_bin_i]);
    		}
    	}
    }

    /**
     * Exchange all necessary bins with the indicated partner.
     *
     * @param partner: Rank of the MPI instance with whom to exchange data.
     * @param recv_bin_map: The map of MPI instances and bins to receive.
     * @param send_bin_map: The map on MPI instances and bins to send.
     * @param send_bins: Vector containing the actual bins to send.
     * @param recv_bins (out): Vector for storing received bins.
     * @param recv_map: Map indicating which received bins should be added to
     *     which of the local bins (in global bin ids).
     */
    void _exchange_send(int partner,
    		inst_bin_map_t& recv_bin_map,
			inst_bin_map_t& send_bin_map,
			const bins_t& send_bins,
			bins_t& recv_bins,
			const std::map<int, int>& recv_map){
    	dbg::trace trace("ea", DBG_HERE);
    	DBO_MPI << "Sending to partner: " << partner << dbe;

    	// These are the bins I am going to send, and in the order in which I
    	// will send them
    	std::vector<int> bins_to_send = send_bin_map[partner];
    	compare::sort(bins_to_send);

    	// These are the bins I am going to receive, and in the order in which
    	// I will receive them
    	std::vector<int> bins_to_receive = recv_bin_map[partner];
    	compare::sort(bins_to_receive);

    	if(_world->rank() == partner){
    		DBO_MPI << "Performing local exchange" << dbe;
    		for(size_t i=0; i<bins_to_send.size(); ++i){
    			int send_bin_i = _get_local_bin_i(bins_to_send[i]);
    			int recv_bin_i = _get_local_bin_i(bins_to_receive[i]);
    			recv_bins[recv_bin_i] = send_bins[send_bin_i];
    		}
    	} else {
    		DBO_MPI << "Sending:" << dbe;
    		for(size_t i=0; i<bins_to_send.size(); ++i){
    			DBO_MPI << "Sending bin: " << bins_to_send[i] << dbe;
    			int local_bin_i = _get_local_bin_i(bins_to_send[i]);
    			DBO_MPI << "Local bin index: " << local_bin_i << dbe;
    			_sendi(partner, send_bins[local_bin_i], int(i));
    		}
    	}
    }

    /**
     * Exchange all necessary bins with the indicated partner.
     *
     * @param partner: Rank of the MPI instance with whom to exchange data.
     * @param recv_bin_map: The map of MPI instances and bins to receive.
     * @param send_bin_map: The map on MPI instances and bins to send.
     * @param send_bins: Vector containing the actual bins to send.
     * @param recv_bins (out): Vector for storing received bins.
     * @param recv_map: Map indicating which received bins should be added to
     *     which of the local bins (in global bin ids).
     */
    void _exchange_recv(int partner,
    		inst_bin_map_t& recv_bin_map,
			inst_bin_map_t& send_bin_map,
			const bins_t& send_bins,
			bins_t& recv_bins,
			const std::map<int, int>& recv_map){
    	dbg::trace trace("ea", DBG_HERE);
    	DBO_MPI << "Receiving from partner: " << partner << dbe;

    	// These are the bins I am going to send, and in the order in which I
    	// will send them
    	std::vector<int> bins_to_send = send_bin_map[partner];
    	compare::sort(bins_to_send);

    	// These are the bins I am going to receive, and in the order in which
    	// I will receive them
    	std::vector<int> bins_to_receive = recv_bin_map[partner];
    	compare::sort(bins_to_receive);

    	if(_world->rank() == partner){
    		DBO_MPI << "Partner is local: exchange already completed" << dbe;
    	} else {
    		DBO_MPI << "Receiving" << dbe;
    		for(size_t i=0; i<bins_to_receive.size(); ++i){
    			DBO_MPI << "Receiving bin: " << bins_to_receive[i] << dbe;
    			int global_bin_i = recv_map.at(bins_to_receive[i]);
    			DBO_MPI << "Adding it to bin: " << global_bin_i << dbe;
    			int local_bin_i = _get_local_bin_i(global_bin_i);
    			DBO_MPI << "Local bin index: " << local_bin_i << dbe;
    			_recv(partner, recv_bins[local_bin_i], int(i));
    		}
    	}
    }

    /**
     * Send the supplied population to the indicated MPI instance.
     *
     * @param partner: The rank of the MPI instance to send to.
     * @param pop: The population to send.
     */
    void _sendi(int partner, const pop_t& pop, int tag=1){
    	dbg::trace trace("ea", DBG_HERE);
    	MPI_Request* req1 = new MPI_Request();
    	MPI_Request* req2 = new MPI_Request();
    	_outstanding_requests.push_back(req1);
    	_outstanding_requests.push_back(req2);
        std::vector<MPI_Request> req(2);
        DBO_MPI << "sending to rank: " << partner << std::endl;

        // Serializing population

        boost::mpi::packed_oarchive* oa = new boost::mpi::packed_oarchive (*_world);
        _send_buffer.push_back(oa);

		//boost::mpi::packed_oarchive oa(*_world);
        (*oa) << pop;
        //_outstanding_requests.push_back(MPI_Request());
        size_t* size_p = new size_t((*oa).size());
        _send_buffer_sizes.push_back(size_p);

        // Send the the number of individuals that will be send
        const void* size_1 = size_p;
        BOOST_MPI_CHECK_RESULT(MPI_Isend,
                (const_cast<void*>(size_1), 1,
                 boost::mpi::get_mpi_datatype<std::size_t>(*size_p),
				 partner, tag, *_world, req1));

        // Actually send the individuals
        BOOST_MPI_CHECK_RESULT(MPI_Isend,
                (const_cast<void*>((*oa).address()),
                (*oa).size(),
                 MPI_PACKED,
				 partner, tag, *_world, req2));

        // Wait until everything is send successfully
        //BOOST_MPI_CHECK_RESULT(MPI_Waitall, (req.siz(), &req[0], MPI_STATUSES_IGNORE));
    }

    /**
     * Send the supplied population to the indicated MPI instance.
@@ -640,29 +748,53 @@ protected:
     */
    void _send(int partner, const pop_t& pop){
       dbg::trace trace("ea", DBG_HERE);
        std::vector<MPI_Request> req(2);
        DBO_MPI << "sending to rank: " << partner << std::endl;
        int tag = 1;

        // Serializing population
        boost::mpi::packed_oarchive oa(*_world);
        oa << pop;

        // Send the the number of individuals that will be send
        const void* size_1 = &oa.size();

        BOOST_MPI_CHECK_RESULT(MPI_Isend,
                (const_cast<void*>(size_1), 1,
                 boost::mpi::get_mpi_datatype<std::size_t>(oa.size()),
                                partner, tag, *_world, &req[0]));

        // Actually send the individuals
        BOOST_MPI_CHECK_RESULT(MPI_Isend,
                (const_cast<void*>(oa.address()), oa.size(),
                 MPI_PACKED, partner, tag, *_world, &req[1]));

        // Wait until everything is send successfully
        BOOST_MPI_CHECK_RESULT(MPI_Waitall, (2, &req[0], MPI_STATUSES_IGNORE));
    }

    void _wait_all(){
    	std::vector<MPI_Request> requests;
    	for(size_t i=0; i<_outstanding_requests.size(); ++i){
    		requests.push_back(*(_outstanding_requests[i]));
    	}
    	BOOST_MPI_CHECK_RESULT(MPI_Waitall, (requests.size(), &requests[0], MPI_STATUSES_IGNORE));
    	for(size_t i=0; i<_outstanding_requests.size(); ++i){
    		delete _outstanding_requests[i];
    	}
    	_outstanding_requests.clear();
    	for(size_t i=0; i<_send_buffer_sizes.size(); ++i){
    		delete _send_buffer_sizes[i];
    	}
    	_send_buffer_sizes.clear();
    	for(size_t i=0; i<_send_buffer.size(); ++i){
    		delete _send_buffer[i];
    	}
    	_send_buffer.clear();
    }

    /**
     * Receive individuals from the indicated MPI instance.
     *
     * @param partner: The rank of the MPI instance to receive from.
     * @param pop (out): Reference to a population in which to store the
     *     received individuals.
     */
    void _recv(int partner, pop_t& pop, int tag = 1){
    	dbg::trace trace("ea", DBG_HERE);
        // Receive data from the root.
        MPI_Status stat;
        std::size_t count;
        int recv = 0;
        using namespace boost::mpi;

        DBO_MPI << "receiving from rank: " << partner << dbe;

        BOOST_MPI_CHECK_RESULT(MPI_Recv,
                (&count, 1,  get_mpi_datatype<std::size_t>(count),
                partner, tag, *_world, &stat));
        MPI_Get_count(&stat,  get_mpi_datatype<std::size_t>(count), &recv);
        DBO_MPI  <<
				" received size_t of size: " << recv <<
				" from " << stat.MPI_SOURCE <<
				" with tag " << stat.MPI_TAG << std::endl;

        // Prepare input buffer and receive the message
        mpi_iarchive_t ia(*_world);
        ia.resize(count);
        BOOST_MPI_CHECK_RESULT(MPI_Recv,
                (ia.address(), ia.size(), MPI_PACKED,
                partner, tag,
                 *_world, &stat));
        MPI_Get_count(&stat, MPI_PACKED, &recv);
        DBO_MPI <<
				" received archive of size: " << recv <<
				" from " << stat.MPI_SOURCE <<
				" with tag " << stat.MPI_TAG << std::endl;

        // De-serializing population
        ia >> pop;
    }

    /**
     * Takes a global bin index and returns the local bin index.
     *
     * @param global_bin_i: The global bin index to convert.
     * @return: The local bin index.
     */
    inline int _get_local_bin_i(int global_bin_i) const{
    	dbg::trace trace("ea", DBG_HERE);
    	dbg::assertion(DBG_ASSERTION(global_bin_i >= 0));
    	dbg::assertion(DBG_ASSERTION(global_bin_i < bins_total));
    	int result = global_bin_i - _first_bin_id;
    	dbg::assertion(DBG_ASSERTION(result >= 0));
    	dbg::assertion(DBG_ASSERTION(result < _bins.size()));
    	return result;
    }

    /**
     * Takes a local bin index and returns its global bin index.
     *
     * @param local_bin_i: The local bin index to convert.
     * @return: The global bin index.
     */
    inline int _get_global_bin_i(int local_bin_i) const{
    	dbg::trace trace("ea", DBG_HERE);
    	dbg::assertion(DBG_ASSERTION(local_bin_i >= 0));
    	dbg::assertion(DBG_ASSERTION(local_bin_i < _bins.size()));
    	int result = local_bin_i + _first_bin_id;
    	dbg::assertion(DBG_ASSERTION(result >= 0));
    	dbg::assertion(DBG_ASSERTION(result < bins_total));
    	return result;
    }

    /**
     * Returns the rank of the instance associated with the supplied bin_id.
     *
     * @param bin_id: The global bin index for which we want to retrieve the MPI
     *     instance.
     * @return: The rank of the MPI instance responsible for the provided bin
     *     index.
     */
    inline int _get_instance_from_bin(int bin_id) const{
    	dbg::trace trace("ea", DBG_HERE);
    	int instance_rank;

		if(bin_id < _overfull_bins){
			// The instance we are looking for is one of the over-full instances
			instance_rank = bin_id / _max_bins_inst;
		} else{
			// The instance we are looking for is one of the other instances
			int other_bins = (bin_id - _overfull_bins);
			int other_instances = other_bins / _min_bins_inst;
			instance_rank = _overfull_inst + other_instances;
		}

		return instance_rank;
    }
};
}
}

#undef DBO
#undef DBO_MPI
#undef DBOF
#undef DBOO

#endif /* MODULES_CMOEA_CMOEA_NSGA2_MPI_HPP_ */


// MIT License
//
// Copyright (c) 2018 degski
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#pragma once

#include <cstdint>
#include <cstdlib>
#include <cmath>

#include <iostream>
#include <random>
#include <unordered_map>
#include <vector>

#include <boost/dynamic_bitset.hpp>
#include <boost/container/static_vector.hpp>

#include <cereal/cereal.hpp>
#include <cereal/types/unordered_map.hpp>
#include <cereal/archives/binary.hpp>

#include "owningptr.hpp"
#include "pool_allocator.hpp"

#include "autotimer.hpp"

#include "Typedefs.hpp"
#include "player.hpp"
#include "flat_search_tree.hpp"

namespace mcts {

    template<typename T>
    class Stack {

        std::vector<T> m_data;

    public:

        explicit Stack ( const T & v_ ) noexcept {
            m_data.push_back ( v_ );
        }

        [[ nodiscard ]] T pop ( ) noexcept {
            const T v { m_data.back ( ) };
            m_data.pop_back ( );
            return v;
        }

        void push ( const T & v_ ) noexcept {
            m_data.push_back ( v_ );
        }

        [[ nodiscard ]] bool not_empty ( ) const noexcept {
            return m_data.size ( );
        }
    };

    template<typename T>
    class Queue {

        boost::container::deque<T> m_data;

    public:

        explicit Queue ( const T & v_ ) noexcept {
            m_data.push_back ( v_ );
        }

        [[ nodiscard ]] const T pop ( ) noexcept {
            const T v = m_data.front ( );
            m_data.pop_front ( );
            return v;
        }

        void push ( const T & v_ ) noexcept {
            m_data.push_back ( v_ );
        }

        [[ nodiscard ]] bool not_empty ( ) const noexcept {
            return m_data.size ( );
        }
    };

    template<typename State>
    struct ArcData { // 1 bytes.

        using state_type = State;
        // float m_score = 0.0f; // 4 bytes.
        // std::int32_t m_visits = 0; // 4 bytes.
        typename State::Move m_move = State::Move::invalid; // 1 bytes.

        ArcData ( ) noexcept {
            // std::cout << "arcdata default constructed\n";
        }
        ArcData ( const State & state_ ) noexcept {
            // std::cout << "arcdata constructed from state\n";
            m_move = state_.lastMove ( );
        }
        ArcData ( const ArcData & ad_ ) noexcept {
            // std::cout << "arcdata copy constructed\n";
            // m_score = nd_.m_score;
            // m_visits = nd_.m_visits;
            m_move = ad_.m_move;
        }
        ArcData ( ArcData && ad_ ) noexcept {
            // std::cout << "arcdata move constructed\n";
            // m_score = std::move ( ad_.m_score );
            // m_visits = std::move ( ad_.m_visits );
            m_move = std::move ( ad_.m_move );
        }

        ~ArcData ( ) noexcept {
        }

        [[ maybe_unused ]] ArcData & operator += ( const ArcData & rhs_ ) noexcept {
            // m_score += rhs_.m_score;
            // m_visits += rhs_.m_visits;
            return * this;
        }

        [[ maybe_unused ]] ArcData & operator = ( const ArcData & ad_ ) noexcept {
            // std::cout << "arcdata copy assigned\n";
            // m_score = nd_.m_score;
            // m_visits = nd_.m_visits;
            m_move = ad_.m_move;
            return * this;
        }

        [[ maybe_unused ]] ArcData & operator = ( ArcData && ad_ ) noexcept {
            // std::cout << "arcdata move assigned\n";
            // m_score = std::move ( ad_.m_score );
            // m_visits = std::move ( ad_.m_visits );
            m_move = std::move ( ad_.m_move );
            return * this;
        }

    private:

        friend class cereal::access;

        template < class Archive >
        void serialize ( Archive & ar_ ) {
            ar_ ( m_move );
        }
    };

    template<typename State>
    struct NodeData { // 17 bytes.

        using state_type = State;
        using Moves = typename State::Moves;
        using Move = typename State::Moves::value_type;
        using MovesPool = pa::pool_allocator<Moves>;
        using MovesPoolPtr = llvm::OwningPtr<MovesPool>;

        Moves * m_moves = nullptr;  // 8 bytes.
        float m_score = 0.0f; // 4 bytes.
        std::int32_t m_visits = 0; // 4 bytes.
        Player m_player_just_moved = Player::Type::invalid; // 1 byte.

        // Constructors.

        NodeData ( ) noexcept {
            // std::cout << "nodedata default constructed\n";
        }
        NodeData ( const State & state_ ) noexcept {
            // std::cout << "nodedata constructed from state\n";
            m_moves = m_moves_pool->new_element ( );
            if ( not ( state_.moves ( m_moves ) ) ) {
                m_moves_pool->delete_element ( m_moves );
                m_moves = nullptr;
            }
            m_player_just_moved = state_.playerJustMoved ( );
        }
        NodeData ( const NodeData & nd_ ) noexcept {
            // std::cout << "nodedata copy constructed\n";
            if ( nd_.m_moves != nullptr ) {
                m_moves = m_moves_pool->new_element ( * nd_.m_moves );
            }
            m_score = nd_.m_score;
            m_visits = nd_.m_visits;
            m_player_just_moved = nd_.m_player_just_moved;
        }
        NodeData ( NodeData && nd_ ) noexcept {
            // std::cout << "nodedata move constructed\n";
            std::swap ( m_moves, nd_.m_moves );
            m_score = std::move ( nd_.m_score );
            m_visits = std::move ( nd_.m_visits );
            m_player_just_moved = std::move ( nd_.m_player_just_moved );
        }

        ~NodeData ( ) noexcept {
            m_moves_pool->delete_element ( m_moves ); // Checked for nullptr in delete_element ( ... ).
        }

        [[ nodiscard ]] Move getUntriedMove ( ) noexcept {
            if ( 1 == m_moves->size ( ) ) {
                // 1 move left, so destroy memory and return that 1 move.
                const Move move = m_moves->front ( );
                m_moves_pool->delete_element ( m_moves );
                m_moves = nullptr;
                return move;
            }
            return m_moves->draw ( );
        }

        [[ maybe_unused ]] NodeData & operator += ( const NodeData & rhs_ ) noexcept {
            m_score += rhs_.m_score;
            m_visits += rhs_.m_visits;
            return * this;
        }

        [[ maybe_unused ]] NodeData & operator = ( const NodeData & nd_ ) noexcept {
            // std::cout << "nodedata copy assigned\n";
            if ( nullptr != nd_.m_moves ) {
                m_moves = m_moves_pool->new_element ( * nd_.m_moves );
            }
            m_score = nd_.m_score;
            m_visits = nd_.m_visits;
            m_player_just_moved = nd_.m_player_just_moved;
            return * this;
        }

        [[ maybe_unused ]] NodeData & operator = ( NodeData && nd_ ) noexcept {
            // std::cout << "nodedata move assigned\n";
            std::swap ( m_moves, nd_.m_moves );
            m_score = std::move ( nd_.m_score );
            m_visits = std::move ( nd_.m_visits );
            m_player_just_moved = std::move ( nd_.m_player_just_moved );
            return * this;
        }

        static MovesPoolPtr m_moves_pool;

    private:

        friend class cereal::access;

        template < class Archive >
        void save ( Archive & ar_ ) const noexcept {
            if ( nullptr != m_moves ) {
                const std::int8_t tmp = 2;
                ar_ ( tmp );
                m_moves->serialize ( ar_ );
            }
            else {
                const std::int8_t tmp = 1;
                ar_ ( tmp );
            }
            ar_ ( m_score, m_visits, m_player_just_moved );
        }

        template < class Archive >
        void load ( Archive & ar_ ) noexcept {
            std::int8_t tmp = -1;
            ar_ ( tmp );
            if ( 2 == tmp ) {
                m_moves = m_moves_pool->new_element ( );
                m_moves->serialize ( ar_ );
            }
            ar_ ( m_score, m_visits, m_player_just_moved );
        }
    };

    template<typename State>
    typename NodeData<State>::MovesPoolPtr NodeData<State>::m_moves_pool ( new MovesPool ( ) );

    template <typename State>
    using Tree = fst::SearchTree<ArcData<State>, NodeData<State>>;

    template<typename State>
    using NodeID = typename Tree<State>::NodeID;

    template<typename State>
    using ArcID = typename Tree<State>::ArcID;

    template < typename State >
    class Mcts {

    public:

        using Tree = Tree<State>;

        using ArcID = typename Tree::ArcID;
        using NodeID = typename Tree::NodeID;

        using ArcData = ArcData<State>;
        using NodeData = NodeData<State>;

        using InIt = typename Tree::in_iterator;
        using OutIt = typename Tree::out_iterator;

        using cInIt = typename Tree::const_in_iterator;
        using cOutIt = typename Tree::const_out_iterator;

        using Move = typename State::Move;
        using Moves = typename State::Moves;

        using Player = typename State::Player;

        using Link = typename Tree::Link;
        using Path = typename Tree::Path;

        using TranspositionTable = std::unordered_map<ZobristHash, NodeID>;
        using InverseTranspositionTable = std::vector < ZobristHash >;
        using TranspositionTablePtr = llvm::OwningPtr<TranspositionTable>;

        // The data.

        Tree m_tree;
        TranspositionTablePtr m_transposition_table;

        bool m_not_initialized = true;

        // The purpose of the m_path is to maintain the path with updates
        // (of the visits/scores) all the way to the original root (start
        // of the game). It's also used as a scrath-pad for the updates
        // after play-out. Each move and players' move gets added to this
        // path.

        Path m_path;
        index_t m_path_size;

        // Init.

        void initialize ( const State & state_ ) noexcept {
            if ( m_transposition_table.get ( ) == nullptr ) {
                m_transposition_table.reset ( new TranspositionTable ( ) );
            }
            // Set root_node data.
            m_tree [ m_tree.root_node ] = NodeData { state_ };
            // Add root_node to transposition_table.
            m_transposition_table->emplace ( state_.zobrist ( ), m_tree.root_node );
            // Has been initialized.
            m_not_initialized = false;
            m_path.reset ( Tree::ArcID::invalid, m_tree.root_node );
            m_path_size = 1;
        }

        [[ nodiscard ]] Link addArc ( const NodeID parent_, const NodeID child_, const State & state_ ) noexcept {
            return { m_tree.addArc ( parent_, child_, state_ ), child_ };
        }

        [[ nodiscard ]] Link addNode ( const NodeID parent_, const State & state_ ) noexcept {
            const Link link_to_child { addArc ( parent_, m_tree.addNode ( state_ ), state_ ) };
            m_transposition_table->emplace ( state_.zobrist ( ), link_to_child.target );
            return link_to_child;
        }

        void printMoves ( const NodeID n_ ) const noexcept {
            std::cout << "moves of " << ( int ) n_ << ": ";
            for ( OutIt a ( m_tree, n_ ); a != OutIt::end ( ); ++a ) {
                std::cout << "[" << ( int ) a.get ( ) << ", " << ( int ) m_tree [ a ].m_move.m_loc << "]";
            }
            putchar ( '\n' );
        }

        [[ nodiscard ]] Move getMove ( const ArcID arc_ ) const noexcept {
            return m_tree [ arc_ ].m_move;
        }

        [[ nodiscard ]] NodeID getNode ( const ZobristHash zobrist_ ) const noexcept {
            const auto it = m_transposition_table->find ( zobrist_ );
            return it == m_transposition_table->cend ( ) ? Tree::NodeID::invalid : it->second;
        }

        [[ nodiscard ]] bool hasChildren ( const NodeID node_ ) const noexcept {
            return m_tree.isInternal ( node_ );
        }

        // Moves.

        [[ nodiscard ]] bool hasNoUntriedMoves ( const NodeID node_ ) const noexcept {
            return m_tree [ node_ ].m_moves == nullptr;
        }

        [[ nodiscard ]] bool hasUntriedMoves ( const NodeID node_ ) const noexcept {
            return m_tree [ node_ ].m_moves != nullptr;
        }

        [[ nodiscard ]] Move getUntriedMove ( const NodeID node_ ) noexcept {
            return m_tree [ node_ ].getUntriedMove ( );
        }

        // Data.

        [[ nodiscard ]] std::int32_t getVisits ( const NodeID node_ ) const noexcept {
            std::int32_t visits = 0;
            for ( InIt a { m_tree, node_ }; a.is_valid ( ); ++a ) {
                visits += m_tree [ a ].m_visits;
            }
            return visits;
        }

        [[ nodiscard ]] float getUCTFromArcs ( const NodeID parent_, const NodeID child_ ) const noexcept {
            ArcData child_data;
            for ( InIt a { m_tree, child_ }; a.is_valid ( ); ++a ) {
                child_data += m_tree [ a ];
            }
            std::int32_t visits = 1;
            for ( InIt a { m_tree, parent_ }; a.is_valid ( ); ++a ) {
                visits += m_tree [ a ].m_visits;
            }
            //                              Exploitation                                                             Exploration
            // Exploitation is the task to select the move that leads to the best results so far.
            // Exploration deals with less promising moves that still have to be examined, due to the uncertainty of the evaluation.
            return ( float ) child_data.m_score / ( float ) child_data.m_visits + sqrtf ( 3.0f * logf ( ( float ) visits ) / ( float ) child_data.m_visits );
        }

        [[ nodiscard ]] float getUCTFromNode ( const NodeID parent_, const NodeID child_ ) const noexcept {
            //                              Exploitation                                                             Exploration
            // Exploitation is the task to select the move that leads to the best results so far.
            // Exploration deals with less promising moves that still have to be examined, due to the uncertainty of the evaluation.
            return ( float ) m_tree [ child_ ].m_score / ( float ) m_tree [ child_ ].m_visits + sqrtf ( 4.0f * logf ( ( float ) ( m_tree [ parent_ ].m_visits + 1 ) ) / ( float ) m_tree [ child_ ].m_visits );
        }

        [[ nodiscard ]] Link selectChildRandom ( const NodeID parent_ ) const noexcept {
            boost::container::static_vector<Link, State::max_no_moves> children;
            for ( OutIt a ( m_tree, parent_ ); a.is_valid ( ); ++a ) {
                children.emplace_back ( a.id ( ), a->target );
            }
            return children [ std::uniform_int_distribution<ptrdiff_t> ( 0, children.size ( ) - 1 ) ( g_rng ) ];
        }

        [[ nodiscard ]] Link selectChildUCT ( const NodeID parent_ ) const noexcept {
            cOutIt a = m_tree.cbeginOut ( parent_ );
            boost::container::static_vector < Link, State::max_no_moves > best_children ( 1, m_tree.link ( a ) );
            float best_UCT_score = getUCTFromNode ( parent_, best_children.back ( ).target );
            ++a;
            for ( ; a.is_valid ( ); ++a ) {
                const Link child = m_tree.link ( a );
                const float UCT_score = getUCTFromNode ( parent_, child.target );
                if ( UCT_score > best_UCT_score ) {
                    best_children.resize ( 1 );
                    best_children.back ( ) = child;
                    best_UCT_score = UCT_score;
                }
                else if ( UCT_score == best_UCT_score ) {
                    best_children.push_back ( child );
                }
            }
            // Ties are broken by fair coin flips.
            return best_children.size ( ) == 1 ? best_children.back ( ) : best_children [ std::uniform_int_distribution < ptrdiff_t > ( 0, best_children.size ( ) - 1 ) ( g_rng ) ];
        }

        // State is updated to reflect move.
        [[ nodiscard ]] Link addChild ( const NodeID parent_, const State & state_ ) noexcept {
            const NodeID child = getNode ( state_.zobrist ( ) );
            return Tree::NodeID::invalid == child ? addNode ( parent_, state_ ) : addArc ( parent_, child, state_ );
        }

        void updateData ( const Link & link_, const State & state_ ) noexcept {
            const float result = state_.result ( m_tree [ link_.target ].m_player_just_moved );
            // ++m_tree [ link_.arc ].m_visits;
            // m_tree [ link_.arc ].m_score += result;
            ++m_tree [ link_.target ].m_visits;
            m_tree [ link_.target ].m_score += result;
        }

        // Find the node (the most robust) with the most visits.
        [[ nodiscard ]] Move getBestMove ( ) noexcept {
            std::int32_t best_child_visits = INT_MIN;
            Move best_child_move = State::Move::none;
            m_path.push ( Tree::ArcID::invalid, Tree::NodeID::invalid );
            ++m_path_size;
            for ( cOutIt a ( m_tree.cbeginOut ( m_tree.root_node ) ); a.is_valid ( ); ++a ) {
                const Link child ( m_tree.link ( a ) );
                const std::int32_t child_visits ( m_tree [ child.target ].m_visits );
                if ( child_visits > best_child_visits ) {
                    best_child_visits = child_visits;
                    best_child_move = m_tree [ child.arc ].m_move;
                    m_path.back ( ) = child;
                }
            }
            return best_child_move;
        }

        // Adding the move of the opponent to the path (and possibly to the tree).
        void connectStatesPath ( const State & state_ ) noexcept {
            const NodeID parent = m_path.back ( ).target; NodeID child = getNode ( state_.zobrist ( ) );
            if ( Tree::NodeID::invalid == child ) {
                child = addNode ( parent, state_ ).target;
            }
            m_path.push ( m_tree.link ( parent, child ) );
            ++m_path_size;
        }

        [[ nodiscard ]] Move compute ( const State & state_, index_t max_iterations_ ) noexcept {

            // constexpr std::int32_t threshold = 5;

            if ( m_not_initialized ) {
                initialize ( state_ );
            }
            else {
                connectStatesPath ( state_ );
            }

            // const Player player = state_.playerToMove ( );
            // if ( player == Player::Type::agent ) {
                // m_path.print ( );
            // }

            // max_iterations_ -= m_tree.nodeNum ( );

            while ( max_iterations_-- > 0 ) {
                NodeID node = m_tree.root_node;
                State state ( state_ );
                // Select a path through the tree to a leaf node.
                while ( hasNoUntriedMoves ( node ) and hasChildren ( node ) ) {
                    // UCT is only applied in nodes of which the visit count
                    // is higher than a certain threshold T
                    Link child = selectChildUCT ( node );
                    state.move_hash ( m_tree [ child.arc ].m_move );
                    m_path.push ( child );
                    node = child.target;
                }

                /*

                static int cnt = 0;

                if ( state != m_tree [ node ].m_state ) {
                    state.print ( );
                    m_tree [ node ].m_state.print ( );
                    ++cnt;
                    if ( cnt == 100 ) exit ( 0 );
                }

                */

                // If we are not already at the final state, expand the tree with a new
                // node and move there.

                // In addition to expanding one node per simulated game, we also expand all the
                // children of a node when a node's visit count equals T

                if ( hasUntriedMoves ( node ) ) {
                    // if ( player == Player::Type::agent and m_tree [ node ].m_visits < threshold )
                    state.move_hash_winner ( getUntriedMove ( node ) ); // State update.
                    m_path.push ( addChild ( node, state ) );
                }

                // The player in back of path is player ( the player to move ).We now play
                // randomly until the game ends.

                // if ( player == Player::Type::human ) {
                    // state.simulate ( );
                    // for ( Link & link : m_path ) {
                        // We have now reached a final state. Backpropagate the result up the
                        // tree to the root node.
                        // updateData ( link, state );
                    // }
                // }

                // else {

                for ( index_t i = 0; i < 3; ++i ) {
                    State sim_state ( state );
                    sim_state.simulate ( );
                    // We have now reached a final state. Backpropagate the result up the
                    // tree to the root node.
                    for ( Link & link : m_path ) {
                        updateData ( link, sim_state );
                    }
                }
                // }
                m_path.resize ( m_path_size );
            }
            return getBestMove ( );
        }

        private:

        void prune_impl ( Mcts * new_mcts_, const State & state_ ) noexcept {
            using Visited = typename Tree::Visited; // New m_nodes by old_index.
            using Stack = typename Tree::Stack;
            // Prune Tree.
            const NodeID old_node = getNode ( state_.zobrist ( ) );
            Tree & new_tree = new_mcts_->m_tree;
            new_tree [ new_tree.root_node ] = std::move ( m_tree [ old_node ] );
            // The Visited-vector stores the new NodeID's indexed by old NodeID's,
            // old NodeID's not present in the new tree have a value of NodeID::invalid.
            static Visited visited;
            visited.clear ( );
            visited.resize ( m_tree.nodesSize ( ), Tree::NodeID::invalid );
            visited [ old_node.value ] = new_tree.root_node;
            static Stack stack;
            stack.clear ( );
            stack.push_back ( old_node );
            while ( stack.size ( ) ) {
                const NodeID parent = stack.back ( ); stack.pop_back ( );
                for ( OutIt a { m_tree, parent }; a.is_valid ( ); ++a ) {
                    const NodeID child { a->target };
                    if ( Tree::NodeID::invalid == visited [ child.value ] ) { // Not visited yet.
                        visited [ child.value ] = new_tree.addNode ( std::move ( m_tree [ child ] ) );
                        stack.push_back ( child );
                    }
                    new_tree.addArc ( visited [ parent.value ], visited [ child.value ], std::move ( m_tree [ a.id ( ) ] ) );
                }
            }
            // Purge TransitionTable.
            auto it = m_transposition_table->begin ( );
            const auto it_cend = m_transposition_table->cend ( );
            while ( it != it_cend ) {
                const auto tmp_it = it;
                ++it;
                const NodeID new_node = visited [ tmp_it->second.value ];
                if ( Tree::NodeID::invalid == new_node ) {
                    m_transposition_table->erase ( tmp_it );
                }
                else {
                    tmp_it->second = new_node;
                }
            }
            // Transfer TranspositionTable.
            new_mcts_->m_transposition_table.reset ( m_transposition_table.take ( ) );
            // Has been initialized.
            new_mcts_->m_not_initialized = false;
            // Reset path.
            new_mcts_->m_path.reset ( new_tree.root_arc, new_tree.root_node );
            new_mcts_->m_path_size = 1;
        }

        public:

        static void prune ( Mcts * & mcts_, const State & state_ ) noexcept {
            Mcts * pruned_mcts = new Mcts ( );
            if ( not ( mcts_->m_not_initialized ) and Mcts::Tree::NodeID::invalid != mcts_->getNode ( state_.zobrist ( ) ) ) {
                // The state exists in the tree and it's not the current
                // root_node, i.e. now prune.
                mcts_->prune_impl ( pruned_mcts, state_ );
            }
            else {
                mcts_->initialize ( state_ );
            }
            std::swap ( mcts_, pruned_mcts );
            delete pruned_mcts;
        }

        static void reset ( Mcts * & mcts_, const State & state_ ) noexcept {
            if ( not ( mcts_->m_not_initialized ) ) {
                const Mcts::NodeID new_root_node = mcts_->getNode ( state_.zobrist ( ) );
                if ( Mcts::Tree::NodeID::invalid != new_root_node ) {
                    // The state exists in the tree and it's not the current
                    // root_node, i.e. re-hang the tree.
                    mcts_->m_tree.root_node = new_root_node;
                }
                else {
                    std::cout << "new tree\n";
                    Mcts * new_mcts = new Mcts ( );
                    new_mcts->initialize ( state_ );
                    std::swap ( mcts_, new_mcts );
                    delete new_mcts;
                }
            }
        }

        [[ nodiscard ]] InverseTranspositionTable invertTranspositionTable ( ) const noexcept {
            InverseTranspositionTable itt ( m_transposition_table->size ( ) );
            for ( auto & e : * m_transposition_table ) {
                itt [ e.second.value ] = e.first;
            }
            return itt;
        }

        static void merge ( Mcts * & t_mcts_, Mcts * & s_mcts_ ) {
            // Same pointer, do nothing.
            if ( t_mcts_ == s_mcts_ ) {
                return;
            }
            // t_mcts_ (target) becomes the largest tree, we're merging the smaller tree (source) into the larger tree.
            if ( t_mcts_->m_tree.nodeNum ( ) < s_mcts_->m_tree.nodeNum ( ) ) {
                std::swap ( t_mcts_, s_mcts_ );
            }
            // Avoid some levels of indirection and make things clearer.
            Tree & t_t = t_mcts_->m_tree, & s_t = s_mcts_->m_tree; // target tree, source tree.
            TranspositionTable & t_tt = * t_mcts_->m_transposition_table; // target transposition table.
            InverseTranspositionTable s_itt { s_mcts_->invertTranspositionTable ( ) }; // source inverse transposition table.
            // bfs help structures.
            using Visited = boost::dynamic_bitset<>;
            using Queue = Queue<NodeID>;
            Visited s_visited { s_t.nodeNum ( ) };
            Queue s_queue { { s_t.root_node } };
            s_visited [ s_t.root_node.value ] = true;
            // Walk the tree, breadth first.
            while ( s_queue.not_empty ( ) ) {
                // The t_source (target parent) does always exist, as we are going at it breadth first.
                const NodeID s_source = s_queue.pop ( ), t_source = t_tt.find ( s_itt [ s_source.value ] )->second;
                // Iterate over children (targets) of the parent (source).
                for ( OutIt soi { s_t, s_source }; soi.is_valid ( ); ++soi ) { // Source Out Iterator (soi).
                    const Link s_link = s_t.link ( soi );
                    if ( not ( s_visited [ s_link.target.value ] ) ) {
                        s_visited [ s_link.target.value ] = true;
                        s_queue.push ( s_link.target );
                        // Now do something. If child in s_mcts_ doesn't exist in t_mcts_, add child.
                        const auto t_it = t_tt.find ( s_itt [ s_link.target.value ] );
                        if ( t_it != t_tt.cend ( ) ) { // Child exists. The arc does or does not exist.
                            // NodeID t_it->second corresponds to NodeID target child.
                            const Link t_link ( t_t.link ( t_source, t_it->second ) );
                            if ( Tree::ArcID::invalid != t_link.arc ) { // The arc does exist.
                                t_t [ t_link.arc.value ] += s_t [ s_link.arc.value ];
                            }
                            else { // The arc does not exist.
                                t_t [ t_t.addArc ( t_source, t_link.target ).arc.value ] = std::move ( s_t [ s_link.arc.value ] );
                            }
                            // Update the values of the target.
                            t_t [ t_link.target.value ] += s_t [ s_link.target.value ];
                        }
                        else { // Child does not exist.
                            const Link t_link = t_t.addNode ( t_source );
                            // m_tree.
                            t_t [ t_link.arc.value    ] = std::move ( s_t [ s_link.arc.value    ] );
                            t_t [ t_link.target.value ] = std::move ( s_t [ s_link.target.value ] );
                            // m_transposition_table.
                            t_tt.emplace ( s_itt [ s_link.target.value ], t_link.target );
                        }
                    }
                }
            }
            t_mcts_->m_path.resize ( 1 );
            t_mcts_->m_path_size = 1;
            delete s_mcts_; // Destruct the smaller tree.
            s_mcts_ = nullptr;
        }

        std::size_t numTranspositions ( ) const noexcept {
            std::size_t nt = 0;
            using Visited = boost::dynamic_bitset<>;
            using Stack = Stack<NodeID>;
            Visited visited { m_tree.nodeNum ( ) };
            Stack stack { m_tree.root_node };
            visited [ m_tree.root_node.value ] = true;
            while ( stack.not_empty ( ) ) {
                const NodeID parent = stack.pop ( );
                for ( OutIt a { m_tree, parent }; a.is_valid ( ); ++a ) {
                    const NodeID child = a->target;
                    if ( not ( visited [ child.value ] ) ) {
                        visited [ child.value ] = true;
                        stack.push ( child );
                        if ( m_tree.inArcNum ( child ) > 1 ) {
                            ++nt;
                        }
                    }
                }
            }
            return nt;
        }

    private:

        friend class cereal::access;

        template < class Archive >
        void save ( Archive & ar_ ) const noexcept {
            ar_ ( m_tree, * m_transposition_table, m_not_initialized );
        }

        template < class Archive >
        void load ( Archive & ar_ ) noexcept {
            m_tree.clearUnsafe ( );
            if ( m_transposition_table.get ( ) == nullptr ) {
                m_transposition_table.reset ( new TranspositionTable ( ) );
            }
            else {
                m_transposition_table->clear ( );
            }
            ar_ ( m_tree, * m_transposition_table, m_not_initialized );
            m_path.reset ( m_tree.root_arc, m_tree.root_node );
            m_path_size = 1;
        }
    };

    template<typename State>
    void compute ( State & state_, index_t max_iterations_ ) noexcept {
        Mcts<State> * mcts = new Mcts<State> ( );
        state_.move_hash_winner ( mcts->compute ( state_, max_iterations_ ) );
        delete mcts;
    }
}


/*
Deterministic Bayesian Sparse Linear Mixed Model (DBSLMM)
Copyright (C) 2019  Sheng Yang and Xiang Zhou

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <algorithm>
#include <vector>
#include <string>
#include <map>
#include <iostream>
#include <sstream>
#include <fstream>
#include <bitset>
#include <numeric>
#include <time.h>
#include <math.h>
#include <sys/time.h>
#include <cctype>
#include <boost/math/distributions/students_t.hpp>

#include <armadillo>

#include "../include/dtpr.hpp"

using namespace std;
using namespace arma;
// input block information
int IO::readBlock(string infile, char *separator, vector <BLOCK> &block){

	int block_num = getRow(infile);
	block.resize(block_num);
	string oneblock, element;
	ifstream file_stream(infile.c_str());
	int block_count = 0;
	while (getline(file_stream, oneblock)) {
		vector<string> block_str;
		stringstream oneblock_stream(oneblock);
		while (getline(oneblock_stream, element, *separator))
			block_str.push_back(element);
		block[block_count].chr = block_str[0];
		block[block_count].start = atol(block_str[1].c_str());
		block[block_count].end = atol(block_str[2].c_str());
		block_count++;
	}
	file_stream.close();
	return 0;
}

// get row number
int IO::getRow(string infile){

	int n_row = 0;
	string row;
	ifstream file_stream(infile.c_str());
	while (getline(file_stream, row))
		n_row++;
	file_stream.close();
	return n_row;
}

// input bim data
int IO::readBim(int n_ref, string ref_str, char *separator, map<string, ALLELE> &bim, bool constr){

	string bed_str = ref_str + ".bed", bim_str = ref_str + ".bim";
	ifstream bim_stream(bim_str.c_str());
	int n_snp = getRow(bim_str);
	vec maf = zeros<vec>(n_snp);
	if (constr == true){
		ifstream bed_stream(bed_str);
		cout << "Calculating MAF of reference panel ..." << endl;
		vector<int> idv(n_ref);
		for (int i = 0; i < n_ref; i++) idv[i] = 1;
		for (int i = 0; i < n_snp; i++){
			vec geno = zeros<vec>(n_ref);
			readSNPIm(i, n_ref, idv, bed_stream, geno, maf(i));
		}
		bed_stream.close();
	} else {
		cout << "[WARNING] Do not consider the difference between reference panel and summary data ..." << endl;
	}

	string snp, elem;
	int count = 0;
	while (getline(bim_stream, snp)) {
		vector<string> snp_vec;
		ALLELE allele;
		stringstream snp_stream(snp);
		while (getline(snp_stream, elem, *separator)) snp_vec.push_back(elem);
		allele.pos = count;
		allele.a1 = snp_vec[4];
		allele.a2 = snp_vec[5];
		allele.maf = maf(count);
		bim.insert(pair<string, ALLELE>(snp_vec[1], allele));
		count++;
	}
	bim_stream.close();
	return 0;
}

// input bim data
int IO::readBim(int n_ref, string ref_str, char *separator, map<string, ALLELEB> &bim, bool constr){

	string bed_str = ref_str + ".bed", bim_str = ref_str + ".bim";
	ifstream bim_stream(bim_str.c_str());
	int n_snp = getRow(bim_str);
	vec maf = zeros<vec>(n_snp);
	if (constr == true){
		ifstream bed_stream(bed_str);
		cout << "Calculating MAF of reference panel ..." << endl;
		vector<int> idv(n_ref);
		for (int i = 0; i < n_ref; i++) idv[i] = 1;
		for (int i = 0; i < n_snp; i++){
			vec geno = zeros<vec>(n_ref);
			readSNPIm(i, n_ref, idv, bed_stream, geno, maf(i));
		}
		bed_stream.close();
	} else {
		cout << "[WARNING] Do not consider the difference between reference panel and external summary data ..." << endl;
	}
	string snp, elem;
	int count = 0;
	while (getline(bim_stream, snp)) {
		vector<string> snp_vec;
		ALLELEB alleleb;
		stringstream snp_stream(snp);
		while (getline(snp_stream, elem, *separator)) snp_vec.push_back(elem);
		alleleb.pos = count;
		alleleb.ps = atoi(snp_vec[3].c_str());
		alleleb.a1 = snp_vec[4];
		alleleb.a2 = snp_vec[5];
		alleleb.maf = maf(count);
		bim.insert(pair<string, ALLELEB>(snp_vec[1], alleleb));
		count++;
	}
	bim_stream.close();
	return 0;
}

// calculate P value from GEMMA output
double IO::calP(double beta, double se, int sampleSize){
	double t = beta / se;
	boost::math::students_t tDis(sampleSize);
	double p = 2 * cdf(complement(tDis, fabs(t)));
	return p;
}

// input summary data
int IO::readSumm(string summ_str, char *separator, vector<SUMM > &summ){

	string snp, element, chr_str;
	ifstream summ_stream(summ_str.c_str());

	int num = getRow(summ_str);
	// read the summary data of specific chromosome
	summ.resize(num);
	int summ_count = 0;
	while (getline(summ_stream, snp)){
		vector<string> snp_summ;
		stringstream summ_stream(snp);
		while (getline(summ_stream, element, *separator))
			snp_summ.push_back(element);
		double P = atof(snp_summ[10].c_str());
		double se, sei;
		if (isdigit(snp_summ[9].c_str()[0])){
			se = atof(snp_summ[9].c_str());
			if (se - 0.0 > 1e-20){
				summ[summ_count].z = atof(snp_summ[8].c_str()) / se;
				summ[summ_count].P = calP(atof(snp_summ[8].c_str()), se, atoi(snp_summ[4].c_str()));
			} else {
				summ[summ_count].z = 0;
				summ[summ_count].P = 1;
			}
		} else {
			summ[summ_count].z = 0;
			summ[summ_count].P = 1;
		}
		summ[summ_count].chr = atoi(snp_summ[0].c_str());
		summ[summ_count].snp = snp_summ[1];
		summ[summ_count].ps = atol(snp_summ[2].c_str());
		summ[summ_count].a1 = snp_summ[5].c_str();
		summ[summ_count].a2 = snp_summ[6].c_str();
		double af = atof(snp_summ[7].c_str());
		summ[summ_count].maf = min(af, 1.0 - af);
		summ_count++;

	}
	summ_stream.close();

	return summ_count;
}

// input DBSLMM result
int IO::readDBSLMM(string summ_str, char *separator, vector<SUMMS> &summ){

	string snp, element, chr_str;
	ifstream summ_stream(summ_str.c_str());

	int num = getRow(summ_str);
	// read the summary data of specific chromosome
	summ.resize(num);
	int summ_count = 0;
	while (getline(summ_stream, snp)){
		vector<string> snp_summ;
		stringstream summ_stream(snp);
		while (getline(summ_stream, element, *separator))
			snp_summ.push_back(element);
		summ[summ_count].snp = snp_summ[0];
		summ[summ_count].a1 = snp_summ[1].c_str();
		summ[summ_count].z = atof(snp_summ[2].c_str());
		summ_count++;
	}
	summ_stream.close();

	return summ_count;
}

// input external result (summary statistics)
int IO::readExt(string summ_str, char *separator, map<string, SUMMS> &summ){

	string snp, element, chr_str;
	ifstream summ_stream(summ_str.c_str());

	// read the summary data of specific chromosome
	int summ_count = 0;
	while (getline(summ_stream, snp)){
		vector<string> snp_summ;
		SUMMS summs;
		stringstream summ_stream(snp);
		while (getline(summ_stream, element, *separator))
			snp_summ.push_back(element);
		summs.snp = snp_summ[0];
		summs.a1 = snp_summ[1];
		summs.maf = atof(snp_summ[2].c_str());
		summs.z = atof(snp_summ[3].c_str());
		summ.insert(pair<string, SUMMS>(snp_summ[0], summs));
	}
	summ_stream.close();

	return summ_count;
}

// input genotype data
// modify from GEMMA, Xiang Zhou et al.
void IO::readSNPIm(const int pos, int ni_test, const vector<int> &indicator_idv, ifstream &infile, vec &geno, double &maf) {

	// debug_msg("entered");
	size_t ni_total = indicator_idv.size(), n_bit;
	if (ni_total % 4 == 0) {
		n_bit = ni_total / 4;
	}
	else {
		n_bit = ni_total / 4 + 1;
	}

	// n_bit, and 3 is the number of magic numbers.
	infile.seekg(pos * n_bit + 3);

	// Read genotypes.
	char ch[1];
	bitset<8> b;

	double geno_mean = 0.0;
	size_t c = 0, c_idv = 0;
	vector<size_t> geno_miss;
	int freq[3];
	freq[0] = freq[1] = freq[2] = 0;
	for (size_t i = 0; i < n_bit; ++i) {
		infile.read(ch, 1);
		b = ch[0];

		// Minor allele homozygous: 2.0; major: 0.0.
		for (size_t j = 0; j < 4; ++j) {
			if ((i == (n_bit - 1)) && c == ni_total) {
				break;
			}
			if (indicator_idv[c] == 0) {
				c++;
				continue;
			}
			c++;

			if (b[2 * j] == 0) {
				if (b[2 * j + 1] == 0) {
					geno(c_idv) = 2.0;
					geno_mean += 2.0;
					freq[2]++;
				}
				else {
					geno(c_idv) = 1.0;
					geno_mean += 1.0;
					freq[1]++;
				}
			}
			else {
				if (b[2 * j + 1] == 1) {
					geno(c_idv) = 0.0;
					geno_mean += 0.0;
					freq[0]++;
				}
				else {
					geno_miss.push_back(c_idv);
				}
			}

			c_idv++;
		}
	}
	// max imputation
	// int imp_val = distance(freq, max_element(freq, freq + 3));
	// mean imputation
	geno_mean /= (double)(c_idv - geno_miss.size());
	for (size_t i = 0; i < geno_miss.size(); ++i)
		geno(geno_miss[i]) = geno_mean;
	double af = 0.5 * sum(geno) / geno.n_elem;
	maf = min(af, 1.0 - af);
	return ;
}

double IO::getWalltime(){
    struct timeval time;
    if (gettimeofday(&time,NULL)){
        //  Handle error
        return 0;
    }
    return (double)time.tv_sec + (double)time.tv_usec * .000001;
}

void SNPPROC::nomalizeVec(vec &x) {

	x -= mean(x);
	x /= stddev(x);
	return;
}

// match summary result and reference panel
int SNPPROC::matchRef(vector<SUMM> summ, map<string, ALLELE> bim, vector<POS> &inter_snp, double mafMax, bool *badsnp_bool) {

	int dis_count = 0, maf_count = 0;
	for (size_t i = 0; i < summ.size(); ++i) {
		bool a1_bool = bim[summ[i].snp].a1 == summ[i].a1;
		bool a2_bool = bim[summ[i].snp].a2 == summ[i].a2;
		bool maf_bool = fabs(bim[summ[i].snp].maf - summ[i].maf) < mafMax;
		if (a1_bool == false || a2_bool == false) dis_count++;
		if (maf_bool == false) maf_count++;
		if (a1_bool == true && a2_bool == true && maf_bool == true && bim.find(summ[i].snp) != bim.end()){
			POS s_summ;
			s_summ.snp = summ[i].snp;
			s_summ.ps = summ[i].ps;
			s_summ.pos = bim[summ[i].snp].pos;
			s_summ.a1 = summ[i].a1;
			s_summ.maf = summ[i].maf;
			s_summ.z = summ[i].z;
			s_summ.P = summ[i].P;
			inter_snp.push_back(s_summ);
			badsnp_bool[i] = true;
		}
	}
	cout << "Number of allele discrepency: " << dis_count << endl;
	cout << "Number of maf discrepency:    " << maf_count << endl;
	return 0;
}

// match external summary statistics to DBSLMM result
int SNPPROC::matchSumm(vector<SUMMS> summ, map<string, SUMMS> summE, vector<SUMMC> &summC) {

	int dis_count = 0;
	for (size_t i = 0; i < summ.size(); ++i) {
		if (summE.find(summ[i].snp) != summE.end()){
			SUMMC summc;
			summc.snp = summE[summ[i].snp].snp;
			summc.a1 = summ[i].a1;
			summc.maf = summE[summ[i].snp].maf;
			summc.z1 = summ[i].z;
			bool a1_bool = summE[summ[i].snp].a1 == summ[i].a1;
			if (a1_bool == false) dis_count++;
			if (a1_bool == true){
				summc.z2 = summE[summ[i].snp].z;
			} else {
				summc.z2 = -summE[summ[i].snp].z;
			}
			summC.push_back(summc);
		}
	}
	cout << "Number of allele discrepency: " << dis_count << endl;
	return 0;
}

// match DBSLMM result, external summary statistics and bim file
int SNPPROC::matchAll(vector <SUMMC> summI, map<string, ALLELEB> bim, double mafMax, vector <SUMMP> &summO) {

	for (size_t i = 0; i < summI.size(); ++i) {
		bool a1_bool = bim[summI[i].snp].a1 == summI[i].a1;
		bool maf_bool = fabs(bim[summI[i].snp].maf - summI[i].maf) < mafMax;
		if (a1_bool == true && maf_bool == true && bim.find(summI[i].snp) != bim.end()){
			SUMMP summp;
			summp.ps = bim[summI[i].snp].ps;
			summp.pos = bim[summI[i].snp].pos;
			summp.snp = summI[i].snp;
			summp.z1 = summI[i].z1;
			summp.z2 = summI[i].z2;
			summO.push_back(summp);
		}
	}
	sort(summO.begin(), summO.end(), sortPS);
	return 0;
}

int SNPPROC::addBlock(vector<POS> summ, vector <BLOCK> block, vector<INFO> &pos_block){

	int num_block = block.size();
	int count = 0;
	pos_block.resize(summ.size());
	for (int i = 0; i < num_block; i++){
		int start = block[i].start;
		int end = block[i].end;
		// cout << i << " " << begin << " " << end << endl;
		for (size_t j = count; j < summ.size(); j++){
			if (summ[j].ps >= start && summ[j].ps < end) {
				pos_block[j].snp = summ[j].snp;
				pos_block[j].block = i;
				pos_block[j].ps = summ[j].ps; // snp position
				pos_block[j].pos = summ[j].pos; // bim file position
				pos_block[j].a1 = summ[j].a1;
				pos_block[j].maf = summ[j].maf;
				pos_block[j].P = summ[j].P;
				pos_block[j].z = summ[j].z;
				count++;
			}else{
				break;
			}
		}
	}
	return num_block;
}

bool sortP(const INFO &snpinfo1, const INFO &snpinfo2){
	return snpinfo1.P < snpinfo2.P;
}

bool sortPOS(const INFO &snpinfo1, const INFO &snpinfo2){
	return snpinfo1.pos < snpinfo2.pos;
}

bool sortPS(const SUMMP &summp1, const SUMMP &summp2)
{
	return summp1.ps < summp2.ps;
}

void printProgBar(int percent) {
	string bar;
	for (int i = 0; i < 50; i++) {
		if (i < (percent / 2)) {
			bar.replace(i, 1, "=");
		}
		else if (i == (percent / 2)) {
			bar.replace(i, 1, ">");
		}
		else {
			bar.replace(i, 1, " ");
		}
	}

	cout << "\r" "[" << bar << "] ";
	cout.width(3);
	cout << percent << "%     " << std::flush;
}


#include <iostream>
#include <cstdio>
#include <fstream>
#include <conio.h>
#include <boost/program_options.hpp>
#include "input_data.h"

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

struct ParamIsNotSet {
    std::string param;
    ParamIsNotSet (std::string param)
        : param(param)
    {}
};

struct ParamIsNotPositive {
    std::string param;
    ParamIsNotPositive (std::string param)
        : param(param)
    {}
};

void print_error (const string msg)
{
    cerr << msg << "\n";
    getch();
    exit(1);
}

void get_double_param (const po::variables_map& vm, const char* param,
    double& val)
{
    if (vm.count(param))
        val = vm[param].as<double>();
    else
        throw ParamIsNotSet(param);
}

void get_pos_double_param (po::variables_map& vm, const char* param,
    double& val)
{
    if (vm.count(param)) {
        val = vm[param].as<double>();
        if (val < 0)
            throw ParamIsNotPositive(param);
    }
    else
        throw ParamIsNotSet(param);
}

void get_int_param (const po::variables_map& vm, const char* param, int& val)
{
    if (vm.count(param)) {
        val = vm[param].as<int>();
        if (val < 0)
            throw ParamIsNotPositive(param);
    }
    else
        throw ParamIsNotSet(param);
}

void get_string_param (const po::variables_map& vm, const char* param,
    string& val)
{
    if (vm.count(param))
        val = vm[param].as<string>();
    else
        throw ParamIsNotSet(param);
}

template <typename T>
std::string toString (T val)
{
    std::ostringstream oss;
    oss << val;
    return oss.str();
}

extern "C" FILE* popen(const char* command, const char* mode);
extern "C" int pclose(FILE* stream);

// call program cmd with arguments len num arg
// len - length of the interval, (num + 1) - number of grid points
// ans - values of the function in the grid points
void calc_fun (string cmd, string arg, double len, int num, vector<double>& ans)
{
    string arg1 = toString(len) + " " + toString(num);
    string full_cmd = cmd + " " + arg1 + " " + arg;
    cout << full_cmd << endl;
    FILE* pf = popen(full_cmd.c_str(), "r");
    ans.resize(num + 1);
    for (int i = 0; i <= num; ++i)
        fscanf(pf, "%lf", &ans[i]);
    pclose(pf);
}

InputData::InputData (string input_file_name)
{
    try {
        po::options_description desc("Input data");
        desc.add_options()
            ("L", po::value<double>(), "Length of the space interval")
            ("T", po::value<double>(), "Length of the time interval")
            ("a", po::value<double>(), "Diffusion coefficient a")
            ("alpha", po::value<double>(), "Diffusion coefficient alpha")
            ("kappa_a", po::value<double>(), "Coefficient kappa_a")
            ("b", po::value<double>(), "Coefficient b")
            ("beta", po::value<double>(), "Coefficient beta (in BC)")
            ("gamma", po::value<double>(), "Coefficient gamma (in BC)")
            ("theta_b1", po::value<double>(), "Boundary temperature theta_b1")
            ("theta_b2", po::value<double>(), "Boundary temperature theta_b2")
            ("mode", po::value<string>(),
                "Mode: whether q(t) or r(t) is given (given_q or given_r)")
            ("q_init_guess", po::value<double>(),
                "Initial guess in the bisection method")
            ("q_init_len", po::value<double>(),
                "Initial length of the interval in the bisection method")
            ("q_tol", po::value<double>(), "Tolerance in the bisection method")
            ("verify_monotonicity", po::value<string>(),
                "Verify monotonicity (on or off)")
            ("monot_ver_q_1", po::value<double>(),
                "Left bound of the range of the monotonicity verification")
            ("monot_ver_q_2", po::value<double>(),
                "Right bound of the range of the monotonicity verification")
            ("monot_ver_q_step", po::value<double>(),
                "Step of the monotonicity verification")
            ("N", po::value<int>(), "Number of space grid points")
            ("M", po::value<int>(), "Number of time grid points")
            ("theta_0_fun_cmd", po::value<string>(), "Command theta_0")
            ("theta_0_fun_arg", po::value<string>(), "Arguments of theta_0")
            ("f_fun_cmd", po::value<string>(), "Command f")
            ("f_fun_arg", po::value<string>(), "Arguments of f")
            ("g_fun_cmd", po::value<string>(), "Command g")
            ("g_fun_arg", po::value<string>(), "Arguments of g")
            ("q_fun_cmd", po::value<string>(), "Command q")
            ("q_fun_arg", po::value<string>(), "Arguments of q")
            ("r_fun_cmd", po::value<string>(), "Command r")
            ("r_fun_arg", po::value<string>(), "Arguments of r")
            ("fd_scheme", po::value<string>(),
                "Finite difference scheme (E - Euler or CN - Crank-Nicolson)")
            ("linear_sys_sol_method", po::value<string>(),
                "Linear system solution method (m - MTL or u - UMFPACK)")
            ("Newton_tol", po::value<double>(), "Tolerance in Newton's method")
            ("linear_sys_tol", po::value<double>(),
                "Tolerance in the linear system solution method")
            ("divided_start_steps", po::value<int>(),
                "Number of start time steps which are to be divided")
            ("start_substeps", po::value<int>(),
                "Number of substeps in the start time steps")
        ;
        po::variables_map vm;
        ifstream ifs(input_file_name.c_str());
        if (!ifs)
            print_error("Can not open input file: " + input_file_name);
        else {
            po::store(parse_config_file(ifs, desc), vm);
            po::notify(vm);
        }
        get_pos_double_param(vm, "L", L);
        get_pos_double_param(vm, "T", T);
        get_pos_double_param(vm, "a", a);
        get_pos_double_param(vm, "alpha", alpha);
        get_pos_double_param(vm, "kappa_a", kappa_a);
        get_pos_double_param(vm, "b", b);
        get_pos_double_param(vm, "beta", beta);
        get_pos_double_param(vm, "gamma", gamma);
        get_pos_double_param(vm, "theta_b1", theta_b1);
        get_pos_double_param(vm, "theta_b2", theta_b2);
        string s;
        get_string_param(vm, "mode", s);
        if (s == "given_q")
            mode = MODE_GIVEN_Q;
        else if (s == "given_r")
            mode = MODE_GIVEN_R;
        else
            print_error("mode doesn't equal either given_q or given_r");
        get_double_param(vm, "q_init_guess", q_init_guess);
        get_pos_double_param(vm, "q_init_len", q_init_len);
        get_pos_double_param(vm, "q_tol", q_tol);
        get_string_param(vm, "verify_monotonicity", s);
        if (s == "on")
            verify_monotonicity = true;
        else if (s == "off")
            verify_monotonicity = false;
        else {
            print_error(
                "verify_monotonicity doesn't equal either on or off");
        }
        if (verify_monotonicity) {
            get_double_param(vm, "monot_ver_q_1", monot_ver_q_1);
            get_double_param(vm, "monot_ver_q_2", monot_ver_q_2);
            get_pos_double_param(vm, "monot_ver_q_step", monot_ver_q_step);
            if (monot_ver_q_1 >= monot_ver_q_2)
                print_error("monot_ver_q_1 >= monot_ver_q_2");
        }
        get_int_param(vm, "N", N);
        get_int_param(vm, "M", M);
        get_pos_double_param(vm, "Newton_tol", Newton_tol);
        get_string_param(vm, "fd_scheme", s);
        if (s == "E")
            fd_scheme = FD_SCHEME_IMPLICIT_EULER;
        else if (s == "CN")
            fd_scheme = FD_SCHEME_CRANK_NICOLSON;
        else
            print_error("fd_scheme doesn't equal either E or CN");
        get_string_param(vm, "linear_sys_sol_method", s);
        if (s == "m")
            linear_sys_sol_method = SOL_METHOD_MTL;
        else if (s == "u")
            linear_sys_sol_method = SOL_METHOD_UMFPACK;
        else {
            print_error(
                "linear_sys_sol_method doesn't equal either m or u");
        }
        if (linear_sys_sol_method == SOL_METHOD_MTL)
            get_pos_double_param(vm, "linear_sys_tol", linear_sys_tol);
        get_int_param(vm, "divided_start_steps", divided_start_steps);
        if (divided_start_steps > 0)
            get_int_param(vm, "start_substeps", start_substeps);

        string theta_0_fun_cmd, theta_0_fun_arg, f_fun_cmd, f_fun_arg,
            g_fun_cmd, g_fun_arg, q_fun_cmd, q_fun_arg, r_fun_cmd, r_fun_arg;
        get_string_param(vm, "theta_0_fun_cmd", theta_0_fun_cmd);
        get_string_param(vm, "theta_0_fun_arg", theta_0_fun_arg);
        calc_fun(theta_0_fun_cmd, theta_0_fun_arg, L, N, theta_0);
        get_string_param(vm, "f_fun_cmd", f_fun_cmd);
        get_string_param(vm, "f_fun_arg", f_fun_arg);
        calc_fun(f_fun_cmd, f_fun_arg, L, N, f);
        get_string_param(vm, "g_fun_cmd", g_fun_cmd);
        get_string_param(vm, "g_fun_arg", g_fun_arg);
        calc_fun(g_fun_cmd, g_fun_arg, L, N, g);
        if (mode == MODE_GIVEN_Q) {
            get_string_param(vm, "q_fun_cmd", q_fun_cmd);
            get_string_param(vm, "q_fun_arg", q_fun_arg);
            calc_fun(q_fun_cmd, q_fun_arg, T, M, q);
        }
        else {  // mode == MODE_GIVEN_R
            get_string_param(vm, "r_fun_cmd", r_fun_cmd);
            get_string_param(vm, "r_fun_arg", r_fun_arg);
            calc_fun(r_fun_cmd, r_fun_arg, T, M, r);
        }
    }
    catch (ParamIsNotSet e) {
        print_error("Parameter \"" + e.param + "\" is not specified");
    }
    catch (ParamIsNotPositive e) {
        print_error("Parameter \"" + e.param + "\" is not positive");
    }
    catch (std::exception& e) {
        print_error(std::string("error: ") + e.what());
    }
    catch (...) {
        print_error("Exception of unknown type!");
    }
    // increase the upper bound a little bit
    // so that the for loop works correctly
    if (monot_ver_q_2 > 0)
        monot_ver_q_2 *= 1 + 1e-5;
    else
        monot_ver_q_2 *= 1 - 1e-5;
}