#include "OpponentPlan.h"

#include "Bases.h"
#include "InformationManager.h"
#include "PlayerSnapshot.h"

using namespace UAlbertaBot;

// Attempt to recognize what the opponent is doing, so we can cope with it.
// For now, only try to recognize a small number of opening situations that require
// different handling.

// This is part of the OpponentModel module. Access should normally be through the OpponentModel instance.

// -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --

bool OpponentPlan::fastPlan(OpeningPlan plan)
{
	return
		plan == OpeningPlan::Proxy ||
		plan == OpeningPlan::WorkerRush ||
		plan == OpeningPlan::FastRush;
}

// NOTE Incomplete test! We don't measure the distance of enemy units from the enemy base,
//      so we don't recognize all the rushes that we should.
bool OpponentPlan::recognizeWorkerRush()
{
	BWAPI::Position myOrigin = InformationManager::Instance().getMyMainBaseLocation()->getPosition();

	int enemyWorkerRushCount = 0;

	for (const auto & kv : InformationManager::Instance().getUnitData(BWAPI::Broodwar->enemy()).getUnits())
	{
		const UnitInfo & ui(kv.second);

		if (ui.type.isWorker() && ui.unit->isVisible() && myOrigin.getDistance(ui.unit->getPosition()) < 1000)
		{
			++enemyWorkerRushCount;
		}
	}

	return enemyWorkerRushCount >= 3;
}

// Factory, possibly with starport, and no sign of many marines intended.
bool OpponentPlan::recognizeFactoryTech()
{
	if (BWAPI::Broodwar->enemy()->getRace() != BWAPI::Races::Terran)
	{
		return false;
	}

	int nMarines = 0;
	int nBarracks = 0;
	int nTechProduction = 0;
	bool tech = false;

	for (const auto & kv : InformationManager::Instance().getUnitData(BWAPI::Broodwar->enemy()).getUnits())
	{
		const UnitInfo & ui(kv.second);

		if (ui.type == BWAPI::UnitTypes::Terran_Marine)
		{
			++nMarines;
		}

		else if (ui.type.whatBuilds().first == BWAPI::UnitTypes::Terran_Barracks)
		{
			return false;			// academy implied, marines seem to be intended
		}

		else if (ui.type == BWAPI::UnitTypes::Terran_Barracks)
		{
			++nBarracks;
		}

		else if (ui.type == BWAPI::UnitTypes::Terran_Academy)
		{
			return false;			// marines seem to be intended
		}

		else if (ui.type == BWAPI::UnitTypes::Terran_Factory ||
			ui.type == BWAPI::UnitTypes::Terran_Starport)
		{
			++nTechProduction;
		}

		else if (ui.type.whatBuilds().first == BWAPI::UnitTypes::Terran_Factory ||
			ui.type.whatBuilds().first == BWAPI::UnitTypes::Terran_Starport ||
			ui.type == BWAPI::UnitTypes::Terran_Armory)
		{
			tech = true;			// indicates intention to rely on tech units
		}
	}

	return (nTechProduction >= 2 || tech) && nMarines <= 6 && nBarracks <= 1;
}

void OpponentPlan::recognize()
{
	//int frame = BWAPI::Broodwar->getFrameCount();

	////table---------------------
	////0=Zerg_Spawning_Pool---1=Zerg_Zergling-----2=Zerg_Hatchery----------3=Zerg_Extractor----------14=Zerg_Sunken_Colony
	////4=Terran_Barracks------5=Terran_Marine-----6=Terran_Refinery--------7=Terran_Command_Center---8=Terran_Bunker
	////9=Protoss_Gateway------10=Protoss_Zealot---11=Protoss_Assimilator---12=Protoss_Nexus----------13=Protoss_Photon_Cannon
	//std::vector<int> numEnemyUnits(20, 0);

	//for (auto unit : BWAPI::Broodwar->enemy()->getUnits())
	//{
	//	if (unit->getType() == BWAPI::UnitTypes::Zerg_Spawning_Pool){ numEnemyUnits[0]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Zerg_Zergling){ numEnemyUnits[1]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Zerg_Hatchery){ numEnemyUnits[2]++; }
	//	else if(unit->getType() == BWAPI::UnitTypes::Zerg_Extractor){ numEnemyUnits[3]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Terran_Barracks){ numEnemyUnits[4]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Terran_Marine){ numEnemyUnits[5]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Terran_Refinery){ numEnemyUnits[6]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Terran_Command_Center){ numEnemyUnits[7]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Terran_Bunker){ numEnemyUnits[8]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Protoss_Gateway){ numEnemyUnits[9]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Protoss_Zealot){ numEnemyUnits[10]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Protoss_Assimilator){ numEnemyUnits[11]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Protoss_Nexus){ numEnemyUnits[12]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Protoss_Photon_Cannon){ numEnemyUnits[13]++; }
	//	else if (unit->getType() == BWAPI::UnitTypes::Zerg_Sunken_Colony){ numEnemyUnits[14]++; }

	//}



	//// Recognize fast rushes.
	//// TODO consider distance and speed: when might units have been produced?
	////      as it stands, 4 pool is unrecognized half the time because lings are seen too late
	//if ((frame < 1600 && numEnemyUnits[0] > 0) ||
	//	(frame < 3200 && numEnemyUnits[1] > 0) ||
	//	(frame < 1750 && numEnemyUnits[9] > 0) ||
	//	(frame < 3300 && numEnemyUnits[10] > 0) ||
	//	(frame < 1400 && numEnemyUnits[4] > 0) ||
	//	(frame < 3000 && numEnemyUnits[5] > 0))
	//{
	//	_openingPlan = OpeningPlan::FastRush;
	//	_planIsFixed = true;
	//	return;
	//}

	// //Plans below here are slow plans. Do not overwrite a fast plan with a slow plan.
	//if (fastPlan(_openingPlan))
	//{
	//	return;
	//}

	//// Recognize slower rushes.
	//// TODO make sure we've seen the bare geyser in the enemy base!
	//// TODO seeing a unit carrying gas also means the enemy has gas
	//if (numEnemyUnits[2] >= 2 &&
	//	numEnemyUnits[0] > 0 &&
	//	numEnemyUnits[3] == 0
	//	||
	//	numEnemyUnits[4] >= 2 &&
	//	numEnemyUnits[6] == 0 &&
	//	numEnemyUnits[7] <= 1
	//	||
	//	numEnemyUnits[9] >= 2 &&
	//	numEnemyUnits[11] == 0 &&
	//	numEnemyUnits[12] <= 1)
	//{
	//	_openingPlan = OpeningPlan::HeavyRush;
	//	_planIsFixed = true;
	//	return;
	//}

	//// Recognize terran factory tech openings.
	//if (recognizeFactoryTech())
	//{
	//	_openingPlan = OpeningPlan::Factory;
	//	return;
	//}

	//// Recognize expansions with pre-placed static defense.
	//// Zerg can't do this.
	//// NOTE Incomplete test! We don't check the location of the static defense
	//if (Bases::Instance().baseCount(BWAPI::Broodwar->enemy()) >= 2)
	//{
	//	if (numEnemyUnits[8] > 0 ||
	//		numEnemyUnits[13] > 0)
	//	{
	//		_openingPlan = OpeningPlan::SafeExpand;
	//		return;
	//	}
	//}

	//// Recognize a naked expansion.
	//// This has to run after the SafeExpand check, since it doesn't check for what's missing.
	//if (Bases::Instance().baseCount(BWAPI::Broodwar->enemy()) >= 2)
	//{
	//	_openingPlan = OpeningPlan::NakedExpand;
	//	return;
	//}

	//// Recognize a turtling enemy.
	//// NOTE Incomplete test! We don't check where the defenses are placed.
	//if (Bases::Instance().baseCount(BWAPI::Broodwar->enemy()) < 2)
	//{
	//	if (numEnemyUnits[8] >= 2 ||
	//		numEnemyUnits[13] >= 2 ||
	//		numEnemyUnits[14] >= 2)
	//	{
	//		_openingPlan = OpeningPlan::Turtle;
	//		return;
	//	}
	//}

	// Don't recognize island plans.
	// The regular plans and reactions do not make sense for island maps.
	if (Bases::Instance().isIslandStart())
	{
		return;
	}

	// Recognize fast plans first, slow plans below.

	// Recognize in-base proxy buildings. Info manager does it for us.
	if (Bases::Instance().getEnemyProxy())
	{
		_openingPlan = OpeningPlan::Proxy;
		_planIsFixed = true;
		return;
	}

	// Recognize worker rushes.
	if (recognizeWorkerRush())
	{
		_openingPlan = OpeningPlan::WorkerRush;
		return;
	}

	int frame = BWAPI::Broodwar->getFrameCount();

	PlayerSnapshot snap;
	snap.takeEnemy();

	BWAPI::Unit enemyUnit = BWAPI::Broodwar->getClosestUnit(
		InformationManager::Instance().getPosition("barracksPoint"),
		BWAPI::Filter::IsEnemy && !BWAPI::Filter::IsWorker, 360);

	// Recognize fast rushes.
	// TODO consider distance and speed: when might units have been produced?
	//      as it stands, 4 pool is unrecognized half the time because lings are seen too late
	if ((frame < 1600 && snap.getCount(BWAPI::UnitTypes::Zerg_Spawning_Pool) > 0) ||
		(frame < 3200 && snap.getCount(BWAPI::UnitTypes::Zerg_Zergling) > 0) ||
		(frame < 1750 && snap.getCount(BWAPI::UnitTypes::Protoss_Gateway) > 0) ||
		(frame < 3300 && snap.getCount(BWAPI::UnitTypes::Protoss_Zealot) > 0) ||
		(frame < 1400 && snap.getCount(BWAPI::UnitTypes::Terran_Barracks) > 0) ||
		(frame < 3000 && snap.getCount(BWAPI::UnitTypes::Terran_Marine) > 0) ||
		(frame < 4500 && enemyUnit))
	{
		_openingPlan = OpeningPlan::FastRush;
		_planIsFixed = true;
		return;
	}

	// Plans below here are slow plans. Do not overwrite a fast plan with a slow plan.
	if (fastPlan(_openingPlan))
	{
		return;
	}

	// Recognize slower rushes.
	// TODO make sure we've seen the bare geyser in the enemy base!
	// TODO seeing a unit carrying gas also means the enemy has gas
	if (snap.getCount(BWAPI::UnitTypes::Zerg_Hatchery) >= 2 &&
		snap.getCount(BWAPI::UnitTypes::Zerg_Spawning_Pool) > 0 &&
		snap.getCount(BWAPI::UnitTypes::Zerg_Extractor) == 0
		||
		snap.getCount(BWAPI::UnitTypes::Terran_Barracks) >= 2 &&
		snap.getCount(BWAPI::UnitTypes::Terran_Refinery) == 0 &&
		snap.getCount(BWAPI::UnitTypes::Terran_Command_Center) <= 1
		||
		snap.getCount(BWAPI::UnitTypes::Protoss_Gateway) >= 2 &&
		snap.getCount(BWAPI::UnitTypes::Protoss_Assimilator) == 0 &&
		snap.getCount(BWAPI::UnitTypes::Protoss_Nexus) <= 1)
	{
		_openingPlan = OpeningPlan::HeavyRush;
		_planIsFixed = true;
		return;
	}

	// Recognize terran factory tech openings.
	if (recognizeFactoryTech())
	{
		_openingPlan = OpeningPlan::Factory;
		return;
	}

	// Recognize expansions with pre-placed static defense.
	// Zerg can't do this.
	// NOTE Incomplete test! We don't check the location of the static defense
	if (Bases::Instance().baseCount(BWAPI::Broodwar->enemy()) >= 2)
	{
		if (snap.getCount(BWAPI::UnitTypes::Terran_Bunker) > 0 ||
			snap.getCount(BWAPI::UnitTypes::Protoss_Photon_Cannon) > 0)
		{
			_openingPlan = OpeningPlan::SafeExpand;
			return;
		}
	}

	// Recognize a naked expansion.
	// This has to run after the SafeExpand check, since it doesn't check for what's missing.
	if (Bases::Instance().baseCount(BWAPI::Broodwar->enemy()) >= 2)
	{
		_openingPlan = OpeningPlan::NakedExpand;
		return;
	}

	// Recognize a turtling enemy.
	// NOTE Incomplete test! We don't check where the defenses are placed.
	if (Bases::Instance().baseCount(BWAPI::Broodwar->enemy()) < 2)
	{
		if (snap.getCount(BWAPI::UnitTypes::Terran_Bunker) >= 2 ||
			snap.getCount(BWAPI::UnitTypes::Protoss_Photon_Cannon) >= 2 ||
			snap.getCount(BWAPI::UnitTypes::Zerg_Sunken_Colony) >= 2)
		{
			_openingPlan = OpeningPlan::Turtle;
			return;
		}
	}

	// Nothing recognized: Opening plan remains unchanged.
}

// -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --

OpponentPlan::OpponentPlan()
	: _openingPlan(OpeningPlan::Unknown)
	, _planIsFixed(false)
{
}

// Update the recognized plan.
// Call this every frame. It will take care of throttling itself down to avoid unnecessary work.
void OpponentPlan::update()
{
	if (!Config::Strategy::UsePlanRecognizer)
	{
		return;
	}

	// The plan is decided. Don't change it any more.
	if (_planIsFixed)
	{
		return;
	}

	int frame = BWAPI::Broodwar->getFrameCount();

	if (frame > 100 && frame < 7200 &&       // only try to recognize openings
		frame % 12 == 7)                     // update interval
	{
		recognize();
	}
}
