#include "QAgent.h"

using namespace BWAPI;

QAgent::QAgent(BWAPI::Unit u, QTable *qt){
	unit = u;
	id = u->getID();
	killCounter = 0;
	qTable = qt;
	accumulatedReward = 0.0;
	followCounter = 0;
	lastHP = u->getHitPoints();
}

/* Getters */
BWAPI::Unit QAgent::getUnit(){
	return unit;
}

int QAgent::getID(){
	return id;
}

double QAgent::getAccumulatedReward(){
	return accumulatedReward;
}

QTable* QAgent::getQTable(){
return qTable;
}

State QAgent::getCurrentState(){
	return currentState;
}

int QAgent::getCurrentAction(){
	return currentAction;
}

int QAgent::getFollowCounter(){
	return followCounter;
}

/* Sets an initial state and takes a first action. */
void QAgent::init(){
	// get the initial state
	currentState = getState(unit);
	// Choose action
	currentAction = getAction(currentState, unit);
	// Take action
	performAction(currentAction, unit);
}

int QAgent::step(bool isFinalState, bool isWinner){
	// Observe the instant reward and the new state.
	nextState = getState(unit);

	// Observe the new instant reward
	double reward = getInstantReward(isWinner); // Observe the instant reward and the new state.
	accumulatedReward += reward;

	if (currentAction != CARRYON) updateQValue(currentState, currentAction, nextState, reward, isFinalState); // Update the QValue

	currentState = nextState;
	if (!isFinalState && unit->exists()){
		currentAction = getAction(currentState, unit); // Choose action
		performAction(currentAction, unit); // Take action
	}

	return reward;
}

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

State QAgent::getState(Unit u){
	if (u->exists()){
		int hp, hpEnemies, allies, enemies, distance, dps, dpsEnemies, coolDown, attacked, speed;
		// WITH THRESHOLDS
		hp = State::getHPThreshold(u->getHitPoints());
		// Gets all the enemies in the unit's sight range
		Unitset foes = u->getUnitsInRadius(u->getType().sightRange(), Filter::IsEnemy);
		// Gets all the alies in the unit's sight range
		Unitset friends = u->getUnitsInRadius(u->getType().sightRange(), Filter::IsAlly);
		hpEnemies = State::getHPEnemiesThreshold(foes, friends.size() + 1); // plus one because the unit is not included in the set of allies
		allies = State::getAlliesThreshold(friends.size() + 1, foes.size());
		enemies = State::getEnemiesThreshold(friends.size() + 1, foes.size());;
		distance = State::getEnemyDistanceThreshold(u->getDistance(u->getClosestUnit(BWAPI::Filter::IsEnemy, unit->getType().sightRange())),u); // Alternative: distance = State::getEnemyDistanceThreshold(State::getClosestEnemyDistance(foes, u), u);
		dps = State::getDPSThreshold(State::calculateDPS(friends));
		dpsEnemies = State::getDPSThreshold(State::calculateDPS(foes));
		if (u->getGroundWeaponCooldown() == 0) coolDown = 0;
		else coolDown = 1;
		if (u->isUnderAttack()) attacked = 1;
		else attacked = 0;
		speed = State::getSpeedThreshold(u);

		return State(hp, hpEnemies, allies, enemies, distance, dps, dpsEnemies, coolDown, attacked, speed);
	}
	State s;
	return s;
}

int QAgent::getAction(State state, BWAPI::Unit unit){
	if (Utils::flipCoin(EPSILON)){  // Probability EPSILON of picking randomly.
		//Broodwar->sendText("Following policy...");
		return getPolicy(state, unit);
	}
	else{
		//Broodwar->sendText("Acting randomly...");
		return getRandomAction(state, unit);
	}
}

int QAgent::getPolicy(State state, BWAPI::Unit unit){
	if (state.getEnemies() == NONE) return EXPLORE;// there's no enemies, only exploring can be performed.

	if (state.getCoolDown() != 0){
		if (unit->getType() == BWAPI::UnitTypes::Protoss_Dragoon){ // The Dragoon has to wait until isAttackFrame returns false
			if (!unit->isAttackFrame()) return FLEE; // the unit's weapon is not ready to shoot, only flee can be performed
		}
		else return FLEE; // the unit's weapon is not ready to shoot, only flee can be performed
	}

	if ((unit->isAttacking() || getCurrentAction() == ATTACK || getCurrentAction() == CARRYON) && followCounter < MAXFOLLOW && unit->getOrderTarget() != NULL) return CARRYON;

	int bestAction = 0; // default: flee TODO: default defend (there shouldn't be a default) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	double utility;
	double maxUtility; // Saves the maximum utility, starts to -Infinity
	maxUtility = -std::numeric_limits<double>::infinity();
	for (int a = 0; a < ACTIONS; a++){
		utility = (*qTable).getQValue(state, a);
		if (utility > maxUtility){
			maxUtility = utility;
			bestAction = a;
		}
		// Breaks ties
		else if (utility == maxUtility){
			if (Utils::flipCoin(0.5)){
				bestAction = a;
			}
		}
	}
	return bestAction;
}

int QAgent::getRandomAction(State state, BWAPI::Unit unit){
	if (state.getEnemies() == NONE) return EXPLORE;// there's no enemies, only exploring can be performed.

	if (state.getCoolDown() != 0){
		if (unit->getType() == BWAPI::UnitTypes::Protoss_Dragoon){ // The Dragoon has to wait until isAttackFrame returns false
			if (!unit->isAttackFrame()) return FLEE; // the unit's weapon is not ready to shoot, only flee can be performed
		}
		else return FLEE; // the unit's weapon is not ready to shoot, only flee can be performed
	}

	if ((unit->isAttacking() || getCurrentAction() == ATTACK || getCurrentAction() == CARRYON) && followCounter < MAXFOLLOW && unit->getOrderTarget() != NULL) return CARRYON;

	srand(time(NULL)); //initialize the random seed
	return rand() % ACTIONS; // there's enemies, attack or flee.
}

void QAgent::performAction(int action, Unit unit){
	switch (action){
		case ATTACK:
			//Broodwar->drawTextMap(unit->getPosition(), "attacking");
			followCounter = 0;
			attack(unit);
			break;
		case FLEE:
			//Broodwar->drawTextMap(unit->getPosition(), "fleeing");
			followCounter = 0;
			flee(unit);
			break;
		case EXPLORE:
			//Broodwar->drawTextMap(unit->getPosition(), "exploring");
			followCounter = 0;
			explore(unit);
			break;
		case CARRYON:
			//Broodwar->drawTextMap(unit->getPosition(), "still attacking");
			carryOn();  // just increments the followCounter
			break;
	}
}

int QAgent::getInstantReward(bool isWinner){
	if (!unit->exists()) return UNITLOST; 
	if (isWinner) return WIN; // If the player wins.
	int instantReward = ALIVE; // Living penalty.
	if (unit->getKillCount() > killCounter){
		killCounter = unit->getKillCount();
		instantReward += ENEMYDESTROYED; // reward for killing an enemy.
	}
	if (unit->getHitPoints() < lastHP){
		lastHP = unit->getHitPoints();
		instantReward += UNITDAMAGED; // Penalty for receiving damage.
	}
	return instantReward;
}

/* ACTIONS */

void QAgent::attack(BWAPI::Unit unit){
	// Ignore the unit if it no longer exists
	// Make sure to include this block when handling any Unit pointer!
	if (!unit->exists())
		return;

	// Attack sum closest weighted 70% plus weakest weighted 30%
	if (unit->getOrderTarget() == NULL){
		BWAPI::Unitset foes = unit->getUnitsInRadius(unit->getType().sightRange(), Filter::IsEnemy);
		if (foes.size() > 0){
			double value = 0;
			BWAPI::Unit bestTarget;
			for (auto &foe : foes){
				if (0.7 / (unit->getDistance(foe) + 1) + 0.3 / (foe->getHitPoints() + 1) > value){
					bestTarget = foe;
					value = 0.7 / (unit->getDistance(foe) + 1) + 0.3 / (foe->getHitPoints() + 1);
				}
			}
			unit->attack(bestTarget);
		}
	}
}

void QAgent::flee(BWAPI::Unit unit){
	// Ignore the unit if it no longer exists
	// Make sure to include this block when handling any Unit pointer!
	if (!unit->exists())
		return;

	fleeMgr.updateBlocks(unit);

	// Gets all the enemies in the unit's sight range
	BWAPI::Unitset foes = unit->getUnitsInRadius(unit->getType().sightRange(), Filter::IsEnemy);

	if (foes.empty())
		return;
	bool moving = false;
	while (!moving && fleeMgr.getBlocks().size() > 0){
		Block *best = fleeMgr.getBestBlock();
		if (unit->hasPath((*best).getCenter().makeValid())){
			unit->move((*best).getCenter().makeValid());
			moving = true;
		}
		else{
			fleeMgr.erase(best);
		}
	}
}

std::set<Block*> QAgent::getFleeGrid(){
	return fleeMgr.getBlocks();
}

Block* QAgent::getFleeBlock(){
	return fleeMgr.getBestBlock();
}

void QAgent::defend(BWAPI::Unit unit){
	unit->holdPosition();
}

void QAgent::explore(BWAPI::Unit unit){
	expMgr.move(unit);
}

void QAgent::carryOn(){
	followCounter++;
}

bool QAgent::operator<(const QAgent &qAgent) const
{
	return id < qAgent.id;
}