import os
from os import sep
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()

    basename='sts_gauge'
    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)
            for j in range(len(x)-1):
                locx=int(x[j])
                locy=int(y[j])
                fid_sts = open(project.boundaries_dir+sep+basename+'_'+str(locx)+'_'+str(locy)+'.csv', 'w')
                s = 'time, stage, xmomentum, ymomentum \n'
                fid_sts.write(s)
                stage = quantities[name][:,j]
                xmomentum = quantities[name][:,j]
                ymomentum = quantities[name][:,j]
                for k in range(len(time)-1):
                    s = '%.2f, %.2f, %.2f, %.2f\n' %(time[k], stage[k], xmomentum[k], ymomentum[k])
                    fid_sts.write(s)
        else:
            msg = 'No gauges found inside polygon'
            raise msg
            fid_sts.close()
    fid.close()

    return inside_points,quantities,elevation,time

polygon=project.poly_all
points,quantities,elevation,time=get_sts_gauge_data(os.path.join(project.boundaries_dir,project.scenario_name),polygon,verbose=False)


assert len(points)==len(elevation)==len(quantities['stage'][0,:])