#include "stdafx.h"
#include "Threat.h"
#include "InfluenceMap.h"
#include "../../SparAIContainer.h"
#include "../PerceptualState.h"
#include "../UnitRecord/UnitRecord.h"
#include "../../Utils/BresenhamCircle.h"
#include <numeric>

// TODO: remove influence of inaccessible enemy units after a certain time?

// TODO: radius should consider dimension of units

// TODO: spider mines
// TODO: bunkers
// TODO: consider minRange() (i.e., Terran Siege Tank cannot attack close to itself)
// TODO: consider unit upgrades (including speed/range)
// TODO: consider type of damage? (explosive/concussive/etc)

/**
 * @brief Functor for Bresenham's circle drawing algorithm, used to update
 * the influence map at certain positions.
 */
struct UpdateCurrentThreatFunctor : std::binary_function<int,int,void>
{
  /**
   * @brief Builds a functor object.
   * @param map a pointer to the influence map to update.
   * @param position the position around which the influence will be updated.
   * @param groundDamage the amount of ground influence to update the map with.
   * @param airDamage the amount of air influence to update the map with.
   */
  UpdateCurrentThreatFunctor(InfluenceMap<Threat>& map, BWAPI::Position position, double groundDamage, double airDamage)
    : _map(map)
    , _x(position.x() / TILE_SIZE), _y(position.y() / TILE_SIZE)
    , _ground(groundDamage)
    , _air(airDamage)
  {
    assert(static_cast<int>(map.getWidth())  == BWAPI::Broodwar->mapWidth()
        && static_cast<int>(map.getHeight()) == BWAPI::Broodwar->mapHeight());
  }

  /**
   * @brief Updates the influence map at the given position.
   * @param x the x coordinate of the tile (relative to the center of the
   * range circle) whose influence will be updated.
   * @param y the y coordinate of the tile (relative to the center of the
   * range circle) whose influence will be updated.
   */
  result_type operator()(first_argument_type x, second_argument_type y)
  {
    x += _x;
    y += _y;

    if (_map.isValid(x, y))
    {
      Threat influence = _map.getData(x, y);
      influence.ground += _ground;
      influence.air += _air;
      _map.setInfluence(x, y, influence);
    }
  }

private:
  InfluenceMap<Threat>& _map;
  const int _x, _y;
  const double _ground, _air;

  // Prevents using the assignement operator.
  UpdateCurrentThreatFunctor& operator=(const UpdateCurrentThreatFunctor&);
};

struct UpdateFutureThreatFunctor : std::binary_function<int,int,void>
{
  /**
   * @brief Builds a functor object.
   * @param map a pointer to the influence map to update.
   * @param position the position around which the influence will be updated.
   * @param futureGroundDamage the amount of ground influence to update the map with.
   * @param futureAirDamage the amount of air influence to update the map with.
   */
  UpdateFutureThreatFunctor(InfluenceMap<Threat>& map, BWAPI::Position position, double futureGroundDamage, double futureAirDamage)
    : _map(map)
    , _x(position.x() / TILE_SIZE), _y(position.y() / TILE_SIZE)
    , _futureGround(futureGroundDamage)
    , _futureAir(futureAirDamage)
  {
    assert(static_cast<int>(map.getWidth())  == BWAPI::Broodwar->mapWidth()
        && static_cast<int>(map.getHeight()) == BWAPI::Broodwar->mapHeight());
  }

  /**
   * @brief Updates the influence map at the given position.
   * @param x the x coordinate of the tile (relative to the center of the
   * range circle) whose influence will be updated.
   * @param y the y coordinate of the tile (relative to the center of the
   * range circle) whose influence will be updated.
   */
  result_type operator()(first_argument_type x, second_argument_type y)
  {
    x += _x;
    y += _y;

    if (_map.isValid(x, y))
    {
      Threat influence = _map.getData(x, y);
      influence.futureGround += _futureGround;
      influence.futureAir += _futureAir;
      _map.setInfluence(x, y, influence);
    }
  }

private:
  InfluenceMap<Threat>& _map;
  const int _x, _y;
  const double _futureGround, _futureAir;

  // Prevents using the assignement operator.
  UpdateFutureThreatFunctor& operator=(const UpdateFutureThreatFunctor&);
};

struct Threat::UnitAttackInfo
{
  UnitAttackInfo(const UnitRecord* unit);

  BWAPI::UnitType _unitType;
  BWAPI::WeaponType _groundWeapon;
  BWAPI::WeaponType _airWeapon;
  
  int _groundDamage;
  int _airDamage;

  int _groundRange;
  int _airRange;

  const UnitRecord::PredictedPaths& _predictedPos;
  double _topSpeed;

  BWAPI::Position _pos;

private:
  UnitAttackInfo & operator=(const UnitAttackInfo &);
};

Threat::UnitAttackInfo::UnitAttackInfo(const UnitRecord* unit)
  : _unitType(unit->getType())
  , _groundWeapon(_unitType.groundWeapon())
  , _airWeapon(_unitType.airWeapon())
  , _topSpeed(_unitType.topSpeed())
  , _predictedPos(unit->getPredictedPaths())
  , _pos(unit->getPosition())
  , _groundDamage(_groundWeapon.damageAmount() + _groundWeapon.damageBonus())
  , _airDamage(_airWeapon.damageAmount() + _airWeapon.damageBonus())
  , _groundRange(_groundWeapon.maxRange())
  , _airRange(_airWeapon.maxRange())
{
  //
  // Upgrades worst case
  //
  if (unit->getPlayer() == BWAPI::Broodwar->enemy() && unit->getType() == BWAPI::UnitTypes::Zerg_Hydralisk)
  {
    _groundRange += 1;
    _airRange += 1;
  }
  if (unit->getPlayer() == BWAPI::Broodwar->enemy() && unit->getType() == BWAPI::UnitTypes::Protoss_Dragoon)
  {
    _groundRange += 2;
    _airRange += 2;
  }
  if (unit->getPlayer() == BWAPI::Broodwar->enemy() && unit->getType() == BWAPI::UnitTypes::Terran_Marine)
  {
    _groundRange += 1;
    _airRange += 1;
  }
  if (unit->getPlayer() == BWAPI::Broodwar->enemy() && unit->getType() == BWAPI::UnitTypes::Terran_Goliath)
  {
    _airRange += 3;
  }
}

void Threat::update(InfluenceMap<Threat>& map)
{
  map.reset();

  std::pair<Map::unit_const_iterator, Map::unit_const_iterator> units;
  if (map.getPlayer() == BWAPI::Broodwar->self())
  {
    units = Spar->getPerceptualState()->getMap()->getOwnUnits();
  }
  else if (map.getPlayer() == BWAPI::Broodwar->enemy())
  {
    units = Spar->getPerceptualState()->getMap()->getEnemyUnits();
  }
  update(map, units.first, units.second);
}

void Threat::update(InfluenceMap<Threat>& map, const UnitRecord* unit)
{
  // TODO: interpolate based on UnitRecord::PredictedPath::getPositionAtFrame(...)

  UnitAttackInfo attackInfo(unit);
  
  if ((unit->getType().isWorker() && (!unit->isAccessible() || !unit->getUnit()->isAttacking())) ||
      (attackInfo._groundDamage == 0 && attackInfo._airDamage == 0))
    return;

  // TODO: would need a more precise update function (that takes the position of the center point inside the tile?)
  const int groundRange = (int)(double(attackInfo._groundRange) / TILE_SIZE) + 2;
  const int airRange = (int)(double(attackInfo._airRange) / TILE_SIZE) + 2;

  //
  // Current influence
  //
  if (attackInfo._groundDamage > 0)
  {
    UpdateCurrentThreatFunctor fCurrentGroundStatic = UpdateCurrentThreatFunctor(map, attackInfo._pos, static_cast<double>(attackInfo._groundDamage), 0);
    BresenhamCircle(fCurrentGroundStatic, groundRange, true);
  }
  if (attackInfo._airDamage > 0)
  {
    UpdateCurrentThreatFunctor fCurrentAirStatic = UpdateCurrentThreatFunctor(map, attackInfo._pos, 0, static_cast<double>(attackInfo._airDamage));
    BresenhamCircle(fCurrentAirStatic, airRange, true);
  }

  //
  // Future influence
  //

  // Consider the next (unweighted) NB_NEXT_FRAMES_THREAT frames for future influence.
  static const double NB_NEXT_FRAMES_THREAT_DBL = double(NB_NEXT_FRAMES_THREAT);

  const int currentFrame = BWAPI::Broodwar->getFrameCount();

  // Sum of probabilities == 1
  assert(std::accumulate(boost::transform_iterator<extract_pair_first_const<UnitRecord::PredictedPaths::value_type>, UnitRecord::PredictedPaths::const_iterator>(attackInfo._predictedPos.begin()),
                         boost::transform_iterator<extract_pair_first_const<UnitRecord::PredictedPaths::value_type>, UnitRecord::PredictedPaths::const_iterator>(attackInfo._predictedPos.end()),
                         0.) <= 1.00001);
  double probabilityLeft = 1.0;

  for (UnitRecord::PredictedPaths::const_iterator itt = attackInfo._predictedPos.begin();
       itt != attackInfo._predictedPos.end();
       ++itt)
  {
    const double probability = itt->first;
    const UnitRecord::PredictedPath& predictedPath = itt->second;
    const std::vector<std::pair<int, BWAPI::Position>>& positions = predictedPath.getPositions();

    // Frame numbers in ascending order
    assert(sequenceRespectsComparator(boost::transform_iterator<extract_pair_first_const<UnitRecord::PredictedPath::const_iterator::value_type>, UnitRecord::PredictedPath::const_iterator>(positions.begin()),
                                      boost::transform_iterator<extract_pair_first_const<UnitRecord::PredictedPath::const_iterator::value_type>, UnitRecord::PredictedPath::const_iterator>(positions.end()),
                                      std::less<UnitRecord::PredictedPath::const_iterator::value_type::first_type>()));

    // Skip empty or irrelevant (last noFrame < currentFrame) path
    if (positions.empty() || positions.rbegin()->first < currentFrame)
      continue;

    probabilityLeft -= probability;

    BWAPI::Position pos = attackInfo._pos; // use current position initially if we can't find anything else
    UnitRecord::PredictedPath::const_iterator curPos = positions.begin();
    while (curPos != positions.end() && curPos->first <= (currentFrame + 1))
    {
      if (!unit->isAccessible())
        pos = curPos->second;
      ++curPos;
    }

    int currentCalcFrame = currentFrame;
    // We consider NB_NEXT_FRAMES_THREAT future frames
    while (currentCalcFrame < currentFrame + NB_NEXT_FRAMES_THREAT)
    {
      int nextCalcFrame;
      if (curPos == positions.end())
      {
        //
        // Since no other position was specified, we assume the unit stays immobile;
        // instead we could perform interpolation on the path.
        //
        nextCalcFrame = currentFrame + NB_NEXT_FRAMES_THREAT;
      }
      else
      {
        nextCalcFrame = curPos->first;
      }

      assert(nextCalcFrame > currentCalcFrame);
      const double weight = probability * double(nextCalcFrame - currentCalcFrame) / NB_NEXT_FRAMES_THREAT_DBL;

      if (attackInfo._groundDamage > 0)
      {
        UpdateFutureThreatFunctor fGroundStatic = UpdateFutureThreatFunctor(map, pos, attackInfo._groundDamage * weight, 0);
	      BresenhamCircle(fGroundStatic, groundRange, true);
      }
      if (attackInfo._airDamage > 0)
      {
        UpdateFutureThreatFunctor fAirStatic = UpdateFutureThreatFunctor(map, pos, 0, attackInfo._airDamage * weight);
        BresenhamCircle(fAirStatic, airRange, true);
      }

      currentCalcFrame = nextCalcFrame;
      if (curPos == positions.end())
      {
        assert(currentCalcFrame == currentFrame + NB_NEXT_FRAMES_THREAT);
      }
      else
      {
        pos = curPos->second;
        ++curPos;
        if (curPos != positions.end())
        {
          if (curPos->first > currentFrame + NB_NEXT_FRAMES_THREAT)
            curPos = positions.end();
        }
      }
    }
  }
  if (probabilityLeft > 0.)
  {
    // Uniform propagation from current position

    const int additionalRange = (int)(ceil(double(NB_NEXT_FRAMES_THREAT * attackInfo._topSpeed) / double(TILE_SIZE)));

    if (attackInfo._groundDamage > 0)
    {
      UpdateFutureThreatFunctor fGroundMobile = UpdateFutureThreatFunctor(map, attackInfo._pos, attackInfo._groundDamage * probabilityLeft, 0);
      BresenhamCircle(fGroundMobile, groundRange + additionalRange, true);
    }
    if (attackInfo._airDamage > 0)
    {
      UpdateFutureThreatFunctor fAirMobile = UpdateFutureThreatFunctor(map, attackInfo._pos, 0, attackInfo._airDamage * probabilityLeft);
      BresenhamCircle(fAirMobile, airRange + additionalRange, true);
    }
  }
}

void Threat::display(int x, int y, int w, int h, BWAPI::Player* player) const
{
  // futureAir and futureGround are not displayed

  if ((ground + air) > 0)
  {
    const int val = std::min(static_cast<int>(ground + air) + 50, 0xff);

    BWAPI::Color colour(val, 0, 0);
    if (player == BWAPI::Broodwar->self())
    {
      colour = BWAPI::Color(0, val, 0);
      x++;
      y++;
      w-=2;
      h-=2;
    }

    BWAPI::Broodwar->drawBoxMap(x, y, x+w, y+h, colour, false);
  }
}

