/*
 * stratified_sampling.h
 *
 *  Created on: Dec 17, 2015
 *      Author: cedric
 */

#ifndef BOOTSTRAP_STRATIFIED_SAMPLING_H
#define BOOTSTRAP_STRATIFIED_SAMPLING_H

#include <random>
#include <vector>

#include "heuristic_wrapper.h"

#include "../state_registry.h"
#include "../successor_generator.h"
#include "../task_proxy.h"


typedef std::pair<State, double> StateWeightPair;
typedef std::vector<std::vector<StateWeightPair> > Frontier;


class StratifiedSampler{
public:
	StratifiedSampler(const std::shared_ptr<AbstractTask> task, const std::shared_ptr<HeuristicWrapper> hw);
	~StratifiedSampler() = default;

	void estimate_max_distance(int dist);

	void set_root(const State &g);
	void add_goal(const State &g);
	int sample();

	std::vector<int> create_type(const State &g);
private:
	const std::shared_ptr<AbstractTask> task;
	TaskProxy task_proxy;
	SuccessorGenerator successor_generator;

	const std::shared_ptr<HeuristicWrapper> hw;

	Frontier A;
	std::vector<Frontier> B;

	int max_distance = 15;


	std::pair<bool, size_t> same_type_in(const std::vector<int> &type, const std::vector<StateWeightPair> &layer);
	void expand_all_children(const StateWeightPair &p, std::vector<StateWeightPair> &next_layer_tmp);
	bool check_overlap_in_layers(int start_a, int start_b, Frontier &tmp_b);


	void expand_frontier(Frontier &F);
	void reset_frontier(Frontier &F, int i);

	int compute_k(int m);
	int advance_only_A(int n, int m);

	std::random_device r_device;
	std::mt19937 generator;

};




#endif /* SRC_SEARCH_BOOTSTRAP_STRATIFIED_SAMPLING_H_ */
