from game import Actions, Directions
from util import raiseNotDefined

from typing import Set, Tuple

class SearchProblem:
    """
    This class outlines the structure of a search problem, but doesn't implement
    any of the methods (in object-oriented terminology: an abstract class).

    You do not need to change anything in this class, ever.
    """

    def getStartState(self):
        """
        Returns the start state for the search problem.
        """
        raiseNotDefined()

    def isGoalState(self, state):
        """
          state: Search state

        Returns True if and only if the state is a valid goal state.
        """
        raiseNotDefined()

    def getSuccessors(self, state):
        """
          state: Search state

        For a given state, this should return a list of triples, (successor,
        action, stepCost), where 'successor' is a successor to the current
        state, 'action' is the action required to get there, and 'stepCost' is
        the incremental cost of expanding to that successor.
        """
        raiseNotDefined()

    def getCostOfActions(self, actions):
        """
         actions: A list of actions to take

        This method returns the total cost of a particular sequence of actions.
        The sequence must be composed of legal moves.
        """
        raiseNotDefined()


class PositionSearchProblem(SearchProblem):
    """
    A search problem defines the state space, start state, goal test, successor
    function and cost function.  This search problem can be used to find paths
    to a particular point on the pacman board.

    The state space consists of (x,y) positions in a pacman game.

    Note: this search problem is fully specified; you should NOT change it.
    """

    def __init__(self, gameState, costFn = lambda x: 1, goal=(1,1), start=None, warn=True, visualize=True):
        """
        Stores the start and goal.

        gameState: A GameState object (pacman.py)
        costFn: A function from a search state (tuple) to a non-negative number
        goal: A position in the gameState
        """
        self.walls = gameState.getWalls()
        self.startState = gameState.getPacmanPosition()
        if start != None: self.startState = start
        self.goal = goal
        self.costFn = costFn
        self.visualize = visualize
        if warn and (gameState.getNumFood() != 1 or not gameState.hasFood(*goal)):
            print('Warning: this does not look like a regular search maze')

        # For display purposes
        (self._visited, self._visitedlist, self._expanded) = ({}, [], 0) # DO NOT CHANGE

    def getStartState(self):
        return self.startState

    def isGoalState(self, state):
        isGoal = state == self.goal

        # For display purposes only
        if isGoal and self.visualize:
            self._visitedlist.append(state)
            import __main__
            if '_display' in dir(__main__):
                if 'drawExpandedCells' in dir(__main__._display): #@UndefinedVariable
                    __main__._display.drawExpandedCells(self._visitedlist) #@UndefinedVariable

        return isGoal

    def getSuccessors(self, state):
        """
        Returns successor states, the actions they require, and a cost of 1.

         As noted in search.py:
             For a given state, this should return a list of triples,
         (successor, action, stepCost), where 'successor' is a
         successor to the current state, 'action' is the action
         required to get there, and 'stepCost' is the incremental
         cost of expanding to that successor
        """

        successors = []
        for action in Directions.LIST:
            x,y = state
            (dx, dy) = Actions.directionToVector(action)
            (nextx, nexty) = (int(x + dx), int(y + dy))
            if not self.walls[nextx][nexty]:
                nextState = (nextx, nexty)
                cost = self.costFn(nextState)
                successors.append( ( nextState, action, cost) )

        # Bookkeeping for display purposes
        self._expanded += 1 # DO NOT CHANGE
        if state not in self._visited:
            self._visited[state] = True
            self._visitedlist.append(state)

        return successors

    def getCostOfActions(self, actions):
        """
        Returns the cost of a particular sequence of actions. If those actions
        include an illegal move, return 999999.
        """
        if actions == None: return 999999
        (x, y)= self.getStartState()
        cost = 0
        for action in actions:
            # Check figure out the next state and see whether its' legal
            (dx, dy) = Actions.directionToVector(action)
            (x, y) = (int(x + dx), int(y + dy))
            if self.walls[x][y]: return 999999
            cost += self.costFn((x,y))
        return cost


class CornersProblem(SearchProblem):
    """
    This search problem finds paths through all four corners of a layout.

    You must select a suitable state space and successor function
    """

    def __init__(self, startingGameState):
        """
        Stores the walls, pacman's starting position and corners.
        """
        self.heuristicInfo = {}
        self.startingGameState = startingGameState
        self.startingPosition = startingGameState.getPacmanPosition()
        self.walls = startingGameState.getWalls()
        (top, right) = self.walls.height-2, self.walls.width-2
        self.corners = ((1,1), (1,top), (right, 1), (right, top))
        for corner in self.corners:
            if not startingGameState.hasFood(*corner):
                print('Warning: no food in corner ' + str(corner))
        self.initial_blob_size = len(self.corners)
        self._expanded = 0 # DO NOT CHANGE; Number of search nodes expanded
        # Please add any code here which you would like to use
        # in initializing the problem
        self.maze_dist_cache = {}

        "*** YOUR CODE HERE ***"
        self.corners_set = frozenset(self.corners)

    def getStartState(self):
        """
        Returns the start state (in your state space, not the full Pacman state
        space)
        """
        "*** YOUR CODE HERE ***"
        # A state is a tuple: (pacman position, set-of-positions-visited)
        start_unvisited_corners = self.corners_set - {self.startingPosition}
        return (self.startingPosition, start_unvisited_corners)
        # raiseNotDefined()

    def isGoalState(self, state):
        """
        Returns whether this search state is a goal state of the problem.
        """
        "*** YOUR CODE HERE ***"
        return len(state[1]) == 0

    def getSuccessors(self, state):
        """
        Returns successor states, the actions they require, and a cost of 1.

         As noted in search.py:
            For a given state, this should return a list of triples, (successor,
            action, stepCost), where 'successor' is a successor to the current
            state, 'action' is the action required to get there, and 'stepCost'
            is the incremental cost of expanding to that successor
        """
        """ YOUR CODE GOES HERE """
        raw_successors = [(self.get_raw_successor(direction, state), direction) for direction in Directions.LIST]
        cost = 1
        successors = [(((nextx, nexty), next_unvisited_corners), action, cost)
                      for (((nextx, nexty), next_unvisited_corners), action) in raw_successors
                      if not self.walls[nextx][nexty]]
        self._expanded += 1 # DO NOT CHANGE
        return successors

    def get_raw_successor(self, action, state) -> (Tuple[int, int], Set[Tuple[int, int]]):
        """ state = ((x, y), unvisited_corners) => ((nextx, nexty), next_unvisited_corners) """
        ((x, y), unvisited_corners) = state
        (dx, dy) = Actions.directionToVector(action)
        (nextx, nexty) = int(x + dx), int(y + dy)
        next_unvisited_corners = (unvisited_corners - {(nextx, nexty)}) if (nextx, nexty) in self.corners_set else \
                                 unvisited_corners
        nextState = ((nextx, nexty), next_unvisited_corners)
        return nextState

    def getCostOfActions(self, actions):
        """
        Returns the cost of a particular sequence of actions.  If those actions
        include an illegal move, return 999999.  This is implemented for you.
        """
        if actions is None: return 999999
        (x, y) = self.startingPosition
        for action in actions:
            (dx, dy) = Actions.directionToVector(action)
            (x, y) = int(x + dx), int(y + dy)
            if self.walls[x][y]: return 999999
        return len(actions)

class FoodSearchProblem:
    """
    A search problem associated with finding the a path that collects all of the
    food (dots) in a Pacman game.

    A search state in this problem is a tuple ( pacmanPosition, foodGrid ) where
      pacmanPosition: a tuple (x,y) of integers specifying Pacman's position
      foodGrid:       a Grid (see game.py) of either True or False, specifying remaining food
    """
    def __init__(self, startingGameState):
        self.start = (startingGameState.getPacmanPosition(), frozenset(startingGameState.getFood().asList()))
        self.walls = startingGameState.getWalls()
        (self.top, self.right) = self.walls.height-2, self.walls.width-2
        self.startingGameState = startingGameState
        self.initial_blob_size = len(startingGameState.getFood().asList())
        self._expanded = 0 # DO NOT CHANGE
        self.maze_dist_cache = {}

    def getStartState(self):
        return self.start

    def isGoalState(self, state):
        return len(state[1]) == 0

    def getSuccessors(self, state):
        """Returns successor states, the actions they require, and a cost of 1."""
        successors = []
        self._expanded += 1 # DO NOT CHANGE
        for direction in Directions.LIST:
            ((x, y), foodSet) = state
            (dx, dy) = Actions.directionToVector(direction)
            (nextx, nexty) = (int(x + dx), int(y + dy))
            if not self.walls[nextx][nexty]:
                nextFood = foodSet - {(nextx, nexty)}
                successor = ( ((nextx, nexty), nextFood), direction, 1)
                successors.append(successor)
        return successors

    def getCostOfActions(self, actions):
        """Returns the cost of a particular sequence of actions.  If those actions
        include an illegal move, return 999999"""
        x,y= self.getStartState()[0]
        cost = 0
        for action in actions:
            # figure out the next state and see whether it's legal
            dx, dy = Actions.directionToVector(action)
            x, y = int(x + dx), int(y + dy)
            if self.walls[x][y]:
                return 999999
            cost += 1
        return cost

class AnyFoodSearchProblem(PositionSearchProblem):
    """
    A search problem for finding a path to any food.

    This search problem is just like the PositionSearchProblem, but has a
    different goal test, which you need to fill in below.  The state space and
    successor function do not need to be changed.

    The class definition above, AnyFoodSearchProblem(PositionSearchProblem),
    inherits the methods of the PositionSearchProblem.

    You can use this search problem to help you fill in the findPathToClosestDot
    method.
    """

    def __init__(self, gameState):
        """Stores information from the gameState.  You don't need to change this."""
        # Store the food for later reference
        self.food = gameState.getFood()

        # Store info for the PositionSearchProblem (no need to change this)
        self.walls = gameState.getWalls()
        self.startState = gameState.getPacmanPosition()
        self.costFn = lambda x: 1
        self._visited, self._visitedlist, self._expanded = {}, [], 0 # DO NOT CHANGE


    def isGoalState(self, state):
        """
        The state is Pacman's position. Fill this in with a goal test that will
        complete the problem definition.
        """
        (x, y) = state

        "*** YOUR CODE HERE ***"
        return self.food[state[0]][state[1]]