#include "QTable.h"

#include <fstream>
#include <iostream>


QTable::QTable(){
	// all weights begin with value 0.
	weights = std::vector<double>(FEATURES);
}

void QTable::insert(QState qState, double value){
	Q.emplace(qState, value);
}

/* Updates the QValue of the current QState. */
double QTable::update(State state, int action, State nextState, double reward, bool isTerminalState){
	double qValue;
	QState qState(state, action);
	auto it = Q.find(qState);
	if (it != Q.end()){
		qValue = it->second + ALPHA * (reward + GAMMA*(getMaxValue(nextState, isTerminalState)) - it->second); // Off-policy one-step Q-learning formula
		it->second = qValue;
	}
	else{
		qValue = 0 + ALPHA * (reward + GAMMA*(getMaxValue(nextState, isTerminalState)) - 0); // Off-policy one-step Q-learning formula (first update)
		insert(qState, qValue);	
	}
	return qValue;
}

/* Updates the weights of the current QState. */
void QTable::updateWeights(State state, int action, State nextState, double reward, bool isTerminalState){
	std::vector<double> features = extractFeatures(state, action);
	for (int i = 0; i < FEATURES; i++){

		// TODO: limites std::numeric_limits<double>::
		/*std::ofstream dataFile(".\\bwapi-data\\AI\\QLearningBotData\\debug.log", std::ios::app); // std::ios::out is implicit and it's combined.
		if (dataFile.is_open()){
			dataFile << "Instant reward :" << reward << "\n";
			dataFile << "Max Value :" << getMaxValue(nextState, isTerminalState) << "\n";
			dataFile << "Q Value :" << getQValue(state, action) << "\n";
			dataFile << "Feature :" << features.at(i) << "\n";
			dataFile << "\n\n\n";
			dataFile.close();
		}*/
		double aux = weights.at(i) + (ALPHA * (reward + GAMMA*getMaxValue(nextState, isTerminalState) - getQValue(state, action)) * features.at(i)); // Off-policy one-step approximate Q-learning formula
		if (std::isfinite(aux) && !std::isnan(aux)) weights.at(i) = aux; // Numbers might go out of range
		/*else{ // FOR DEBUG!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
			std::ofstream dataFile(".\\bwapi-data\\AI\\QLearningBotData\\error_inf_or_nan.log", std::ios::app); // std::ios::out is implicit and it's combined.
			if (dataFile.is_open()){
				dataFile << "infinite or nan update\n";
				dataFile.close();
			}
		}*/
		//BWAPI::Broodwar->drawTextScreen(550, 140 + (i*10), "Weight %d -> %f", i, weights.at(i)); // FOR DEBUG
		//BWAPI::Broodwar->drawTextScreen(10, 140 + (i * 10), "Feature %d -> %.2f", i, features.at(i)); // FOR DEBUG
	}
}

/* Returns the expected value of performing the best action in a given state. */
double QTable::getMaxValue(State state, bool terminalState){
	if (terminalState)
		return 0;

	double utility;
	double maxUtility = -std::numeric_limits<double>::infinity(); // Saves the maximum utility, starts to -Infinity
	for (int a = 0; a < ACTIONS; a++){ // For all actions
		utility = getQValue(state, a);
		if (utility > maxUtility)
			maxUtility = utility;
	}
	return maxUtility;
}

/* Returns the dot product of weights and features. */
double QTable::getQValue(State state, int action){
	std::vector<double> features = extractFeatures(state, action);
	double value = 0;
	for (int i = 0; i < FEATURES; i++){
		value += weights.at(i) * features.at(i);
	}
	return value;
}

double QTable::getQValue(QState qState){
	return Q[qState];
}

QState QTable::getFirstQState(){
	return Q.begin()->first;
}

boost::unordered_map<QState, double, StateActionHash> QTable::getQ(){
	return Q;
}

std::vector<double> QTable::getWeights(){
	return weights;
}

void QTable::setWeights(std::vector<double> w){
	weights = w;
}

std::vector<double> QTable::extractFeatures(State state, int action){
	std::vector<double> vector(FEATURES);

	vector.at(0) = feature0(state, action); // hp of our own units in state s.
	vector.at(1) = feature1(state, action); // hp of the enemy units in state s.
	vector.at(2) = feature2(state, action); // dps of our own units in state s.
	vector.at(3) = feature3(state, action); // dps of the enemy units in state s.
	vector.at(4) = feature4(state, action); // number of allies in state s.
	vector.at(5) = feature5(state, action); // number of enemies in state s.
	vector.at(6) = feature6(state, action); // distance to closest enemy in state s.
	vector.at(7) = feature7(state, action); // weapon's cool down frames.
	vector.at(8) = feature8(state, action); // action a takes us most likely closer to an enemy.
	vector.at(9) = feature9(state, action); // action a takes us most likely away from an enemy.
	vector.at(10) = feature10(state, action); // action a will most likely hurt the enemy.
	vector.at(11) = feature11(state, action); // action a will most likely kill the enemy.
	vector.at(12) = feature12(state, action); // action a will most likely cause us damage from the enemy.
	vector.at(13) = feature13(state, action); // action a will most likely cause us death from the enemy.
	vector.at(14) = feature14(state, action); // Whether the unit is under attack or not.
	// TODO: speed feature
	return vector;
}

// Quantitative features
int QTable::feature0(State state, int action){ return state.getHP(); } // hp of our own units in state s.
int QTable::feature1(State state, int action){ return state.getHPEnemies(); } // hp of the enemy units in state s.
int QTable::feature2(State state, int action){ return state.getDPS(); } // dps of our own units in state s.
int QTable::feature3(State state, int action){ return state.getDPSEnemies(); } // dps of the enemy units in state s.
int QTable::feature4(State state, int action){ return state.getAllies(); } // number of allies in state s.
int QTable::feature5(State state, int action){ return state.getEnemies(); } // number of enemies in state s.
int QTable::feature6(State state, int action){ return 1/(state.getDistance()+1); } // distance to closest enemy in state s.
int QTable::feature7(State state, int action){ return state.getCoolDown(); } // weapon's cool down frames.

// Boolean features
int QTable::feature8(State state, int action){ // action a takes us most likely closer to an enemy.
	if (state.getEnemies() == 0) return 0; // If there's no enemies this feature is 0.
	if (action == FLEE) return 0;
	return 1;
}
int QTable::feature9(State state, int action){ // action a takes us most likely away from an enemy.
	if (state.getEnemies() == 0) return 0; // If there's no enemies this feature is 0.
	if (action == FLEE) return 1;
	return 0;
}
int QTable::feature10(State state, int action){ // action a will most likely hurt the enemy.
	if (state.getEnemies() == 0) return 0; // If there's no enemies this feature is 0.
	if (action == ATTACK || action == CARRYON) return 1;
	return 0;
}
int QTable::feature11(State state, int action){ // action a will most likely kill the enemy.
	if (state.getEnemies() == 0) return 0; // If there's no enemies this feature is 0.
	if (state.getDPS() >= state.getHPEnemies() && (action == ATTACK || action == CARRYON) ) return 1;
	return 0;
}
int QTable::feature12(State state, int action){ // action a will most likely cause us damage from the enemy.
	if (state.getEnemies() == 0) return 0; // If there's no enemies this feature is 0.
	if (state.getDPSEnemies() > 0 && action != FLEE ) return 1;
	return 0;
}
int QTable::feature13(State state, int action){ // action a will most likely cause us death from the enemy.
	if (state.getEnemies() == 0) return 0; // If there's no enemies this feature is 0.
	if (state.getDPSEnemies() >= state.getHP() && action != FLEE) return 1;
	return 0;
}
int QTable::feature14(State state, int action){ return state.getAttacked(); } // Whether the unit is under attack or not.



std::ostream& operator<<(std::ostream &strm, const QTable &qTable){
	std::cout << "The QTable's buckets contain:\n";
	for (unsigned i = 0; i < qTable.Q.bucket_count(); ++i) {
		std::cout << "bucket #" << i << " contains:";
		for (auto local_it = qTable.Q.begin(i); local_it != qTable.Q.end(i); ++local_it)
			std::cout << " " << local_it->first << ":" << local_it->second;
		std::cout << std::endl;
	}

	return strm << "";
}