#include "stdafx.h"
#include "Utils/Utils.h"
#include "BWAPI_Addon.h"

double BWAPI::detail::mapDiagonalLength;
DWORD BWAPI::detail::mainThreadId;

void BWAPI::detail::AdditionalUnitTypeData::addUnitType(BWAPI::UnitType unitType)
{
  morphTypes.insert(unitType);
}

void addDerivedMorphTypes(std::set<BWAPI::UnitType>& derivedMorphTypes, BWAPI::UnitType unitType)
{
  const std::set<BWAPI::UnitType>& morphTypes = BWAPI::detail::singleton.additionalUnitTypeData[unitType].morphTypes;
  for (std::set<BWAPI::UnitType>::const_iterator it = morphTypes.begin(); it != morphTypes.end(); ++it)
  {
    BWAPI::UnitType derivedMorphType = *it;
    if (derivedMorphTypes.insert(derivedMorphType).second)
    {
      addDerivedMorphTypes(derivedMorphTypes, derivedMorphType);
    }
  }
}

BWAPI::detail::Singleton::Singleton()
  : unitDiscoverEvent(*this)
  , unitEvadeEvent(*this)
  , unitShowEvent(*this)
  , unitHideEvent(*this)
  , unitCreateEvent(*this)
  , unitDestroyEvent(*this)
  , unitMorphEvent(*this)
  , unitRenegadeEvent(*this)
{
}

void BWAPI::Addon_Init()
{
  BWAPI::detail::mainThreadId = ::GetCurrentThreadId();

  BWAPI::detail::Singleton& singleton = BWAPI::detail::singleton;

  BWAPI::detail::mapDiagonalLength = sqrt(static_cast<double>(
    BWAPI::Broodwar->mapWidth() * BWAPI::Broodwar->mapWidth() + 
    BWAPI::Broodwar->mapHeight() * BWAPI::Broodwar->mapHeight())) * TILE_SIZE;

  // Note that neutral units are not included here, but feel free to add them.

  // TODO: formalize as "actions"?: training / morphing -> integrate cost (resources,etc), dependencies, time, units produced?

  //
  // Add immediate morph transitions
  //

  // TODO: add morph types when cancelling something? (like Lurker egg or Devourer cocoon)

  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Resource_Vespene_Geyser].addUnitType(BWAPI::UnitTypes::Protoss_Assimilator);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Resource_Vespene_Geyser].addUnitType(BWAPI::UnitTypes::Terran_Refinery);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Resource_Vespene_Geyser].addUnitType(BWAPI::UnitTypes::Zerg_Extractor);

  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Protoss_Assimilator].addUnitType(BWAPI::UnitTypes::Resource_Vespene_Geyser);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Protoss_Dark_Templar].addUnitType(BWAPI::UnitTypes::Protoss_Dark_Archon);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Protoss_High_Templar].addUnitType(BWAPI::UnitTypes::Protoss_Archon);

  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Terran_Command_Center].addUnitType(BWAPI::UnitTypes::Zerg_Infested_Command_Center);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Terran_Refinery].addUnitType(BWAPI::UnitTypes::Resource_Vespene_Geyser);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Terran_Siege_Tank_Siege_Mode].addUnitType(BWAPI::UnitTypes::Terran_Siege_Tank_Tank_Mode);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Terran_Siege_Tank_Tank_Mode].addUnitType(BWAPI::UnitTypes::Terran_Siege_Tank_Siege_Mode);
  
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Cocoon].addUnitType(BWAPI::UnitTypes::Zerg_Devourer);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Cocoon].addUnitType(BWAPI::UnitTypes::Zerg_Guardian);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Creep_Colony].addUnitType(BWAPI::UnitTypes::Zerg_Spore_Colony);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Creep_Colony].addUnitType(BWAPI::UnitTypes::Zerg_Sunken_Colony);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Drone].addUnitType(BWAPI::UnitTypes::Zerg_Creep_Colony);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Drone].addUnitType(BWAPI::UnitTypes::Zerg_Defiler_Mound);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Drone].addUnitType(BWAPI::UnitTypes::Zerg_Evolution_Chamber);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Drone].addUnitType(BWAPI::UnitTypes::Zerg_Extractor);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Drone].addUnitType(BWAPI::UnitTypes::Zerg_Hatchery);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Drone].addUnitType(BWAPI::UnitTypes::Zerg_Hydralisk_Den);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Drone].addUnitType(BWAPI::UnitTypes::Zerg_Nydus_Canal);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Drone].addUnitType(BWAPI::UnitTypes::Zerg_Queens_Nest);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Drone].addUnitType(BWAPI::UnitTypes::Zerg_Spawning_Pool);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Drone].addUnitType(BWAPI::UnitTypes::Zerg_Spire);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Drone].addUnitType(BWAPI::UnitTypes::Zerg_Ultralisk_Cavern);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Egg].addUnitType(BWAPI::UnitTypes::Zerg_Defiler);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Egg].addUnitType(BWAPI::UnitTypes::Zerg_Drone);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Egg].addUnitType(BWAPI::UnitTypes::Zerg_Hydralisk);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Egg].addUnitType(BWAPI::UnitTypes::Zerg_Mutalisk);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Egg].addUnitType(BWAPI::UnitTypes::Zerg_Overlord);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Egg].addUnitType(BWAPI::UnitTypes::Zerg_Queen);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Egg].addUnitType(BWAPI::UnitTypes::Zerg_Scourge); // Two scourges in one egg
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Egg].addUnitType(BWAPI::UnitTypes::Zerg_Ultralisk);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Egg].addUnitType(BWAPI::UnitTypes::Zerg_Zergling); // Two zerglings in one egg
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Extractor].addUnitType(BWAPI::UnitTypes::Resource_Vespene_Geyser);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Hatchery].addUnitType(BWAPI::UnitTypes::Zerg_Lair);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Hydralisk].addUnitType(BWAPI::UnitTypes::Zerg_Lurker_Egg);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Lair].addUnitType(BWAPI::UnitTypes::Zerg_Hive);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Larva].addUnitType(BWAPI::UnitTypes::Zerg_Egg);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Lurker_Egg].addUnitType(BWAPI::UnitTypes::Zerg_Lurker);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Mutalisk].addUnitType(BWAPI::UnitTypes::Zerg_Cocoon);
  singleton.additionalUnitTypeData[BWAPI::UnitTypes::Zerg_Spire].addUnitType(BWAPI::UnitTypes::Zerg_Greater_Spire);

  // Zerg broodlings
  const std::set<BWAPI::UnitType>& allUnitTypes = BWAPI::UnitTypes::allUnitTypes();
  for (std::set<BWAPI::UnitType>::const_iterator it = allUnitTypes.begin(); it != allUnitTypes.end(); ++it)
  {
    BWAPI::UnitType unitType = *it;
    if (!unitType.isBuilding() && unitType.isFlyer() &&
      (unitType != BWAPI::UnitTypes::Protoss_Probe && unitType != BWAPI::UnitTypes::Protoss_Reaver &&
       unitType != BWAPI::UnitTypes::Protoss_Archon && unitType != BWAPI::UnitTypes::Protoss_Dark_Archon))
    {
      singleton.additionalUnitTypeData[unitType].addUnitType(BWAPI::UnitTypes::Zerg_Broodling); // Two broodlings per successful spell
    }
  }

  //
  // Compute all possible "derived" morph transitions (i.e. Mutalisk -> Devourer through Cocoon)
  //

  for (size_t i = 0; i < array_size(singleton.additionalUnitTypeData); ++i)
  {
    addDerivedMorphTypes(singleton.additionalUnitTypeData[i].allDerivedMorphTypes, BWAPI::UnitType(i));
    for (std::set<BWAPI::UnitType>::const_iterator it = singleton.additionalUnitTypeData[i].allDerivedMorphTypes.begin();
         it != singleton.additionalUnitTypeData[i].allDerivedMorphTypes.end();
         ++it)
    {
      BWAPI::UnitType derivedMorphType = *it;
      if (isMobile(derivedMorphType))
        singleton.additionalUnitTypeData[i].allDerivedMobileMorphTypes.insert(derivedMorphType);
    }
  }
}