#include "solver.hpp"

void Solver::generate_set(vector<int> carry) {
	if(carry.size() == N) {
		possible.insert(carry);
		return;
	}

	for(int col = 0; col < M; col++) {
		carry.push_back(col);
		generate_set(carry);
		carry.pop_back();
	}
}
Solver::Solver(int _N, int _M) : N(_N), M(_M) {
	generate_set({});
}

vector<int> Solver::guess() {
	// Optimal first pick is always the same (at least for 5x8)
	if(first_pick) {
		first_pick = false;
		vector<int> pick = {0};

		int times = (N-1) / M + 1;
		for(int i = 0; i < times; i++)
			for(int j = 0; j < M && pick.size() < N; j++)
				pick.push_back(j);

		return pick;
	}

	return choose_possible().guess;
}
void Solver::learn(vector<int> guess, Response response) {
	// Eliminating impossible sequences
	set<vector<int>> next_possible;
	for(auto sequence : possible)
		if(validate(sequence, guess) == response)
			next_possible.insert(sequence);

	possible = next_possible;
}

int Solver::get_weight(vector<int> guess) {
	// Bucketing possible sequences by responses
	vector<int> response_count(N*N+N+1, 0);
	for(auto sequence : possible) {
		Response response = validate(sequence, guess);
		response_count[(N+1)*response.somewhere + response.correct]++;
	}

	// Get size of the fullest bucket
	int max = 0;
	for(int count : response_count)
		if(count > max)
			max = count;
	// Possible guesses have higher priority
	return 2*max - possible.count(guess);
}
// Choosing next guess
Weighed_guess Solver::minimax(vector<int> carry) {
	if(carry.size() == N)
		return {get_weight(carry), carry};

	// Pick the best of next picks
	Weighed_guess best = {-1, {}};
	for(int col = 0; col < M; col++) {
		carry.push_back(col);

		Weighed_guess next = minimax(carry);
		if(best.weight == -1 || next.weight < best.weight)
			best = next;

		carry.pop_back();
	}
	return best;
}
Weighed_guess Solver::choose_possible() {
	Weighed_guess best = {-1, {}};
	for(auto sequence : possible) {
		Weighed_guess next = {get_weight(sequence), sequence};
		if(best.weight == -1 || next.weight < best.weight)
			best = next;
	}
	return best;
}