#include "stdafx.h"
#include "../../../../PerceptualState/PerceptualState.h"
#include "../Layer3Script.h"
#include "../L3BaseLocation.h"
#include "ManageResourcesProcess.h"

// Replaced the const bool by a #define to fix the warning... might consider using a policy instead
#define MOVE_PROBES_AWAY_IF_UNDER_ATTACK
const size_t MAX_TOTAL_PROBES = 70;

void ManageResourcesProcess::onStart()
{
  Layer3Process::onStart();

  std::pair<Map::base_location_iterator, Map::base_location_iterator> baseLocations = Spar->getPerceptualState()->getMap()->getBaseLocations();
  for (; baseLocations.first != baseLocations.second; ++baseLocations.first)
  {
    // Why use a dynamic_cast if you never test the return value? SC: to offset the address if necessary for multiple inheritance
    dynamic_cast<L3BaseLocation*>(*baseLocations.first)->instantiateCollectActions(this);
  }

  // Set up mining probes at main base
  L3BaseLocation* mainBase = Layer3Script::getInstance().getMainBase();
  const std::set<BWAPI::Unit*>& units=BWAPI::Broodwar->self()->getUnits();
  for(std::set<BWAPI::Unit*>::const_iterator i = units.begin(); i != units.end(); i++) {
    if((*i)->getType() == BWAPI::UnitTypes::Protoss_Probe) {
      mainBase->getCollectMinerals()->addUnit(*i);
    }
  }
}

void ManageResourcesProcess::run()
{
  // TODO: island locations... or situations where ground transfer is not feasible
  // TODO: depleted geysers
  // TODO: destroyed nexus or mineral line under attack...

  const ResourcesSaturation saturation = Spar->getPerceptualState()->getResourceManager()->calculateProbesSaturationRatio();
  const double probesPerCrystalRatio = double(saturation.getNbProbesOnMinerals()) / double(saturation.getNbCrystals());

  const size_t totalNbProbes = saturation.getTotalNbProbes();
  const size_t targetNbProbesMinerals = size_t(totalNbProbes * m_targetResourcesAcquisitionPercentage.getPctMinerals() + 0.5);
  const size_t targetNbProbesGas = totalNbProbes - targetNbProbesMinerals;
  size_t probesTransfertGas = targetNbProbesGas - saturation.getNbProbesOnGas();

  const bool stopRecruitingProbes = (saturation.getMineralsSaturationRatio() >= 1.0 && saturation.getGasSaturationRatio() >= 1.0)
                                 || (saturation.getTotalNbProbes() >= MAX_TOTAL_PROBES);

  struct BaseManagement
  {
    BaseManagement(L3BaseLocation* p_baseLocation, int p_probesDiff, bool p_isUnderAttack)
      : baseLocation(p_baseLocation)
      , probesDiff(p_probesDiff)
      , isUnderAttack(p_isUnderAttack)
    {}
    L3BaseLocation* baseLocation;
    int probesDiff;
    bool isUnderAttack;
  };
  std::list<BaseManagement> basesSaturation;

  //
  // "Sort" mining bases by probes saturation
  //
  std::pair<Map::base_location_iterator, Map::base_location_iterator> baseLocations = Spar->getPerceptualState()->getMap()->getBaseLocations();
  for (Map::base_location_iterator it = baseLocations.first; it != baseLocations.second; ++it)
  {
    L3BaseLocation* baseLocation = dynamic_cast<L3BaseLocation*>(*it);
    if (baseLocation->getNbCompletedOwnUnits(BWAPI::UnitTypes::Protoss_Nexus) > 0)
    {
      baseLocation->executeCollectResources(this);

      if (stopRecruitingProbes)
        baseLocation->cancelRecruitProbes();

      bool isUnderAttack = baseLocation->hasEnemyArmyUnits();
      if (isUnderAttack)
      {
#ifdef MOVE_PROBES_AWAY_IF_UNDER_ATTACK
          baseLocation->cancelRecruitProbes();

          // Ghetto fix: transfer gas probes to minerals to send them elsewhere
          GroupTacticalAction::transfer(baseLocation->getCollectGas(), baseLocation->getCollectMinerals(), baseLocation->getCollectGas()->getGroup()->size());

          basesSaturation.insert(basesSaturation.end(), BaseManagement(baseLocation, baseLocation->getCollectMinerals()->getGroup()->size(), isUnderAttack));
#endif
      }
      else
      {
        if (!stopRecruitingProbes)
          baseLocation->executeRecruitProbes(this);

        const int probesDiff = baseLocation->getNbProbesOnMinerals() - int(std::ceil(probesPerCrystalRatio * double(baseLocation->getLastKnownNbCrystalsLeft())));
        if (probesDiff > 0)
          basesSaturation.insert(basesSaturation.begin(), BaseManagement(baseLocation, probesDiff, isUnderAttack));
        else
          basesSaturation.insert(basesSaturation.end(), BaseManagement(baseLocation, probesDiff, isUnderAttack));
      }
    }
  }

  //
  // Transfer probes to gas or build assimilators
  //
  typedef general_iterator<std::list<BaseManagement>::iterator, std::list<BaseManagement>::reverse_iterator> list_iterator;
  list_iterator it = probesTransfertGas >= 0 ? list_iterator(basesSaturation.begin()) : list_iterator(basesSaturation.rbegin());
  list_iterator end = probesTransfertGas >= 0 ? list_iterator(basesSaturation.end()) : list_iterator(basesSaturation.rend());
  for (; probesTransfertGas != 0 && it != end; ++it)
  {
    BaseManagement& overSaturation = *it;
    if (!overSaturation.isUnderAttack)
    {
      const size_t probesOnGas = overSaturation.baseLocation->getNbProbesOnGas();
      const size_t maxProbesOnGas = overSaturation.baseLocation->probesGasSaturation();
      const size_t toTransferToGas = std::min(std::min(overSaturation.baseLocation->getNbProbesOnMinerals(), maxProbesOnGas - probesOnGas), probesTransfertGas);

      const size_t nbGeysersNeeded = size_t(std::ceil((probesOnGas + toTransferToGas) / 3.));
      for (size_t n = overSaturation.baseLocation->getNbOwnUnits(BWAPI::UnitTypes::Protoss_Assimilator); n < nbGeysersNeeded; ++n)
      {
        overSaturation.baseLocation->buildAssimilator(this);
      }

      const size_t maxCurrentProbesOnGas = 3 * overSaturation.baseLocation->getNbCompletedOwnUnits(BWAPI::UnitTypes::Protoss_Assimilator);
      const size_t toTransferCurrently = std::min(toTransferToGas, maxCurrentProbesOnGas - probesOnGas);
      GroupTacticalAction::transfer(overSaturation.baseLocation->getCollectMinerals(), overSaturation.baseLocation->getCollectGas(), toTransferCurrently);

      overSaturation.probesDiff -= toTransferToGas;

      probesTransfertGas -= toTransferToGas;
    }
  }

  //
  // Transfer (minerals) probes to under-saturated bases
  //
  for (std::list<BaseManagement>::iterator it = basesSaturation.begin();
       it != basesSaturation.end();
       ++it)
  {
    BaseManagement& overSaturation = *it;

    if (overSaturation.probesDiff > 1) // > 1 (instead of > 0) to avoid "flickering" probes between bases
    {
      for (std::list<BaseManagement>::reverse_iterator rit = basesSaturation.rbegin();
           overSaturation.probesDiff > 0 && rit != basesSaturation.rend();
           ++rit)
      {
        BaseManagement& underSaturation = *rit;
        if (underSaturation.probesDiff > overSaturation.probesDiff)
          break;
        const int nbToTransfer = std::min(overSaturation.probesDiff, -underSaturation.probesDiff);
        const int nbTransferred = GroupTacticalAction::transfer(overSaturation.baseLocation->getCollectMinerals(), underSaturation.baseLocation->getCollectMinerals(), nbToTransfer);
        overSaturation.probesDiff -= nbTransferred;
        underSaturation.probesDiff += nbTransferred;
      }
    }
  }
}

void ManageResourcesProcess::updateResourcesAcquisitionPercentage(ResourcesAcquisitionPercentage targetResourcesAcquisitionPercentage)
{
  m_targetResourcesAcquisitionPercentage = targetResourcesAcquisitionPercentage;
}