from zipfile import *
from PIL import Image
import StringIO
import os
from math import sqrt,cos,sin,atan2

# The representation of the graph.
# A graph is made of:
# 1. an array of nodes.  Each node with id n is at the nth index of this array. Each node contains a list of adjacencies.
# An adjacency is a 2-tuple, with element 0 being the node it is adjacent to, and element 0 being the cost to that node.
# 2. A dictionary, which maps coordinates to nodes.  If you are not sure the requested coordinate is valid, use the get() method
# of the dictionary class.
class PathGraph:

    # The nodes that make up the graph.  Contain minimal information, namely an id and the elements in the "coordinates".
    class PathGraphNode:
        def __init__(self, nid, coords = [], adjs = [], state = []):
            self.nid = nid # Node id
            self.coords = coords # The coordinates of the node.  Typically posX, posY
            self.adjs = adjs # List of adjacencies.  (adjacent node id, cost)
            self.state = state # A list of the states of the node.

        # Return the node's id
        def getID(self):
            return self.nid

        # Get a particular coordinate of the node
        def getCoordinate(self, coord=0):
            return self.coords[coord]

        # Get all the coordinates of the node
        def getCoordinates(self):
            return self.coords

        # Get how many coordinates there are in the node
        def getNumCoordinates(self):
            return len(self.coords)

        # Get the list of Adjacencies
        def getAdjacencies(self):
            return self.adjs

        # Return if a given node with id nid is adjacent to this node
        def isAdjacent(self, nid):
            for (n, d) in self.adjs:
                if n == nid:
                    return True
            return False

        # get the state stateNum.  Returns None if not found.
        def getState(self, stateNum=0):
            if self.hasState(stateNum):
                return self.state[stateNum]
            else:
                return None
				
        # Set the state stateNum to the stateVal.
        # If stateNum doesn't exist, create it.
        def setState(self, stateVal, stateNum=0):
            if self.hasState(stateNum):
                self.state[stateNum] = stateVal
            else:
                while len(self.state) < stateNum -1:
                    self.state.append(None)
                self.state.append(stateVal)

        # Add a new state to the statelist, and assign it value stateVal.
        # Returns the state number created.
        def addState(self, stateVal):
            self.state.append(stateVal)
            return len(self.state) - 1

        # Return True if the stateNum exists, False otherwise.
        def hasState(self, stateNum=0):
            return len(self.state) > stateNum

        # Set a particular coordinate of the node
        def setCoordinate(self, val, coord=0):
            coords[coord]=val
            return self.coords[coord]

        # Set all the coordinates of the node
        def setCoordinates(self,somecoords):
            self.coords = somecoords
            return self.coords

        # Returns a string representation of the node
        def coordinateString(self):
            retStr = ""
            for i in self.coords:
                retStr += str(i) + " "
            return retStr.rstrip(' ')

    # Creates a trajectory on the specified node ids.
    # Converts it to a path if possible.  Currently, trajectory can only consist of two node ids.
    # Trajectories only work for integer graphs
    class Trajectory:
        def __init__(self, graph, trajectory = []):
            self.graph = graph
            self.trajectory = trajectory
            self.path = None
            self.completePath = None

        # Converts the trajectory to a node path as far as it can along a valid path.
        def ConvertToPath(self):
            self.completePath = False
            self.path = [self.trajectory[0]]
            if self.trajectory[0] == self.trajectory[-1]:
                return 
            for nodeFromID, nodeToID in zip(self.trajectory[:-1],self.trajectory[1:]):
                #print "ConvertToPath ", nodeFromID, nodeToID #jrs
                nodeTo = self.graph.nodeList[nodeToID]
                nodeFrom = self.graph.nodeList[nodeFromID]

                toVals = self.graph.nodeList[nodeToID].coords
                fromVals = self.graph.nodeList[nodeFromID].coords

                distVals = tuple(toVal - fromVal for toVal, fromVal in zip(toVals, fromVals))
                gran = max(tuple(abs(dVal) for dVal in distVals)) # maximum distance cooordinate, used to figure how many steps need to be taken

                deltaVals = tuple(distVal / float(gran) for distVal in distVals) # a vector that should take us from current node to next node on traj (if no obstacles)
                #jrs
		#print "toVals: ", toVals #jrs
		#print "fromVals: ", fromVals #jrs
                #print "distVals: ", distVals #jrs
                #print "gran: ", gran #jrs
                #print "deltaVals: ", deltaVals #jrs
                prevNode = nodeFrom
                for i in range(1,gran+1):
                    node = self.graph.nodeDict.get(tuple(int(fromVal + i*deltaVal + 0.5) for fromVal, deltaVal in zip(fromVals, deltaVals)))
                    if node and prevNode.isAdjacent(node): # The step we wanted to take is valid
                        self.path.append(node)
                        prevNode = self.graph.nodeList[node]
                        if node == nodeToID:
                            self.completePath = True
                            #print "Converted traj to path: ", map(self.graph.getNode,self.path)
                            break
                    else:
                        break

    # Creates a new graph
    def __init__(self, coordCount, nodeList=[], nodeDict = {}):
        self.nodeList = nodeList
        self.coords = coordCount
        self.nodeDict = nodeDict

    # Get the node with coordinates nodeCoords
    def getNodeID(self, nodeCoords):
        return self.nodeDict[nodeCoords]

    # Get the adjacency list of node ID nid
    def getAdjacencies(self, nid):
        return self.nodeList[nid].getAdjacencies()

    # Gets the cost of the adjacency between node id nid1 and node id nid2.
    # Returns None if no such cost.
    def getAdjacencyCost(self, nid1, nid2):
        adjs = self.getAdjacencies(nid1)
        for (n, d) in adjs:
            if n == nid2:
                return d
        return None

    # Get the node list
    def getNodes(self):
        return self.nodeList

    # Get the node with id nid
    def getNode(self, nid):
        return self.nodeList[nid]

    # Get a node by the coords.  Returns None is there is no such node
    def getNodeByCoords(self, coords):
        return self.nodeDict[coords]

    # Get the number of nodes.
    def getNodeCount(self):
        return len(self.nodeList)

    # Get the coordinates of node with id nid
    def getNodeCoords(self, nid):
        return self.nodeList[nid].coords

    # Get the node dictionary
    def getNodeDict(self):
        return self.nodeDict

    # Get the id of a node that is adjacent in the specified direction to the node id nid.
    # Returns None if there is no adjacency.
    def getAdjNID(self, nid, direction):
        node = self.nodeList[nid]
        tgt = self.nodeDict(node.getCoordinates() + direcion)
        if tgt != None and node.isAdjacent(tgt):
            return tgt
        else:
            return None
            

    # Print out the graph.
    def printGraph(self):
        for i in range(len(self.nodeList)):
            print "Node " + str(self.nodeList[i].nid) + " at " + self.nodeList[i].coordinateString()
            for j in range(len(self.nodeList[i].adjs)):
                print "\tAdjacent to " + str(self.nodeList[i].adjs[j][0]) + " with cost " + str(self.nodeList[i].adjs[j][1])


    # Add a new node to the list, with the given coordinates and ajacencies.
    # Returns the new node's id
    def addNode(self, coordinates, adjacencies):
        nid = len(self.nodeList)
        self.nodeList.append(PathGraph.PathGraphNode(nid,coordinates, adjacencies))
        for adj in adjacencies:
            nbr = adj[0]
            edge = ( nid, adj[1] )
            self.nodeList[nbr].adjs.append(edge)
        self.nodeDict[coordinates] = nid
        return nid

# "Static" functions

#Creates and returns a graph object for the mapFile passed in.
def createGraph(mapFile):
        
    nodeList = []
    nodeDict = {}


    coordCount = int(mapFile.readline())      
    graphData = mapFile.readline().split()

    while graphData != []:
        coords = []
        for i in range(coordCount):
            coords.append(int(graphData[i]))

        coords = tuple(coords)
        nodeDict[coords] = len(nodeList)

        adjs = []
        i = coordCount
        while i < len(graphData):
            adjs.append( (int(graphData[i]), float(graphData[i+1]) ) )
            i += 2

        nodeList.append(PathGraph.PathGraphNode(len(nodeList), coords, adjs, []))
                
        graphData = mapFile.readline().split()

    return PathGraph(coordCount, nodeList, nodeDict)

# Create a graph object from the bitmap image im.
# Returns a tuple.  Entry 1 is the graphe object, entry 2 is the opened bitmap.
def loadImageFile(im):
    graphIm = Image.open(im)
    pixels = graphIm.load()
    #print graphIm.size

    nodeList = []
    nodeDict = {}

    for i in range(graphIm.size[0]):
        for j in range(graphIm.size[1]):
            if pixels[i,j] == ( 255,255,255 ): # Only white nodes are pathfindable.
                nid = len(nodeList)
                coords = (i,j)

                nodeDict[coords] = nid
                adjs = []

                #look at neighbors already seen.  Guaranteed to only be north, west
                if (i-1,j) in nodeDict : #north
                    adjs.append( (nodeDict[ (i-1,j) ], 1) ) #Add ajacency for this node
                    nodeList[ nodeDict[ (i-1,j) ] ].adjs.append( (nid, 1) ) #Add adjacency for other node
                if (i,j-1) in nodeDict: #west
                    adjs.append( (nodeDict[ (i,j-1) ], 1) )
                    nodeList[ nodeDict[ (i,j-1) ] ].adjs.append( (nid, 1) ) #Add adjacency for other node
                if (i-1,j-1) in nodeDict: #nw
                    adjs.append( (nodeDict[ (i-1,j-1) ], 1.4143) )
                    nodeList[ nodeDict[ (i-1,j-1) ] ].adjs.append( (nid, 1.4143) )
                if (i-1,j+1) in nodeDict: #sw
                    adjs.append( (nodeDict[ (i-1,j+1) ], 1.4143) )
                    nodeList[ nodeDict[ (i-1,j+1) ] ].adjs.append( (nid, 1.4143) )

                nodeList.append(PathGraph.PathGraphNode(nid, coords, adjs))

    graph = PathGraph(2, nodeList, nodeDict)
    return (graph, graphIm)
    
# Load a .map file
# Returns a tuple.  Entry 1 is the graphe object, entry 2 is the opened bitmap.
def loadMapFile(mapFile):
    if mapFile == None:
        return None
    else:
        unzipMapFile = ZipFile(mapFile, mode='r')
            
        for f in unzipMapFile.namelist():
            if f.endswith('.graph'):
                graph = createGraph(unzipMapFile.open(f))
            else:
                image = Image.open(StringIO.StringIO(unzipMapFile.read(f)))

        unzipMapFile.close()
        return (graph, image)
        
# Create a .map file with image im, saved at mapFile
# OVERWRITES mapFile!
def createMapFile(im, mapFile):
    zf = ZipFile(mapFile, mode='w', compression=ZIP_DEFLATED)

    base = os.path.basename(im)
    baseWOExt = os.path.splitext(base)[0]

    graph, image = loadImageFile(im)

    mapDataStr = ""
    
    mapDataStr += (str(graph.coords))
    for node in graph.getNodes():
        mapDataStr += ('\n' + str(node.coordinateString()))
        for adj in graph.getAdjacencies(node.getID()):
            mapDataStr += (" " + str(adj[0]) + " " + str(adj[1]))

    try:
        zf.writestr(baseWOExt + ".graph",mapDataStr)
        zf.write(im, base)
    finally:
        zf.close()