#include "Squad.h"
#include "UnitUtil.h"
#include "ProductionManager.h"
#include "UAlbertaBotModule.h"
#include <random>

using namespace UAlbertaBot;

Squad::Squad()
	: _name("Default")
	, _combatSquad(false)
	, _combatSimRadius(Config::Micro::CombatRegroupRadius)
	, _fightVisibleOnly(false)
	, _enemyHasAir(false)
	, _enemyHasGround(false)
	, _hasAir(false)
	, _hasGround(false)
	, _canAttackAir(false)
	, _canAttackGround(false)
	, _attackAtMax(false)
    , _lastRetreatSwitch(0)
    , _lastRetreatSwitchVal(true)
    , _priority(0)
{
    int a = 10;
}

Squad::Squad(const std::string & name, SquadOrder order, size_t priority) 
	: _name(name)
	, _order(order)
	, _combatSquad(name != "Idle")
	, _combatSimRadius(Config::Micro::CombatRegroupRadius)
	, _fightVisibleOnly(false)
	, _enemyHasAir(false)
	, _enemyHasGround(false)
	, _hasAir(false)
	, _hasGround(false)
	, _canAttackAir(false)
	, _canAttackGround(false)
	, _attackAtMax(false)
    , _lastRetreatSwitch(0)
    , _lastRetreatSwitchVal(true)
    , _priority(priority)
{
	setSquadOrder(order);
}

Squad::~Squad()
{
    clear();
}

void Squad::update()
{
	// Profile debug
	//PROFILE_FUNCTION();

	// update all necessary unit information within this squad
	updateUnits();

	if (_units.empty())
	{
		_regroupStatus = std::string("\x04 No attackers available");
		return;
	}

	// execute surveyer micro
	_surveyerManager.execute(_order);

	SquadOrder squad_order = _order;
	bool _goAggressive = ProductionManager::Instance().getAggression();

	// determine whether or not we should regroup
	bool needToRegroup = needsToRegroup() && _goAggressive;

	// draw some debug info
	//if (Config::Debug::DrawSquadInfo && 
	//	_order.getType() == SquadOrderTypes::Attack || 
	//	_order.getType() == SquadOrderTypes::Regroup ||
	//	_order.getType() == SquadOrderTypes::DestroyNeutral)
	//{
	//	BWAPI::Broodwar->drawTextScreen(200, 260, "%s", _regroupStatus.c_str());
	//}

	// if we do need to regroup, do it
	if (needToRegroup)
	{
		BWAPI::Position regroupPosition = calcRegroupPosition();

		if (Config::Debug::DrawSquadInfo)
		{
			BWAPI::Broodwar->drawCircleMap(regroupPosition.x, regroupPosition.y, 10, BWAPI::Colors::Purple, true);
		}
        
		squad_order = SquadOrder(SquadOrderTypes::Regroup, regroupPosition, 600, "Regroup");
		setSquadOrder(squad_order);

		//_meleeManager.regroup(regroupPosition);
		//_rangedManager.regroup(regroupPosition);
		//_lurkerManager.regroup(regroupPosition);
        //_tankManager.regroup(regroupPosition);
		// Note: Medics, detectors and transports do not regroup.
		
	}
	else // otherwise, execute micro
	{
		_lastOrderPosition = _order.getPosition();

		//_meleeManager.execute(_order);
		//_rangedManager.execute(_order);
		//_lurkerManager.execute(_order);
		//_tankManager.execute(_order);
	}

	_meleeManager.execute(_order);
	_rangedManager.execute(_order);
	_lurkerManager.execute(_order);
	_tankManager.execute(_order);

	_casterManager.setUnitClosestToEnemy(_vanguard);
	_casterManager.execute(_order);

	_medicManager.setUnitClosestToEnemy(_vanguard);
	_medicManager.setCenter(_center);
	_medicManager.execute(_order);

	_detectorManager.setUnitClosestToEnemy(_vanguard);
	_detectorManager.setCenter(_center);
	_detectorManager.execute(_order);

	// For now, we don't use this
	//_transportManager.update();
}

bool Squad::isEmpty() const
{
    return _units.empty();
}

size_t Squad::getPriority() const
{
    return _priority;
}

void Squad::setPriority(const size_t & priority)
{
    _priority = priority;
}

void Squad::updateUnits()
{
	// Profile debug
	//PROFILE_FUNCTION();

	setAllUnits();
	setNearEnemyUnits();
	addUnitsToMicroManagers();

	_vanguard = unitClosestToEnemy();
	_center = calcCenter();
}

// Clean up the _units vector.
// Also notice and remember a few facts about the members of the squad.
// Note: Some units may be loaded in a bunker or transport and cannot accept orders.
//       Check unit->isLoaded() before issuing orders.
void Squad::setAllUnits()
{
	_hasAir = false;
	_hasGround = false;
	_canAttackAir = false;
	_canAttackGround = false;

	BWAPI::Unitset goodUnits;
	for (const auto unit : _units)
	{
		if (UnitUtil::IsValidUnit(unit))
		{
			goodUnits.insert(unit);

			if (unit->isFlying())
			{
				if (!unit->getType().isDetector())    // mobile detectors don't count
				{
					_hasAir = true;
				}
			}
			else
			{
				_hasGround = true;
			}
			if (UnitUtil::CanAttackAir(unit))
			{
				_canAttackAir = true;
			}
			if (UnitUtil::CanAttackGround(unit))
			{
				_canAttackGround = true;
			}
		}
	}
	_units = goodUnits;
}


// Notice and remember a few facts about the enemies nearby the squad.
// Note: Some units may be loaded in a bunker or transport and cannot accept orders.
//       Check unit->isLoaded() before issuing orders.
void Squad::setAllUnitsEnemy()
{
	_enemyHasAir = false;
	_enemyHasGround = false;

	std::vector<UnitInfo> enemyCombatUnits;
	InformationManager::Instance().getNearbyForce(enemyCombatUnits, _center, BWAPI::Broodwar->enemy(), _combatSimRadius);

	for (const auto ui : enemyCombatUnits)
	{
		if (!ui.unit->exists() || ui.unit->isLoaded())
		{
			continue;
		}

		if (ui.unit->isFlying())
		{
			if (!ui.unit->getType().isDetector())    // mobile detectors don't count
			{
				_enemyHasAir = true;
			}
		}
		else
		{
			_enemyHasGround = true;
		}
	}
}

void Squad::setNearEnemyUnits()
{
	_nearEnemy.clear();

	for (const auto unit : _units)
	{
		if (!unit->getPosition().isValid())   // excludes loaded units
		{
			continue;
		}

		_nearEnemy[unit] = unitNearEnemy(unit);

		if (Config::Debug::DrawSquadInfo) {
			int left = unit->getType().dimensionLeft();
			int right = unit->getType().dimensionRight();
			int top = unit->getType().dimensionUp();
			int bottom = unit->getType().dimensionDown();

			int x = unit->getPosition().x;
			int y = unit->getPosition().y;

			BWAPI::Broodwar->drawBoxMap(x - left, y - top, x + right, y + bottom,
				(_nearEnemy[unit]) ? Config::Debug::ColorUnitNearEnemy : Config::Debug::ColorUnitNotNearEnemy);
		}
	}
}

void Squad::addUnitsToMicroManagers()
{
	BWAPI::Unitset meleeUnits;
	BWAPI::Unitset rangedUnits;
	BWAPI::Unitset detectorUnits;
	BWAPI::Unitset transportUnits;
	BWAPI::Unitset lurkerUnits;
    BWAPI::Unitset tankUnits;
    BWAPI::Unitset medicUnits;
	BWAPI::Unitset surveyerUnits;
	BWAPI::Unitset casterUnits;

	// add _units to micro managers
	for (auto & unit : _units)
	{
		if (unit->isCompleted() && unit->exists() && unit->getHitPoints() > 0 && !unit->isLoaded())
		{
			if (_name == "Survey" && unit->getType() == BWAPI::UnitTypes::Zerg_Overlord)
			{
				surveyerUnits.insert(unit);
			}
            else if (unit->getType() == BWAPI::UnitTypes::Terran_Medic)
            {
                medicUnits.insert(unit);
            }
			else if (unit->getType().isSpellcaster() ||
				unit->getType() == BWAPI::UnitTypes::Zerg_Infested_Terran) {
				casterUnits.insert(unit);
			}
			else if (unit->getType() == BWAPI::UnitTypes::Zerg_Lurker)
			{
				lurkerUnits.insert(unit);
			}
			else if (unit->getType() == BWAPI::UnitTypes::Terran_Siege_Tank_Siege_Mode ||
				unit->getType() == BWAPI::UnitTypes::Terran_Siege_Tank_Tank_Mode)
            {
                tankUnits.insert(unit);
            }   
			// select detector _units
			else if (unit->getType().isDetector() && !unit->getType().isBuilding())
			{
				detectorUnits.insert(unit);
			}
			// select transport _units
			else if (unit->getType() == BWAPI::UnitTypes::Protoss_Shuttle || 
				unit->getType() == BWAPI::UnitTypes::Terran_Dropship)
			{
				transportUnits.insert(unit);
			}
			// select ranged _units
			else if ((unit->getType().groundWeapon().maxRange() > 32) || 
				unit->getType() == BWAPI::UnitTypes::Zerg_Scourge ||
				unit->getType() == BWAPI::UnitTypes::Protoss_Reaver ||
				unit->getType() == BWAPI::UnitTypes::Protoss_Carrier)
			{
				rangedUnits.insert(unit);
			}
			// select melee _units
			else if (unit->getType().isWorker() && _combatSquad)
			{
				// If this is a combat squad, then workers are melee units like any other,
				// but we have to tell WorkerManager about them.
				// If it's not a combat squad, WorkerManager owns them; don't add them to a micromanager.
				WorkerManager::Instance().setCombatWorker(unit);
				meleeUnits.insert(unit);
			}
			// select melee _units
			else if (unit->getType().groundWeapon().maxRange() <= 32)
			{
				meleeUnits.insert(unit);
			}
		}
	}

	_meleeManager.setUnits(meleeUnits);
	_rangedManager.setUnits(rangedUnits);
	_detectorManager.setUnits(detectorUnits);
	_transportManager.setUnits(transportUnits);
	_lurkerManager.setUnits(lurkerUnits);
	_tankManager.setUnits(tankUnits);
    _medicManager.setUnits(medicUnits);
	_surveyerManager.setUnits(surveyerUnits);
	_casterManager.setUnits(casterUnits);
}

// calculates whether or not to regroup
bool Squad::needsToRegroup()
{
	// Profile debug
	//PROFILE_FUNCTION();

	// if we have no units, never regroup
	if (_units.empty())
	{
		_regroupStatus = std::string("\x04 No attackers available");
		return false;
	}

	// Our order may not allow us to regroup.
	if (!_order.isRegroupableOrder())
	{
		_regroupStatus = std::string("\x04 Never retreat!");
		return false;
	}

	// If we're nearly maxed and have good income and cash, don't regroup.
	if (BWAPI::Broodwar->self()->supplyUsed() >= 380 &&
		BWAPI::Broodwar->self()->minerals() > 1000 && 
		BWAPI::Broodwar->self()->gas() > 1000)
	{
		_attackAtMax = true;
	}
	if (_attackAtMax)
	{
		if (BWAPI::Broodwar->self()->supplyUsed() < 320)
		{
			_attackAtMax = false;
		}
		else
		{
			_regroupStatus = std::string("\x04 Maxed. Banzai!");
			return false;
		}
	}
	
	// No closest unit to enemy, don't regroup
	BWAPI::Unit vanguard = _vanguard;
	BWAPI::Position center = _center;
	if (!vanguard)
	{
		_regroupStatus = std::string("\x04 No vanguard");
		return false;
	}
	
	// Is there static defense nearby that we should take into account?
	// vanguard is known to be set thanks to the test immediately above.
	BWAPI::Unit nearestStaticDefense = nearbyStaticDefense(vanguard->getPosition());
	const BWAPI::Position finalPosition = finalRegroupPosition();
	if (nearestStaticDefense)
	{
		// Don't retreat if we are in range of static defense that is attacking.
		if (nearestStaticDefense->getOrder() == BWAPI::Orders::AttackUnit)
		{
			_regroupStatus = std::string("\x04 Go static defense!");
			return false;
		}

		// Don't retreat if we are behind static defense.
		// Assume that the static defense is between the final regroup position and the enemy.
		/*if (vanguard->getPosition().getApproxDistance(nearestStaticDefense->getPosition()) < 192 &&
			nearestStaticDefense->getPosition().getApproxDistance(finalPosition) - vanguard->getPosition().getApproxDistance(finalPosition) > 96)
		{
			_regroupStatus = std::string("\x04 Behind static defense");
			return false;
		}*/
	}
	//else
	//{
		// There is no static defense to retreat to.
		if (vanguard->getPosition().getApproxDistance(finalPosition) < 224)
		{
			_regroupStatus = std::string("\x04 Back to the wall");
			return false;
		}
	//}


	// if none of our units are in range of any enemy units, don't regroup
    std::vector<UnitInfo> enemyCombatUnits;
    const auto & enemyUnitInfo = InformationManager::Instance().getUnitInfo(BWAPI::Broodwar->enemy());

    bool anyInRange = false;
    for (const auto & eui : enemyUnitInfo)
    {
        for (const auto & u : _units)
        {
			if (!u->exists() || u->isLoaded())
			{
				continue;
			}

			// Max of weapon range and vision range. Vision range is as long or longer, except for tanks.
			// We assume that the tanks may siege, and check the siege range of unsieged tanks.
			if (UnitUtil::CanAttack(eui.second.type, u->getType()))
			{
				int range = 0;     // range of enemy unit

				if (eui.second.type == BWAPI::UnitTypes::Terran_Siege_Tank_Tank_Mode ||
					eui.second.type == BWAPI::UnitTypes::Terran_Siege_Tank_Siege_Mode)
				{
					range = (BWAPI::UnitTypes::Terran_Siege_Tank_Siege_Mode).groundWeapon().maxRange() + 128;  // plus safety fudge
				}
				else
				{
					// Sight range is >= weapon range, so we can stop here.
					range = eui.second.type.sightRange();
				}

				range += 128;    // plus some to account for our squad spreading out

				if (range >= u->getPosition().getApproxDistance(eui.second.lastPosition))
				{
					anyInRange = true;
					break;   // break out of inner loop
				}
			}
        }

		if (anyInRange)
        {
            break;   // break out of outer loop
        }
    }

    if (!anyInRange)
    {
        _regroupStatus = std::string("\x04 No enemy units in range");
        return false;
    }
	
	// ---
	// All other checks are done. Finally do the expensive combat simulation.
	// ---

	// If we most recently retreated, don't attack again until retreatDuration frames have passed.
	const int retreatDuration = 5 * 24;
	bool retreat = _lastRetreatSwitchVal && (BWAPI::Broodwar->getFrameCount() - _lastRetreatSwitch < retreatDuration);

	if (!retreat)
	{
		// All checks are done. Finally do the expensive combat simulation.
		CombatSimulation sim;

		// FAP
		//sim.setCombatUnits(unitClosest->getPosition(), Config::Micro::CombatRegroupRadius);
		int score = sim.simulateCombatFAP(vanguard->getPosition(), _combatSimRadius, _canAttackGround, _canAttackAir);

		retreat = score < 0;
		_lastRetreatSwitch = BWAPI::Broodwar->getFrameCount();
		_lastRetreatSwitchVal = retreat;
	}
	
	if (retreat)
	{
		_regroupStatus = std::string("\x04 Retreat");
	}
	else
	{
		_regroupStatus = std::string("\x04 Attack");
	}

	return retreat;
}

void Squad::setSquadOrder(const SquadOrder & so)
{
	_order = so;
}

bool Squad::containsUnit(BWAPI::Unit u) const
{
    return _units.contains(u);
}

bool Squad::containsUnitType(BWAPI::UnitType t) const
{
	for (const auto u : _units)
	{
		if (u->getType() == t)
		{
			return true;
		}
	}
	return false;
}

void Squad::clear()
{
	// If game has ended WorkerManager could be deconstrcuted and result in crash
	// Also SC will quit, so no need
	if (gameEnded) return;

    for (const auto unit : getUnits())
    {
        if (unit->getType().isWorker())
        {
            WorkerManager::Instance().finishedWithWorker(unit);
        }
    }

    _units.clear();
}

bool Squad::unitNearEnemy(BWAPI::Unit unit)
{
	UAB_ASSERT(unit, "missing unit");

	BWAPI::Unitset enemyNear;
	BWAPI::Unitset enemyNearFiltered;

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

	for (const auto u : enemyNear)
	{
		if (!u->exists() || u->isLoaded())
		{
			continue;
		}

		if (u->isFlying())
		{
			if (_canAttackAir)    // only add flying units if we can attack air
			{
				enemyNearFiltered.insert(u);
			}
		}
		else
		{
			enemyNearFiltered.insert(u);
		}
	}

	return enemyNearFiltered.size() > 0;
}

BWAPI::Position Squad::calcCenter() const
{
	// Profile debug
	//PROFILE_FUNCTION();

    if (_units.empty())
    {
        //if (Config::Debug::DrawSquadInfo)
        //{
        //	BWAPI::Broodwar->printf("Squad::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());
}

BWAPI::Position Squad::calcRegroupPosition()
{
	// Profile debug
	//PROFILE_FUNCTION();

	// 1. Retreat toward static defense, if any is near.
	BWAPI::Unit vanguard = _vanguard;  // squad vanguard
	if (vanguard)
	{
		BWAPI::Unit nearest = nearbyStaticDefense(vanguard->getPosition());
		if (nearest)
		{
			BWAPI::Position behind = DistanceAndDirection(nearest->getPosition(), vanguard->getPosition(), -128);
			return behind;
		}
	}
		
	// 2. Retreat to the location of the unit not near the enemy which is
	// closest to the order position.
	// NOTE May retreat somewhere silly if the chosen unit was newly produced.
	//      Sometimes retreats back and forth through the enemy when new
	//      units are produced in bases on opposite sides.
	BWAPI::Position base_position = finalRegroupPosition();
	BWAPI::Position regroup(0, 0);

	if (vanguard)
	{
		BWAPI::Position order_position(_order.getPosition());
		BWAPI::TilePosition order_location(_order.getPosition());
		BWAPI::Position enemy_position(vanguard->getPosition());
		int minDistEnemy = 100000;

		for (const auto unit : _units)
		{
			// Don't return the position of a detector, which may be in a weird place.
			// That means science vessel, protoss observer, or overlord.
			// Bug fix thanks to AIL!
			if (unit->exists() && 
				!_nearEnemy[unit] &&
				!unit->getType().isDetector() &&
				unit->getType() != BWAPI::UnitTypes::Terran_Medic &&
				unit->getPosition().isValid())    // excludes loaded units
			{
				//int distOrder = unit->getPosition().getApproxDistance(order_position);
				int distBase = unit->getPosition().getApproxDistance(base_position);
				int distEnemy = unit->getPosition().getApproxDistance(enemy_position);
				int distBaseEnemy = base_position.getApproxDistance(enemy_position);

				if (distEnemy < minDistEnemy && distBase < distBaseEnemy)
				{
					// If the squad has any ground units, don't try to retreat to the position of an air unit
					// which is flying in a place that a ground unit cannot reach.
					if (!_hasGround || MapTools::Instance().getGroundTileDistance(unit->getTilePosition(), order_location) >= 0)
					{
						minDistEnemy = distEnemy;
						regroup = unit->getPosition();
					}
				}
			}
		}
	}
	if (regroup != BWAPI::Position(0, 0))
	{
		return regroup;
	}

	// 3. Failing that, retreat to a base we own.
	return base_position;
}

// Return the rearmost position we should retreat to, which puts our "back to the wall".
BWAPI::Position Squad::finalRegroupPosition()
{
	// Retreat to the main base, unless we change our mind below.
	const BWEM::Base * base = InformationManager::Instance().getMyMainBaseLocation();

	// If the natural has been taken, retreat there instead.
	const BWEM::Base * natural = InformationManager::Instance().getMyNaturalLocation();
	if (natural && BWAPI::Broodwar->self() == InformationManager::Instance().getBaseOwner(natural))
	{
		base = natural;
	}

	return base->Center();
}

BWAPI::Unit Squad::nearbyStaticDefense(const BWAPI::Position & pos) const
{
	BWAPI::Unit nearest = nullptr;

	// NOTE What matters is whether the enemy has ground or air units.
	if (_enemyHasGround)
	{
		nearest = InformationManager::Instance().nearestGroundStaticDefense(pos);
	}
	else
	{
		nearest = InformationManager::Instance().nearestAirStaticDefense(pos);
	}
	if (nearest && nearest->getPosition().getApproxDistance(pos) < 192)
	{
		return nearest;
	}
	return nullptr;
}

// Return the unit closest to the order position (not actually closest to the enemy).
BWAPI::Unit Squad::unitClosestToEnemy() const
{
	// Profile debug
	//PROFILE_FUNCTION();

	BWAPI::Position order_position(_order.getPosition());
	BWAPI::TilePosition order_location(_order.getPosition());
	const BWEM::Base * enemyBase = InformationManager::Instance().getEnemyMainBaseLocation();
	BWAPI::Unit closest = nullptr;
	int closestDist = 100000;

	UAB_ASSERT(_order.getPosition().isValid(), "bad order position");

	for (const auto unit : _units)
	{
		// Non-combat units should be ignored for this calculation.
		if (!unit->getPosition().isValid() ||   // includes units loaded into bunkers or transports
			unit->getType().isDetector() ||
			unit->getType().isSpellcaster() ||  // ignore queens and defilers
			unit->getType() == BWAPI::UnitTypes::Terran_Medic || 
			unit->getType() == BWAPI::UnitTypes::Zerg_Infested_Terran ||
			unit->getType() == BWAPI::UnitTypes::Zerg_Broodling)
		{
			continue;
		}

		int dist;
		if (_order.getType() == SquadOrderTypes::Regroup)
		{
			dist = unit->getPosition().getApproxDistance(_lastOrderPosition);
		}
		else
		{
			dist = unit->getPosition().getApproxDistance(order_position);
		}
		/*if (_hasGround)
		{
			// A ground or air-ground squad. Use ground distance.
			// It is -1 if no ground path exists.
			dist = MapTools::Instance().getGroundTileDistance(unit->getTilePosition(), order_location);
		}
		else
		{
			// An all-air squad. Use air distance (which is what unit->getPosition().getApproxDistance() gives).
			dist = unit->getPosition().getApproxDistance(order_position);
		}*/

		if (dist < closestDist && dist > 0)
		{
			closest = unit;
			closestDist = dist;
		}
	}

	if (closest && Config::Debug::DrawSquadInfo) {
		BWAPI::Broodwar->drawCircleMap(closest->getPosition(), 5, BWAPI::Colors::White, true);
	}

	return closest;
}

int Squad::squadUnitsNear(BWAPI::Position p)
{
	int numUnits = 0;

	for (auto & unit : _units)
	{
		if (unit->getPosition().getApproxDistance(p) < 600)
		{
			numUnits++;
		}
	}

	return numUnits;
}

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

const SquadOrder & Squad::getSquadOrder()	const			
{ 
	return _order; 
}

void Squad::addUnit(BWAPI::Unit u)
{
	_units.insert(u);
}

void Squad::removeUnit(BWAPI::Unit u)
{
	if (_combatSquad && u->getType().isWorker())
	{
		WorkerManager::Instance().finishedWithWorker(u);
	}
	_units.erase(u);
}

// Remove all workers from the squad, releasing them back to WorkerManager.
void Squad::releaseWorkers()
{
	for (auto it = _units.begin(); it != _units.end();)
	{
		if (_combatSquad && (*it)->getType().isWorker())
		{
			WorkerManager::Instance().finishedWithWorker(*it);
			it = _units.erase(it);
		}
		else
		{
			++it;
		}
	}
}

const std::string & Squad::getName() const
{
    return _name;
}

const bool Squad::hasCombatUnits() const
{
	// If the only units we have are detectors, then we have no combat units.
	return !(_units.empty() || _units.size() == _detectorManager.getUnits().size());
}
