#include "BananaBrain.h"

void ResultStore::init()
{
	init_filenames();
	read_prepared_results();
	read_results();
}

void ResultStore::init_filenames()
{
	sprintf_s(ai_filename_, sizeof(ai_filename_), "bwapi-data\\AI\\Results_%s.txt", determine_name().c_str());
	sprintf_s(read_filename_, sizeof(read_filename_), "bwapi-data\\read\\Results_%s.txt", determine_name().c_str());
	sprintf_s(write_filename_, sizeof(write_filename_), "bwapi-data\\write\\Results_%s.txt", determine_name().c_str());
}

void ResultStore::read_prepared_results()
{
	FILE *f;
	int err = fopen_s(&f, ai_filename_, "r");
	if (err == 0) read_file(f, prepared_results_);
}

void ResultStore::read_results()
{
	FILE *f;
	int err = fopen_s(&f, read_filename_, "r");
	if (err != 0) err = fopen_s(&f, write_filename_, "r");
	if (err == 0) read_file(f, results_);
}

void ResultStore::read_file(FILE *f,std::vector<Result>& results)
{
	while (true) {
		char timestamp[512];
		int start_positions;
		char map[512];
		char strategy[512];
		char opponent_strategy[512];
		double duration;
		int is_win;
		int ret = fscanf_s(f, "%[^,],%d,%[^,],%[^,],%[^,],%lf,%d\n",
						   timestamp, sizeof(timestamp),
						   &start_positions,
						   map, sizeof(map),
						   strategy, sizeof(strategy),
						   opponent_strategy, sizeof(opponent_strategy),
						   &duration, &is_win);
		if (ret == EOF) break;
		if (ret == 7) {
			Result result;
			result.timestamp = timestamp;
			result.start_positions = start_positions;
			result.map = map;
			result.strategy = strategy;
			result.opponent_strategy = opponent_strategy;
			result.duration = duration;
			result.is_win = (is_win != 0);
			results.push_back(std::move(result));
		}
	}
	fclose(f);
}

std::string ResultStore::pick_strategy(std::vector<std::string> strategies)
{
	if (configuration.human_opponent() || configuration.ucb1()) {
		return pick_strategy_ucb1(strategies);
	} else {
		return pick_strategy_greedy(strategies);
	}
}

std::string ResultStore::pick_strategy_greedy(const std::vector<std::string>& strategies)
{
	std::set<std::string> strategies_set(strategies.begin(), strategies.end());
	std::map<std::string,std::vector<bool>> result_map;
	for (auto& result : prepared_results_) {
		if (contains(strategies_set, result.strategy)) {
			result_map[result.strategy].push_back(result.is_win);
		}
	}
	for (auto& result : results_) {
		if (contains(strategies_set, result.strategy)) {
			result_map[result.strategy].push_back(result.is_win);
		}
	}
	
	const double decay_factor = configuration.tournament() ? kDecayFactorTournament : kDecayFactor;
	std::map<std::string,double> estimator_map;
	for (auto& strategy : strategies) {
		const std::vector<bool>& results = result_map[strategy];
		double weight_sum = 0.0;
		double won_weight_sum = 0.0;
		for (size_t i = 0; i < results.size(); i++) {
			size_t index = results.size() - i;
			double weight = 1.0 / (1.0 + (index / decay_factor));
			weight_sum += weight;
			if (results[i]) won_weight_sum += weight;
		}
		estimator_map[strategy] = (kPriorGames * kTargetWinRate + won_weight_sum) / (kPriorGames + weight_sum);
	}
	
	auto it = std::max_element(estimator_map.begin(), estimator_map.end(), [](auto& entry1,auto& entry2){
		return entry1.second < entry2.second;
	});
	double max_estimator = (*it).second;
	
	std::vector<std::string> potential_strategies;
	for (auto& entry : estimator_map) {
		if (entry.second == max_estimator) potential_strategies.push_back(entry.first);
	}
	return pick_at_random(potential_strategies);
}

std::string ResultStore::pick_strategy_ucb1(const std::vector<std::string>& strategies)
{
	struct StrategyRecord
	{
		int win_count = 0;
		int loss_count = 0;
		double score = 0.0;
		
		bool empty() { return win_count == 0 && loss_count == 0; }
		int count() { return win_count + loss_count; }
	};
	
	std::map<std::string,StrategyRecord> map;
	for (auto& strategy : strategies) {
		StrategyRecord& record = map[strategy];
		for (auto& result : prepared_results_) {
			if (result.strategy == strategy) {
				if (result.is_win) record.win_count++; else record.loss_count++;
			}
		}
		for (auto& result : results_) {
			if (result.strategy == strategy) {
				if (result.is_win) record.win_count++; else record.loss_count++;
			}
		}
	}
	
	std::vector<std::string> unplayed_strategies;
	for (auto& strategy : strategies) {
		if (map[strategy].empty()) unplayed_strategies.push_back(strategy);
	}
	if (!unplayed_strategies.empty()) return pick_at_random(unplayed_strategies);
	
	int n = 0;
	for (auto& strategy : strategies) {
		StrategyRecord& result = map[strategy];
		n += result.count();
	}
	double twologn = 2.0 * std::log(n);
	for (auto& strategy : strategies) {
		StrategyRecord& record = map[strategy];
		double wins = record.win_count;
		double attempts = record.count();
		double average_reward = wins / attempts;
		record.score = average_reward + std::sqrt(twologn / attempts);
	}
	
	auto it = std::max_element(map.begin(), map.end(), [](auto& entry1,auto &entry2) {
		return entry1.second.score < entry2.second.score;
	});
	double max_score = (*it).second.score;
	
	std::vector<std::string> potential_strategies;
	for (auto& strategy : strategies) if (map[strategy].score == max_score) potential_strategies.push_back(strategy);
	return pick_at_random(potential_strategies);
}

std::string ResultStore::pick_at_random(const std::vector<std::string>& strategies)
{
	if (strategies.size() > 1) {
		std::uniform_int_distribution<size_t> dist(0, strategies.size() - 1);
		int r = dist(random_generator());
		return strategies[r];
	} else {
		return strategies[0];
	}
}

void ResultStore::apply_result(const std::string& strategy,const std::string& opponent_strategy,bool win)
{
	Result result;
	
	char timestamp[80];
	time_t rawtime;
	time(&rawtime);
	struct tm timeinfo;
	gmtime_s(&timeinfo, &rawtime);
	strftime(timestamp, sizeof(timestamp), "%Y-%m-%d %H:%M:%S", &timeinfo);
	
	result.timestamp = timestamp;
	result.start_positions = determine_start_positions();
	result.map = determine_map_name();
	result.strategy = strategy;
	result.opponent_strategy = opponent_strategy;
	result.duration = determine_duration();
	result.is_win = win;
	
	results_.push_back(std::move(result));
}

void ResultStore::store()
{
	FILE *f;
	int err = fopen_s(&f, write_filename_, "w");
	if (err == 0) {
		for (auto& result : results_) {
			fprintf(f, "%s,%d,%s,%s,%s,%.1lf,%d\n",
					result.timestamp.c_str(), result.start_positions, result.map.c_str(),
					result.strategy.c_str(), result.opponent_strategy.c_str(),
					result.duration, result.is_win ? 1 : 0);
		}
		fclose(f);
	}
}

std::string ResultStore::determine_name()
{
	if (Broodwar->enemy()->getType() != PlayerTypes::Computer) {
		return Broodwar->enemy()->getName();
	} else {
		return kComputerPlayerName;
	}
}

int ResultStore::determine_start_positions()
{
	return base_state.start_base_count();
}

std::string ResultStore::determine_map_name()
{
	return Broodwar->mapFileName();
}

double ResultStore::determine_duration()
{
	return Broodwar->getFrameCount() / ((1000.0 / 42.0) * 60.0);
}
