import os, glob
from gee.display.seis import SeisDisplayGroup, SeisDisplay
from gee.display.parmo import ParMoDisplay
from gee.display.borders import SimpleTimeBorder, TimeBorder
from gee.display.widget import VBox, HBox, Image, InfoBox
from gee.display.chart import Contour
from gee.pdf import PDF
from gee.units import Time, TimeRange, Unit
from gee.files.reader import SeisReader

for station in os.listdir('stations'):
    stationdir = 'stations/' + station
    orig = glob.glob(stationdir + '/orig/*')
    rt = glob.glob(stationdir + '/rt/*')
    norm = glob.glob(stationdir + '/norm/*')
    normrt = glob.glob(stationdir + '/normrt/*')
    displays = [SeisDisplay(seismograms=SeisReader(seis)) for seis in orig, rt, norm, normrt]
    for display in displays:
        display.topBorder = None
        display.bottomBorder = SimpleTimeBorder()
    displays[3].addLabel('NORMALIZED FOR VIEWING')
    displays[4].addLabel('CORRECTED TRACES')
    displays[4].bottomBorder = TimeBorder(label='Travel-time (sec)')
    displayGroup = SeisDisplayGroup(title='Analysis Window at %s' % station, displays)
    displayGroup.setTime(TimeRange(Time(1490, Unit.SECOND), Time(1550, Unit.SECOND)))
    
    origParMo = ParMoDisplay(vertical=SacReader(stationdir + '/rt/radial.sac'),
                             horizontal=SacReader(stationdir + '/rt/tang.sac'),
                             label='Original R-T System')
    origParMo.bottomBorder.addLabel('Tangential Motion')
    
    corrParMo = ParMoDisplay(vertical=SeisReader(stationdir + '/normrt/radial.sac'),
                             horizontal=SeisReader(stationdir + '/normrt/tang.sac'),
                             label='Corrected R-T System')
    corrParMo.leftBorder.addLabel('Radial Motion')
    corrParMo.bottomBorder.addLabel('Tangential Motion')
        
    leftside = VBox(displays,
                    HBox(origParMo, corrParMo))
                    
    ems = Contour('ems.csv', 
                   title='Energy Minimization Space')
    ems.leftBorder.addTitle('Azimuth(deg)')
    ems.bottomBorder.addTitle('Lag(sec)')   
                   
    rightside = VBox(ems,
                     InfoBox(title='Shear Wave Splitting Analysis',
                             contents=open(stationdir +'/info.txt').read())
    PDF(HBox(leftside, rightside)).write(stationdir + '/analysis.pdf')   