# Dead-end avoidance team for CS470 Assignment 3.

from captureAgents import CaptureAgent
import random
import util
from game import Directions
from util import nearestPoint


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


def computeDeadEndDepths(gameState):
  """Return a map from legal cells to their peeled dead-end depth."""
  walls = gameState.getWalls()
  width = walls.width
  height = walls.height
  legal = []
  degree = {}
  neighbors = {}

  for x in range(width):
    for y in range(height):
      if walls[x][y]:
        continue
      pos = (x, y)
      legal.append(pos)
      nbrs = []
      for dx, dy in ((1, 0), (-1, 0), (0, 1), (0, -1)):
        nx, ny = x + dx, y + dy
        if 0 <= nx < width and 0 <= ny < height and not walls[nx][ny]:
          nbrs.append((nx, ny))
      neighbors[pos] = nbrs
      degree[pos] = len(nbrs)

  queue = []
  depths = {}
  head = 0
  for pos in legal:
    if degree[pos] <= 1:
      depths[pos] = 1
      queue.append(pos)

  removed = set()
  while head < len(queue):
    pos = queue[head]
    head += 1
    if pos in removed:
      continue
    removed.add(pos)
    for nbr in neighbors[pos]:
      if nbr in removed:
        continue
      degree[nbr] -= 1
      if degree[nbr] <= 1 and nbr not in depths:
        depths[nbr] = depths[pos] + 1
        queue.append(nbr)

  return depths


class ReflexCaptureAgent(CaptureAgent):
  def registerInitialState(self, gameState):
    self.start = gameState.getAgentPosition(self.index)
    self.width = gameState.data.layout.width
    self.height = gameState.data.layout.height
    CaptureAgent.registerInitialState(self, gameState)
    self.homeEntries = self.getHomeEntries(gameState)
    self.deadEndDepths = computeDeadEndDepths(gameState)

  def chooseAction(self, gameState):
    actions = gameState.getLegalActions(self.index)
    values = [self.evaluate(gameState, action) for action in actions]
    bestValue = max(values)
    bestActions = [a for a, v in zip(actions, values) if v == bestValue]
    if Directions.STOP in bestActions and len(bestActions) > 1:
      bestActions.remove(Directions.STOP)
    return random.choice(bestActions)

  def getSuccessor(self, gameState, action):
    successor = gameState.generateSuccessor(self.index, action)
    pos = successor.getAgentState(self.index).getPosition()
    if pos != nearestPoint(pos):
      return successor.generateSuccessor(self.index, action)
    return successor

  def evaluate(self, gameState, action):
    return self.getFeatures(gameState, action) * self.getWeights(gameState, action)

  def getFeatures(self, gameState, action):
    features = util.Counter()
    successor = self.getSuccessor(gameState, action)
    features['successorScore'] = self.getScore(successor)
    return features

  def getWeights(self, gameState, action):
    return {'successorScore': 1.0}

  def getHomeEntries(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 minDistance(self, pos, targets):
    if not pos or not targets:
      return 0
    return min(self.getMazeDistance(pos, target) for target in targets)

  def activeEnemyGhostDistances(self, gameState, myPos):
    distances = []
    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:
        distances.append(self.getMazeDistance(myPos, enemyPos))
    return distances

  def allVisibleGhostDistances(self, gameState, myPos):
    active = []
    scared = []
    for opponent in self.getOpponents(gameState):
      enemy = gameState.getAgentState(opponent)
      enemyPos = enemy.getPosition()
      if enemyPos is None or enemy.isPacman:
        continue
      dist = self.getMazeDistance(myPos, enemyPos)
      if enemy.scaredTimer <= 1:
        active.append(dist)
      else:
        scared.append(dist)
    return active, scared

  def deadEndDepth(self, pos):
    if pos is None:
      return 0
    return self.deadEndDepths.get((int(pos[0]), int(pos[1])), 0)


class DeadEndOffensiveAgent(ReflexCaptureAgent):
  def getFeatures(self, gameState, action):
    features = util.Counter()
    successor = self.getSuccessor(gameState, action)
    myState = successor.getAgentState(self.index)
    myPos = myState.getPosition()

    foodList = self.getFood(successor).asList()
    capsules = self.getCapsules(successor)
    carrying = myState.numCarrying
    activeGhostDistances, scaredGhostDistances = self.allVisibleGhostDistances(successor, myPos)
    closestGhost = min(activeGhostDistances) if activeGhostDistances else None
    closestScared = min(scaredGhostDistances) if scaredGhostDistances else None

    features['successorScore'] = self.getScore(successor)
    features['foodRemaining'] = len(foodList)
    features['distanceToFood'] = self.adjustedFoodDistance(myPos, foodList, closestGhost, closestScared)
    features['distanceHome'] = self.minDistance(myPos, self.homeEntries)
    features['distanceToCapsule'] = self.minDistance(myPos, capsules)
    features['deadEndDepth'] = self.deadEndDepth(myPos)

    if action == Directions.STOP:
      features['stop'] = 1
    reverse = Directions.REVERSE[gameState.getAgentState(self.index).configuration.direction]
    if action == reverse:
      features['reverse'] = 1

    if closestGhost is not None:
      features['ghostDistance'] = closestGhost
      if closestGhost <= 2:
        features['immediateDanger'] = 1
      elif closestGhost <= 5:
        features['nearGhost'] = 1

    if self.isEnteringRiskyDeadEnd(myPos, closestGhost, closestScared):
      features['trapRisk'] = features['deadEndDepth']

    if carrying >= 3 or len(foodList) <= 2:
      features['shouldReturn'] = 1
    if carrying > 0 and closestGhost is not None and closestGhost <= 6:
      features['shouldReturn'] = 1
    if closestGhost is not None and closestGhost <= 3 and features['deadEndDepth'] > 0:
      features['shouldReturn'] = 1

    if gameState.data.timeleft < features['distanceHome'] + 20:
      features['shouldReturn'] = 1

    return features

  def adjustedFoodDistance(self, myPos, foodList, closestGhost, closestScared):
    if not myPos or not foodList:
      return 0
    deadEndWeight = self.deadEndWeight(closestGhost, closestScared)
    bestCost = None
    for food in foodList:
      cost = self.getMazeDistance(myPos, food) + self.deadEndDepth(food) * deadEndWeight
      if bestCost is None or cost < bestCost:
        bestCost = cost
    return bestCost

  def deadEndWeight(self, closestGhost, closestScared):
    if closestGhost is None:
      return 0.0
    if closestGhost <= 2:
      return 10.0
    if closestGhost <= 4:
      return 6.0
    if closestGhost <= 6:
      return 3.5
    if closestScared is not None and closestScared <= 5:
      return 0.2
    return 1.0

  def isEnteringRiskyDeadEnd(self, myPos, closestGhost, closestScared):
    depth = self.deadEndDepth(myPos)
    if depth <= 0:
      return False
    if closestGhost is None:
      return False
    if closestScared is not None and closestGhost > 6:
      return False
    return closestGhost <= 6

  def getWeights(self, gameState, action):
    features = self.getFeatures(gameState, action)
    weights = {
      'successorScore': 200,
      'foodRemaining': -100,
      'distanceToFood': -3,
      'distanceToCapsule': -2,
      'ghostDistance': 2,
      'immediateDanger': -1000,
      'nearGhost': -180,
      'deadEndDepth': 0,
      'trapRisk': -45,
      'stop': -100,
      'reverse': -3,
      'shouldReturn': 0,
      'distanceHome': 0,
    }
    if features['shouldReturn']:
      weights['distanceHome'] = -16
      weights['distanceToFood'] = -1
      weights['distanceToCapsule'] = -1
      weights['deadEndDepth'] = -3
    if features['nearGhost'] or features['immediateDanger']:
      weights['distanceToCapsule'] = -8
      weights['deadEndDepth'] = -12
    return weights


class PatrolDefensiveAgent(ReflexCaptureAgent):
  def registerInitialState(self, gameState):
    ReflexCaptureAgent.registerInitialState(self, gameState)
    self.patrolTarget = self.choosePatrolTarget(gameState)
    self.lastEatenFood = None

  def chooseAction(self, gameState):
    self.updateLastEatenFood(gameState)
    invaders = self.visibleInvaders(gameState)
    if invaders:
      myPos = gameState.getAgentPosition(self.index)
      self.patrolTarget = min(
        [a.getPosition() for a in invaders],
        key=lambda p: self.getMazeDistance(myPos, p)
      )
    elif self.lastEatenFood is not None:
      self.patrolTarget = self.lastEatenFood
    elif self.patrolTarget is None:
      self.patrolTarget = self.choosePatrolTarget(gameState)
    return ReflexCaptureAgent.chooseAction(self, gameState)

  def getFeatures(self, gameState, action):
    features = util.Counter()
    successor = self.getSuccessor(gameState, action)
    myState = successor.getAgentState(self.index)
    myPos = myState.getPosition()

    invaders = self.visibleInvaders(successor)
    features['onDefense'] = 1
    if myState.isPacman:
      features['onDefense'] = 0
    features['numInvaders'] = len(invaders)
    if invaders:
      features['invaderDistance'] = min(
        self.getMazeDistance(myPos, invader.getPosition()) for invader in invaders
      )
    else:
      features['distanceToPatrol'] = self.getMazeDistance(myPos, self.patrolTarget)

    if action == Directions.STOP:
      features['stop'] = 1
    reverse = Directions.REVERSE[gameState.getAgentState(self.index).configuration.direction]
    if action == reverse:
      features['reverse'] = 1
    return features

  def getWeights(self, gameState, action):
    return {
      'numInvaders': -1000,
      'onDefense': 120,
      'invaderDistance': -20,
      'distanceToPatrol': -4,
      'stop': -100,
      'reverse': -2,
    }

  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 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)
      self.lastEatenFood = min(eaten, key=lambda p: self.getMazeDistance(myPos, p))

  def choosePatrolTarget(self, gameState):
    defendingFood = self.getFoodYouAreDefending(gameState).asList()
    if not defendingFood:
      return random.choice(self.homeEntries)
    foodCenter = (
      sum(x for x, y in defendingFood) / float(len(defendingFood)),
      sum(y for x, y in defendingFood) / float(len(defendingFood))
    )
    return min(
      self.homeEntries,
      key=lambda p: abs(p[0] - foodCenter[0]) + abs(p[1] - foodCenter[1])
    )