#include "Macro.h"

std::vector<CBase> Macro::bases;
std::vector<int> Macro::enemyIDs;
std::vector<Location> Macro::locations;
std::vector<CPlayer> Macro::players;
std::vector<QueueEntry> Macro::queue;
std::vector<CUnit> Macro::scouts;
std::vector<Squad> Macro::squads;
int Macro::lastAttack = 0;
int Macro::selfID;
TilePosition Macro::lastAttackLocation;
TilePosition Macro::lastBuildLocation;
bool Macro::macroHatch;
bool Macro::newBase;
bool Macro::stopProduction;
bool Macro::workerScout = false;
namespace { auto & map = BWEM::Map::Instance(); }

void Macro::addBase(TilePosition tp)
{
	bases.push_back(CBase(tp));
}

void Macro::addLocation(TilePosition tp, bool StartLocation)
{
	locations.push_back(Location(tp, StartLocation));
}

void Macro::addPlayer(BWAPI::Player player)
{
	stopProduction = false; // needed to initialize it somewhere
	players.push_back(CPlayer(player));
}

void Macro::addSquad()
{
	squads.push_back(Squad());
}

// does top-level analysis of what to build
//@TODO clean this up
void Macro::analyzeQueue()
{
	std::string entry = "";
	bool status = false;
	double score = 0;
	int index = -1;
	int count = 0;

	for (auto e : queue)
	{
		if ((!e.unit.isBuilding() && e.unit.supplyProvided() == 0) || !e.inProgress)
		{
			if (e.isTech()) e.score = Analysis::analyze(e.tech);
			else if (e.isUnit()) e.score = Analysis::analyze(e.unit);
			else if (e.isUpgrade()) e.score = Analysis::analyze(e.upgrade);
			if (e.score > score)
			{
				score = e.score;
				index = count;
			}
		}
		count++;
	}

	if (index > -1)
	{
		if (queue.at(index).isUnit())
		{
			if (!queue.at(index).unit.isBuilding())
			{
				if (!stopProduction || Broodwar->self()->minerals() > 300 + queue.at(index).unit.mineralPrice())
				{
					queue.at(index).inProgress = getTrainingBase().unit->train(queue.at(index).unit);
					if (queue.at(index).inProgress) queue.erase(queue.begin() + index);
				}
			}
			else if (queue.at(index).unit.isBuilding())
			{
				if (queue.at(index).unit == UnitTypes::Zerg_Lair || queue.at(index).unit == UnitTypes::Zerg_Hive)
				{
					queue.at(index).inProgress = players.at(selfID).bases.front().base->morph(queue.at(index).unit);
				}
				else
				{
					queue.at(index).inProgress = true;
					// find a drone and send them off
					int workerIndex = Util::buildWorkerIndex();

					if (workerIndex > -1)
					{
						players.at(selfID).bases.at(players.at(selfID).units.at(workerIndex).miningBase).subtractWorker("minerals");
						if (newBase && queue.at(index).unit == UnitTypes::Zerg_Hatchery && !stopProduction)
						{
							players.at(selfID).units.at(workerIndex).action = "expand";
							CBase base = Util::getNextExpand(selfID);
							players.at(selfID).units.at(workerIndex).target = base.tilePosition;
							players.at(selfID).units.at(workerIndex).unit->move(Position(players.at(selfID).units.at(workerIndex).target));
							stopProduction = true;
						}
						else
						{
							players.at(selfID).units.at(workerIndex).action = "build";
							TilePosition target = BuildingPlacement::getPosition(players.at(selfID).units.at(workerIndex).targetUnit);
							if (queue.at(index).unit == UnitTypes::Zerg_Sunken_Colony
								|| queue.at(index).unit == UnitTypes::Zerg_Spore_Colony)
							{
								players.at(selfID).units.at(workerIndex).targetUnit = UnitTypes::Zerg_Creep_Colony;
								target = BuildingPlacement::getPosition(queue.at(index).unit);
							}
							else
							{
								players.at(selfID).units.at(workerIndex).targetUnit = queue.at(index).unit;
							}
							players.at(selfID).units.at(workerIndex).target = target;
						}
					}
				}
			}
		}
		else if (queue.at(index).isUpgrade())
		{
			if (!stopProduction || Broodwar->self()->minerals() > 300 + queue.at(index).upgrade.mineralPrice())
			{
				if (Util::findUpgradeBuilding(queue.at(index).upgrade) > -1)
				{
					players.at(selfID).units.at(Util::findUpgradeBuilding(queue.at(index).upgrade)).unit->upgrade(queue.at(index).upgrade);
					queue.erase(queue.begin() + index);
				}
			}
		}
	}
}

// assigns a base to a player
void Macro::assignBase(int playerID, CBase base)
{
	players.at(playerID).bases.push_back(base);
}

// find the base that's closest to the unit and also unsaturated
int Macro::findMiningBase(BWAPI::Unit unit, std::string type)
{
	bool check = false;
	int distance = 10000;
	int index = -1;
	index = Util::getClosestOwnedBase(unit);

	if (players.at(selfID).bases.at(index).isSaturated(type)) index = -1;

	if (index == -1)
	{
		for (size_t i = 0; i < players.at(selfID).bases.size(); i++)
		{
			if (players.at(selfID).bases.at(i).complete)
			{
				int d = unit->getDistance(players.at(selfID).bases.at(i).base);

				if (d < distance &&	!players.at(selfID).bases.at(i).isSaturated(type))
				{
					distance = d;
					index = i;
				}
			}
		}
	}
	return index;
}

// stops a drone from mining minerals or gas
void Macro::freeUnit(int index)
{
	if (players.at(selfID).units.at(index).isMining())
	{
		int base = players.at(selfID).units.at(index).miningBase;
		if (base > -1)
			players.at(selfID).bases.at(base).subtractWorker(players.at(selfID).units.at(index).action);
	}
	players.at(selfID).units.at(index).action = "";
	players.at(selfID).units.at(index).miningBase = -1;
	players.at(selfID).units.at(index).target = TilePosition(-1, -1);
	players.at(selfID).units.at(index).unit->move(players.at(selfID).units.at(index).unit->getPosition());
	//players.at(selfID).units.at(index).unit->stop();
}

// find a base with larva to morph a unit at
CUnit Macro::getTrainingBase()
{
	int id = 0;
	int count = 0;
	for (auto u : players.at(selfID).units)
	{
		if (u.unit->getType() == UnitTypes::Zerg_Hatchery) id = count;
		if (u.unit->getLarva().size() > 0)
		{
			return u;
		}
		count++;
	}

	return players.at(selfID).units.at(id);
}

void Macro::morphCreepColonies()
{
	for (size_t i = 0; i < players.at(selfID).units.size(); i++)
	{
		if (players.at(selfID).units.at(i).unit->getType() == UnitTypes::Zerg_Creep_Colony)
		{
			int base = Util::getClosestOwnedBase(players.at(selfID).units.at(i).unit);
			if(players.at(selfID).bases.at(base).sunkenCount < players.at(selfID).bases.at(base).desiredSunkenCount())
			{
				players.at(selfID).units.at(i).unit->morph(UnitTypes::Zerg_Sunken_Colony);
			}
			else players.at(selfID).units.at(i).unit->morph(UnitTypes::Zerg_Spore_Colony);
		}
	}
}

void Macro::onUnitComplete(Unit unit)
{
	if (int id = unit->getPlayer()->getID() == selfID 
		&& unit->getType() != UnitTypes::Zerg_Sunken_Colony
		&& unit->getType() != UnitTypes::Zerg_Spore_Colony)
	{
		int index = -1;
		int count = 0;
		for (auto e : queue)
		{
			if (e.unit == unit->getType() && e.inProgress)
			{
				index = count;
				break;
			}
			count++;
		}
		if(index > -1) queue.erase(queue.begin()+index);
		if (unit->getType() == UnitTypes::Zerg_Hatchery)
		{
			bool chk = false;

			for (auto b : bases)
			{
				if (Util::isSameTilePosition(unit->getTilePosition(), b.tilePosition))
				{
					chk = true;
					break;
				}
			}

			if (Util::isBase(unit) && players.at(selfID).bases.size() > 0) players.at(selfID).bases.back().complete = true;
		}
		int unitIndex = Util::getUnitIndex(unit);
		if (unitIndex > -1)
		{
			players.at(selfID).units.at(unitIndex).unitType = unit->getType();
		}
	}
}

// assigns a unit or base to the respective player
void Macro::onUnitCreate(Unit unit)
{
	if (unit->getPlayer() != Broodwar->neutral())
	{
		int id = unit->getPlayer()->getID();
		if (unit->getType().isResourceDepot())
		{
			for (auto b : bases)
			{
				if (Util::isSameTilePosition(unit->getTilePosition(), b.tilePosition))
				{
					CBase base = CBase(unit);
					base.complete = true;
					assignBase(id, base);
				}
			}
		}
		players.at(id).units.push_back(CUnit(unit));
	}
}

//@TODO clean this up
void Macro::onUnitDestroy(Unit unit)
{
	if (unit->getPlayer() != Broodwar->neutral())
	{
		int id = unit->getPlayer()->getID();
		int index = Util::getUnitIndex(unit);
		players.at(id).gasLost += unit->getType().gasPrice();
		players.at(id).mineralsLost += unit->getType().mineralPrice();

		if (index > -1)
		{
			if (unit->getPlayer() == Broodwar->self())
			{
				if (players.at(id).units.at(index).action == "build"
					|| players.at(id).units.at(index).action == "expand")
				{
					//TODO newbase isn't being cleared properly when a drone dies
					// on the way to build a hatchery
					if (players.at(id).units.at(index).action == "expand")
					{
						newBase = false;
						stopProduction = false;
					}
					int queueIndex = -1;
					for (size_t i = 0; i < queue.size(); i++)
					{
						if (queue.at(i).unit == players.at(id).units.at(index).targetUnit)
						{
							queueIndex = i;
							break;
						}
					}
					if (queueIndex > -1) queue.erase(queue.begin() + queueIndex);
				}

				else if (players.at(id).units.at(index).action == "minerals")
				{
					int baseIndex = players.at(id).units.at(index).miningBase;
					if (baseIndex > -1)
					{
						players.at(id).bases.at(baseIndex).mineralWorkers--;
					}
				}

				else if (players.at(id).units.at(index).action == "gas")
				{
					int baseIndex = players.at(id).units.at(index).miningBase;
					if (baseIndex > -1)
					{
						players.at(id).bases.at(baseIndex).gasWorkers--;
					}
				}

				int squadindex = -1;
				int unitindex = -1;
				for (size_t j = 0; j < squads.size(); j++)
				{
					for (size_t k = 0; k < squads.at(j).units.size(); k++)
					{
						if (squads.at(j).units.at(k).id == unit->getID())
						{
							squadindex = j;
							unitindex = k;
							break;
						}
					}
				}
				if (unitindex > -1)
				{
					squads.at(squadindex).units.erase(squads.at(squadindex).units.begin() + unitindex);
				}
				int scoutindex = -1;
				for (size_t j = 0; j < scouts.size(); j++)
				{
					if (scouts.at(j).unit->getID() == unit->getID())
					{
						scoutindex = j;
						break;
					}
				}
				if (scoutindex > -1) scouts.erase(scouts.begin() + scoutindex);
				
			}
			players.at(id).units.erase(players.at(id).units.begin() + index);
		}

		if (unit->getType().isResourceDepot())
		{
			index = Util::getBaseIndex(unit);

			if (index > -1) players.at(id).bases.erase(players.at(id).bases.begin() + index);
		}
	}
}

void Macro::onUnitDiscover(Unit unit)
{
	if (unit->getPlayer() != Broodwar->neutral() && unit->getPlayer() != Broodwar->self() && !Util::unitCounted(unit))
	{
		int id = unit->getPlayer()->getID();
		players.at(id).gasSpent += unit->getType().gasPrice();
		players.at(id).mineralsSpent += unit->getType().mineralPrice();
		if (unit->getType().isResourceDepot())
		{
			for (auto b : bases)
			{
				if (unit->getTilePosition().x == b.tilePosition.x && unit->getTilePosition().y == b.tilePosition.y)
				{
					CBase base = CBase(unit);
					base.complete = unit->isCompleted();
					assignBase(id, base);
				}
			}
		}
		CUnit u = CUnit(unit);
		u.lastSeen = u.unit->getTilePosition();
		players.at(id).units.push_back(u);
	}
}

void Macro::onUnitMorph(Unit unit)
{
	if (unit->getPlayer()->getID() == selfID)
	{
		if (unit->getType().isRefinery())
		{
			players.at(selfID).units.push_back(CUnit(unit));
		}
		int index = Util::getQueueIndex(unit->getType());

		if (index > -1)
		{
			queue.erase(queue.begin() + index);
		}

		if (unit->getType() == UnitTypes::Zerg_Hatchery && newBase
			&& Util::isBasePosition(unit->getTilePosition()))
		{
			players.at(selfID).bases.push_back(CBase(unit));
			newBase = false;
			stopProduction = false;
		}
		else if (unit->getType() == UnitTypes::Zerg_Hatchery && macroHatch)
		{
			macroHatch = false;
		}

		int unitIndex = Util::getUnitIndex(unit);
		if (unitIndex > -1)
		{
			players.at(selfID).units.at(unitIndex).unitType = unit->getType();
		}
	}
}

void Macro::onUnitShow(Unit unit)
{
	if (!unit->getPlayer()->isNeutral() && !unit->getPlayer()->getID() == Broodwar->self()->getID())
	{
		int id = unit->getPlayer()->getID();
		for (size_t i = 0; i < players.at(id).units.size(); i++)
		{
			if (players.at(id).units.at(i).unit->getID() == unit->getID())
			{
				players.at(id).units.at(i).lastSeen = unit->getTilePosition();
				players.at(id).units.at(i).unitType = unit->getType();
			}
		}
	}
}

int Macro::queueCount(BWAPI::UnitType unit)
{
	int count = 0;
	for (auto e : queue)
	{
		if (e.unit == unit) count++;
	}
	return count;
}

bool Macro::queueHas(BWAPI::TechType tech)
{
	for (auto e : queue)
	{
		if (e.tech == tech) return true;
	}
	return false;
}

bool Macro::queueHas(BWAPI::UnitType unit)
{
	for (auto e : queue)
	{
		if (e.unit == unit) return true;
	}
	return false;
}

bool Macro::queueHas(BWAPI::UpgradeType upgrade)
{
	for (auto e : queue)
	{
		if (e.upgrade == upgrade) return true;
	}
	return false;
}

void Macro::setEnemyID(int id)
{
	enemyIDs.push_back(id);
}

void Macro::setLastAttack(int frame)
{
	lastAttack = frame;
}

void Macro::setLastBuildLocation(TilePosition tp)
{
	lastBuildLocation = tp;
}

void Macro::setMacroHatch(bool status)
{
	macroHatch = status;
}

void Macro::setNewBase(bool status)
{
	newBase = status;
}

void Macro::setSelfID(int id)
{
	selfID = id;
}

void Macro::setStopProduction(bool status)
{
	stopProduction = status;
}