# myTeam.py
# ---------
# Licensing Information:  You are free to use or extend these projects for
# educational purposes provided that (1) you do not distribute or publish
# solutions, (2) you retain this notice, and (3) you provide clear
# attribution to UC Berkeley, including a link to http://ai.berkeley.edu.
#
# Attribution Information: The Pacman AI projects were developed at UC Berkeley.
# The core projects and autograders were primarily created by John DeNero
# (denero@cs.berkeley.edu) and Dan Klein (klein@cs.berkeley.edu).

from collections import deque

from captureAgents import CaptureAgent
from game import Directions


def createTeam(firstIndex, secondIndex, isRed,
               first='PlanningOffensiveAgent',
               second='PlanningPortalDefender'):
  return [eval(first)(firstIndex), eval(second)(secondIndex)]


class PlanningCaptureAgent(CaptureAgent):
  actionOrder = [
    Directions.NORTH,
    Directions.SOUTH,
    Directions.EAST,
    Directions.WEST,
  ]
  vectors = {
    Directions.NORTH: (0, 1),
    Directions.SOUTH: (0, -1),
    Directions.EAST: (1, 0),
    Directions.WEST: (-1, 0),
  }

  def registerInitialState(self, gameState):
    self.start = gameState.getAgentPosition(self.index)
    CaptureAgent.registerInitialState(self, gameState)
    self.width = gameState.data.layout.width
    self.height = gameState.data.layout.height
    self.walls = gameState.getWalls()
    self.portals = self.computePortals(gameState)

  def computePortals(self, gameState):
    walls = gameState.getWalls()
    if self.red:
      homeX = self.width // 2 - 1
      enemyX = self.width // 2
    else:
      homeX = self.width // 2
      enemyX = self.width // 2 - 1

    portals = []
    fallback = []
    for y in range(1, self.height - 1):
      if not walls[homeX][y]:
        fallback.append((homeX, y))
        if not walls[enemyX][y]:
          portals.append((homeX, y))
    return portals or fallback or [self.start]

  def cleanPosition(self, pos):
    if pos is None:
      return None
    return (int(pos[0]), int(pos[1]))

  def isHomeSide(self, pos):
    if pos is None:
      return True
    if self.red:
      return pos[0] < self.width // 2
    return pos[0] >= self.width // 2

  def nextPosition(self, pos, action):
    dx, dy = self.vectors[action]
    return (int(pos[0] + dx), int(pos[1] + dy))

  def isLegalPosition(self, pos):
    x, y = pos
    return 0 <= x < self.width and 0 <= y < self.height and not self.walls[x][y]

  def orderedLegalActions(self, gameState):
    legal = [action for action in gameState.getLegalActions(self.index)
             if action != Directions.STOP]
    current = gameState.getAgentState(self.index).configuration.direction
    reverse = Directions.REVERSE[current]

    ordered = []
    if current in legal:
      ordered.append(current)
    for action in self.actionOrder:
      if action in legal and action != current and action != reverse:
        ordered.append(action)
    if reverse in legal:
      ordered.append(reverse)
    return ordered or gameState.getLegalActions(self.index)

  def minDistance(self, pos, targets):
    if pos is None or not targets:
      return 0
    return min(self.getMazeDistance(pos, target) for target in targets)

  def activeEnemyGhostPositions(self, gameState):
    positions = []
    for opponent in self.getOpponents(gameState):
      enemy = gameState.getAgentState(opponent)
      pos = gameState.getAgentPosition(opponent)
      if pos is None:
        continue
      if not enemy.isPacman and enemy.scaredTimer <= 1:
        positions.append(pos)
    return positions

  def visibleInvaderPositions(self, gameState):
    positions = []
    for opponent in self.getOpponents(gameState):
      enemy = gameState.getAgentState(opponent)
      pos = gameState.getAgentPosition(opponent)
      if enemy.isPacman and pos is not None:
        positions.append(pos)
    return positions

  def dangerZones(self, gameState, radius):
    zones = set()
    for ghost in self.activeEnemyGhostPositions(gameState):
      queue = deque([(ghost, 0)])
      seen = set([ghost])
      while queue:
        pos, depth = queue.popleft()
        zones.add(pos)
        if depth >= radius:
          continue
        for action in self.actionOrder:
          nxt = self.nextPosition(pos, action)
          if nxt in seen or not self.isLegalPosition(nxt):
            continue
          seen.add(nxt)
          queue.append((nxt, depth + 1))
    return zones

  def choosePathAction(self, gameState, targets, avoid=None, allowStop=False):
    start = self.cleanPosition(gameState.getAgentPosition(self.index))
    targetSet = set(self.cleanPosition(target) for target in targets if target is not None)
    avoid = avoid or set()

    if start is None or not targetSet:
      return self.safeFallbackAction(gameState, targetSet, avoid)
    if start in targetSet and allowStop:
      return Directions.STOP

    queue = deque()
    seen = set([start])
    legal = self.orderedLegalActions(gameState)
    for action in legal:
      if action == Directions.STOP:
        continue
      nxt = self.nextPosition(start, action)
      if not self.isLegalPosition(nxt) or nxt in seen or nxt in avoid:
        continue
      if nxt in targetSet:
        return action
      seen.add(nxt)
      queue.append((nxt, action))

    while queue:
      pos, firstAction = queue.popleft()
      for action in self.actionOrder:
        nxt = self.nextPosition(pos, action)
        if not self.isLegalPosition(nxt) or nxt in seen or nxt in avoid:
          continue
        if nxt in targetSet:
          return firstAction
        seen.add(nxt)
        queue.append((nxt, firstAction))

    return self.safeFallbackAction(gameState, targetSet, avoid)

  def safeFallbackAction(self, gameState, targets, avoid):
    myPos = self.cleanPosition(gameState.getAgentPosition(self.index))
    ghosts = self.activeEnemyGhostPositions(gameState)
    reverse = Directions.REVERSE[gameState.getAgentState(self.index).configuration.direction]
    bestAction = Directions.STOP
    bestScore = None

    for action in gameState.getLegalActions(self.index):
      nxt = myPos if action == Directions.STOP else self.nextPosition(myPos, action)
      if not self.isLegalPosition(nxt):
        continue
      score = 0
      if action == Directions.STOP:
        score -= 8
      if action == reverse:
        score -= 1
      if nxt in avoid:
        score -= 10000
      if targets:
        score -= self.minDistance(nxt, targets)
      if ghosts:
        score += 8 * self.minDistance(nxt, ghosts)

      if bestScore is None or score > bestScore:
        bestAction = action
        bestScore = score

    return bestAction

  def nearestTargets(self, pos, targets, limit=None):
    ordered = sorted(targets, key=lambda target: self.getMazeDistance(pos, target))
    if limit is None:
      return ordered
    return ordered[:limit]


class PlanningOffensiveAgent(PlanningCaptureAgent):
  def chooseAction(self, gameState):
    myState = gameState.getAgentState(self.index)
    myPos = self.cleanPosition(gameState.getAgentPosition(self.index))
    foodList = self.getFood(gameState).asList()
    capsules = self.getCapsules(gameState)
    ghosts = self.activeEnemyGhostPositions(gameState)
    ghostDistance = self.minDistance(myPos, ghosts) if ghosts else 999
    homeDistance = self.minDistance(myPos, self.portals)
    carrying = myState.numCarrying

    dangerRadius = 2 if carrying > 0 or not self.isHomeSide(myPos) else 1
    avoidWide = self.dangerZones(gameState, dangerRadius)
    avoidTight = set(ghosts)

    if carrying > 0 and gameState.data.timeleft < homeDistance + 25:
      return self.moveToTargets(gameState, self.portals, avoidTight)
    if carrying >= 4 or len(foodList) <= 2:
      return self.moveToTargets(gameState, self.portals, avoidTight)
    if carrying >= 3 and ghostDistance <= 8:
      return self.moveToTargets(gameState, self.portals, avoidWide)
    if carrying > 0 and ghostDistance <= 7:
      return self.moveToTargets(gameState, self.portals, avoidWide)
    if ghostDistance <= 3:
      if capsules:
        return self.moveToTargets(gameState, capsules, avoidTight)
      return self.moveToTargets(gameState, self.portals, avoidTight)
    if ghostDistance <= 4 and capsules:
      return self.moveToTargets(gameState, capsules, avoidTight)

    if foodList:
      targets = self.chooseFoodTargets(myPos, foodList)
      return self.moveToTargets(gameState, targets, avoidWide)

    return self.moveToTargets(gameState, self.portals, avoidTight)

  def moveToTargets(self, gameState, targets, avoid):
    action = self.choosePathAction(gameState, targets, avoid=avoid, allowStop=False)
    if action == Directions.STOP and avoid:
      action = self.choosePathAction(gameState, targets, avoid=set(), allowStop=False)
    return action

  def chooseFoodTargets(self, myPos, foodList):
    def foodScore(food):
      return self.getMazeDistance(myPos, food) + 0.8 * self.minDistance(food, self.portals)

    return sorted(foodList, key=foodScore)[:6]


class PlanningPortalDefender(PlanningCaptureAgent):
  def registerInitialState(self, gameState):
    PlanningCaptureAgent.registerInitialState(self, gameState)
    self.currentTarget = self.chooseFoodClusterPortal(gameState)
    self.lastMissingFood = None

  def chooseAction(self, gameState):
    self.updateMissingFood(gameState)
    myState = gameState.getAgentState(self.index)
    invaders = self.visibleInvaderPositions(gameState)
    self.currentTarget = self.chooseDefensiveTarget(gameState, invaders)

    avoid = set()
    allowStop = self.currentTarget in self.portals and not invaders and not myState.isPacman
    if myState.scaredTimer > 0:
      avoid = set(invaders)

    return self.choosePathAction(
      gameState,
      [self.currentTarget],
      avoid=avoid,
      allowStop=allowStop
    )

  def updateMissingFood(self, gameState):
    previous = self.getPreviousObservation()
    if previous is None:
      return
    previousFood = set(self.getFoodYouAreDefending(previous).asList())
    currentFood = set(self.getFoodYouAreDefending(gameState).asList())
    missing = list(previousFood - currentFood)
    if missing:
      myPos = gameState.getAgentPosition(self.index)
      self.lastMissingFood = min(missing, key=lambda food: self.getMazeDistance(myPos, food))

  def chooseDefensiveTarget(self, gameState, invaders):
    myPos = self.cleanPosition(gameState.getAgentPosition(self.index))
    myState = gameState.getAgentState(self.index)

    if myState.isPacman:
      return min(self.portals, key=lambda portal: self.getMazeDistance(myPos, portal))
    if invaders:
      return self.chooseInvaderTarget(gameState, invaders, myPos)
    if self.lastMissingFood is not None:
      if myPos is not None and self.getMazeDistance(myPos, self.lastMissingFood) <= 1:
        self.lastMissingFood = None
      else:
        return self.lastMissingFood
    return self.chooseFoodClusterPortal(gameState)

  def chooseInvaderTarget(self, gameState, invaders, myPos):
    myState = gameState.getAgentState(self.index)
    bestTarget = None
    bestScore = None

    for invaderPos in invaders:
      nearestPortal = min(
        self.portals,
        key=lambda portal: self.getMazeDistance(invaderPos, portal)
      )
      myToInvader = self.getMazeDistance(myPos, invaderPos)
      myToPortal = self.getMazeDistance(myPos, nearestPortal)
      invaderToPortal = self.getMazeDistance(invaderPos, nearestPortal)

      if myState.scaredTimer > 0:
        target = nearestPortal
        score = myToPortal
      elif myToInvader <= 4 or myToInvader <= invaderToPortal + 1:
        target = invaderPos
        score = myToInvader
      else:
        target = nearestPortal
        score = myToPortal + 0.5 * invaderToPortal

      if bestScore is None or score < bestScore:
        bestTarget = target
        bestScore = score

    return bestTarget

  def chooseFoodClusterPortal(self, gameState):
    return self.sortedFoodClusterPortals(gameState)[0]

  def sortedFoodClusterPortals(self, gameState):
    defendingFood = self.getFoodYouAreDefending(gameState).asList()
    if not defendingFood:
      return sorted(self.portals, key=lambda portal: self.getMazeDistance(self.start, portal))

    clusterSize = max(1, min(8, len(defendingFood) // 2))

    def portalFoodScore(portal):
      foodDistances = sorted(
        self.getMazeDistance(portal, food)
        for food in defendingFood
      )
      clusterDistance = sum(foodDistances[:clusterSize]) / float(clusterSize)
      startDistance = self.getMazeDistance(self.start, portal)
      return (clusterDistance, startDistance)

    return sorted(self.portals, key=portalFoodScore)