# 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).
# Student side autograding was added by Brad Miller, Nick Hay, and
# Pieter Abbeel (pabbeel@cs.berkeley.edu).


from captureAgents import CaptureAgent
import heapq
import random
from game import Directions


#################
# Team creation #
#################

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


##########
# Agents #
##########

class PlanningRoleAgent(CaptureAgent):
  """Dynamic role agent that chooses targets first, then plans by BFS."""

  DIRECTION_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)
    self.width = gameState.data.layout.width
    self.height = gameState.data.layout.height
    self.useRiskAwarePlanning = self.width >= 36
    self.lastEatenFood = None
    CaptureAgent.registerInitialState(self, gameState)
    self.homeEntries = self.computeHomeEntries(gameState)
    self.chokeTargets = self.computeChokeTargets()

  def chooseAction(self, gameState):
    self.updateLastEatenFood(gameState)
    mode = self.decideTeamMode(gameState)
    role = self.decideRole(gameState, mode)
    targets = self.chooseTargets(gameState, mode, role)

    avoidDanger = role in ('offense', 'conservative')
    action = self.firstActionToTargets(gameState, targets, avoidDanger)
    if action is not None:
      return action

    if avoidDanger:
      action = self.firstActionToTargets(gameState, targets, False)
      if action is not None:
        return action

    return self.fallbackAction(gameState, targets, role)

  def decideTeamMode(self, gameState):
    invaders = self.visibleInvaders(gameState)
    maxInvaderCarrying = 0
    if invaders:
      maxInvaderCarrying = max(invader.numCarrying for invader in invaders)

    score = self.getScore(gameState)
    timeLeft = self.getTimeLeft(gameState)

    if len(invaders) >= 2:
      return 'defense'
    if maxInvaderCarrying >= 3:
      return 'defense'
    if score < -5 and not invaders:
      return 'offense'
    if score > 5 or timeLeft < 150:
      return 'conservative'
    return 'split'

  def decideRole(self, gameState, mode):
    if mode == 'split':
      team = sorted(self.getTeam(gameState))
      if self.index == team[0]:
        return 'offense'
      return 'defense'
    return mode

  def chooseTargets(self, gameState, mode, role):
    if role == 'defense':
      return self.defenseTargets(gameState)
    if role == 'conservative':
      return self.conservativeTargets(gameState)
    return self.offenseTargets(gameState, mode)

  def offenseTargets(self, gameState, mode):
    myState = gameState.getAgentState(self.index)
    myPos = gameState.getAgentPosition(self.index)
    food = self.getFood(gameState).asList()
    capsules = self.getCapsules(gameState)
    carrying = myState.numCarrying
    homeDistance = self.minDistance(myPos, self.homeEntries)
    ghostDistances = self.activeEnemyGhostDistances(gameState, myPos)
    closestGhost = min(ghostDistances) if ghostDistances else None

    returnThreshold = 5 if mode == 'offense' else 4
    shouldReturn = (
      carrying >= returnThreshold or
      len(food) <= 2 or
      (carrying > 0 and closestGhost is not None and closestGhost <= 5) or
      self.getTimeLeft(gameState) < homeDistance + 24
    )
    if shouldReturn:
      return self.sortedTargets(myPos, self.homeEntries)

    if capsules and closestGhost is not None and closestGhost <= 5:
      return self.sortedTargets(myPos, capsules)

    if food:
      return self.sortedTargets(myPos, food)[:6]
    return self.sortedTargets(myPos, self.homeEntries)

  def defenseTargets(self, gameState):
    myPos = gameState.getAgentPosition(self.index)
    invaders = self.visibleInvaders(gameState)
    if invaders:
      positions = [invader.getPosition() for invader in invaders]
      positions.sort(
        key=lambda p: (
          -gameState.getAgentState(self.nearestOpponentAt(gameState, p)).numCarrying,
          self.getMazeDistance(myPos, p)
        )
      )
      return positions

    if self.lastEatenFood is not None:
      return [self.lastEatenFood]
    return [self.defensivePatrolTarget(gameState)]

  def conservativeTargets(self, gameState):
    myState = gameState.getAgentState(self.index)
    invaders = self.visibleInvaders(gameState)
    if invaders:
      return self.defenseTargets(gameState)
    if myState.numCarrying > 0 or myState.isPacman:
      return self.sortedTargets(gameState.getAgentPosition(self.index),
                                self.homeEntries)
    return [self.defensivePatrolTarget(gameState)]

  def firstActionToTargets(self, gameState, targets, avoidDanger):
    if not targets:
      return None

    start = gameState.getAgentPosition(self.index)
    if start is None:
      return None
    start = self.normalizePosition(start)

    targetSet = set(self.normalizePosition(t) for t in targets if t is not None)
    if not targetSet:
      return None

    danger = self.visibleActiveGhostPositions(gameState)
    if avoidDanger and danger and self.useRiskAwarePlanning:
      action = self.riskAwareActionToTargets(gameState, start, targetSet, danger)
      if action is not None:
        return action

    legalActions = [a for a in gameState.getLegalActions(self.index)
                    if a != Directions.STOP]
    random.shuffle(legalActions)

    walls = gameState.getWalls()
    queue = []
    visited = set([start])

    for action in legalActions:
      nextPos = self.nextPosition(start, action)
      if not self.isOpen(nextPos, walls):
        continue
      if avoidDanger and self.isDangerous(nextPos, danger):
        continue
      if nextPos in targetSet:
        return action
      queue.append((nextPos, action))
      visited.add(nextPos)

    head = 0
    while head < len(queue):
      pos, firstAction = queue[head]
      head += 1
      for action in (Directions.NORTH, Directions.SOUTH,
                     Directions.EAST, Directions.WEST):
        nextPos = self.nextPosition(pos, action)
        if nextPos in visited or not self.isOpen(nextPos, walls):
          continue
        if avoidDanger and self.isDangerous(nextPos, danger):
          continue
        if nextPos in targetSet:
          return firstAction
        visited.add(nextPos)
        queue.append((nextPos, firstAction))

    return None

  def riskAwareActionToTargets(self, gameState, start, targetSet, danger):
    walls = gameState.getWalls()
    frontier = []
    bestCost = {start: 0}
    counter = 0
    limit = self.width * self.height * 4
    heapq.heappush(frontier, (
      self.targetHeuristic(start, targetSet), 0, counter, start, None
    ))

    while frontier and counter < limit:
      priority, cost, _, pos, firstAction = heapq.heappop(frontier)
      if cost != bestCost.get(pos):
        continue
      if pos in targetSet and firstAction is not None:
        return firstAction

      for action in (Directions.NORTH, Directions.SOUTH,
                     Directions.EAST, Directions.WEST):
        nextPos = self.nextPosition(pos, action)
        if not self.isOpen(nextPos, walls):
          continue
        if self.isDangerous(nextPos, danger):
          continue
        stepCost = 1 + self.positionRisk(nextPos, danger)
        nextCost = cost + stepCost
        if nextCost >= bestCost.get(nextPos, 999999):
          continue
        counter += 1
        first = action if firstAction is None else firstAction
        bestCost[nextPos] = nextCost
        heapq.heappush(frontier, (
          nextCost + self.targetHeuristic(nextPos, targetSet),
          nextCost,
          counter,
          nextPos,
          first,
        ))

    return None

  def positionRisk(self, pos, dangerPositions):
    risk = 0
    for danger in dangerPositions:
      distance = abs(pos[0] - danger[0]) + abs(pos[1] - danger[1])
      if distance <= 2:
        risk += 18
      elif distance <= 5:
        risk += 2 * (6 - distance)
    return risk

  def targetHeuristic(self, pos, targetSet):
    return min(abs(pos[0] - target[0]) + abs(pos[1] - target[1])
               for target in targetSet)

  def fallbackAction(self, gameState, targets, role):
    actions = gameState.getLegalActions(self.index)
    if Directions.STOP in actions and len(actions) > 1:
      actions.remove(Directions.STOP)

    values = []
    for action in actions:
      successor = gameState.generateSuccessor(self.index, action)
      pos = successor.getAgentPosition(self.index)
      value = self.getScore(successor) * 100
      if targets and pos is not None:
        value -= self.minDistance(pos, targets)
      if role in ('offense', 'conservative'):
        ghostDistances = self.activeEnemyGhostDistances(successor, pos)
        if ghostDistances:
          closest = min(ghostDistances)
          if closest <= 1:
            value -= 1000
          else:
            value += min(closest, 5)
      if action == Directions.STOP:
        value -= 100
      values.append((value, action))

    bestValue = max(value for value, action in values)
    bestActions = [action for value, action in values if value == bestValue]
    return random.choice(bestActions)

  def computeHomeEntries(self, gameState):
    walls = gameState.getWalls()
    x = self.width // 2 - 1 if self.red else self.width // 2
    entries = []
    for y in range(1, self.height - 1):
      if not walls[x][y]:
        entries.append((x, y))
    return entries or [self.start]

  def computeChokeTargets(self):
    entries = sorted(self.homeEntries, key=lambda p: p[1])
    if len(entries) <= 2:
      return entries

    groups = []
    current = [entries[0]]
    for entry in entries[1:]:
      if entry[1] == current[-1][1] + 1:
        current.append(entry)
      else:
        groups.append(current)
        current = [entry]
    groups.append(current)

    if len(groups) >= 2:
      lower = groups[0][len(groups[0]) // 2]
      upper = groups[-1][len(groups[-1]) // 2]
      return [lower, upper]

    lowerIndex = max(0, len(entries) // 3)
    upperIndex = min(len(entries) - 1, (2 * len(entries)) // 3)
    if lowerIndex == upperIndex and upperIndex + 1 < len(entries):
      upperIndex += 1
    return [entries[lowerIndex], entries[upperIndex]]

  def defensivePatrolTarget(self, gameState):
    myPos = gameState.getAgentPosition(self.index)
    if self.lastEatenFood is not None:
      if myPos is None or self.getMazeDistance(myPos, self.lastEatenFood) > 1:
        return self.lastEatenFood
      self.lastEatenFood = None

    team = sorted(self.getTeam(gameState))
    rank = team.index(self.index) if self.index in team else 0
    if self.chokeTargets:
      return self.chokeTargets[rank % len(self.chokeTargets)]

    targetY = self.height * (rank + 1) / float(len(team) + 1)
    return min(self.homeEntries, key=lambda p: abs(p[1] - targetY))

  def visibleInvaders(self, gameState):
    enemies = [gameState.getAgentState(i) for i in self.getOpponents(gameState)]
    return [enemy for enemy in enemies
            if enemy.isPacman and enemy.getPosition() is not None]

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

  def activeEnemyGhostDistances(self, gameState, myPos):
    distances = []
    if myPos is None:
      return distances
    for ghostPos in self.visibleActiveGhostPositions(gameState):
      distances.append(self.getMazeDistance(myPos, ghostPos))
    return distances

  def nearestOpponentAt(self, gameState, pos):
    bestOpponent = self.getOpponents(gameState)[0]
    bestDistance = 9999
    for opponent in self.getOpponents(gameState):
      enemyPos = gameState.getAgentState(opponent).getPosition()
      if enemyPos is None:
        continue
      distance = self.getMazeDistance(pos, enemyPos)
      if distance < bestDistance:
        bestDistance = distance
        bestOpponent = opponent
    return bestOpponent

  def updateLastEatenFood(self, gameState):
    previous = self.getPreviousObservation()
    if previous is None:
      return
    previousFood = set(self.getFoodYouAreDefending(previous).asList())
    currentFood = set(self.getFoodYouAreDefending(gameState).asList())
    eaten = list(previousFood - currentFood)
    if eaten:
      myPos = gameState.getAgentPosition(self.index)
      if myPos is None:
        self.lastEatenFood = eaten[0]
      else:
        self.lastEatenFood = min(
          eaten,
          key=lambda p: self.getMazeDistance(myPos, p)
        )

  def sortedTargets(self, pos, targets):
    if pos is None:
      return list(targets)
    return sorted(targets, key=lambda target: self.getMazeDistance(pos, target))

  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 normalizePosition(self, pos):
    return (int(pos[0]), int(pos[1]))

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

  def isOpen(self, pos, walls):
    x, y = pos
    if x < 0 or x >= self.width or y < 0 or y >= self.height:
      return False
    return not walls[x][y]

  def isDangerous(self, pos, dangerPositions):
    for danger in dangerPositions:
      if abs(pos[0] - danger[0]) + abs(pos[1] - danger[1]) <= 1:
        return True
    return False

  def getTimeLeft(self, gameState):
    return getattr(gameState.data, 'timeleft', 9999)