#include "AgentManager.h"

AgentManager::AgentManager(){}

/* Getters */
QTable AgentManager::getQTable(){
	return qTable;
}

std::set<QAgent*> AgentManager::getQAgents(){
	return qAgents;
}

void AgentManager::init(QTable qt){
	qTable = qt;
	totalReward = 0;
	initAgents();
}

void AgentManager::initAgents(){
	// Iterate through all the units that we own
	for (auto &u : BWAPI::Broodwar->self()->getUnits())
	{
		// Ignore the unit if it no longer exists
		// Make sure to include this block when handling any Unit pointer!
		if (!u->exists())
			continue;

		QAgent *qa = new QAgent(u, &qTable);
		qAgents.emplace(qa);
		qa->init();
	}
}

void AgentManager::step(bool isFinalState, bool isWinner){

	if (isFinalState && !isWinner){ // We lost
		int reward = getInstantRewardOnDefeat();// Observe the instant reward and the new state.
		totalReward += reward;
		// Define a final state
		State s;
		update(s, EXIT, s, reward, isFinalState);
	}
	else{
		for (auto& qAgent : qAgents){
			totalReward += (*qAgent).step(isFinalState, isWinner);
		}
	}
}

void AgentManager::update(State state, int action, State nextState, double reward, bool isTerminalState){
	qTable.update(state, action, nextState, reward, isTerminalState);
}

int AgentManager::getInstantRewardOnDefeat(){
	return LOSE; // If the player loses.
}

int AgentManager::getTotalReward(){
	return totalReward;
}

void AgentManager::terminate(BWAPI::Unit u){
	for (auto& qAgent : qAgents){
		if (u->getID() == (*qAgent).getID()){
			totalReward += (*qAgent).step(false, false);
			qAgents.erase(qAgent);
		}
	}
}