'''
 Learned A* through Dubrovnik
 Z. Butler, 4/2015
 note xml includes all nodes for all partially-present ways
 use a bounding box to ignore nodes outside the region, should be safe

 All but compute_params, make_cost is the Lab 1 solution
'''

from Tkinter import *
import struct
import xml.etree.ElementTree as ET
from Queue import *
import math
from sys import *
import random

LEFTLON = 18.055
RIGHTLON = 18.125
TOPLAT = 42.675
BOTLAT = 42.635
WIDTH = RIGHTLON-LEFTLON
HEIGHT = TOPLAT-BOTLAT
LONRATIO = math.cos(TOPLAT*3.1415/180)
WINWID = 800
WINHGT = (int)((WINWID/LONRATIO)*HEIGHT/WIDTH)
TOXPIX = WINWID/WIDTH
TOYPIX = WINHGT/HEIGHT
EPIX = 3601 #width,height of elevation array
MPERLAT = 111000
MPERLON = MPERLAT*LONRATIO

def node_dist_3_cost(edge):
    return node_dist_3(edge.src,edge.dest)

def node_dist_2(n1, n2):
    dx = (n2.pos[0]-n1.pos[0])*MPERLON
    dy = (n2.pos[1]-n1.pos[1])*MPERLAT
    return math.sqrt(dx*dx+dy*dy) # in meters

def node_dist_3(n1, n2):
    dx = (n2.pos[0]-n1.pos[0])*MPERLON
    dy = (n2.pos[1]-n1.pos[1])*MPERLAT
    dz = n2.elev-n1.elev
    if dz < 0: dz = 0
    return math.sqrt(dx*dx+dy*dy+dz*dz) # in meters
    
class Path(): #recorded path, not calculated
    __slots__ = ('nodes','time','creator', 'dist', 'uphill','allnodes')
    def __init__(self,line,planner):
        self.nodes = []
        parts = line.split(",")
        #print parts
        timeind = len(parts)-4
        for n in range(timeind):
            self.nodes.append(long(parts[n]))
        self.time = float(parts[timeind]) + float(parts[timeind+1])/60.0
        self.creator = parts[timeind+2]

        self.dist = 0
        self.uphill = 0
        self.allnodes = []
        for n in range(len(self.nodes)-1):
            (segnodes,segways),_ = planner.plan(self.nodes[n],self.nodes[n+1])
            for sn in range(len(segnodes)-1):
                self.allnodes.append(segnodes[sn].id) # all but last node, which will be first node of next segment
                self.dist += node_dist_2(segnodes[sn],segnodes[sn+1])
                ediff = segnodes[sn+1].elev - segnodes[sn].elev
                if ediff > 0:
                    self.uphill += ediff
        self.allnodes.append(self.nodes[-1]) # very last node
        self.dist /= 1000.0
        self.uphill /= 100.0

    def __str__(self):
        return self.creator + ": " + str(self.dist) + " time " + str(self.time)
        
class Node():
    __slots__ = ('id', 'pos', 'ways', 'elev')
    def __init__(self,id,p,e=0):
        self.id = id
        self.pos = p
        self.ways = []
        self.elev = e
        self.waystr = None
    def __str__(self):
        if self.waystr is None:
            self.waystr = self.get_waystr()
        return str(self.pos) + ": " + self.waystr
    def get_waystr(self):
        if self.waystr is None:
            self.waystr = ""
            self.wayset = set()
            for w in self.ways:
                self.wayset.add(w.way.name)
            for w in self.wayset:
                self.waystr += w.encode("utf-8") + " "
        return self.waystr
        

class Edge():
    __slots__ = ('way','src','dest')
    def __init__(self, w, src, d):
        self.way = w
        self.dest = d
        self.src = src
        self.len = node_dist_2(src,d)

class Way():
    __slots__ = ('name','type','nodes')
    # nodes here for ease of drawing only
    def __init__(self,n,t):
        self.name = n
        self.type = t
        self.nodes = []

class Planner():
    __slots__ = ('nodes', 'ways', 'costfn')
    def __init__(self,n,w,c):
        self.nodes = n
        self.ways = w
        self.costfn = c

    def heur(self,node,gnode):
        return node_dist_2(node,gnode)
    
    def plan(self,sid,gid):
        if not isinstance(sid,Node):
            s = nodes[sid]
            g = nodes[gid]
        else:
            s = sid
            g = gid
        parents = {}
        costs = {}
        q = PriorityQueue()
        q.put((self.heur(s,g),s))
        parents[s] = None
        costs[s] = 0
        while not q.empty():
            cf, cnode = q.get()
            if cnode == g:
                #print ("Path found, time will be",costs[g])
                return self.make_path(parents,g), costs[g]
            for edge in cnode.ways:
                newcost = costs[cnode] + self.costfn(edge)
                #print (edge.way.name, edge.src.id, edge.dest.id, \
                #       node_dist_2(edge.src,edge.dest), self.costfn(edge))
                if edge.dest not in parents or newcost < costs[edge.dest]:
                    parents[edge.dest] = (cnode, edge.way)
                    costs[edge.dest] = newcost
                    q.put((self.heur(edge.dest,g)+newcost,edge.dest))

    def make_path(self,par,g):
        nodes = []
        ways = []
        curr = g
        nodes.append(curr)
        while par[curr] is not None:
            prev, way = par[curr]
            ways.append(way.name)
            nodes.append(prev)
            curr = prev
        nodes.reverse()
        ways.reverse()
        return nodes,ways

class DispWin(Frame):
    
    __slots__ = ('whatis', 'nodes', 'ways', 'elevs', 'nodelab', 'elab', \
                 'planner', 'lastnode', 'startnode', 'goalnode', 'knownpaths', \
                 'lastpath')
    
    def lat_lon_to_pix(self,latlon):
        x = (latlon[1]-LEFTLON)*(TOXPIX)
        y = (TOPLAT-latlon[0])*(TOYPIX)
        return x,y

    def pix_to_elev(self,x,y):
        return self.lat_lon_to_elev(((TOPLAT-(y/TOYPIX)),((x/TOXPIX)+LEFTLON)))

    def lat_lon_to_elev(self,latlon):
        # row is 0 for 43N, 1201 (EPIX) for 42N
        row = (int)((43 - latlon[0]) * EPIX)
        # col is 0 for 18 E, 1201 for 19 E
        col = (int)((latlon[1]-18) * EPIX)
        return self.elevs[row*EPIX+col]

    def maphover(self,event):
        #print event.x,event.y
        self.elab.configure(text = str(self.pix_to_elev(event.x,event.y)))
        for (dx,dy) in [(0,0),(-1,0),(0,-1),(1,0),(0,1),(-1,-1),(-1,1),(1,-1),(1,1)]:
            ckpos = (event.x+dx,event.y+dy)
            if ckpos in self.whatis:
                self.lastnode = self.whatis[ckpos]
                lnpos = self.lat_lon_to_pix(self.nodes[self.lastnode].pos)
                self.canvas.coords('lastdot',(lnpos[0]-2,lnpos[1]-2,lnpos[0]+2,lnpos[1]+2))
                nstr = str(self.lastnode)
                nstr += " "
                nstr += str(self.nodes[self.whatis[ckpos]].get_waystr())
                self.nodelab.configure(text=nstr)
                return
            #self.nodelab.configure(text="None")

    def mapclick(self,event):
        if self.lastnode is None:
            return
        #print "Clicked on "+str(event.x)+","+str(event.y)+" last node "+str(self.lastnode)
        if self.startnode is None:
            self.startnode = self.nodes[self.lastnode]
            self.snpix = self.lat_lon_to_pix(self.startnode.pos)
            self.canvas.coords('startdot',(self.snpix[0]-2,self.snpix[1]-2,self.snpix[0]+2,self.snpix[1]+2))
        elif self.goalnode is None:
            self.goalnode = self.nodes[self.lastnode]
            self.snpix = self.lat_lon_to_pix(self.goalnode.pos)
            self.canvas.coords('goaldot',(self.snpix[0]-2,self.snpix[1]-2,self.snpix[0]+2,self.snpix[1]+2))

    def clear(self):
        self.lastnode = None
        self.goalnode = None
        self.startnode = None
        self.canvas.coords('startdot',(0,0,0,0))
        self.canvas.coords('goaldot',(0,0,0,0))
        self.canvas.coords('path',(0,0,0,0))
            
    def draw_path(self):
        #print "Drawing!"
        pcoords = []
        self.lastpath += 1
        path = self.knownpaths[self.lastpath]
        for nodeid in path.allnodes:
            nodepos = self.lat_lon_to_pix(self.nodes[(long)(nodeid)].pos)
            pcoords.append(nodepos[0])
            pcoords.append(nodepos[1])
        #print path.allnodes
        #print [nodes[n].elev for n in path.allnodes]
        self.canvas.coords('path',*pcoords)
        pstr = "Dist: " + str(path.dist) + " up: " + str(path.uphill)
        self.nodelab.configure(text=pstr)
        
            
    def plan_path(self):
        print "Planning!"
        if self.startnode is None or self.goalnode is None:
            print "Sorry, not enough info."
            return
        print ("From", self.startnode.id, "to", self.goalnode.id)
        (nodes,ways),cost = self.planner.plan(self.startnode, self.goalnode)
        print ("Expected time:",cost)
        lastway = ""
        for wayname in ways:
            if wayname != lastway:
                print wayname
                lastway = wayname
        coords = []
        for node in nodes:
            npos = self.lat_lon_to_pix(node.pos)
            coords.append(npos[0])
            coords.append(npos[1])
            #print node.id
        self.canvas.coords('path',*coords)

    def __init__(self,master,nodes,ways,coastnodes,elevs,paths,planner):
        self.whatis = {}
        self.nodes = nodes
        self.ways = ways
        self.elevs = elevs
        self.startnode = None
        self.goalnode = None
        self.knownpaths = paths
        self.lastpath = -1
        self.planner = planner
        thewin = Frame(master)
        w = Canvas(thewin, width=WINWID, height=WINHGT)#, cursor="crosshair")
        w.bind("<Button-1>", self.mapclick)
        w.bind("<Motion>", self.maphover)
        for waynum in self.ways:
            nlist = self.ways[waynum].nodes
            thispix = self.lat_lon_to_pix(self.nodes[nlist[0]].pos)
            if len(self.nodes[nlist[0]].ways) > 2:
                #w.create_oval(thispix[0],thispix[1],thispix[0]+1,thispix[1]+1,outline='blue')
                self.whatis[((int)(thispix[0]),(int)(thispix[1]))] = nlist[0]
            for n in range(len(nlist)-1):
                nextpix = self.lat_lon_to_pix(self.nodes[nlist[n+1]].pos)
                self.whatis[((int)(nextpix[0]),(int)(nextpix[1]))] = nlist[n+1]
                w.create_line(thispix[0],thispix[1],nextpix[0],nextpix[1])
                #if len(self.nodes[nlist[n+1]].ways) > 2:
                #    w.create_oval(nextpix[0],nextpix[1],nextpix[0]+1,nextpix[1]+1,outline='blue')
                thispix = nextpix
        thispix = self.lat_lon_to_pix(self.nodes[coastnodes[0]].pos)
        for n in range(len(coastnodes)-1):
            nextpix = self.lat_lon_to_pix(self.nodes[coastnodes[n+1]].pos)
            w.create_line(thispix[0],thispix[1],nextpix[0],nextpix[1],fill="blue")
            thispix = nextpix

        w.create_line(0,0,0,0,fill='orange',width=3,tag='path')

        w.create_oval(0,0,0,0,outline='green',fill='green',tag='startdot')
        w.create_oval(0,0,0,0,outline='red',fill='red',tag='goaldot')
        w.create_oval(0,0,0,0,outline='blue',fill='blue',tag='lastdot')
        w.pack(fill=BOTH)
        self.canvas = w

        cb = Button(thewin, text="Clear", command=self.clear)
        cb.pack(side=RIGHT,pady=5)

        pb = Button(thewin, text="Plan!", command=self.plan_path)
        pb.pack(side=RIGHT,pady=5)

        sb = Button(thewin, text="Next", command=self.draw_path)
        sb.pack(side=RIGHT,pady=5)

        nodelablab = Label(thewin, text="Node:")
        nodelablab.pack(side=LEFT, padx = 5)
        
        self.nodelab = Label(thewin,text="None")
        self.nodelab.pack(side=LEFT,padx = 5)

        elablab = Label(thewin, text="Elev:")
        elablab.pack(side=LEFT, padx = 5)

        self.elab = Label(thewin, text = "0")
        self.elab.pack(side=LEFT, padx = 5)
        
        thewin.pack()

def read_paths(filename,planner):
    paths = []
    #print "pathdata = ["
    for line in open(filename):
        onepath = Path(line,planner)
        paths.append(onepath)
        #print onepath.dist, onepath.uphill, onepath.time
        #print ";"
    #print "];"
    #tograde = sorted(paths, key=lambda p: p.creator)
    #for p in tograde:
    #    print p
    return paths

def build_elevs(efilename):
    efile = open(efilename)
    estr = efile.read()
    elevs = []
    for spot in range(0,len(estr),2):
        elevs.append(struct.unpack('>h',estr[spot:spot+2])[0])

    #elevs = array.array('h') #signed shorts
    #elevs.fromfile(efile,EPIX*EPIX)
    #swapping endian-ness here:
    #for e in range(len(elevs)):
    #    elevs[e] = 256*(elevs[e]%256)+(elevs[e]/256)
    return elevs

def build_graph(elevs):
    tree = ET.parse('dbv.osm')
    root = tree.getroot()

    nodes = dict()
    ways = dict()
    waytypes = set()
    coastnodes = []
    for item in root:
        if item.tag == 'node':
            coords = ((float)(item.get('lat')),(float)(item.get('lon')))
            # row is 0 for 43N, 1201 (EPIX) for 42N
            erow = (int)((43 - coords[0]) * EPIX)
            # col is 0 for 18 E, 1201 for 19 E
            ecol = (int)((coords[1]-18) * EPIX)
            try:
                el = elevs[erow*EPIX+ecol]
            except IndexError:
                el = 0
            nodes[(long)(item.get('id'))] = Node((long)(item.get('id')),coords,el)            
            '''
            if (coords[0] < south):
            south = coords[0]
            if (coords[0] > north):
            north = coords[0]
            if (coords[1] < west):
            west = coords[1]
            if (coords[1] > east):
            east = coords[1]            
            '''
        elif item.tag == 'way':
            if item.get('id') == '157161112': #main coastline
                for thing in item:
                    if thing.tag == 'nd':
                        coastnodes.append((long)(thing.get('ref')))
                continue
            useme = False
            oneway = False
            myname = 'unnamed way'
            for thing in item:
                if thing.tag == 'tag' and thing.get('k') == 'highway':
                    useme = True
                    mytype = thing.get('v')
                if thing.tag == 'tag' and thing.get('k') == 'name':
                    myname = thing.get('v')
                if thing.tag == 'tag' and thing.get('k') == 'omeway':
                    if thing.get('v') == 'yes':
                        oneway = True
            if useme:
                wayid = (long)(item.get('id'))
                ways[wayid] = Way(myname,mytype)
                nlist = []
                for thing in item:
                    if thing.tag == 'nd':
                        nlist.append((long)(thing.get('ref')))
                thisn = nlist[0]
                for n in range(len(nlist)-1):
                    nextn = nlist[n+1]
                    nodes[thisn].ways.append(Edge(ways[wayid],nodes[thisn],nodes[nextn]))
                    thisn = nextn
                if not oneway:
                    thisn = nlist[-1]
                    for n in range(len(nlist)-2,-1,-1):
                        nextn = nlist[n]
                        nodes[thisn].ways.append(Edge(ways[wayid],nodes[thisn],nodes[nextn]))
                        thisn = nextn                
                ways[wayid].nodes = nlist
    #print len(coastnodes)
    #print coastnodes[0]
    #print nodes[coastnodes[0]]
    return nodes, ways, coastnodes

def compute_params(paths):
    # training vs testing
    numtest = len(paths)//5
    testp = []
    trainp = paths[:]
    for _ in range(numtest):
        testp.append(trainp.pop(random.randint(0,len(trainp)-1)))
        
    w = [.1,.1, .01]
    alpha = 0.01
    for epoch in range(1000):
        delta = [0,0,0]
        testerr = 0
        trainerr = 0
        for p in trainp:
            err = p.time - (w[0]*p.dist + w[1]*p.uphill + w[2])
            delta[0] += p.dist*err#/p.time
            delta[1] += p.uphill*err#/p.time
            delta[2] += err#/p.time
            trainerr = err*err
            #print p.dist, p.uphill, p.time, err, delta
        for p in testp:
            err = p.time - (w[0]*p.dist + w[1]*p.uphill + w[2])
            testerr += err*err
        w[0] += alpha*delta[0]
        w[1] += alpha*delta[1]
        w[2] += alpha*delta[2]
        #print trainerr, testerr
        #print w
    print trainerr, testerr
    for p in trainp:
        predict = (w[0]*p.dist + w[1]*p.uphill + w[2])
        print predict, p.time
    print "----"
    for p in testp:
        predict = (w[0]*p.dist + w[1]*p.uphill + w[2])
        print predict, p.time
    return w

def make_cost(params):
    ''' Construct a function that can be passed into the graph maker! '''
    def thefn(edge):
        val = params[0]*node_dist_2(edge.src,edge.dest)/1000.0
        if (edge.dest.elev > edge.src.elev):
            val += params[1]*(edge.dest.elev - edge.src.elev)/100.0
        #print (edge.src.id, edge.dest.id, val)
        return val
    #print thefn
    return thefn

#[8.7199461257862563, 3.2382728168494443, 0.36016872363089847]

elevs = build_elevs("N42E018.HGT")
nodes, ways, coastnodes = build_graph(elevs)
planner = Planner(nodes,ways,node_dist_3_cost)
paths = read_paths(argv[1],planner)
print "Computing parameters"
params = compute_params(paths)
print params
myfun = make_cost(params)
#print myfun(nodes[436448505].ways[0])
    
learnedplanner = Planner(nodes,ways,myfun)

master = Tk()
thewin = DispWin(master,nodes,ways,coastnodes,elevs,paths,learnedplanner)
mainloop()