# 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 random
import util
from game import Directions
from util import nearestPoint


def createTeam(firstIndex, secondIndex, isRed,
               first='RiskAwareOffensiveAgent',
               second='InvaderInterceptDefensiveAgent'):
  return [globals()[first](firstIndex), globals()[second](secondIndex)]


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)

  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()
    if self.red:
      x = self.width // 2 - 1
    else:
      x = 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


class DirectPlanningAgent(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 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 chooseBestAction(self, gameState, scorer):
    actions = gameState.getLegalActions(self.index)
    values = [scorer(action) for action in actions]
    bestValue = max(values)
    bestActions = [action for action, value in zip(actions, values) if value == bestValue]
    if Directions.STOP in bestActions and len(bestActions) > 1:
      bestActions.remove(Directions.STOP)
    return random.choice(bestActions)

  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 nearestTarget(self, pos, targets):
    if not targets:
      return None
    return min(targets, key=lambda target: self.getMazeDistance(pos, target))

  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 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 reversePenalty(self, gameState, action):
    reverse = Directions.REVERSE[gameState.getAgentState(self.index).configuration.direction]
    return 3 if action == reverse else 0


class DirectRiskOffensiveAgent(DirectPlanningAgent):
  def chooseAction(self, gameState):
    return self.chooseBestAction(gameState, lambda action: self.scoreOffenseAction(gameState, action))

  def scoreOffenseAction(self, gameState, action):
    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

    ghostDistances = self.activeEnemyGhostDistances(successor, myPos)
    closestGhost = min(ghostDistances) if ghostDistances else None
    homeDistance = self.minDistance(myPos, self.homeEntries)
    foodDistance = self.minDistance(myPos, foodList)
    capsuleDistance = self.minDistance(myPos, capsules)
    returnMode = (
      carrying >= 3
      or len(foodList) <= 2
      or (carrying > 0 and closestGhost is not None and closestGhost <= 5)
      or gameState.data.timeleft < homeDistance + 20
    )

    score = 200 * self.getScore(successor)
    score -= 100 * len(foodList)
    if returnMode:
      score -= 15 * homeDistance
      score -= foodDistance
      score -= capsuleDistance
    else:
      score -= 3 * foodDistance
      score -= 2 * capsuleDistance

    if closestGhost is not None:
      score += 2 * closestGhost
      if closestGhost <= 2:
        score -= 1000
      elif closestGhost <= 5:
        score -= 180
      if closestGhost <= 5:
        score -= 6 * capsuleDistance

    if action == Directions.STOP:
      score -= 100
    score -= self.reversePenalty(gameState, action)
    return score


class DirectInterceptDefensiveAgent(DirectPlanningAgent):
  def registerInitialState(self, gameState):
    DirectPlanningAgent.registerInitialState(self, gameState)
    self.patrolTarget = self.choosePatrolTarget(gameState)
    self.lastEatenFood = None
    self.currentTarget = self.patrolTarget
    self.targetMode = 'patrol'

  def chooseAction(self, gameState):
    self.updateLastEatenFood(gameState)
    self.currentTarget, self.targetMode = self.chooseDefenseTarget(gameState)
    return self.chooseBestAction(gameState, lambda action: self.scoreDefenseAction(gameState, action))

  def chooseDefenseTarget(self, gameState):
    invader = self.primaryInvader(gameState)
    if invader is not None:
      myPos = gameState.getAgentPosition(self.index)
      invaderPos = invader.getPosition()
      portal = self.nearestHomeEntry(invaderPos)
      distToInvader = self.getMazeDistance(myPos, invaderPos)
      invaderToPortal = self.getMazeDistance(invaderPos, portal)
      nearEscapeRoute = invaderToPortal <= max(4, self.width // 6)
      invaderCloserToExit = invaderToPortal < distToInvader
      highValueInvader = invader.numCarrying >= 2
      catchSoon = distToInvader <= 2 or distToInvader <= invaderToPortal
      atEscapePortal = self.getMazeDistance(myPos, portal) <= 1
      tightEscapeBoard = self.width <= 22 and self.height <= 8 and len(self.homeEntries) <= 3
      releaseChase = invader.numCarrying < 2 and atEscapePortal and distToInvader <= invaderToPortal + 1
      narrowReleaseChase = (
        tightEscapeBoard
        and invader.numCarrying < 2
        and atEscapePortal
        and distToInvader <= invaderToPortal + 2
      )

      if catchSoon or releaseChase or narrowReleaseChase:
        return invaderPos, 'chase'
      if highValueInvader or nearEscapeRoute or invaderCloserToExit:
        return portal, 'hold'
      return invaderPos, 'chase'

    if self.lastEatenFood is not None:
      myPos = gameState.getAgentPosition(self.index)
      if myPos is not None and self.getMazeDistance(myPos, self.lastEatenFood) <= 1:
        self.lastEatenFood = None
      else:
        return self.lastEatenFood, 'inspect'
    if self.patrolTarget is None:
      self.patrolTarget = self.choosePatrolTarget(gameState)
    return self.patrolTarget, 'patrol'

  def scoreDefenseAction(self, gameState, action):
    successor = self.getSuccessor(gameState, action)
    myState = successor.getAgentState(self.index)
    myPos = myState.getPosition()
    target = self.currentTarget or self.choosePatrolTarget(successor)
    targetDistance = self.getMazeDistance(myPos, target)

    score = -12 * targetDistance
    if myState.isPacman:
      score -= 1200

    invaders = self.visibleInvaders(successor)
    if invaders:
      closestInvaderDistance = min(
        self.getMazeDistance(myPos, invader.getPosition()) for invader in invaders
      )
      if self.targetMode == 'chase':
        score -= 10 * closestInvaderDistance
      if closestInvaderDistance <= 1 and myState.scaredTimer == 0:
        score += 500
      if myState.scaredTimer > 0 and closestInvaderDistance <= 2:
        score -= 250

    if action == Directions.STOP:
      if self.targetMode == 'hold' and targetDistance == 0:
        score -= 1
      else:
        score -= 100
    reverse = Directions.REVERSE[gameState.getAgentState(self.index).configuration.direction]
    if action == reverse:
      score -= 2
    return score

  def primaryInvader(self, gameState):
    invaders = self.visibleInvaders(gameState)
    if not invaders:
      return None
    myPos = gameState.getAgentPosition(self.index)
    return max(
      invaders,
      key=lambda invader: (
        invader.numCarrying,
        -self.getMazeDistance(myPos, invader.getPosition())
      )
    )

  def nearestHomeEntry(self, pos):
    return min(self.homeEntries, key=lambda entry: self.getMazeDistance(pos, entry))

  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)
    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 pos: abs(pos[0] - foodCenter[0]) + abs(pos[1] - foodCenter[1])
    )


class RiskAwareOffensiveAgent(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
    ghostDistances = self.activeEnemyGhostDistances(successor, myPos)
    closestGhost = min(ghostDistances) if ghostDistances else None

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

    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 carrying >= 3 or len(foodList) <= 2:
      features['shouldReturn'] = 1
    if carrying > 0 and closestGhost is not None and closestGhost <= 5:
      features['shouldReturn'] = 1
    if gameState.data.timeleft < features['distanceHome'] + 20:
      features['shouldReturn'] = 1

    return features

  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,
      'stop': -100,
      'reverse': -3,
      'shouldReturn': 0,
      'distanceHome': 0,
    }
    if features['shouldReturn']:
      weights['distanceHome'] = -15
      weights['distanceToFood'] = -1
      weights['distanceToCapsule'] = -1
    if features['nearGhost'] or features['immediateDanger']:
      weights['distanceToCapsule'] = -8
    return weights


class InvaderInterceptDefensiveAgent(ReflexCaptureAgent):
  def registerInitialState(self, gameState):
    ReflexCaptureAgent.registerInitialState(self, gameState)
    self.patrolTarget = self.choosePatrolTarget(gameState)
    self.lastEatenFood = None
    self.currentTarget = self.patrolTarget
    self.holdingPortal = False
    self.chasingInvader = False

  def chooseAction(self, gameState):
    self.updateLastEatenFood(gameState)
    self.currentTarget = self.chooseDefenseTarget(gameState)
    actions = gameState.getLegalActions(self.index)
    values = [self.evaluateDefenseAction(gameState, action) for action in actions]
    bestValue = max(values)
    bestActions = [a for a, v in zip(actions, values) if v == bestValue]
    return random.choice(bestActions)

  def chooseDefenseTarget(self, gameState):
    invader = self.primaryInvader(gameState)
    self.holdingPortal = False
    self.chasingInvader = False

    if invader is not None:
      myPos = gameState.getAgentPosition(self.index)
      invaderPos = invader.getPosition()
      portal = self.nearestHomeEntry(invaderPos)
      distToInvader = self.getMazeDistance(myPos, invaderPos)
      invaderToPortal = self.getMazeDistance(invaderPos, portal)
      nearEscapeRoute = invaderToPortal <= max(4, self.width // 6)
      invaderCloserToExit = invaderToPortal < distToInvader
      highValueInvader = invader.numCarrying >= 2
      catchSoon = distToInvader <= 2 or distToInvader <= invaderToPortal
      atEscapePortal = self.getMazeDistance(myPos, portal) <= 1
      tightEscapeBoard = self.width <= 22 and self.height <= 8 and len(self.homeEntries) <= 3
      releaseChase = invader.numCarrying < 2 and atEscapePortal and distToInvader <= invaderToPortal + 1
      narrowReleaseChase = (
        tightEscapeBoard
        and invader.numCarrying < 2
        and atEscapePortal
        and distToInvader <= invaderToPortal + 2
      )

      if catchSoon or releaseChase or narrowReleaseChase:
        self.chasingInvader = True
        return invaderPos
      if highValueInvader or nearEscapeRoute or invaderCloserToExit:
        self.holdingPortal = True
        return portal
      self.chasingInvader = True
      return invaderPos

    if self.lastEatenFood is not None:
      myPos = gameState.getAgentPosition(self.index)
      if myPos is not None and self.getMazeDistance(myPos, self.lastEatenFood) <= 1:
        self.lastEatenFood = None
      else:
        return self.lastEatenFood
    self.patrolTarget = self.choosePatrolTarget(gameState)
    return self.patrolTarget

  def evaluateDefenseAction(self, gameState, action):
    successor = self.getSuccessor(gameState, action)
    myState = successor.getAgentState(self.index)
    myPos = myState.getPosition()
    target = self.currentTarget or self.choosePatrolTarget(successor)
    score = 0

    if myState.isPacman:
      score -= 1200

    targetDistance = self.getMazeDistance(myPos, target)
    score -= 12 * targetDistance

    invaders = self.visibleInvaders(successor)
    if invaders:
      closestInvaderDistance = min(
        self.getMazeDistance(myPos, invader.getPosition()) for invader in invaders
      )
      if self.chasingInvader:
        score -= 10 * closestInvaderDistance
      if closestInvaderDistance <= 1 and myState.scaredTimer == 0:
        score += 500
      if myState.scaredTimer > 0 and closestInvaderDistance <= 2:
        score -= 250

    if action == Directions.STOP:
      if self.holdingPortal and targetDistance == 0:
        score -= 1
      else:
        score -= 100
    reverse = Directions.REVERSE[gameState.getAgentState(self.index).configuration.direction]
    if action == reverse:
      score -= 2

    return score

  def primaryInvader(self, gameState):
    invaders = self.visibleInvaders(gameState)
    if not invaders:
      return None
    myPos = gameState.getAgentPosition(self.index)
    return max(
      invaders,
      key=lambda invader: (
        invader.numCarrying,
        -self.getMazeDistance(myPos, invader.getPosition())
      )
    )

  def nearestHomeEntry(self, pos):
    return min(self.homeEntries, key=lambda entry: self.getMazeDistance(pos, entry))

  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])
    )