#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));

	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) {
	int final_col, count = 0;
	for(int col = 0; col < M; col++)
		if(possible[n][col]) {
			final_col = col;
			count++;
		}

	if(count == 1)
		return final_col;
	return -1;
}
bool Game::is_empty(int col) {
	for(int n = 0; n < N; n++)
		if(possible[n][col])
			return false;
	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
void Game::if_not_here_then_nowhere(vector<int> guess) {
	auto positions_of_colors = get_positions_of_colors(guess);

	// 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)
			empty_color(col);
	}
}
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> p_guess, vector<int> p_response) {
	guess = p_guess;
	response = p_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[0] != hist.response[0] || response[1] != hist.response[1])
			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;
}

// Guessing
vector<int> Solver::guess() {
	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] = -2;

	// 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[1] -= 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[0] -= 1;
				already_used_in_cleaning[i] = 1;
				break;
			}
		}
	}

	return hist;
}


// Here there be learning
bool Solver::extract_info(Historic_guess hist) {
	bool something_to_learn = true;

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

	// 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[0] == 0 && hist.response[1] == 0) {
		known.empty(hist.guess);
		something_to_learn = false;
	}

	// None at the right spot
	else if(response[1] == 0)
		known.not_here(guess);

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

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

	return something_to_learn;
}
void Solver::learn(vector<int> p_guess, vector<int> p_response) {
	// All guessed colors are in the sequence
	if(p_response[0] + p_response[1] == N)
		known.all_are_here(p_guess);

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

	// Write to history
	history.push_back({p_guess, p_response});

	// Repeat multiple times, if new information turned out
	for(auto _ : history)
		// Learn from previous guesses
		for(int i = 0; i < history.size(); i++)
			if(!extract_info(history[i])) {
				// If there is nothing left to learn from the guess
				history.erase(history.begin()+i);
				i--;
			}
}