from gtfs import Schedule
from gtfs.entity import Trip, StopTime, Stop
from sqlalchemy.orm import eagerload
import csv
import sys

from prender import processing

def cons(ary):
  for i in range(len(ary)-1):
    yield ary[i],ary[i+1]

def find_freqs(sched):
  counts = {}

  n_trips = sched.session.query(Trip).count()

  for i in range(0,n_trips,500):
      print (i/float(n_trips))*100
      for trip in sched.session.query(Trip).options(eagerload("stop_times")).limit(500).offset(i):
          for st1,st2 in cons(sorted(trip.stop_times,key=lambda x:x.stop_sequence)):
              counts[(st1.stop_id,st2.stop_id)] = counts.get( (st1.stop_id,st2.stop_id), 0 ) + 1

  fp = open("freqs","w")
  for (st1,st2),count in counts.items():
    fp.write( "%s,%s,%s\n"%(st1,st2,count) )
  fp.close()

def map_line_freqs():
    sched = Schedule( "davepeck_20100806_2257.db" )

    ll,bb,rr,tt = sched.session.execute("SELECT min(stop_lon),min(stop_lat),max(stop_lon),max(stop_lat) from Stops").fetchone()


    mr = processing.MapRenderer()
    mr.start(ll,bb,rr,tt,4000) #left,bottom,right,top,width
    mr.background(255,255,255)
    mr.smooth()

    stops = dict([(stop.stop_id,stop) for stop in sched.stops])

    for st1_id, st2_id, count in csv.reader( open( "freqs" ) ):
        st1 = stops[st1_id]
        st2 = stops[st2_id]
        print st1.stop_lat,st1.stop_lon, st2.stop_lat,st2.stop_lon, count
        mr.strokeWeight( (float(count))*0.000008 )
        mr.line(st1.stop_lon,st1.stop_lat,st2.stop_lon,st2.stop_lat)

    mr.saveLocal("map.png")
    mr.stop()

def map_stop_freqs():
    sched = Schedule( "trimet.db" )

    ll,bb,rr,tt = sched.session.execute("SELECT min(stop_lon),min(stop_lat),max(stop_lon),max(stop_lat) from Stops").fetchone()

    mr = processing.MapRenderer()
    mr.start(ll,bb,rr,tt,4000) #left,bottom,right,top,width
    mr.background(255,255,255)
    mr.smooth()
    mr.strokeWeight( 0.0001 )

    stops = dict([(stop.stop_id,stop) for stop in sched.stops])

    counts = sched.session.execute( "select stop_id, count(*) from stop_times group by stop_id" )
    for stop_id, count in counts:
        print stop_id
        stop = stops[stop_id]
        mr.ellipse( stop.stop_lon, stop.stop_lat, 0.0001*count**0.5, 0.0001*count**0.5 )

    mr.saveLocal("trimet_stop_freqs.png")
    mr.stop()

def hyper_loader():
    from gtfs.feed import Feed

    feed = Feed( "/home/brandon/Desktop/wmata_20100822_0217.zip" )
    stop_times_table = feed.get_table( "stop_times.txt" )
    print stop_times_table

    for i, record in enumerate( stop_times_table ):
        if i%500==0: sys.stdout.write("."); sys.stdout.flush()
        pass
    sys.stdout.write("\n")

if __name__=='__main__':
    map_stop_freqs()