#ifndef BOOTSTRAP_H
#define BOOTSTRAP_H

#include "training_set.h"

#include "ann/neural_net.h"

#include "feature_extractor/feature_extractor.h"
#include "heuristic_wrapper.h"

#include "../blind_search_heuristic.h"
#include "../heuristic.h"
#include "../operator_cost.h"
#include "../option_parser_util.h"
#include "../search_engine.h"
#include "../state_registry.h"
#include "../successor_generator.h"
#include "../task_proxy.h"


#include <fstream>
#include <iostream>
#include <random>
#include <vector>

typedef std::pair<State, int> StateValuePair;
typedef std::vector<StateValuePair > StateValuePairVector;

class Bootstrap{
public:
	Bootstrap(const std::shared_ptr<AbstractTask> task,
			const Options opts);
	~Bootstrap() = default;
	int get_heuristic_from_neural_net(const std::vector<int> &features);
	/* This is the neural network, which is trained by the bootstrap instance.
	 */
	NeuralNet *nn;

private:
	const std::shared_ptr<AbstractTask> task;
	TaskProxy task_proxy;
	SuccessorGenerator successor_generator;

	// Instances is a set of states, which bootstrap will try to solve. Solved instances are then used for training.
	std::vector<State> instances;

	// This is a set of states which are on the path to the solution. Use this set to train the neural network.
	TrainingSet learning_set;

	// Use search engine to solve instances in 'instances'. If a solution is found. Add its path to the learning_set.
	SearchEngine *engine = 0;

	// t_max is the time limit for one search instance in seconds. If an instance cannot be solved in the given time, we just ignore it.
	double t_max = 10;

	// This flag is used to decide if we should backup the g_state_registry, or just override the existing.
	bool first_run = true;
	// This vector stores the initial state data. It will be used to restore the initial state after solving instances.
	std::vector<int> set_back_to_initial_state;
	int number_of_random_walks;
	int random_walk_length;
	std::string ann_topology;
	std::string initial_heuristic_name;
	std::string method;
	double eta;
	int epochs;

	std::streambuf* cout_backup = std::cout.rdbuf();
	std::ofstream sup_out;

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


	const std::shared_ptr<HeuristicWrapper> hw;

    void initialize();

    /* Creates a neural network, which is used later to learn
     * the heuristic function.
     */
    void initialize_neural_network();

    /* Fill the set of instances, with instances which bootstrap will
       try to solve to learn a stronger heuristic.
       If search type is none, do not fill the instances, just perform random walk from the goal state
       and put every state with its distance to the goal into the trainingset.
     */
	void create_instances();

	/* Try to solve all instances in the set instances. Add each solved
	   instance to the learning set and remove it from the set of instances.
	 */
	void solve_instances();

	void train_learner();

	/* Initialize the search engine that will be used to solve instances.
	 */
	void set_search_engine(const State &start);


	void perform_random_walk(const State &in_state, int rw_length, StateValuePairVector &rw_set);
	StateValuePairVector perform_random_walk(const State &in_state, int rw_length);


	void suppress_output();
	void restore_output();

	State create_goal();
	State create_start();
	State create_random_state();
	void extract_states_on_path(std::vector<State>::iterator &instance_iterator);

	void sample_state_space(const State &start_state);
	void sample_state_space_random();


	int generate_rw_length();

	int call_biss(State s);
	void run_bootstrap();
	void initialize_training_set();
};

#endif
