#!/usr/bin/python

"""

Create illustrations for Amit's game programming web site---in
particular, to illustrate pathfinding and movement.  These
illustrations typically show something on top of a (square) grid.

"""

# Comparisons to make:
# 1. Diagonal vs. non-diagonal
# 2. Grid distance vs. Grid+tiebreaking vs. Euclidean vs. squared Euclidean
# 3. Regular grid vs. grid with skip links vs. hierarchal

__author__ = 'amitp@cs.stanford.edu (Amit J Patel)'
__copyright__ = '2003 Amit J Patel'

import pygame
import math

from norvig import search

# Grid concepts: a grid class represents a set of locations.  A
# 'location' is a location on the map.  A 'position' is a position on
# the screen.  The grid class can convert locations to positions and
# back.  It can also draw the grid locations.

START_POS = (3, 3)
BG_COLOR = (140, 140, 128)
EMPTY_COLOR = (224, 224, 216)
START_COLOR = (255, 128, 128)
GOAL_COLOR =  (128, 128, 255)
LINE_COLOR =  (255, 255, 255)
SHADOW_COLOR =(  0,   0,   0)

class SquareGrid:
    """
    A square grid has two orthogonal axes, I and J; total of I*J locations
    """
    SIZE_X = 10 # pixels
    SIZE_Y = 10 # pixels
    GRID_I = 45
    GRID_J = 45
    PIXELS_X = 1 + GRID_I*SIZE_X
    PIXELS_Y = 1 + GRID_J*SIZE_Y

    def locations(self):
        """Return a list of all locations in this grid"""
        locations = []
        for i in range(self.GRID_I):
            for j in range(self.GRID_J):
                locations.append((i, j))
        return locations

    def neighbors(self, (i, j)):
        """Return the locations adjacent to a location"""
        neighbors = [(i+1, j), (i-1, j), (i, j+1), (i, j-1)]
        if flags.diagonals:
            neighbors.extend([(i+1, j+1), (i-1, j+1), (i+1, j-1), (i-1, j-1)])
        return neighbors
    
    def loc_to_pos(self, (i, j)):
        """Convert any location to a position (x, y)"""
        return (1 + i*self.SIZE_X, 1 + j*self.SIZE_Y)

    def loc_to_textpos(self, (i, j)):
        return self.loc_to_pos((i+0.5, j+0.5))
    
    def pos_to_fractional_loc(self, (x, y)):
        """Convert a position (x, y) to a fractional location"""
        return ((x-1) / self.SIZE_X, (y-1) / self.SIZE_Y)
    
    def pos_to_loc(self, (x, y)):
        """Convert a position (x, y) to a grid-aligned location"""
        return ((x-1) // self.SIZE_X, (y-1) // self.SIZE_Y)

    def draw(self, surface, loc, color=EMPTY_COLOR):
        """Draw a single grid location in any color"""
        (x, y) = self.loc_to_pos(loc)
        pygame.draw.rect(surface, color,
                         (x, y, self.SIZE_X-1, self.SIZE_Y-1), 0)
        
class TriangleGrid:
    """
    A triangle has two non-orthogonal axes, I and J.  In each row J
    the I values are in range(0, 2*(J+1))).  One way to think
    of this is a parallelogram grid, where each of the parallelograms
    is two triangles (one pointed up and one pointed down).
    """
    GRID_SIZE = 20
    SIZE_X = 48 # width of bounding box, should be even
    SIZE_Y = int(SIZE_X*math.sqrt(3)/2) # height of bounding box
    PIXELS_X = 2 + SIZE_X * GRID_SIZE
    PIXELS_Y = 2 + SIZE_Y * GRID_SIZE
    
    def locations(self):
        locations = []
        for j in range(self.GRID_SIZE):
            for i in range(0, 2*j+1):
                locations.append((i, j))
        return locations

    def neighbors(self, (i, j)):
        neighbors = [(i-1, j), (i+1, j)]
        if i%2 == 0:
            # Triangle points up, so the neighbor is down
            neighbors.append((i+1, j+1))
        else:
            # Triangle points down, the neighbor is up
            neighbors.append((i-1, j-1))
        return neighbors
    
    def loc_to_pos(self, (i, j)):
        return (1 + (i+(self.GRID_SIZE-j))*(self.SIZE_X/2),
                (j+0.5)*self.SIZE_Y)
    def loc_to_textpos(self, loc):
        (x, y) = self.loc_to_pos(loc)
        if loc[0] % 2 == 1: # downward triangle
            return (x, y-0.25*self.SIZE_Y)
        else: # upward triangle
            return (x, y+0.25*self.SIZE_Y)
        
    def pos_to_fractional_loc(self, (x, y)):
        # TODO: this is wrong!!
        return ((x-1) / (self.SIZE_X/2), (y-1) / self.SIZE_Y)
    def pos_to_loc(self, (x, y)):
        # TODO: this is wrong!!
        return ((x-1) // (self.SIZE_X/2), (y-1) // self.SIZE_Y)

    def draw(self, surface, loc, color=EMPTY_COLOR):
        (x, y) = self.loc_to_pos(loc)
        if loc[0] % 2 == 1:
            points = [(x, y+0.5*self.SIZE_Y), # bottom middle
                      (x-0.5*self.SIZE_X, y-0.5*self.SIZE_Y), # top left
                      (x+0.5*self.SIZE_X, y-0.5*self.SIZE_Y), # top right
                      ]
        else:
            points = [(x, y-0.5*self.SIZE_Y), # top middle
                      (x-0.5*self.SIZE_X, y+0.5*self.SIZE_Y), # bottom left
                      (x+0.5*self.SIZE_X, y+0.5*self.SIZE_Y), # bottom right
                      ]
        pygame.draw.polygon(surface, color, points, 0)
        pygame.draw.polygon(surface, (0, 0, 0), points, 1)
                            
grid = SquareGrid()
# grid = TriangleGrid()

OBSTACLES = [
    'None',
    'Line',
    'Trap',
    ]
(OBSTACLE_NONE,
 OBSTACLE_LINE,
 OBSTACLE_TRAP) = range(len(OBSTACLES))

COLORS = [
    'None',
    'f values',
    'g values',
    'h values',
    'blend',
    ]
(COLOR_NONE,
 COLOR_F_VALUES,
 COLOR_G_VALUES,
 COLOR_H_VALUES,
 COLOR_BLEND) = range(len(COLORS))

ANNOTATIONS = [
    'None',
    'Coordinates',
    'f values',
    'g/h values',
    ]
(ANNOTATE_NONE,
 ANNOTATE_COORDINATES,
 ANNOTATE_F_VALUES,
 ANNOTATE_G_H_VALUES) = range(len(ANNOTATIONS))

HEURISTICS = [
    'None',
    'Grid-optimized',
    'Grid-optimized with scale-based tiebreaker',
    'Grid-optimized with cross-product tiebreaker',
    'Grid-optimized with both scale-based and cross-product tiebreakers',
    'Euclidean distance',
    'Euclidean distance squared',
    ]
(HEURISTIC_NONE,
 HEURISTIC_GRID,
 HEURISTIC_GRID_TIEBREAKING_SCALE,
 HEURISTIC_GRID_TIEBREAKING_CROSS,
 HEURISTIC_GRID_TIEBREAKING_BOTH,
 HEURISTIC_EUCLIDEAN,
 HEURISTIC_EUCLIDEAN_SQUARED) = range(len(HEURISTICS))

class flags:
    diagonals = 0
    draw_edges = 0
    pathfind_to_pointer = 0
    color = COLOR_BLEND
    annotation = ANNOTATE_NONE
    obstacles = OBSTACLE_NONE
    heuristic = HEURISTIC_GRID

class MapPathProblem(search.GraphProblem):
    """The problem of finding a path on a grid map"""
    def __init__(self, *args, **kwargs):
        search.GraphProblem.__init__(self, *args, **kwargs)
        self.probed = {} # {location: Node}

    def cross_product_adjustment(self, location):
        """Minor adjustment to nudge simple paths into staying on the
        straight line between initial position and goal"""
        di1 = location[0] - self.goal[0]
        dj1 = location[1] - self.goal[1]
        di2 = self.initial[0] - self.goal[0]
        dj2 = self.initial[1] - self.goal[1]
        cross_product = abs(di1*dj2 - di2*dj1)
        return cross_product * 1e-5
    
    def heuristic(self, location):
        """Multiple heuristic functions, controlled by flags.*"""
        di = self.goal[0] - location[0]
        dj = self.goal[1] - location[1]
        if flags.heuristic == HEURISTIC_NONE:
            h = 0
        elif flags.heuristic == HEURISTIC_EUCLIDEAN:
            (x0, y0) = grid.loc_to_pos(location)
            (x1, y1) = grid.loc_to_pos(self.goal)
            di = (x1-x0) / grid.SIZE_X
            dj = (y1-y0) / grid.SIZE_Y
            h = math.sqrt(di**2 + dj**2)
        elif flags.heuristic == HEURISTIC_EUCLIDEAN_SQUARED:
            (x0, y0) = grid.loc_to_pos(location)
            (x1, y1) = grid.loc_to_pos(self.goal)
            di = (x1-x0) / grid.SIZE_X
            dj = (y1-y0) / grid.SIZE_Y
            h = di**2 + dj**2
        else:
            if flags.diagonals:
                h = max(abs(di), abs(dj))
            else:
                h = abs(di) + abs(dj)
                
            if flags.heuristic in (HEURISTIC_GRID_TIEBREAKING_SCALE,
                                   HEURISTIC_GRID_TIEBREAKING_BOTH):
                h *= 1.001
            if flags.heuristic in (HEURISTIC_GRID_TIEBREAKING_CROSS,
                                   HEURISTIC_GRID_TIEBREAKING_BOTH):
                h += self.cross_product_adjustment(location)
        return h
        
    def h(self, node):
        """Heuristic that remembers which nodes are accessed"""
        # NOTE: this is the interface that the base class calls
        if self.probed.has_key(node.state):
            if node.path_cost < self.probed[node.state].path_cost:
                self.probed[node.state] = node
        else:
            self.probed[node.state] = node
            
        return self.heuristic(node.state)
    
class Map(search.Graph):
    """The undirected graph representing a grid map"""
    def __init__(self):
        nodes = {} # {position: {}}
        self.locations = {} # {position: position}
        for location in grid.locations():
            (i, j) = location
            if flags.obstacles == OBSTACLE_LINE and i in (10, 11, 12, 13, 20, 21) and 2 <= j <= 15:
                continue
            if flags.obstacles == OBSTACLE_TRAP and (
                (i in (20, 21) and 5 <= j <= 15)
                or (j in (6, 20) and 2 <= i <= 15)
                ):
                continue
                
            nodes[location] = {}
            self.locations[location] = location
        search.Graph.__init__(self, nodes, False)

        for location1 in grid.locations():
            if not nodes.has_key(location1):
                continue

            for location2 in grid.neighbors(location1):
                if nodes.has_key(location2):
                    if flags.heuristic in (
                        HEURISTIC_EUCLIDEAN,
                        HEURISTIC_EUCLIDEAN_SQUARED):
                        # TODO: this is square specific
                        di = location1[0] - location2[0]
                        dj = location1[1] - location2[1]
                        cost = math.sqrt(di**2 + dj**2)
                    else:
                        cost = 1
                            
                    self.connect(location1, location2, cost)

class Display:
    def __init__(self):
        pygame.init()
        self.screen = pygame.display.set_mode((grid.PIXELS_X, grid.PIXELS_Y))
        pygame.display.set_caption('Illustrations')
        self.font = pygame.font.SysFont('Tahoma', 11)

        self.background = pygame.Surface(self.screen.get_size()).convert()
        self.background.fill(BG_COLOR)

        self.map = Map()
        self.problem = None
        self.path = None
        
    def compute_path(self, destination):
        if self.path is None or self.path[0].state != destination:
            self.problem = MapPathProblem(START_POS, destination, self.map)
            results = search.astar_search(self.problem)
            self.path = results and results.path()
            assert self.path is None or self.path[0].state == destination
        
    def draw_line_segment(self, loc1, loc2, color,
                          intrusion=0.6, xoffset=0.0, yoffset=0.0):
        p = 0.5 * (1.0 - intrusion)
        (x1, y1) = grid.loc_to_textpos(loc1)
        (x2, y2) = grid.loc_to_textpos(loc2)
        # Adjust the positions by intrusion
        dx, dy = (x2-x1), (y2-y1)
        x1 += dx * intrusion*0.5
        y1 += dy * intrusion*0.5
        x2 -= dx * intrusion*0.5
        y2 -= dy * intrusion*0.5
        pygame.draw.line(self.screen, color,
                         (x1+xoffset, y1+yoffset), (x2+xoffset, y2+yoffset), 1)

    def draw_annotation(self, loc, annotation):
        text = self.font.render(annotation, 1, (0, 0, 0))
        text_x, text_y = grid.loc_to_textpos(loc)
        textpos = text.get_rect()
        textpos.centerx = text_x
        textpos.centery = text_y
        self.screen.blit(text, textpos)
        
    def main(self):
        done = 0
        while not done:
            for event in pygame.event.get():
                if event.type == QUIT:
                    done = 1
                elif event.type == KEYDOWN and event.key == K_ESCAPE:
                    done = 1
                elif event.type == KEYDOWN and event.key == K_a:
                    flags.annotation = (1 + flags.annotation) % len(ANNOTATIONS)
                    print 'Annotate:', ANNOTATIONS[flags.annotation]
                elif event.type == KEYDOWN and event.key == K_c:
                    flags.color = (1 + flags.color) % len(COLORS)
                    print 'Color:', COLORS[flags.color]
                elif event.type == KEYDOWN and event.key == K_d:
                    flags.diagonals = not flags.diagonals
                    self.path = None
                    self.problem = None
                    self.map = Map()
                    print 'Move along diagonals:', flags.diagonals
                elif event.type == KEYDOWN and event.key == K_h:
                    flags.heuristic = (1 + flags.heuristic) % len(HEURISTICS)
                    self.path = None
                    self.problem = None
                    self.map = Map()
                    print 'Heuristic:', HEURISTICS[flags.heuristic]
                elif event.type == KEYDOWN and event.key == K_o:
                    flags.obstacles = (1 + flags.obstacles) % len(OBSTACLES)
                    self.path = None
                    self.problem = None
                    self.map = Map()
                    print 'Obstacles:', OBSTACLES[flags.obstacles]
                elif event.type == KEYDOWN and event.key == K_f:
                    flags.pathfind_to_pointer = not flags.pathfind_to_pointer
                    print 'Pathfind to mouse pointer:', flags.pathfind_to_pointer
                elif event.type == MOUSEBUTTONDOWN:
                    self.path = None # Force recomputation
                    self.compute_path(grid.pos_to_loc(event.pos))
                    pass
                elif event.type == MOUSEBUTTONUP:
                    pass

            if pygame.mouse.get_focused() and flags.pathfind_to_pointer:
                self.compute_path(grid.pos_to_loc(pygame.mouse.get_pos()))
                
            self.screen.blit(self.background, (0, 0))

            for loc in grid.locations():
                if self.map.locations.has_key(loc):
                    grid.draw(self.screen, loc)

            if self.problem is not None:
                max_cost = 1
                min_fcost = 1e9
                max_fcost = 1
                for loc, node in self.problem.probed.items():
                    gcost = node.path_cost
                    hcost = self.problem.heuristic(loc)
                    max_cost = max(max_cost, gcost, hcost)
                    min_fcost = min(min_fcost, gcost+hcost)
                    max_fcost = max(max_fcost, gcost+hcost)
                    
                for loc, node in self.problem.probed.items():
                    gcost = node.path_cost
                    hcost = self.problem.heuristic(loc)
                    if flags.color == COLOR_F_VALUES:
                        color = (50,
                                 min(255, 100 + 200*(gcost+hcost-min_fcost)/(max_fcost-min_fcost)),
                                 50)
                    elif flags.color == COLOR_G_VALUES:
                        color = (50,
                                 50,
                                 min(255, 50 + 200*gcost/max_cost))
                    elif flags.color == COLOR_H_VALUES:
                        color = (min(255, 50 + 200*hcost/max_cost),
                                 50,
                                 50)
                    elif flags.color == COLOR_BLEND:
                        color = (min(255, 50 + 200*hcost/max_cost),
                                 min(255, 50 + 200*(gcost+hcost)/max_cost),
                                 min(255, 50 + 200*gcost/max_cost))
                    else: # COLOR_NONE
                        color = EMPTY_COLOR
                    grid.draw(self.screen, loc, color)

            grid.draw(self.screen, START_POS, START_COLOR)
            if self.path is not None:
                grid.draw(self.screen, self.path[0].state, GOAL_COLOR)
                grid.draw(self.screen, self.path[-1].state, START_COLOR)
                for k in range(len(self.path)-1):
                    loc1 = self.path[k].state
                    loc2 = self.path[k+1].state
                    self.draw_line_segment(loc1, loc2, SHADOW_COLOR,
                                           yoffset=1.0)
                    if flags.heuristic in (
                        HEURISTIC_NONE,
                        HEURISTIC_EUCLIDEAN_SQUARED
                        ):
                        # HACK: for dark grid colors
                        self.draw_line_segment(loc1, loc2, LINE_COLOR)

                if flags.annotation in (ANNOTATE_F_VALUES,
                                        ANNOTATE_G_H_VALUES):
                    for p in self.path:
                        annotation = '??'
                        if self.problem.probed.has_key(p.state):
                            gcost = self.problem.probed[p.state].path_cost
                            hcost = self.problem.heuristic(p.state)
                            if flags.annotation == ANNOTATE_F_VALUES:
                                annotation = '%.1f' % (gcost + hcost)
                            else:
                                annotation = '%.0f+%.0f' % (gcost, hcost)
                                
                        self.draw_annotation(p.state, annotation)
                        
            if flags.annotation == ANNOTATE_COORDINATES:
                for loc in grid.locations():
                    self.draw_annotation(loc, '%d,%d' % loc)
                    
            if flags.draw_edges:
                # Draw all edges
                for node1 in self.map.nodes():
                    for node2 in self.map.get(node1).keys():
                        self.draw_line_segment(node1, node2, SHADOW_COLOR, 0.2)
                        
            pygame.display.update()
            pygame.time.delay(100)
        
if __name__ == '__main__':
    from pygame.locals import *
    Display().main()