#include "stdafx.h"
#include "GroupAttackBehavior.h"
#include "UnitAttack/UnitAttackBehavior.h"
#include "UnitAttack/DragoonUnitAttackBehavior.h"
#include "UnitAttack/HighTemplarUnitAttackBehavior.h"
#include "UnitAttack/DefaultUnitAttackBehavior.h"
#include "../../../../../../../PerceptualState/PerceptualState.h"
#include "../../../../../../../PerceptualState/World/CurrentWorld.h"
#include "../../../../Commands/FollowUnitCommand.h"
#include "../../../../Commands/AttackMoveCommand.h"
#include "../../../../Commands/HighTemplarCommands/UseArchonWarpCommand.h"
#include "../../../General/Move/GroupMoveBehavior.h"

// TODO: KEEP A FORMATION

static const int MAX_NB_TEMPLARS = 6;
static const double MAX_TEMPLARS_RATIO = 0.5;
static const int MAX_TEMPLARS_DISTANCE_FOR_ARCHON_WARP = 20 * TILE_SIZE;

GroupAttackBehavior::GroupAttackBehavior(const Process* parent, unsigned int initialPriority, const PersistentUnitGroup* group, Location* location)
  : EngageEnemiesBehaviorT(parent, initialPriority, group)
  , EngageEnemiesBehavior(parent, initialPriority, group)
  , GroupBehavior(parent, initialPriority, group)
  , Behavior(parent, initialPriority)
  , FSM(parent)
  , Process(parent)
  , m_engagingState(*this)
  , m_goToState(*this)
  , m_terminateState(*this)
  , m_states(m_goToState, m_engagingState)
  , m_location(*this, location)
  , m_unitMorphEventHandler(*this)
  , m_observerGroup("GroupAttackBehavior Observers")
  , m_highTemplars("GroupAttackBehavior High Templars")
{
}

void GroupAttackBehavior::onStart()
{
  EngageEnemiesBehaviorT::onStart();
  m_goToState.enter();
  m_unitMorphID = BWAPI::getUnitMorphEvent().subscribe(&m_unitMorphEventHandler, NULL);

  for (PersistentUnitGroup::const_iterator it = getGroup()->begin(); it != getGroup()->end(); ++it)
  {
    BWAPI::Unit* unit = *it;
    if(unit->getType() == BWAPI::UnitTypes::Protoss_Observer)
      m_observerGroup.insert(unit);
    else if(unit->getType() == BWAPI::UnitTypes::Protoss_High_Templar)
      m_highTemplars.insert(unit);    
  }
}

void GroupAttackBehavior::onEnd()
{
  BWAPI::getUnitMorphEvent().unsubscribe(m_unitMorphID);
  EngageEnemiesBehaviorT::onEnd();
}

void GroupAttackBehavior::onUnitAddedToGroup(void*, BWAPI::Unit* unit)
{
  if(unit->getType() == BWAPI::UnitTypes::Protoss_Observer)
    m_observerGroup.insert(unit);
  else if(unit->getType() == BWAPI::UnitTypes::Protoss_High_Templar)
    m_highTemplars.insert(unit);

  if(getCurrentState() == static_cast<FSMState*>(&m_goToState))
    m_goToState.onUnitAdded(unit);
  else if(getCurrentState() == static_cast<FSMState*>(&m_engagingState))
    m_engagingState.onUnitAdded(unit);
  else
    assert(false);

  if(unit->getType() == BWAPI::UnitTypes::Protoss_High_Templar)
    manageUnits();
}

void GroupAttackBehavior::onUnitRemovedFromGroup(void*, BWAPI::Unit* unit)
{
  stdext::hash_map<BWAPI::Unit*, BehaviorWrapper>::iterator pos = m_subBehaviors.find(unit);
  if(pos != m_subBehaviors.end())
    pos->second.behavior->externalTerminate();
  if(unit->getType() == BWAPI::UnitTypes::Protoss_Observer)
    m_observerGroup.remove(unit);
  else if(unit->getType() == BWAPI::UnitTypes::Protoss_High_Templar)
    m_highTemplars.remove(unit);

  if(unit->getType() == BWAPI::UnitTypes::Protoss_High_Templar)
    manageUnits();
}

void GroupAttackBehavior::onUnitMorph(void*, BWAPI::Unit* unit, BWAPI::UnitType previousType)
{
  if (previousType == BWAPI::UnitTypes::Protoss_High_Templar)
  {
    // The unit that just morphed may have been a templar in this group
    m_highTemplars.remove(unit);
  }

  if (getGroup()->find(unit) != getGroup()->end())
  {
    onUnitRemovedFromGroup(NULL, unit);
    onUnitAddedToGroup(NULL, unit);
  }
}

BWAPI::Unit* GroupAttackBehavior::getUnitToFollow(BWAPI::Unit* unit) const
{
  // TODO: use the unit, so that for example templars are spread out in the formation (when they follow different units)

  unused(unit);

  return getUnitToFollow();
}

BWAPI::Unit* GroupAttackBehavior::getUnitToFollow() const
{
  BWAPI::Unit* closest = NULL;
  int closestDistance = boost::numeric::bounds<int>::highest();    
  for(PersistentUnitGroup::const_iterator it=getGroup()->begin(); it != getGroup()->end(); ++it)
  {
    if((*it)->getType() != BWAPI::UnitTypes::Protoss_Observer &&
       (*it)->getType() != BWAPI::UnitTypes::Protoss_High_Templar)
    {
      int distance = (*it)->getPosition().getApproxDistance(m_location->getCurrentCenter());
      if(distance < closestDistance)
      {
        closestDistance = distance;
        closest = (*it);
      }
    }
  }
  return closest;
}

struct SortByArchonWarpPotential : std::binary_function<BWAPI::Unit*, BWAPI::Unit*, bool>
{
  bool operator()(BWAPI::Unit* templar1, BWAPI::Unit* templar2) const
  {
    // min energy, min hit points, min shields
    if (templar1->getEnergy() != templar2->getEnergy())
      return templar1->getEnergy() < templar2->getEnergy();
    if (templar1->getHitPoints() != templar2->getHitPoints())
      return templar1->getHitPoints() < templar2->getHitPoints();
    if (templar1->getShields() != templar2->getShields())
      return templar1->getShields() < templar2->getShields();
    return templar1 < templar2;
  }
};

void GroupAttackBehavior::manageUnits()
{
  if (m_highTemplars.empty())
    return;

  std::set<BWAPI::Unit*, SortByArchonWarpPotential> candidates;
  for (PersistentUnitGroup::const_iterator it = m_highTemplars.begin(); it != m_highTemplars.end(); ++it)
  {
    BWAPI::Unit* templar = *it;
    if (templar->getOrder() != BWAPI::Orders::CastPsionicStorm &&
        templar->getOrder() != BWAPI::Orders::ArchonWarp)
    {
      candidates.insert(templar);
    }
  }

  bool found;
  do
  {
    found = false;
    std::set<BWAPI::Unit*, SortByArchonWarpPotential>::iterator it1,it2;
    for (it1 = candidates.begin(); it1 != candidates.end(); ++it1)
    {
      for (it2 = ++(std::set<BWAPI::Unit*, SortByArchonWarpPotential>::iterator(it1)); it2 != candidates.end(); ++it2)
      {
        if (!CurrentWorld::PsionicStorm->isAcquired() ||                                    // psionic storm not researched yet
            candidates.size() > MAX_NB_TEMPLARS ||                                          // too many templars
            double(candidates.size()) / double(getGroup()->size()) > MAX_TEMPLARS_RATIO ||  // templars ratio too high
            ((*it1)->getEnergy() < 70 && (*it2)->getEnergy() < 70))                         // both low-energy templars
        {
          if ((*it1)->getPosition().getApproxDistance((*it2)->getPosition()) < MAX_TEMPLARS_DISTANCE_FOR_ARCHON_WARP && 
              BWTA::isConnected((*it1)->getTilePosition(), (*it1)->getTilePosition()))      // templars not too far apart
          {
            found = true;
            break;
          }
        }
      }
      if (found)
        break;
    }
    if (found)
    {
#ifdef _DEBUG
      if (!(*it1)->canIssueCommand(BWAPI::UnitCommand::useTech(*it1, BWAPI::TechTypes::Archon_Warp, *it2)))
      {
        SPAR_LOG(LogTypes.GENERAL_WARNING, "Assertion failed GroupAttackBehavior::manageUnits(): \"(*it1)->canIssueCommand(BWAPI::UnitCommand::useTech(*it1, BWAPI::TechTypes::Archon_Warp, *it2))\"");
      }
#endif
      BehaviorState* currentState = dynamic_cast<BehaviorState*>(getCurrentState());
      Layer1FSM::getInstance().getArbitrator().
        submitCommand(*it1, UseArchonWarpCommand(*currentState, *it2), Arbitration::ActionImportance(*this));
      candidates.erase(it1);
      candidates.erase(it2);
    }
  } while (found);
}

UnitAttackBehavior* GroupAttackBehavior::instantiateUnitAttackBehavior(BWAPI::Unit* unit) const
{
  const unsigned int priority = getCurrentPriority() * 2 + 1;

  if (unit->getType() == BWAPI::UnitTypes::Protoss_Dragoon)
    return new DragoonUnitAttackBehavior(this, priority, unit);
  else if (unit->getType() == BWAPI::UnitTypes::Protoss_High_Templar)
    return new HighTemplarUnitAttackBehavior(this, priority, unit, boost::bind(&GroupAttackBehavior::getUnitToFollow, this, unit));
  else
    return new DefaultUnitAttackBehavior(this, priority, unit);
}

// GoToLocation state //////////////////////////////////////////////////////////

GroupAttackBehavior::GoToLocation::GoToLocation(GroupAttackBehavior& behavior)
  : BehaviorStateT(behavior)
  , m_transitionTimerEvent(*this, Scheduler::FSM_SCHEDULING_PRIORITY, 10, 5)
  , m_orderTimerEvent(*this, Scheduler::FSM_SCHEDULING_PRIORITY, 30, 5)
  , m_transitionToEngaging(*this,
                           GroupIsAttacking(behavior.getGroup()),
                           NoTransitionAction(),
                           behavior.m_engagingState)
  , m_transitionToTerminate(*this,
                            GroupReachedLocation<PersistentUnitGroup>(*behavior.getGroup(), behavior.m_location, DEFAULT_GROUP_MAX_DISTANCE), 
                            NoTransitionAction(),
                            behavior.m_terminateState)
  , m_events(m_transitionTimerEvent, m_orderTimerEvent)
  , m_transitions(m_transitionToEngaging,m_transitionToTerminate)
{
}

void GroupAttackBehavior::GoToLocation::orderUnit(BWAPI::Unit* unit, BWAPI::Unit* closest)
{
  if(unit->getType() == BWAPI::UnitTypes::Protoss_Observer)
  {
    if (closest != NULL)
    {
      assert(unit->canIssueCommand(BWAPI::UnitCommand::follow(unit, closest)));
      Layer1FSM::getInstance().getArbitrator().
        submitCommand(unit, FollowUnitCommand(*this, closest), Arbitration::ActionImportance(m_context));
    }
  }
  else if(unit->getType() == BWAPI::UnitTypes::Protoss_High_Templar)
  {
    if (closest != NULL)
    {
      assert(unit->canIssueCommand(BWAPI::UnitCommand::follow(unit, closest)));
      if (unit->getOrder() != BWAPI::Orders::ArchonWarp)
      {
        Layer1FSM::getInstance().getArbitrator().
          submitCommand(unit, FollowUnitCommand(*this, closest), Arbitration::ActionImportance(m_context));
      }
    }
  }
  else
  {
#ifdef _DEBUG
    if (!unit->canIssueCommand(BWAPI::UnitCommand::attack(unit, m_context.m_location->getCurrentCenter())))
    {
      SPAR_LOG(LogTypes.GENERAL_WARNING, "Assertion failed GroupAttackBehavior::GotoLocation::orderUnit(): \"unit->canIssueCommand(BWAPI::UnitCommand::attack(unit, m_location->getCurrentCenter()))\"");
    }
#endif
    Layer1FSM::getInstance().getArbitrator().
      submitCommand(unit, AttackMoveCommand(*this, m_context.m_location->getCurrentCenter()), Arbitration::ActionImportance(m_context));
  }
}

void GroupAttackBehavior::GoToLocation::onEnter()
{
  BWAPI::Unit* closest = m_context.getUnitToFollow();
  for(PersistentUnitGroup::const_iterator it=m_context.getGroup()->begin(); it != m_context.getGroup()->end(); ++it)
  {
    orderUnit(*it, closest);
  }
  m_context.manageUnits();
}

void GroupAttackBehavior::GoToLocation::onTransitionTimer()
{
  m_transitionToEngaging.testTransition(boost::make_tuple());
  m_transitionToTerminate.testTransition(boost::make_tuple());
}

void GroupAttackBehavior::GoToLocation::onOrderTimer()
{
  BWAPI::Unit* closest = m_context.getUnitToFollow();
  for(PersistentUnitGroup::const_iterator it=m_context.getGroup()->begin(); it != m_context.getGroup()->end(); ++it)
  {
    orderUnit(*it, closest);
  }
}

void GroupAttackBehavior::GoToLocation::onUnitAdded(BWAPI::Unit* unit)
{
  BWAPI::Unit* closest = m_context.getUnitToFollow();
  orderUnit(unit, closest);
}

// Engaging state //////////////////////////////////////////////////////////////

GroupAttackBehavior::Engaging::Engaging(GroupAttackBehavior& behavior)
: BehaviorStateT(behavior)
, m_engagingTimerEvent(*this, Scheduler::FSM_SCHEDULING_PRIORITY, 10, 5)
, m_onSubBehaviorTerminated(*this)
, m_transitionToGoTo(*this,
                     NoCondition(), 
                     NoTransitionAction(),
                     behavior.m_goToState)
, m_events(m_engagingTimerEvent)
, m_transitions(m_transitionToGoTo)
{
}

void GroupAttackBehavior::Engaging::onEnter()
{
  for (PersistentUnitGroup::const_iterator i=m_context.getGroup()->begin(); i!=m_context.getGroup()->end(); ++i)
  {
    BWAPI::Unit* unit = *i;
    if(unit->getType() != BWAPI::UnitTypes::Protoss_Observer)
    {
      UnitAttackBehavior* behavior = m_context.instantiateUnitAttackBehavior(unit);
      behavior->zombifyOnTerminate();

      BehaviorWrapper wrapper;
      wrapper.behavior = behavior;
      wrapper.subscriberId = behavior->getTerminateEvent().subscribe(&m_onSubBehaviorTerminated, unit);

      m_context.m_subBehaviors.insert(std::make_pair(unit, wrapper));

      behavior->start();
    }
  }
  m_context.manageUnits();
  selectTarget();
}

void GroupAttackBehavior::Engaging::onTimer()
{
  BWAPI::Unit* closest = m_context.getUnitToFollow();
  if (closest != NULL)
  {
    for(PersistentUnitGroup::const_iterator it = m_context.m_observerGroup.begin(); it != m_context.m_observerGroup.end(); ++it)
    {
      Layer1FSM::getInstance().getArbitrator().
        submitCommand(*it, FollowUnitCommand(*this, closest), Arbitration::ActionImportance(m_context));
      //orderUnit(*it, closest);
    }
  }

  m_context.manageUnits();
  if(!selectTarget())
    m_transitionToGoTo.testTransition(boost::make_tuple());
}

void GroupAttackBehavior::Engaging::onUnitAdded(BWAPI::Unit* unit)
{
  if(unit->getType() != BWAPI::UnitTypes::Protoss_Observer)
  {
    UnitAttackBehavior* behavior = m_context.instantiateUnitAttackBehavior(unit);
    behavior->zombifyOnTerminate();

    BehaviorWrapper wrapper;
    wrapper.behavior = behavior;
    wrapper.subscriberId = behavior->getTerminateEvent().subscribe(&m_onSubBehaviorTerminated, unit);

    m_context.m_subBehaviors.insert(std::make_pair(unit, wrapper));

    behavior->start();
  }
}

void GroupAttackBehavior::Engaging::onSubBehaviorTerminated(void* data)
{
  BWAPI::Unit* unit = reinterpret_cast<BWAPI::Unit*>(data);
  stdext::hash_map<BWAPI::Unit*, BehaviorWrapper>::iterator pos = m_context.m_subBehaviors.find(unit);

  pos->second.behavior->getTerminateEvent().unsubscribe(pos->second.subscriberId);

  m_context.m_subBehaviors.erase(pos);
}

bool GroupAttackBehavior::Engaging::selectTarget()
{
  std::set<BWAPI::Unit*> enemies = Spar->getPerceptualState()->getEnemyUnits(GroupRectangularLocation(m_context.getGroup()));
  bool atLeastOneTargetFound = false;
  for(stdext::hash_map<BWAPI::Unit*, BehaviorWrapper>::iterator it = m_context.m_subBehaviors.begin(); it != m_context.m_subBehaviors.end(); ++it)
  {
    if(it->second.behavior->selectTarget(enemies))
      atLeastOneTargetFound = true;
  }
  return atLeastOneTargetFound;
}
