#include "QTable.h"

#include <fstream>
#include <iostream>


QTable::QTable(){}

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

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;
}

/* 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;
}

double QTable::getQValue(State state, int action){
	QState qState(state, action);
	return Q[qState];
}

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 << "";
}