#include "stdafx.h"
#include "ThreatAwareUnitMoveBehavior.h"

#include <BWAPI.h>
#include <boost/bind.hpp>
#include <ctime>
#include "../../../../Pathfinding.h"
#include "../../../../../../../BWAPI_Addon.h"

// TODO: break up processing on several frames (or a separate thread) to correctly determine when no path is available?

ThreatAwareUnitMoveBehavior::ThreatAwareUnitMoveBehavior(const Process* parent, unsigned int initialPriority, BWAPI::Unit* unit,
                    BWAPI::Position destination, const MoveOptions& moveOptions)
  : UnitMoveBehaviorT(parent, initialPriority, unit, destination, moveOptions)
  , UnitMoveBehavior(parent, initialPriority, unit, destination, moveOptions)
  , UnitBehavior(parent, initialPriority, unit)
  , Behavior(parent, initialPriority)
  , FSM(parent)
  , Process(parent)
  , m_gotoState(*this, moveOptions.m_refreshRate, moveOptions.m_moveWeight)
  , m_states(m_gotoState)
{
  assert(m_moveOptions.m_threatAware);
}

///////////////////////////////////////////////////
// Goto state
///////////////////////////////////////////////////

ThreatAwareUnitMoveBehavior::Goto::Goto(ThreatAwareUnitMoveBehavior& behavior, size_t refreshDelay, const double moveWeight)
  : UnitMoveBehaviorT::Goto(behavior, refreshDelay)
  , m_events(m_timerEvent
#ifdef NON_FINAL_VERSION
    , m_frameDebugEvent
#endif
  )
  , m_transitions()
  , m_h(moveWeight)
  , m_updateInProgress(false)
  , m_isRetreatPath(false)
{
  m_currentTarget = m_smoothPath.end();
}

void ThreatAwareUnitMoveBehavior::Goto::onEnter()
{
  UnitMoveBehaviorT::Goto::onEnter();
}

void ThreatAwareUnitMoveBehavior::Goto::onExit()
{
  UnitMoveBehaviorT::Goto::onExit();
}

ThreatAwareUnitMoveBehavior::Goto::WeightedPath ThreatAwareUnitMoveBehavior::Goto::generateWeightedPath(const Path& path) const
{
  WeightedPath weightedPath;
  weightedPath.reserve(path.size());
  for(Path::const_iterator it = path.begin(); it != path.end(); ++it)
    weightedPath.push_back(WeightedPathElement(*it, m_context.m_moveOptions.m_threatEvaluator(*it)));
  return weightedPath;
}

Path ThreatAwareUnitMoveBehavior::Goto::calculatePath(const BWAPI::TilePosition current)
{
  lookahead_t lookahead = m_context.m_moveOptions.m_lookahead;
  pathfinding::Closed closed;

  pathfinding::Path path = boundedAStar(
    current,                                    // initial state
    m_context.m_destinationTile,                // goal state
    closed,                                     // closed set
    lookahead,                                  // lookahead
    m_h,                                        // heuristic evaluator
    m_context.m_moveOptions.m_isObstacle,       // obstacle evaluator
    m_context.m_moveOptions.m_threatEvaluator,  // threat evaluator
    m_context.m_moveOptions.m_moveWeight);      // move weight
    
  if (!path.path.empty())
  {
    // Update heuristic for all nodes in the closed set (RTAA* algorithm)
    for (pathfinding::Closed::const_iterator i=closed.begin(); i!=closed.end(); ++i)
    {
      const double g = i->second;
      m_h.update(i->first, path.f - g);
    }
  }

  return path.path;
}

bool ThreatAwareUnitMoveBehavior::Goto::hasGreaterThreatAlongPath()
{
  // TODO: should check only the leftover portion of the path

  // Invalidate the current path if the danger level has risen somewhere along the path
  for(WeightedPath::const_iterator it = m_fullPath.begin(); it != m_fullPath.end(); ++it)
  {
    if(m_context.m_moveOptions.m_threatEvaluator(it->m_tile) > it->m_threat)
      return true;
  }
  return false;
}

void ThreatAwareUnitMoveBehavior::Goto::onNewPathFound(const Path& path)
{
  if (!m_context.isTerminated()) // Avoid the asynchronous callback when the behavior is already terminated! (and the unit is potentially dead)
  {
    onPathPlanning(path.reachGoal);

    if (path.empty())
    {
      //if (m_currentTarget == m_smoothPath.end())
      //{
      //  // Retreat if threatened
      //  if ((!m_context.m_unit->isAttacking() && !m_context.m_unit->isAttackFrame()) && 
      //      m_context.m_moveOptions.m_threatEvaluator(m_context.m_unit->getTilePosition()) > 0)
      //  {
      //    Path retreatPath = calculateRetreatPath(
      //      m_context.m_unit,
      //      pathfinding::GeneralObstacleEvaluator(m_context.m_unit),
      //      m_context.m_moveOptions.m_threatEvaluator).path;

      //    m_fullPath = generateWeightedPath(retreatPath);
      //    m_smoothPath = smoothPath(m_fullPath.begin(), m_fullPath.end());
      //    m_isRetreatPath = true;
      //    SPAR_LOG(LogTypes.GENERAL_ERROR, "Retreat path");
      //  }
      //  else
      //  {
      //    // Will try again on next iteration
      //    m_fullPath.clear();
      //    m_smoothPath.clear();
      //  }
      //}
    }
    else if (m_currentTarget == m_smoothPath.end() || !m_isRetreatPath)
    {
      m_fullPath = generateWeightedPath(path);
      m_smoothPath = smoothPath(m_fullPath.begin(), m_fullPath.end());
      m_isRetreatPath = false;
      //SPAR_LOG(LogTypes.GENERAL_ERROR, "Threat-aware path");
    }

    m_currentTarget = m_smoothPath.begin();
    m_updateInProgress = false;
  }
}

void ThreatAwareUnitMoveBehavior::Goto::onTimer()
{
  UnitMoveBehaviorT::Goto::onTimer();
  if (m_context.isTerminated())
    return;

  // We should probably have some way to invalidate retreat paths as well
  if(!m_isRetreatPath && hasGreaterThreatAlongPath())
    m_currentTarget = m_smoothPath.end();

  if(m_currentTarget == m_smoothPath.end())
  {
    // Retreat if threatened
    if ((!m_context.m_unit->isAttacking() && !m_context.m_unit->isAttackFrame()) && 
        m_context.m_moveOptions.m_threatEvaluator(m_context.m_unit->getTilePosition()) > 0)
    {
      Path retreatPath = calculateRetreatPath(
        m_context.m_unit,
        pathfinding::GeneralObstacleEvaluator(m_context.m_unit),
        m_context.m_moveOptions.m_threatEvaluator).path;

      m_fullPath = generateWeightedPath(retreatPath);
      m_smoothPath = smoothPath(m_fullPath.begin(), m_fullPath.end());
      m_isRetreatPath = true;
      m_currentTarget = m_smoothPath.begin();
      //SPAR_LOG(LogTypes.GENERAL_ERROR, "Retreat path");
    }

    // Replan
    if(!m_updateInProgress)
    {
      // If the path is empty, try planning again
      BWAPI::TilePosition current = BWAPI::TilePosition(m_context.m_unit->getPosition());
      m_updateInProgress = true;
      m_context.scheduleParallelTask<Path>(
        boost::bind(&ThreatAwareUnitMoveBehavior::Goto::calculatePath, this, current), // operation
        boost::bind(&ThreatAwareUnitMoveBehavior::Goto::onNewPathFound, this, _1),  // callback
        "Path planning" // name
        );
    }
  }

  // If the path is not empty, keep following it
  if(m_currentTarget != m_smoothPath.end())
  {
    BWAPI::Position step = getCenter(*m_currentTarget);
    submitMoveCommand(step);
    if(step.getDistance(m_context.m_unit->getPosition()) <=  m_context.m_moveOptions.m_maxDistance)
      ++m_currentTarget;
  }
}

void ThreatAwareUnitMoveBehavior::Goto::onFrameDebug()
{
  UnitMoveBehaviorT::Goto::onFrameDebug();

  const std::set<BWAPI::Unit*>& selectedUnits = BWAPI::Broodwar->getSelectedUnits();
  if (Config::getInstance().displayOwnUnitsActions() && selectedUnits.find(m_context.m_unit) != selectedUnits.end() /*m_context.m_unit->isSelected()*/)
  {
    BWAPI::Position unitPosition = m_context.m_unit->getPosition();

    // Display unit path
    {
      BWAPI::Broodwar->drawCircle(BWAPI::CoordinateType::Map, unitPosition.x(), unitPosition.y(), 15, BWAPI::Colors::Red, false);
      BWAPI::Position previous = unitPosition;
      for(std::vector<BWAPI::TilePosition>::const_iterator itt = m_currentTarget; itt != m_smoothPath.end(); ++itt) 
      {
        const BWAPI::Position position = getCenter(*itt);
        BWAPI::Broodwar->drawCircleMap(position.x() , position.y(), 15, BWAPI::Colors::Red, false);
        BWAPI::Broodwar->drawLineMap(previous.x(), previous.y(), position.x(), position.y(), BWAPI::Colors::Red);
        previous = position;
      }
    }

    {
      // Display threat level for each tile
      const int maxDistance = 50;
      const double threatThreshold = 20;
      const int colors[] = {0x07 /*green*/,0x17 /*yellow*/,0x11 /*orange*/,0x06 /*red*/,0x08 /*bright red*/};
      const BWAPI::TilePosition currentTile(unitPosition);
      const BWAPI::TilePosition t(currentTile.x() - maxDistance/2, currentTile.y() - maxDistance/2); //Top-left corner
      for(int x = 0; x < maxDistance; ++x)
      {
        for(int y = 0; y < maxDistance; ++y)
        {
          BWAPI::TilePosition current = t + BWAPI::TilePosition(x, y);
          if(current.isValid())
          {
            BWAPI::Position p = getCenter(current);
            if(m_context.m_moveOptions.m_isObstacle(current, current))
            {
              BWAPI::Broodwar->drawTextMap(p.x(), p.y(), "%c %c", 0x04 /* white */, 'X');
            }
            else
            {
              double value = m_context.m_moveOptions.m_threatEvaluator(current);
              int c = colors[0];
              for(int i = 1; i<sizeof(colors)/sizeof(colors[0]); ++i)
              {
                if(value > threatThreshold * i)
                  c = colors[i];
              }
              BWAPI::Broodwar->drawTextMap(p.x(), p.y(), "%c %g", c, value);
            }
          }
        }
      }
    }
  }
}