import os
from os import sep
#import create_boundary
import project
try:
    from pylab import plot,show,xlabel,ylabel,legend,title,savefig,hold
except:
    print 'Cannot import pylab plotting will not work. Csv files are still created'
from anuga.utilities.polygon import inside_polygon
from anuga.utilities.numerical_tools import ensure_numeric
from Scientific.IO.NetCDF import NetCDFFile
def get_sts_gauge_data(filename,polygon,verbose=False):
    from Numeric import asarray,transpose,sqrt,argmax,argmin,arange,Float,\
        compress,zeros,fabs,take
    fid = NetCDFFile(filename+'.sts', 'r')    #Open existing file for read
    x = fid.variables['x'][:]+fid.xllcorner   #x-coordinates of vertices
    y = fid.variables['y'][:]+fid.yllcorner   #y-coordinates of vertices
    points=transpose(asarray([x.tolist(),y.tolist()]))
    time=fid.variables['time'][:]
    elevation=fid.variables['elevation'][:]
    indices=inside_polygon(points, polygon, closed=True, verbose=False)
    inside_points=take(points,indices,0)
    elevation=take(elevation,indices,0)
    if verbose:
        inside_points_array=ensure_numeric(inside_points)#for plotting only
        plot(points[:,0],points[:,1],'o')
        plot(inside_points_array[:,0][::5],inside_points_array[:,1][::5],'o')
        polygon=ensure_numeric(polygon)
        plot(polygon[:,0],polygon[:,1])
        show()
        
    quantity_names=['stage','xmomentum','ymomentum']
    quantities = {}
    for i, name in enumerate(quantity_names):
        quantities[name] = fid.variables[name][:]
        if inside_points is not None:
            quantities[name] = take(quantities[name],indices,1)
        else:
            msg = 'No gauges found inside polygon'
            raise msg
    fid.close()

    return inside_points,quantities,elevation,time

from anuga.shallow_water.data_manager import urs2sts
base_name='tgs/gauges'
if os.path.exists(base_name+'.sts'):
    print 'sts file already created'
else:
    print 'creating sts file'
    urs2sts(base_name,
            basename_out=base_name+'.sts',
            mean_stage=0.0, verbose=False)

#polygon=create_boundary.bounding_polygon
polygon=project.bounding_polygon
points,quantities,elevation,time=get_sts_gauge_data(base_name,polygon,verbose=False)
assert len(points)==len(elevation)==len(quantities['stage'][0,:])

def plot_gauge(time,quantity,quantity_name,elevation,gauge_id,pos):
    figname=figname = '%s_%s%s.png' %('tgs/urs_gauge',
                                      quantity_name,
                                      str(gauge_id))
    hold(False)
    plot(time/60.0,quantity,label='%s (elevation = %sm)'%('urs',str(elevation)))
    title('Gauge %d at (%f,%f)' %(gauge_id,pos[0],pos[1]))
    xlabel('Time (minutes)')
    ylabel(quantity_name)
    legend()
    print 'saving figure here %s' %figname
    savefig(figname)
    
points=points[::5]
elevation=elevation[::5]
gauge_filename='urs_gauges'
d=','
if os.path.exists(gauge_filename+ '.txt'):
    print 'file with list of gauges already created'
else:
    fid = open(gauge_filename+ '.txt', 'w')
    fid.write('easting,northing,name,elevation\n')
    for i in range(len(points)):
        fid.write(str(points[i][0])+d+str(points[i][1])+d+str(i)+d+str(elevation[i])+'\n')
    fid.close()

def get_sww_gauge_data(sww_filename,gauge_filename):
    from anuga.abstract_2d_finite_volumes.util import sww2csv_gauges, csv2timeseries_graphs
    dir_name, base = os.path.split(sww_filename)
    if os.path.exists(dir_name+sep+'gauge_0.csv'):
        print 'csv gauge files exist already'
    else:
        print 'writing gauge csv files to '+dir_name
        sww2csv_gauges(sww_file=sww_filename+'.sww',
                   quantities = ['stage','elevation','xmomentum','ymomentum'],
                   gauge_file=gauge_filename+'.txt')
        
    print 'generating graphs'
    csv2timeseries_graphs(directories_dic={dir_name:[dir_name,0, 0]},
                          base_name='gauge_',
                          plot_numbers='',#['1'],
                          quantities=['stage','xmomentum','ymomentum'],
                          extra_plot_name='',
                          assess_all_csv_files=True,
                          create_latex=False,
                          verbose=True)

sww_filename='good_simulation/good_polyline20080721_130608'
#sww_filename='polyline/good_polyline20080607_205151'
#sww_filename='most/good_most20080607_205101'
#sww_filename='cloud/good_urs20080607_205134'
#get_sww_gauge_data(sww_filename,gauge_filename)

print 'plotting from sww file', sww_filename

for i in range(len(points)):
    quantity=quantities['stage'][:,i]
    plot_gauge(time,quantity,'stage',elevation[i],i,points[i])
    quantity=quantities['xmomentum'][:,i]
    plot_gauge(time,quantity,'xmomentum',elevation[i],i,points[i])
    #quantity=quantities['ymomentum'][:,i]
    #plot_gauge(time,quantity,'ymomentum',elevation[i],i,points[i])