# Risk-aware courier team for CS470 Assignment 3.

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


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


class TargetAgent(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)

  def chooseAction(self, gameState):
    actions = gameState.getLegalActions(self.index)
    values = [self.evaluateAction(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 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 pos is None or not targets:
      return 0
    return min(self.getMazeDistance(pos, target) for target in targets)

  def closestTarget(self, pos, targets):
    if not targets:
      return None
    return min(targets, key=lambda target: self.getMazeDistance(pos, target))

  def activeGhostDistances(self, gameState, pos):
    distances = []
    if pos is None:
      return 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(pos, enemyPos))
    return distances

  def evaluateAction(self, gameState, action):
    return 0


class RiskAwareCourierAgent(TargetAgent):
  carryHomeWeight = -5
  ghostClosePenalty = -2000
  ghostDangerRadius = 5
  scoreWeight = 10000
  targetDistanceWeight = -20

  def evaluateAction(self, gameState, action):
    successor = self.getSuccessor(gameState, action)
    oldState = gameState.getAgentState(self.index)
    newState = successor.getAgentState(self.index)
    myPos = newState.getPosition()

    target, returning = self.chooseTarget(gameState)
    if target is None:
      target = self.closestTarget(myPos, self.homeEntries)

    value = self.scoreWeight * self.getScore(successor)
    value += self.targetDistanceWeight * self.getMazeDistance(myPos, target)

    homeDistance = self.minDistance(myPos, self.homeEntries)
    if newState.numCarrying > 0:
      value += self.carryHomeWeight * newState.numCarrying * homeDistance
    if returning:
      value += -12 * homeDistance

    ghostDistances = self.activeGhostDistances(successor, myPos)
    if ghostDistances:
      closestGhost = min(ghostDistances)
      if closestGhost <= self.ghostDangerRadius:
        value += self.ghostClosePenalty * (self.ghostDangerRadius + 1 - closestGhost)

    if action == Directions.STOP:
      value -= 500
    reverse = Directions.REVERSE[oldState.configuration.direction]
    if action == reverse:
      value -= 5

    if self.likelyDied(gameState, successor):
      value -= 8000
    return value

  def chooseTarget(self, gameState):
    myState = gameState.getAgentState(self.index)
    myPos = myState.getPosition()
    food = self.getFood(gameState).asList()
    home = self.closestTarget(myPos, self.homeEntries)

    if self.shouldReturnHome(gameState):
      return home, True
    if len(food) <= 2:
      return home, True
    if not food:
      return home, True

    return self.bestRoundTripFood(myPos, food), False

  def shouldReturnHome(self, gameState):
    myState = gameState.getAgentState(self.index)
    myPos = myState.getPosition()
    carrying = myState.numCarrying
    if carrying <= 0:
      return False

    carryThreshold = 2 if self.getScore(gameState) >= 0 else 3
    if carrying >= carryThreshold:
      return True

    ghostDistances = self.activeGhostDistances(gameState, myPos)
    if myState.isPacman and ghostDistances and min(ghostDistances) <= self.ghostDangerRadius:
      return True

    homeDistance = self.minDistance(myPos, self.homeEntries)
    timeToBank = 4 * (homeDistance + 5)
    if gameState.data.timeleft <= max(60, timeToBank):
      return True
    return False

  def bestRoundTripFood(self, pos, food):
    return min(
      food,
      key=lambda target: (
        self.getMazeDistance(pos, target) + self.minDistance(target, self.homeEntries),
        self.getMazeDistance(pos, target)
      )
    )

  def likelyDied(self, gameState, successor):
    oldState = gameState.getAgentState(self.index)
    newState = successor.getAgentState(self.index)
    oldPos = oldState.getPosition()
    newPos = newState.getPosition()
    if oldState.isPacman and newPos == self.start and oldPos != self.start:
      return True
    if oldState.numCarrying > 0 and newState.numCarrying == 0:
      return self.getScore(successor) <= self.getScore(gameState)
    return False


class CourierDefensiveAgent(TargetAgent):
  def registerInitialState(self, gameState):
    TargetAgent.registerInitialState(self, gameState)
    self.patrolTarget = self.choosePatrolTarget(gameState)
    self.lastEatenFood = None

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

  def evaluateAction(self, gameState, action):
    successor = self.getSuccessor(gameState, action)
    myState = successor.getAgentState(self.index)
    myPos = myState.getPosition()
    invaders = self.visibleInvaders(successor)

    value = 0
    if not myState.isPacman:
      value += 120
    else:
      value -= 250

    value -= 1000 * len(invaders)
    if invaders:
      value -= 25 * min(self.getMazeDistance(myPos, invader.getPosition()) for invader in invaders)
    elif self.patrolTarget is not None:
      value -= 4 * self.getMazeDistance(myPos, self.patrolTarget)

    if action == Directions.STOP:
      value -= 100
    reverse = Directions.REVERSE[gameState.getAgentState(self.index).configuration.direction]
    if action == reverse:
      value -= 2
    return value

  def visibleInvaders(self, gameState):
    enemies = [gameState.getAgentState(index) for index 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 pos: self.getMazeDistance(myPos, pos))

  def choosePatrolTarget(self, gameState):
    defendingFood = self.getFoodYouAreDefending(gameState).asList()
    if not defendingFood:
      return random.choice(self.homeEntries)
    center = (
      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 pos: abs(pos[0] - center[0]) + abs(pos[1] - center[1])
    )