#include "ArmyStructureInferenceManager.h"
#include "Overseer.h"

ArmyStructureInferenceManager::ArmyStructureInferenceManager( Overseer* ov )
{
	overseer = ov;
}

void ArmyStructureInferenceManager::calculateArmies()
{
	if(Broodwar->getFrameCount()%8 == 0) {
	knownArmies.clear();
	assignedUnits.clear();
	// first get all the enemy units
	std::vector<UnitModel*> units = overseer->getUnitManager()->getEnemyMilitaryUnits();
	if(units.empty()) {
		////////Broodwar->sendText("no enemy units");
		return;
	}
	for (std::vector<UnitModel*>::iterator i = units.begin(); i != units.end(); ++i)
	{
		UnitModel* cur = *i;


		std::vector<UnitModel*> neighbours = overseer->getUnitManager()->getEnemyMilitaryUnitsInRange(cur->getPosition(), 128);
		std::set<UnitModel*> currentArmy;
		while(!neighbours.empty()) {
			if(currentArmy.count(*(neighbours.begin())) == 1 || assignedUnits.count(*(neighbours.begin())) == 1 ) {
				neighbours.erase(neighbours.begin());
			} else {
				currentArmy.insert(*(neighbours.begin()));
				assignedUnits.insert(*(neighbours.begin()));
				std::vector<UnitModel*> subNeighbours = overseer->getUnitManager()->getEnemyMilitaryUnitsInRange((*(neighbours.begin()))->getPosition(), 128);
				neighbours.insert(neighbours.end(), subNeighbours.begin(), subNeighbours.end());
			}
		}
		collaboratorMap[cur] = neighbours;
		knownArmies.push_back(currentArmy);
	}
	}
}

void ArmyStructureInferenceManager::drawArmies()
{
	if(knownArmies.empty()) {
		return;
	}

	int col = 0;
	for (vector<std::set<UnitModel*>>::iterator i = knownArmies.begin(); i != knownArmies.end(); ++i)
	{
		std::set<UnitModel*> cur = *i;
		BWAPI::Color drawCol = BWAPI::Colors::White;

			if(col == 0)
				drawCol = BWAPI::Colors::Red;
	
			if(col == 1)
				drawCol = BWAPI::Colors::Green;
		
			if(col == 2)
				drawCol = BWAPI::Colors::Blue;
		
			if(col == 3)
				drawCol = BWAPI::Colors::Yellow;
			
			if(col == 4)
				drawCol = BWAPI::Colors::Orange;
			
			if(col == 5)
				drawCol = BWAPI::Colors::Purple;
		
		for (std::set<UnitModel*>::iterator it = cur.begin(); it != cur.end(); ++it)
		{
			UnitModel* curUnit = *it;
			////Broodwar->drawEllipseMap(curUnit->getPosition().x(), curUnit->getPosition().y(), 32, 32, drawCol, false);
			////Broodwar->drawEllipseMap(curUnit->getPosition().x(), curUnit->getPosition().y(), 26, 26, drawCol, false);
			//Broodwar->drawEllipseMap(curUnit->getPosition().x(), curUnit->getPosition().y(), 20, 20, drawCol, false);
			//Broodwar->drawEllipseMap(curUnit->getPosition().x(), curUnit->getPosition().y(), 14, 14, drawCol, false);
			//Broodwar->drawEllipseMap(curUnit->getPosition().x(), curUnit->getPosition().y(), 8, 8, drawCol, false);
		}
		col++;
	}

}

std::vector<UnitModel*> ArmyStructureInferenceManager::getCollaboratorsOf( UnitModel* u )
{
	std::vector<UnitModel*> collaborators = collaboratorMap[u];
	return collaborators;
}

std::vector<UnitModel*> ArmyStructureInferenceManager::getCollaboratorsOf( Unit* u )
{
	UnitModel* um = overseer->getUnitManager()->getUnitModelFromMap(u);
	return getCollaboratorsOf(um);
}
