#include "MicroManager.h"
#include "UnitUtil.h"
#include "ProductionManager.h"

using namespace UAlbertaBot;

MicroManager::MicroManager() 
{
}

void MicroManager::setUnits(const BWAPI::Unitset & u) 
{ 
	_units = u; 
	/*for (auto & unit : u)
	{
		_units.insert(unit);
	}*/
}

BWAPI::Position MicroManager::calcCenter() const
{
    if (_units.empty())
    {
        //if (Config::Debug::DrawSquadInfo)
        //{
        //    BWAPI::Broodwar->printf("calcCenter() called on empty squad");
        //}
        return BWAPI::Position(0,0);
    }

	BWAPI::Position accum(0,0);
	for (auto & unit : _units)
	{
		if (unit->getPosition().isValid())
		{
			accum += unit->getPosition();
		}
	}
	return BWAPI::Position(accum.x / _units.size(), accum.y / _units.size());
}

void MicroManager::execute(const SquadOrder & inputOrder)
{
	order = inputOrder;
	drawOrderText();

	// Nothing to do if we have no units
	if (_units.empty())
	{
		return;
	}

	// No order
	if (!(order.getType() == SquadOrderTypes::Attack || 
		order.getType() == SquadOrderTypes::Defend || 
		order.getType() == SquadOrderTypes::Regroup ||
		order.getType() == SquadOrderTypes::Survey ||
		order.getType() == SquadOrderTypes::DestroyNeutral))
	{
		return;
	}

	// The targets that micro managers have available to shoot at
	BWAPI::Unitset nearbyEnemies;

	// Always include enemies in the radius of the order.
	MapGrid::Instance().GetUnits(nearbyEnemies, order.getPosition(), order.getRadius(), false, true);

	// If the order is to destroy neutral units at a given location.
	if (order.getType() == SquadOrderTypes::DestroyNeutral)
	{
		for (BWAPI::Unit unit : BWAPI::Broodwar->getStaticNeutralUnits())
		{
			if (unit &&
				unit->getPlayer() == BWAPI::Broodwar->neutral() &&
				!unit->getInitialType().isInvincible() &&
				!unit->getType().canMove() &&
				!unit->isFlying() &&
				order.getPosition().getDistance(unit->getInitialPosition()) < 4.5 * 32)
			{
				nearbyEnemies.insert(unit);
			}
		}

		// Add enemy units
		for (auto & unit : _units)
		{
			if (!unit->getType().isFlyer())
			{
				MapGrid::Instance().GetUnits(nearbyEnemies, unit->getPosition(), unit->getType().sightRange(), false, true);
			}
		}

		// Allow micromanager to handle enemies
		destroyNeutralTargets(nearbyEnemies);
		//executeMicro(nearbyEnemies);
		return;
	}

	// if the order is to attack
	else if (order.getType() == SquadOrderTypes::Attack)
	{
		// Add enemy units
		auto ourMain = InformationManager::Instance().getMyMainBaseLocation();
		auto natural = InformationManager::Instance().getMyNaturalLocation();
		auto enemyMain = InformationManager::Instance().getEnemyMainBaseLocation();
		bool _goAggressive = ProductionManager::Instance().getAggression();
		for (auto & unit : _units)
		{
			if (!_goAggressive && unit->getDistance(order.getPosition()) > 800)
			{
				MapGrid::Instance().GetUnits(nearbyEnemies, unit->getPosition(), 2 * 32, false, true);
			}
			else
			{
				MapGrid::Instance().GetUnits(nearbyEnemies, unit->getPosition(), 1 * unit->getType().sightRange(), false, true);
			}
		}

		// if this is an attack squad
		BWAPI::Unitset workersRemoved;

		for (auto & enemyUnit : nearbyEnemies)
		{
			// if it's not a worker, or if it is but we don't like it, add it to the targets
			if (!enemyUnit->getType().isWorker() ||
				enemyUnit->isRepairing() || enemyUnit->isConstructing() || unitNearChokepoint(enemyUnit) ||
				enemyUnit->isBraking() || !enemyUnit->isMoving() || enemyUnit->isGatheringMinerals() || enemyUnit->isGatheringGas())
			{
				workersRemoved.insert(enemyUnit);
			}
			// if it is a worker
			else
			{
				for (auto& enemyRegion : InformationManager::Instance().getOccupiedRegions(BWAPI::Broodwar->enemy()))
				{
					// only add it if it's in their region
					if (BWEMmap.GetArea(BWAPI::TilePosition(enemyUnit->getPosition())) == enemyRegion)
					{
						workersRemoved.insert(enemyUnit);
						break;
					}
				}
			}
		}

		// Allow micromanager to handle enemies
		executeMicro(workersRemoved);
		return;
	}

	// if the order is to regroup
	else if (order.getType() == SquadOrderTypes::Regroup)
	{
		// Add enemy units
		for (auto & unit : _units)
		{
			MapGrid::Instance().GetUnits(nearbyEnemies, unit->getPosition(), 2 * 32, false, true);
		}

		BWAPI::Unitset nearbyEnemiesFiltered;
		for (auto & unit : _units)
		{
			for (auto & enemyUnit : nearbyEnemies)
			{
				if (UnitUtil::CanAttack(enemyUnit, unit) && 
					UnitUtil::CanAttack(unit, enemyUnit))
				{
					nearbyEnemiesFiltered.insert(enemyUnit);
				}
			}
		}

		executeMicro(nearbyEnemiesFiltered);
		return;
	}

	// if the order is to defend
	else if (order.getType() == SquadOrderTypes::Defend)
	{
		// Add enemy units
		for (auto & unit : _units)
		{
			if (unit->getDistance(order.getPosition()) > 800)
			{
				MapGrid::Instance().GetUnits(nearbyEnemies, unit->getPosition(), 2 * 32, false, true);
			}
			else
			{
				MapGrid::Instance().GetUnits(nearbyEnemies, unit->getPosition(), 1 * unit->getType().sightRange(), false, true);
			}
		}
	}

	// if the order is to survey
	else if (order.getType() == SquadOrderTypes::Survey)
	{
		// Add enemy units
		for (auto & unit : _units)
		{
			MapGrid::Instance().GetUnits(nearbyEnemies, unit->getPosition(), 1 * unit->getType().sightRange(), false, true);
		}
	}

	executeMicro(nearbyEnemies);
	return;

	/*
	// do survey micro
	if (order.getType() == SquadOrderTypes::Survey)
	{
		executeMicro(nearbyEnemies);
	}

	// the following block of code attacks all units on the way to the order position
	// we want to do this if the order is attack, defend, or harass
	if (order.getType() == SquadOrderTypes::Attack || 
		order.getType() == SquadOrderTypes::Defend ||
		order.getType() == SquadOrderTypes::Regroup)
	{
		// if this is a worker defense force
		if (_units.size() == 1 && (*_units.begin())->getType().isWorker())
		{
			executeMicro(nearbyEnemies);
		}
		// otherwise it is a normal attack force
		else
		{
			// if this is a defense squad then we care about all units in the area
			if (order.getType() == SquadOrderTypes::Defend)
			{
				executeMicro(nearbyEnemies);
			}
			// otherwise we only care about workers if they are in their own region
			// or are building something
			// Idea: Don't goose chase the enemy scout.
			// Unfortunately there are bad side effects.
			else
			{
				// if this is an attack squad
				BWAPI::Unitset workersRemoved;

				for (auto & enemyUnit : nearbyEnemies)
				{
					// if it's not a worker, or if it is but we don't like it, add it to the targets
					if (!enemyUnit->getType().isWorker() ||
						enemyUnit->isRepairing() || enemyUnit->isConstructing() || unitNearChokepoint(enemyUnit) ||
						enemyUnit->isBraking() || !enemyUnit->isMoving() || enemyUnit->isGatheringMinerals() || enemyUnit->isGatheringGas())
					{
						workersRemoved.insert(enemyUnit);
					}
					// if it is a worker
					else
					{
						for (BWTA::Region * enemyRegion : InformationManager::Instance().getOccupiedRegions(BWAPI::Broodwar->enemy()))
						{
							// only add it if it's in their region
							if (BWTA::getRegion(BWAPI::TilePosition(enemyUnit->getPosition())) == enemyRegion)
							{
								workersRemoved.insert(enemyUnit);
								break;
							}
						}
					}
				}

				// Allow micromanager to handle enemies
				executeMicro(workersRemoved);
			}
		}
	}
	*/
	
}

// The order is DestroyNeutral. Carry it out.
void MicroManager::destroyNeutralTargets(const BWAPI::Unitset & targets)
{
	// Is any target in sight? We only need one.
	BWAPI::Unitset visibleTargets;
	for (const auto target : targets)
	{
		if (target->exists() &&
			target->isTargetable() &&
			target->isDetected())			// not e.g. a neutral egg under a neutral arbiter
		{
			visibleTargets.insert(target);
		}
	}

	int count = 0;
	for (const auto unit : _units)
	{
		count++;

		// pick closest visible target
		auto visibleTarget = closestUnit(unit, visibleTargets);

		if (visibleTarget)
		{
			// We see a target, so we can issue attack orders to units that can attack.
			if (UnitUtil::CanAttackGround(unit) && unit->canAttack())
			{
				// There are targets. Attack it.
				//Micro::CatchAndAttackUnit(unit, visibleTarget);
				Micro::SmartAttackUnit(unit, visibleTarget);
			}
		}
		else if (unit->getDistance(order.getPosition()) > 96)
		{
			// There are no targets. Move to the order position if not already close.
			//Micro::SmartMovePath(unit, order.getPosition());
			Micro::SmartMove(unit, order.getPosition());
		}
	}
}

BWAPI::Unit MicroManager::closestUnit(BWAPI::Unit unit, const BWAPI::Unitset & targets)
{
	int minDistance = 0;
	BWAPI::Unit closest = nullptr;

	for (auto & target : targets)
	{
		int distance = -1;
		BWEMmap.GetPath(unit->getPosition(), target->getPosition(), &distance);
		//int distance = target->getDistance(unit);
		if (distance >= 0 && (!closest || distance < minDistance))
		{
			minDistance = distance;
			closest = target;
		}
	}

	return closest;
}

BWAPI::Unit MicroManager::closestThreat(BWAPI::Unit unit, const BWAPI::Unitset & targets)
{
	int minDistance = 0;
	BWAPI::Unit closest = nullptr;

	for (auto & target : targets)
	{
		if (!UnitUtil::IsThreat(unit, target)){
			continue;
		}

		int distance = -1;
		BWEMmap.GetPath(unit->getPosition(), target->getPosition(), &distance);
		//int distance = target->getDistance(unit);
		if (distance >= 0 && (!closest || distance < minDistance))
		{
			minDistance = distance;
			closest = target;
		}
	}

	return closest;
}

const BWAPI::Unitset & MicroManager::getUnits() const 
{ 
    return _units; 
}

void MicroManager::regroup(const BWAPI::Position & regroupPosition) const
{
	// If we have our natural regroup to it
	// Else regroup to mainbase
	BWAPI::Position ourBasePosition;
	BWAPI::Unit ourNatural = nullptr;
	BWAPI::Unit ourMain = nullptr;
	ourNatural = InformationManager::Instance().getBaseDepot(InformationManager::Instance().getMyNaturalLocation());

	if (UnitUtil::IsValidUnit(ourNatural))
	{
		ourBasePosition = InformationManager::Instance().getMyNaturalLocation()->Center();
	}
	else
	{
		ourBasePosition = InformationManager::Instance().getMyMainBaseLocation()->Center();
	}

	int regroupDistanceFromBase = MapTools::Instance().getGroundDistance(regroupPosition, ourBasePosition);

	// for each of the units we have
	for (auto & unit : _units)
	{
        int unitDistanceFromBase = MapTools::Instance().getGroundDistance(unit->getPosition(), ourBasePosition);
		int unitDistanceFromOrder = MapTools::Instance().getGroundDistance(unit->getPosition(), order.getPosition());

		// 1. A broodling or scourge should never retreat, but attack as long as it lives.
		// 2. If none of its kind has died yet, a dark templar or lurker should not retreat.
		// 3. A ground unit next to an enemy sieged tank should not move away.
		// TODO 4. A unit in stay-home mode should stay home, not "regroup" away from home.
		// TODO 5. A unit whose retreat path is blocked by enemies should do something else, at least attack-move.
		if (unit->getType() == BWAPI::UnitTypes::Zerg_Broodling || 
			(BWAPI::Broodwar->enemy()->getRace() == BWAPI::Races::Terran &&
			!unit->isFlying() &&
			BWAPI::Broodwar->getClosestUnit(unit->getPosition(),
			BWAPI::Filter::IsEnemy && BWAPI::Filter::GetType == BWAPI::UnitTypes::Terran_Siege_Tank_Siege_Mode,
			64)))
		{
			if (!immobilizeUnit(unit))
			{
				Micro::SmartAttackMove(unit, unit->getPosition());
			}
		}
		
		else if (unit->getType() == BWAPI::UnitTypes::Zerg_Scourge ||
			/*(unit->getType() == BWAPI::UnitTypes::Zerg_Lurker && BWAPI::Broodwar->self()->deadUnitCount(BWAPI::UnitTypes::Zerg_Lurker) == 0) ||*/
			(unit->getType() == BWAPI::UnitTypes::Protoss_Dark_Templar && BWAPI::Broodwar->self()->deadUnitCount(BWAPI::UnitTypes::Protoss_Dark_Templar) == 0))
		{
			if (unitDistanceFromBase > regroupDistanceFromBase)
			{
				if (!immobilizeUnit(unit))
				{
					Micro::SmartAttackMove(unit, unit->getPosition());
				}
			}
			else {
				if (!immobilizeUnit(unit))
				{
					Micro::SmartAttackMove(unit, regroupPosition);
				}
			}
		}
		/*
		else if (unitDistanceFromBase > regroupDistanceFromBase)
		{
			if (!mobilizeUnit(unit))
			{
				Micro::SmartMove(unit, ourBasePosition);
			}
		}
		*/
		else if (unit->getDistance(regroupPosition) > 96)
		{
			if (!mobilizeUnit(unit))
			{
				Micro::SmartMove(unit, regroupPosition);
			}
		}
		else
		{
			// We have retreated to a good position.
			if (!immobilizeUnit(unit))
			{
				Micro::SmartAttackMove(unit, unit->getPosition());
			}
		}
	}
}

bool MicroManager::unitNearEnemy(BWAPI::Unit unit)
{
	assert(unit);

	BWAPI::Unitset enemyNear;

	MapGrid::Instance().GetUnits(enemyNear, unit->getPosition(), 800, false, true);

	return enemyNear.size() > 0;
}

// returns true if position:
// a) is walkable
// b) doesn't have buildings on it
// c) doesn't have a unit on it that can attack ground
bool MicroManager::checkPositionWalkable(BWAPI::Position pos) 
{
	// get x and y from the position
	int x(pos.x), y(pos.y);

	// walkable tiles exist every 8 pixels
	bool good = BWAPI::Broodwar->isWalkable(x/8, y/8);
	
	// if it's not walkable throw it out
	if (!good) return false;
	
	// for each of those units, if it's a building or an attacking enemy unit we don't want to go there
	for (auto & unit : BWAPI::Broodwar->getUnitsOnTile(x/32, y/32)) 
	{
		if	(unit->getType().isBuilding() || unit->getType().isResourceContainer() || 
			(unit->getPlayer() != BWAPI::Broodwar->self() && unit->getType().groundWeapon() != BWAPI::WeaponTypes::None)) 
		{		
				return false;
		}
	}

	// otherwise it's okay
	return true;
}

void MicroManager::trainSubUnits(BWAPI::Unit unit) const
{
	if (unit->getType() == BWAPI::UnitTypes::Protoss_Reaver)
	{
		if (!unit->isTraining() && unit->canTrain(BWAPI::UnitTypes::Protoss_Scarab))
		{
			unit->train(BWAPI::UnitTypes::Protoss_Scarab);
		}
	}
	else if (unit->getType() == BWAPI::UnitTypes::Protoss_Carrier)
	{
		if (!unit->isTraining() && unit->canTrain(BWAPI::UnitTypes::Protoss_Interceptor))
		{
			unit->train(BWAPI::UnitTypes::Protoss_Interceptor);
		}
	}
}

bool MicroManager::unitNearChokepoint(BWAPI::Unit unit) const
{
	for (const auto & area : BWEMmap.Areas())
	{
		for (const BWEM::ChokePoint * choke : area.ChokePoints())
		{
			if (unit->getDistance(BWAPI::Position(choke->Center())) < 5 * 32)
			{
				return true;
			}
		}
	}

	return false;
}

// Mobilize the unit if it is immobile: A sieged tank or a burrowed zerg unit.
// Return whether any action was taken.
bool MicroManager::mobilizeUnit(BWAPI::Unit unit) const
{
	if (unit->getType() == BWAPI::UnitTypes::Terran_Siege_Tank_Siege_Mode && unit->canUnsiege())
	{
		return unit->unsiege();
	}
	if (unit->isBurrowed() && unit->canUnburrow() &&
		!unit->isIrradiated() &&
		(double(unit->getType().maxHitPoints()) / double(unit->getHitPoints()) < 6.25))  // very weak units stay burrowed
	{
		return unit->unburrow();
	}
	return false;
}

// Immobilixe the unit: Siege a tank, burrow a lurker. Otherwise do nothing.
// Return whether any action was taken.
bool MicroManager::immobilizeUnit(BWAPI::Unit unit) const
{
	if (unit->getType() == BWAPI::UnitTypes::Terran_Siege_Tank_Tank_Mode && unit->canSiege())
	{
		return unit->siege();
	}
	if (!unit->isBurrowed() && unit->canBurrow() &&
		(unit->getType() == BWAPI::UnitTypes::Zerg_Lurker || unit->isIrradiated()))
	{
		return unit->burrow();
	}
	return false;
}

// Sometimes a unit on ground attack-move freezes in place.
// Luckily it's easy to recognize, though units may be on PlayerGuard for other reasons.
// Return whether any action was taken.
// This solves stuck zerglings, but doesn't always prevent other units from getting stuck.
bool MicroManager::unstickStuckUnit(BWAPI::Unit unit)
{
	if (!unit->isMoving() && !unit->getType().isFlyer() && !unit->isBurrowed() &&
		unit->getOrder() == BWAPI::Orders::PlayerGuard &&
		BWAPI::Broodwar->getFrameCount() % 4 == 0)
	{
		Micro::SmartStop(unit);
		return true;
	}

	return false;
}

void MicroManager::initializeVectors() {
	kiteVector = BWAPI::Position(0, 0);
	tangentVector = BWAPI::Position(0, 0);
	accumulatedTangents = 0;
}

void MicroManager::updateVectors(BWAPI::Unit unit, BWAPI::Unit target, int inclusionScale, BWAPI::Position heading) {
	int distance = unit->getDistance(target);
	if (UnitUtil::IsThreat(unit, target, false)) {
		int range = UnitUtil::GetAttackRange(target, unit);
		if (distance < range + inclusionScale * (32 + 3 * (target->getType().topSpeed()))) {
			kiteVector += Micro::GetKiteVector(target, unit);
		}
		//difficult to say what to do with it right now until other things have been set up
		else if (heading.isValid() && distance < range + inclusionScale*(2 * 32) && target->getDistance(heading) < unit->getDistance(heading)) {
			accumulatedTangents++;
			//int scale = (target->getDistance(unit))/64; //let's just do short range for now...
			tangentVector += Micro::GetTangentVector(target, unit, heading);
		}

	}
}

BWAPI::Position	MicroManager::computeAttractionVector(BWAPI::Unit unit) {
	auto allies = BWAPI::Broodwar->getUnitsInRadius(unit->getPosition(), 48, BWAPI::Filter::GetType == unit->getType());
	return (allies.getPosition() - unit->getPosition());
}

void MicroManager::normalizeVectors() {
	if (kiteVector.getLength() >= 64) {  //normalize
		kiteVector = BWAPI::Position(int(64 * kiteVector.x / kiteVector.getLength()), int(64 * kiteVector.y / kiteVector.getLength()));
	}
	if (accumulatedTangents > 0) {
		tangentVector /= accumulatedTangents;
	}
}

void MicroManager::drawOrderText()
{
	if (Config::Debug::DrawUnitTargetInfo)
	{
		for (auto & unit : _units)
		{
			BWAPI::Broodwar->drawTextMap(unit->getPosition().x, unit->getPosition().y, "%s", order.getStatus().c_str());
		}
	}
}
