//////////////////////////////////////////////////////////////////////////
//
// This file is part of Iron's source files.
// Iron is free software, licensed under the MIT/X11 License. 
// A copy of the license is provided with the library in the LICENSE file.
// Copyright (c) 2016, Igor Dimitrijevic
//
//////////////////////////////////////////////////////////////////////////


#include "army.h"
#include "factory.h"
#include "../Iron.h"

namespace { auto & bw = Broodwar; }


namespace iron
{



//////////////////////////////////////////////////////////////////////////////////////////////
//                                                                                          //
//                                  class Army
//                                                                                          //
//////////////////////////////////////////////////////////////////////////////////////////////


Army::Army()
{
}


void Army::Update()
{
	UpdateHisArmyStats();
	UpdateGroundLead();
	UpdateRatios();
}


void Army::UpdateHisArmyStats()
{
	m_hisSoldiers = 0;
	m_hisMarines = 0;
	m_hisMedics = 0;
	m_hisSoldiersFearingAir = 0;
	m_hisInvisibleUnits = 0;

	for (const auto & info : him().UnitTrace())
		if (!info.second.Type().isWorker())
		{
			++m_hisSoldiers;
			if (info.second.Type() == Terran_Marine) ++m_hisMarines;
			if (info.second.Type() == Terran_Medic) ++m_hisMedics;

			switch (info.second.Type())
			{
				case Terran_Medic:
				case Terran_Firebat:
				case Terran_Vulture:
				case Terran_Siege_Tank_Tank_Mode:
				case Terran_Siege_Tank_Siege_Mode:
				case Terran_Dropship:
				case Terran_Science_Vessel:
				case Protoss_Zealot:
				case Protoss_Dark_Templar:
				case Protoss_Dark_Archon:
				case Protoss_Reaver:
				case Protoss_High_Templar:
				case Protoss_Observer:
				case Protoss_Shuttle:
				case Zerg_Zergling:
				case Zerg_Lurker:
				case Zerg_Defiler:
				case Zerg_Guardian:
				case Zerg_Overlord:
				case Zerg_Queen:
				case Zerg_Ultralisk:
					++m_hisSoldiersFearingAir;
			}

			switch (info.second.Type())
			{
				case Terran_Ghost:
				case Terran_Wraith:
				case Protoss_Dark_Templar:
				case Protoss_Observer:
				case Protoss_Arbiter:
				case Zerg_Lurker:
					++m_hisInvisibleUnits;
			}
		}
}


void Army::UpdateGroundLead()
{
	m_groundLeadStatus_vultures = m_groundLeadStatus_mines = "-";

	if (me().SupplyUsed() >= 175)
	{
		m_groundLead = true;
		return;
	}

	double vulturesNeeded = 0.0;
	double vulturesWithMinesNeeded = 0.0;
	int hisUnitsFearingAir = 0;

	for (const auto & info : him().UnitTrace())
		if (!info.second.Type().isWorker())
		{
			switch (info.second.Type())
			{
			case Terran_Marine:					vulturesNeeded += 1/3.0; break;
			case Terran_Medic:					vulturesNeeded += 2/3.0; break;
			case Terran_Vulture:				vulturesNeeded += 1.4;								   ++hisUnitsFearingAir; break;
			case Terran_Siege_Tank_Tank_Mode:
			case Terran_Siege_Tank_Siege_Mode:	vulturesNeeded += 3.0; vulturesWithMinesNeeded += 2.0; ++hisUnitsFearingAir; break;
			case Terran_Goliath:				vulturesNeeded += 2.5; vulturesWithMinesNeeded += 1.5; break;

			case Protoss_Zealot:				vulturesNeeded += 0.4; break;
			case Protoss_Dragoon:				vulturesNeeded += 2.0; vulturesWithMinesNeeded += 2.0; break;
			case Protoss_Dark_Templar:			vulturesNeeded += 2.0; vulturesWithMinesNeeded += 2.0; break;
			case Protoss_Reaver:				vulturesNeeded += 4.0; vulturesWithMinesNeeded += 2.0; break;

			case Zerg_Zergling:					vulturesNeeded += 0.25; break;
			case Zerg_Hydralisk:				vulturesNeeded += 1.5; vulturesWithMinesNeeded += 1.0; break;
			case Zerg_Lurker:					vulturesNeeded += 2.0; vulturesWithMinesNeeded += 2.0; break;
			}
		}

	double vulturesAvailable = 0.0;
	double vulturesWithMinesAvailable = 0.0;

	const int wraithBonus = min(hisUnitsFearingAir, me().CompletedUnits(Terran_Wraith));
	vulturesNeeded = max(0.0, vulturesNeeded - 1.5 * wraithBonus);
	vulturesWithMinesNeeded = max(0.0, vulturesWithMinesNeeded - 2.0 * wraithBonus);

	for (const auto & u : me().Units(Terran_Vulture))
		if (u->Completed())
		{
			vulturesAvailable += 1.0 * (0.25 + 0.75 * u->Life() / double(u->MaxLife()));
			if (u->IsMy<Terran_Vulture>()->RemainingMines() >= 1)
				vulturesWithMinesAvailable += 1.0;
		}

	for (const auto & u : me().Units(Terran_Siege_Tank_Tank_Mode))
		if (u->Completed())
		{
			vulturesAvailable += 2.0 * (0.25 + 0.75 * u->Life() / double(u->MaxLife()));
		}

	for (const auto & u : me().Units(Terran_Goliath))
		if (u->Completed())
		{
			vulturesAvailable += 1.5 * (0.25 + 0.75 * u->Life() / double(u->MaxLife()));
		}

	for (const auto & u : me().Units(Terran_Vulture_Spider_Mine))
		if (u->Completed())
		{
			vulturesWithMinesAvailable += 0.5 * u->Life() / double(u->MaxLife());
		}


	vulturesAvailable += min(3*me().CompletedUnits(Terran_Wraith), m_hisSoldiersFearingAir/2);


	if (him().StartingBase())
	{
		m_hisGroundUnitsAhead = count_if(him().Units().begin(), him().Units().end(),
				[](const unique_ptr<HisUnit> & u)
				{
					return u->Completed() && !u->Type().isWorker() && u->GroundAttack() && !u->Flying() &&
						(groundDist(u->Pos(), me().GetBase(0)->Center()) < groundDist(u->Pos(), him().StartingBase()->BWEMPart()->Center()));
				});

		m_myGroundUnitsAhead = count_if(me().Units().begin(), me().Units().end(),
				[](const MyUnit * u)
				{
					return u->Completed() && !u->Type().isWorker() && u->GroundAttack() && !u->Flying() &&
						(groundDist(u->Pos(), me().GetBase(0)->Center()) > groundDist(u->Pos(), him().StartingBase()->BWEMPart()->Center()));
				});
	}


	m_vulturesCond = vulturesAvailable + 0.0001 > vulturesNeeded*1.1 + 3.0;
	m_minesCond = vulturesWithMinesAvailable + 0.0001 > vulturesWithMinesNeeded;

	m_groundLeadStatus_vultures = "vultures : " + to_string(int(vulturesAvailable + 0.5)) + "(need " + to_string(int(vulturesNeeded*1.1 + 3 + 0.5)) + ") --> " + (m_vulturesCond ? "YES" : "NO");
	m_groundLeadStatus_mines = "mines : " + to_string(int(vulturesWithMinesAvailable + 0.5)) + "(need " + to_string(int(vulturesWithMinesNeeded + 0.5)) + ") --> " + (m_minesCond ? "YES" : "NO");

	m_groundLead = m_vulturesCond && m_minesCond;
}


void Army::UpdateRatios_vsTerran()
{
	const int hisAirUnits = count_if(him().UnitTrace().begin(), him().UnitTrace().end(), [](const pair<Unit, HisUnitTrace> & p)
									{ return p.second.Type().isFlyer() && !p.second.Type().isBuilding(); });
	const int hisGoliaths = count_if(him().UnitTrace().begin(), him().UnitTrace().end(), [](const pair<Unit, HisUnitTrace> & p)
									{ return (p.second.Type() == Terran_Goliath); });

	m_tankRatioWanted = (me().SupplyUsed() < 100) ? 2 : 3;
	if		(hisGoliaths >= 10) m_tankRatioWanted += 3;
	else if (hisGoliaths >= 5) m_tankRatioWanted += 2;
	else if (hisGoliaths >= 2) m_tankRatioWanted += 1;


	m_goliathRatioWanted = 0;
	if (me().CompletedBuildings(Terran_Armory) >= 1)
		if		(hisAirUnits >= 10) m_goliathRatioWanted = 3;
		else if (hisAirUnits >= 5) m_goliathRatioWanted = 2;
		else if (hisAirUnits >= 2) m_goliathRatioWanted = 1;

//	m_goliathRatioWanted = 4;
//	m_tankRatioWanted = 4;

}


void Army::UpdateRatios_vsProtoss()
{
	int dragoons = 0;
	int reavers = 0;
	int cannons = 0;
	int arbiters = 0;
	int scouts = 0;
	int carriers = 0;
	int others = 0;

	for (const auto & info : him().UnitTrace())
		switch (info.second.Type())
		{
		case Protoss_Dragoon: ++dragoons;		break;
		case Protoss_Reaver: ++reavers;			break;
		case Protoss_Arbiter: ++arbiters;		break;
		case Protoss_Scout: ++scouts;			break;
		case Protoss_Carrier: ++carriers;		break;
		case Protoss_Photon_Cannon: ++cannons;	break;
		default:
			if (!info.second.Type().isFlyer())
			if (!info.second.Type().isBuilding())
			if (!info.second.Type().isWorker())
				++others;
		}

	int prosVultures = max(1, others);
	int prosTanks = dragoons + cannons + 2*reavers;
	int prosGoliaths = arbiters + scouts + 2*carriers;

	m_tankRatioWanted = int(0.5 + 10 * prosTanks / double(prosVultures + prosTanks + prosGoliaths));
	m_tankRatioWanted = max(1, min(5, m_tankRatioWanted));

	m_goliathRatioWanted = int(0.5 + 10 * prosGoliaths / double(prosVultures + prosTanks + prosGoliaths));
	m_goliathRatioWanted = max(0, min(4, m_goliathRatioWanted));
}


void Army::UpdateRatios_vsZerg()
{
	int prosVultures = 0;
	int prosTanks = 0;
	int prosGoliaths = 0;

	for (const auto & info : him().UnitTrace())
		switch (info.second.Type())
		{
		case Zerg_Zergling:		prosVultures += 1;								break;
		case Zerg_Hydralisk:	prosVultures += 2; ++prosTanks;					break;
		case Zerg_Mutalisk:										++prosGoliaths;	break;
		case Zerg_Lurker:		prosVultures += 2; ++prosTanks;					break;
		case Zerg_Defiler:		prosVultures += 2; ++prosTanks;					break;
		case Zerg_Guardian:										++prosGoliaths;	break;
		case Zerg_Overlord:														break;
		case Zerg_Queen:														break;
		case Zerg_Ultralisk:	prosVultures += 6;								break;
		case Zerg_Devourer:										++prosGoliaths;	break;
		case Zerg_Scourge:										++prosGoliaths;	break;
		case Zerg_Sunken_Colony:					 ++prosTanks;				break;
		}

	prosVultures /= 2;

	m_tankRatioWanted = int(0.5 + 10 * prosTanks / double(prosVultures + prosTanks + prosGoliaths));
	m_tankRatioWanted = max(1, min(3, m_tankRatioWanted));

	m_goliathRatioWanted = int(0.5 + 10 * prosGoliaths / double(prosVultures + prosTanks + prosGoliaths));
	m_goliathRatioWanted = max(0, min(4, m_goliathRatioWanted));
}


void Army::UpdateRatios()
{
	static bool started = false;
	
	if (!me().Buildings(Terran_Armory).empty() ||
		(him().Race() != Races::Zerg) && (me().Units(Terran_Siege_Tank_Tank_Mode).size() >= 1))
		started = true;


	if (started)
	{
		if		(him().Race() == Races::Terran)  UpdateRatios_vsTerran();
		else if (him().Race() == Races::Protoss) UpdateRatios_vsProtoss();
		else if (him().Race() == Races::Zerg)    UpdateRatios_vsZerg();
	}
}


} // namespace iron



