#include "solver.hpp"


Game::Game(int p_N, int p_M) : N(p_N), M(p_M) {
	possible = vector<vector<bool>>(N, vector<bool>(M, 1));
	empty_colors = vector<bool>(M, 0);
	final = vector<int>(N, -1);

	unsigned int seed = std::chrono::system_clock::now().time_since_epoch().count();
	random_engine = std::default_random_engine(seed);
}

// Getting known information
bool Game::can(int n, int col) {
	return possible[n][col];
}
vector<vector<int>> Game::get_positions_of_colors(vector<int> guess) {
	auto positions_of_colors = vector<vector<int>>(M, vector<int>(0));
	for(int n = 0; n < N; n++)
		if(guess[n] > -1)
			positions_of_colors[guess[n]].push_back(n);
	return positions_of_colors;
}
int Game::final_color(int n) {
	if(final[n] > -1)
		return final[n];
	int final_col, count = 0;
	for(int col = 0; col < M; col++)
		if(possible[n][col]) {
			final_col = col;
			count++;
		}

	if(count == 1) {
		final[n] = final_col;
		return final_col;
	}
	return -1;
}
bool Game::is_empty(int col) {
	if(empty_colors[col])
		return true;

	for(int n = 0; n < N; n++)
		if(possible[n][col])
			return false;
	empty_colors[col] = true;
	return true;
}
vector<vector<int>> Game::list_all_possibilities() {
	auto r = vector<vector<int>>(N, vector<int>(0));
	for(int col = 0; col < M; col++)
		for(int n = 0; n < N; n++)
			if(possible[n][col])
				r[n].push_back(col);

	for(int n = 0; n < N; n++)
		std::shuffle(r[n].begin(), r[n].end(), random_engine);
	return r;
}
void Game::print() {
	cout << " ";
	for(int col = 0; col < M; col++)
		cout << col;
	cout << std::endl;
	for(int i = 0; i < N; i++) {
		cout << i;
		for(auto col : possible[i])
			cout << col;
		cout << std::endl;;
	}
	cout << std::endl;
}

// Learning functions
void Game::cannot_be(int n, int col) {
	possible[n][col] = false;
}
void Game::must_be(int n, int must_col) {
	for(int col = 0; col < M; col++)
		if(col != must_col)
			possible[n][col] = 0;
}
void Game::empty_color(int col) {
	for(int n = 0; n < N; n++)
		possible[n][col] = 0;
}

// Specific reactions
bool Game::if_not_here_then_nowhere(vector<int> guess) {
	auto positions_of_colors = get_positions_of_colors(guess);
	bool learned_something = false;

	// If color isn't here, it can't be in the sequence
	for(int col = 0; col < M; col++) {
		int possible_count = 0;
		for(int n : positions_of_colors[col])
			if(possible[n][col])
				possible_count++;
		if(possible_count == 0 && positions_of_colors[col].size() > 0 && !is_empty(col)) {
			empty_color(col);
			learned_something = true;
		}
	}

	return learned_something;
}
void Game::here(vector<int> guess) {
	for(int n = 0; n < N; n++)
		if(guess[n] > -1 && possible[n][guess[n]])
			must_be(n, guess[n]);
}
void Game::not_here(vector<int> guess) {
	for(int n = 0; n < N; n++)
		if(guess[n] > -1)
			cannot_be(n, guess[n]);
}
void Game::empty(vector<int> guess) {
	for(int col : guess)
		if(col > -1)
			empty_color(col);
}
void Game::all_are_here(vector<int> guess) {
	auto positions_of_colors = get_positions_of_colors(guess);

	for(int col = 0; col < M; col++) {
		if(positions_of_colors[col].size() == 0)
			empty_color(col);
	}
}


// For remembering guesses with their responses
Historic_guess::Historic_guess(vector<int> _guess, Response _response) {
	guess = _guess;
	response = _response;
}


// Solver
Solver::Solver(int p_N, int p_M) : N(p_N), M(p_M), known({p_N, p_M}) {}

// Check, if it could have been this
bool Solver::all_are_consistent(vector<int> supposed_sequence) {
	for(auto hist : history) {
		auto response = validate(supposed_sequence, hist.guess);
		if(response.somewhere != hist.response.somewhere || response.correct != hist.response.correct)
			return false;
	}
	return true;
}

// Try all remaining sequences
vector<int> Solver::brute_force(vector<vector<int>> *possibilities, vector<int> *chosen, int index) {
	vector<int> r = vector<int>(N, -1);
	if(index == N) {
		if(all_are_consistent(*chosen))
			r = *chosen;
		return r;
	}
	for(int col : (*possibilities)[index]) {
		chosen->push_back(col);
		r = brute_force(possibilities, chosen, index+1);
		if(r[0] != -1)
			return r;
		chosen->pop_back();
	}
	return r;
}

int Solver::get_weight(vector<int> guess) {
	if(!all_are_consistent(guess))
		return -1;

	// Get weight
	for(auto hist : history) {
		// TODO get worst-case weight
	}

	return 1;
}
// Now featuring: minimax pick
Weighed_guess Solver::minimax(vector<vector<int>> *possibilities, vector<int> *chosen, int index) {
	// If complete guess, get weight and return
	if(index == N)
		return {get_weight(*chosen), *chosen};
	// Get max-weighted child
	Weighed_guess r = {-2, {}};
	for(int col : (*possibilities)[index]) {
		chosen->push_back(col);
		Weighed_guess r2 = minimax(possibilities, chosen, index+1);
		if(r2.weight > r.weight || r.weight == -2)
			r = r2;
		chosen->pop_back();
	}
	return r;
}

// Guessing
vector<int> Solver::guess() {
	// TODO make it smart
	auto possibilities = known.list_all_possibilities();
	auto chosen = vector<int>(0);
	return brute_force(&possibilities, &chosen, 0);
}

// Prints what the solver deduced
void Solver::print() {
	known.print();
}

// Clean guess and response from info we know
Historic_guess Solver::clean(Historic_guess hist) {
	vector<bool> already_used_in_cleaning(N, 0);

	// Clean empty colors
	for(int n = 0; n < N; n++)
		if(known.is_empty(hist.guess[n]))
			hist.guess[n] = -1;

	// The in-place colors we know
	for(int n = 0; n < N; n++) {
		if(hist.guess[n] <= -1)
			continue;
		if(known.final_color(n) == hist.guess[n]) {
			hist.guess[n] = -1;
			hist.response.correct -= 1;
			already_used_in_cleaning[n] = 1;
		}
	}

	// The out-of-place colors we know
	for(int n = 0; n < N; n++) {
		if(hist.guess[n] <= -1 || known.can(n, hist.guess[n]))
			continue;
		for(int i = 0; i < N; i++) {
			if(i == n || hist.guess[i] <= -1 || already_used_in_cleaning[i])
				continue;
			if(known.final_color(i) == hist.guess[n]) {
				hist.guess[n] = -1;
				hist.response.somewhere -= 1;
				already_used_in_cleaning[i] = 1;
				break;
			}
		}
	}

	return hist;
}


// Here there be learning
vector<bool> Solver::extract_info(Historic_guess hist) {
	bool something_to_learn = true, learned_something = false;

	// A bit of cleaning
	auto cleaned = clean(hist);
	auto guess = cleaned.guess;
	auto response = cleaned.response;

	bool something = false;
	for(int n = 0; n < N; n++)
		if(guess[n] > -1)
			something = true;

	if(!something)
		return {false, false};

	// Get number of colors, that can be on their positions
	int possible_count = 0;
	for(int n = 0; n < N; n++)
		if(guess[n] > -1 && known.can(n, guess[n]))
			possible_count++;

	// None of these colors are there
	if(hist.response.somewhere == 0 && hist.response.correct == 0) {
		known.empty(hist.guess);
		something_to_learn = false;
		learned_something = true;
	}

	// None at the right spot
	else if(response.correct == 0) {
		if(possible_count > 0) {
			known.not_here(guess);
			learned_something = true;
		}
	}

	// At least only on the right spot
	else if(response.somewhere == 0) {
		// Only colors that can be on these positions are left
		if(response.correct == possible_count) {
			known.here(guess);
			something_to_learn = false;
			learned_something = true;
		}
		else if(hist.response.somewhere == 0){
			if(known.if_not_here_then_nowhere(hist.guess)) {
				learned_something = true;
			}
		}
	}

	// Nonzero / nonzero
	else {
		// Only colors that can be on these positions are left
		if(response.correct == possible_count) {
			known.here(guess);
			learned_something = true;
		}
	}

	return {something_to_learn, learned_something};
}
void Solver::learn(vector<int> _guess, Response _response) {
	// All guessed colors are in the sequence
	if(_response.somewhere + _response.correct == N)
		known.all_are_here(_guess);

	if(_guess[0] == -1)
		return;

	// Write to history
	history.push_back({_guess, _response});

	// Repeat multiple times, if new information turned out
	bool learned_something = true;
	while(learned_something) {
		learned_something = false;

		// Learn from previous guesses
		for(int i = 0; i < history.size(); i++) {
			auto info = extract_info(history[i]);
			if(!info[0]) {
				// If there is nothing left to learn from the guess
				history.erase(history.begin()+i);
				i--;
			}
			if(info[1])
				learned_something = true;
		}
	}
}