#include "Common.h"
#include "UnitData.h"

using namespace MyBot;

UnitData::UnitData() 
	: mineralsLost(0)
	, gasLost(0)
{
	int maxTypeID(0);
	for (const BWAPI::UnitType & t : BWAPI::UnitTypes::allUnitTypes())
	{
		maxTypeID = maxTypeID > t.getID() ? maxTypeID : t.getID();
	}

	numDeadUnits	    = std::vector<int>(maxTypeID + 1, 0);
	numCreatedUnits = std::vector<int>(maxTypeID + 1, 0);
	numUnits = std::vector<int>(maxTypeID + 1, 0);
}

void UnitData::updateUnitInfo(BWAPI::Unit unit)
{
	if (!unit) { return; }

    bool firstSeen = false;
    auto & it = unitAndUnitInfoMap.find(unit);
    if (it == unitAndUnitInfoMap.end())
    {
        firstSeen = true;
        unitAndUnitInfoMap[unit] = UnitInfo();
    }
    
	UnitInfo & ui   = unitAndUnitInfoMap[unit];
    ui.unit         = unit;
    ui.player       = unit->getPlayer();
	ui.lastPosition = unit->getPosition();
	ui.lastHealth   = unit->getHitPoints();
    ui.lastShields  = unit->getShields();
	ui.unitID       = unit->getID();
	ui.type         = unit->getType();
    ui.completed    = unit->isCompleted();

	// 여기서부터 추가
	ui.lastPositionRecordFrame = BWAPI::Broodwar->getFrameCount();

	if (unit->exists()) {
		ui.isLastPositionValid = true;
	}

	// ui.positionBefore 저장
	if (ui.positionBeforeRecordFrame + 10 <= BWAPI::Broodwar->getFrameCount()
		|| ui.positionBefore == BWAPI::Positions::None)
	{
		ui.positionBefore = ui.lastPosition;
		ui.positionBeforeRecordFrame = BWAPI::Broodwar->getFrameCount();
	}

    if (firstSeen)
    {
		numCreatedUnits[unit->getType().getID()]++;
		numUnits[unit->getType().getID()]++;

		if (unit->getType() == BWAPI::UnitTypes::Zerg_Hydralisk) {
			std::cout << "Zerg_Hydralisk added to " << numUnits[unit->getType().getID()] << std::endl;
		}
	}
}

void UnitData::removeUnit(BWAPI::Unit unit)
{
	if (!unit) { return; }

	mineralsLost += unit->getType().mineralPrice();
	gasLost += unit->getType().gasPrice();
	numUnits[unit->getType().getID()]--;
	numDeadUnits[unit->getType().getID()]++;
		
	unitAndUnitInfoMap.erase(unit);
}

void UnitData::removeBadUnits()
{
	for (auto iter(unitAndUnitInfoMap.begin()); iter != unitAndUnitInfoMap.end();)
	{
		// 추가 수정
		BWAPI::Unit unit = iter->first;
		UnitInfo & unitInfo = iter->second;

		if (BWAPI::Broodwar->isVisible(BWAPI::TilePosition(unitInfo.lastPosition))) {
			if (unit->exists() == false) {
				unitInfo.isLastPositionValid = false;
			}
		}

		if (isBadUnitInfo(iter->second))
		{
			numUnits[iter->second.type.getID()]--;

			iter = unitAndUnitInfoMap.erase(iter);
		}
		else
		{
			iter++;
		}
	}
}

const bool UnitData::isBadUnitInfo(const UnitInfo & ui) const
{
    if (!ui.unit)
    {
        return false;
    }

	// Cull away any refineries/assimilators/extractors that were destroyed and reverted to vespene geysers
	if (ui.unit->getType() == BWAPI::UnitTypes::Resource_Vespene_Geyser)
	{ 
		return true;
	}

	// If the unit is a building and we can currently see its position and it is not there
	// 추가 : 건물이 있었던 lastPosition 이 visible 하지만, 거기에 건물이 없는 경우는
	// 테란 건물이 이동했거나, (여전히 존재하지만, 어디있는지 알 수 없음)
	// 테란 건물이 에너지가 다 닳아서 파괴되었거나 (없어짐)
	// 건물 건설을 취소했거나, (없어짐)
	// 건물이 파괴된 경우 (플레이어가 스스로 건물을 파괴한 경우. 없어짐) 이다.
	// 이동이 가능한 테란 건물의 경우, 어딘가에 존재할 것이라고 생각할 수 있으므로, 없애지않는다
	if (ui.type.isBuilding() && BWAPI::Broodwar->isVisible(ui.lastPosition.x/32, ui.lastPosition.y/32) && !ui.unit->isVisible())
	{
		if (ui.type.isFlyingBuilding()) {
			return false;
		}
		else  {
			return true;
		}
	}

	return false;
}

int UnitData::getGasLost() const 
{ 
    return gasLost; 
}

int UnitData::getMineralsLost() const 
{ 
    return mineralsLost; 
}

int UnitData::getNumUnits(BWAPI::UnitType t) const 
{ 
	//return numUnits[t.getID()]; 

	// 추가 수정
	int count = 0;
	for (auto & it : unitAndUnitInfoMap) {
		if (it.second.type == t) {
			count++;
		}
	}

	return count;
}

int UnitData::getNumDeadUnits(BWAPI::UnitType t) const 
{ 
    return numDeadUnits[t.getID()]; 
}

int UnitData::getNumCreatedUnits(BWAPI::UnitType t) const
{
	return numCreatedUnits[t.getID()];
}

std::map<BWAPI::Unit, UnitInfo> & UnitData::getUnitAndUnitInfoMap()
{
	return unitAndUnitInfoMap;
}