#include "stdafx.h"
#include "UpdateTargetArmyAndTechsProcess.h"
#include "../Layer4Heuristic.h"
#include "../../../Layer3Tactical/Interface/ArmyComposition.h"
#include "../../../../PerceptualState/World/CurrentWorld.h"
#include "../../../../PerceptualState/PerceptualState.h"

//
// Warning: ugly ad-hoc code follows
//

// TODO: build army before teching straight away!

static const int PSI_TO_TRIGGER_FORGE = 40 * 2;
static const int PSI_TO_TRIGGER_OBSERVERS = 60 * 2;

static const int PSI_ZERG_TO_TRIGGER_TEMPLARS = 70 * 2;

ArmyComposition UpdateTargetArmyAndTechsProcess::estimateBestCounter() const
{
  const Map* map = Spar->getPerceptualState()->getMap();

  ArmyComposition counter;

  // Input: own army, current resources, enemy army

  BWAPI::Race enemyRace = BWAPI::Broodwar->enemy()->getRace();

  if (enemyRace == BWAPI::Races::Protoss)
  {
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_Zealot, 3);
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_Dragoon, 2);
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_High_Templar, 1);
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_Dark_Templar, 1);
  }
  else if (enemyRace == BWAPI::Races::Terran)
  {
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_Zealot, 2);
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_Dragoon, 3);
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_High_Templar, 1);
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_Dark_Templar, 1);
  }
  else if (enemyRace == BWAPI::Races::Zerg)
  {
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_Zealot, 3);
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_Dragoon, 1);
    //if (CurrentWorld::TemplarArchives->isCompleted() || BWAPI::Broodwar->self()->supplyUsed() >= PSI_ZERG_TO_TRIGGER_TEMPLARS)
    //{
      counter.addUnlimited(BWAPI::UnitTypes::Protoss_High_Templar, 1);
    //}
    if (map->getNbOwnUnits(BWAPI::UnitTypes::Protoss_Corsair) < 4)
      counter.addUnlimited(BWAPI::UnitTypes::Protoss_Corsair, 1);
  }
  else
  {
    // 1a2a3a for the win!
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_Zealot, 1);
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_Dragoon, 1);
  }

  if((double)BWAPI::Broodwar->self()->minerals() / (double)BWAPI::Broodwar->self()->gas() < 1 && BWAPI::Broodwar->self()->gas() > 500)
  {
    counter.addUnlimited(BWAPI::UnitTypes::Protoss_High_Templar, 1);
  }

  if (CurrentWorld::Observer->isCreatable())
  {
    const size_t nbObservers = map->getNbOwnUnits(BWAPI::UnitTypes::Protoss_Observer);
    if (nbObservers < 4)
    {
      counter.addLimited(BWAPI::UnitTypes::Protoss_Observer, 4 - nbObservers);
    }
  }

  return counter;
}

TechSet UpdateTargetArmyAndTechsProcess::getTargetTechs(const ArmyComposition& targetArmyComposition) const
{
  const EnemyArmyComposition* enemyArmyComposition = Spar->getPerceptualState()->getEnemyArmyComposition();

  TechSet techs;

  std::map<BWAPI::UnitType, size_t> infiniteComposition = targetArmyComposition.getUnlimitedComposition();
  for (std::map<BWAPI::UnitType, size_t>::const_iterator it = infiniteComposition.begin(); it != infiniteComposition.end(); ++it)
  {
    std::pair<BWAPI::UnitType, size_t> units = *it;
    techs.add(CurrentWorld::getUnitTypeObject(units.first), 0);
  }

  if (!CurrentWorld::Observer->isCreatable() &&
      (enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Protoss_Dark_Templar) > 0 || enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Protoss_High_Templar) > 0 || 
       enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Protoss_Archon) > 0 || enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Protoss_Dark_Archon) > 0 || enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Protoss_Templar_Archives) > 0 ||
       enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Protoss_Arbiter) > 0  || enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Protoss_Arbiter_Tribunal) > 0 ||
       enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Terran_Ghost) > 0  || enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Terran_Covert_Ops) > 0 ||
       enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Terran_Wraith) > 0  || (enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Terran_Starport) > 0 && enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Terran_Control_Tower) > 0) ||
       enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Zerg_Lurker) > 0  || (enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Zerg_Lair) > 0 && enemyArmyComposition->getNbEnemyUnits(BWAPI::UnitTypes::Zerg_Hydralisk_Den) > 0) ||
       BWAPI::Broodwar->self()->supplyUsed() >= PSI_TO_TRIGGER_OBSERVERS))
  {
    techs.add(CurrentWorld::Observer, 0);
  }

  if (!CurrentWorld::Forge->isCompleted() &&
      BWAPI::Broodwar->self()->supplyUsed() >= PSI_TO_TRIGGER_FORGE)
  {
    techs.add(CurrentWorld::Forge, 0);
  }

  addImprovementTechs(targetArmyComposition, techs);

  return techs;
}

void UpdateTargetArmyAndTechsProcess::addImprovementTechs(const ArmyComposition&, TechSet& techs) const
{
  const Map* map = Spar->getPerceptualState()->getMap();

  const size_t initialNbTechs = techs.size();

  if (!CurrentWorld::PsionicStorm->isAcquired()
   && CurrentWorld::PsionicStorm->allDependenciesAcquired()
   //&& map->getNbOwnUnits(BWAPI::UnitTypes::Protoss_High_Templar) >= 1
   )
  {
    techs.add(CurrentWorld::PsionicStorm, 0);
  }
  if (!CurrentWorld::SingularityCharge->isAcquired() &&
      CurrentWorld::SingularityCharge->allDependenciesAcquired() &&
      map->getNbOwnUnits(BWAPI::UnitTypes::Protoss_Dragoon) >= 4)
  {
    techs.add(CurrentWorld::SingularityCharge, 0);
  }
  if (!CurrentWorld::LegEnhancements->isAcquired() &&
      CurrentWorld::LegEnhancements->allDependenciesAcquired() &&
      map->getNbOwnUnits(BWAPI::UnitTypes::Protoss_Zealot) >= 4)
  {
    techs.add(CurrentWorld::LegEnhancements, 0);
  }
  if (!CurrentWorld::CarrierCapacity->isAcquired() &&
      CurrentWorld::CarrierCapacity->allDependenciesAcquired() &&
      map->getNbOwnUnits(BWAPI::UnitTypes::Protoss_Carrier) >= 1)
  {
    techs.add(CurrentWorld::CarrierCapacity, 0);
  }

  if (map->getNbCompletedOwnUnits(BWAPI::UnitTypes::Protoss_Forge) > 0)
  {
    size_t maxTechs = (techs.size() == initialNbTechs) ? 3 : 1;
    for (size_t i = 0; maxTechs > 0 && i < array_size(CurrentWorld::GroundWeapons); ++i) {
      if (!CurrentWorld::GroundWeapons[i]->isAcquired() &&
          CurrentWorld::GroundWeapons[i]->allDependenciesAcquired())
      {
        --maxTechs;
        techs.add(CurrentWorld::GroundWeapons[i], 0);
        break;
      }
    }
    for (size_t i = 0; maxTechs > 0 && i < array_size(CurrentWorld::GroundArmor); ++i) {
      if (!CurrentWorld::GroundArmor[i]->isAcquired() &&
          CurrentWorld::GroundArmor[i]->allDependenciesAcquired())
      {
        --maxTechs;
        techs.add(CurrentWorld::GroundArmor[i], 0);
        break;
      }
    }
    for (size_t i = 0; maxTechs > 0 && i < array_size(CurrentWorld::PlasmaShields); ++i) {
      if (!CurrentWorld::PlasmaShields[i]->isAcquired() &&
          CurrentWorld::PlasmaShields[i]->allDependenciesAcquired())
      {
        --maxTechs;
        techs.add(CurrentWorld::PlasmaShields[i], 0);
        break;
      }
    }
  }
}

void UpdateTargetArmyAndTechsProcess::run()
{
  ArmyComposition counter = estimateBestCounter();

  TechSet techs = getTargetTechs(counter);

  Layer4Heuristic::getInstance().getLayer3().setTargetArmyComposition(counter);

  Layer4Heuristic::getInstance().getLayer3().setTargetTechs(techs);
}
