"""
Created: 11/01/2006
Author: John Jakeman

Program Visualises .sww files
.sww: Netcdf format for storing model output f(t,x,y)
"""
import time as systime
def update_viewer(time_interp, yeildstep, t):
    stage = interpolated_quantity(fid.variables['stage'][:],time_interp)
    domain.set_quantity('stage',stage)
    
    ratio = time_interp[1]
    domain.set_time(t+ratio*yieldstep)
    
    #domain.visualiser.update_quantity('stage')
    #domain.visualiser.update_timer()
    if domain.visualise is True:
        domain.visualiser.redraw_ready.set()
        domain.visualiser.idle.wait()
        domain.visualiser.idle.clear()
        domain.visualiser.unpaused.wait()    
    
    print 'time = ',domain.write_time()
    #systime.sleep(0.5)

import project
from anuga.pyvolution.data_manager import get_time_interp, interpolated_quantity
from caching import cache
filename = project.filename[:-3] + '.sww' #filename is of form filename.py
#filename = 'run_new_meribula.sww'
#filename = 'results29_01_2006.sww'
filename = 'cairns250m.sww'
verbose = True
very_verbose = False
boundary = None
t = 0.0
fail_if_NaN=False
NaN_filler=0

"""
The following is borrowed from anuga.pyvolution.data_manager.sww2domain
"""
NaN=9.969209968386869e+036
#initialise NaN.

from Scientific.IO.NetCDF import NetCDFFile
from anuga.pyvolution.shallow_water_vtk import Domain
from Numeric import asarray, transpose, resize

if verbose: print 'Reading from ', filename
fid = NetCDFFile(filename, 'r')    #Open existing file for read
time = fid.variables['time']       #Timesteps

#if t is None:
#    t = time[-1]
time_interp = 0, 0.0
#time_interp = get_time_interp(time,t)

# Get the variables as Numeric arrays
x = fid.variables['x'][:]             #x-coordinates of vertices
y = fid.variables['y'][:]             #y-coordinates of vertices
elevation = fid.variables['elevation']     #Elevation
stage = fid.variables['stage']     #Water level
xmomentum = fid.variables['xmomentum']   #Momentum in the x-direction
ymomentum = fid.variables['ymomentum']   #Momentum in the y-direction

starttime = fid.starttime[0]
volumes = fid.variables['volumes'][:] #Connectivity
coordinates=transpose(asarray([x.tolist(),y.tolist()]))

conserved_quantities = []
interpolated_quantities = {}
other_quantities = []

# get geo_reference
#sww files don't have to have a geo_ref
try:
    geo_reference = Geo_reference(NetCDFObject=fid)
except: #AttributeError, e:
    geo_reference = None

if verbose: print '    getting quantities'
for quantity in fid.variables.keys():
    dimensions = fid.variables[quantity].dimensions
    if 'number_of_timesteps' in dimensions:
        conserved_quantities.append(quantity)
        interpolated_quantities[quantity]=\
              interpolated_quantity(fid.variables[quantity][:],time_interp)
    else:
        other_quantities.append(quantity)

other_quantities.remove('x')
other_quantities.remove('y')
other_quantities.remove('z')
other_quantities.remove('volumes')

conserved_quantities.remove('time')

if verbose: print '    building domain'
#    From domain.Domain:
#    domain = Domain(coordinates, volumes,\
#                    conserved_quantities = conserved_quantities,\
#                    other_quantities = other_quantities,zone=zone,\
#                    xllcorner=xllcorner, yllcorner=yllcorner)

#   From shallow_water.Domain:
coordinates=coordinates.tolist()
volumes=volumes.tolist()
#FIXME:should this be in mesh?(peter row)
if fid.smoothing == 'Yes': unique = False
else: unique = True
if unique:
    coordinates,volumes,boundary=weed(coordinates,volumes,boundary)


domain = cache(Domain, (coordinates, volumes, boundary))

if not boundary is None:
    domain.boundary = boundary

domain.geo_reference = geo_reference

domain.starttime=float(starttime)+float(t)
domain.time=0.0

for quantity in other_quantities:
    try:
        NaN = fid.variables[quantity].missing_value
    except:
        pass #quantity has no missing_value number
    X = fid.variables[quantity][:]
    if very_verbose:
        print '       ',quantity
        print '        NaN =',NaN
        print '        max(X)'
        print '       ',max(X)
        print '        max(X)==NaN'
        print '       ',max(X)==NaN
        print ''
    if (max(X)==NaN) or (min(X)==NaN):
        if fail_if_NaN:
            msg = 'quantity "%s" contains no_data entry'%quantity
            raise msg
        else:
            data = (X<>NaN)
            X = (X*data)+(data==0)*NaN_filler
    if unique:
        X = resize(X,(len(X)/3,3))
    domain.set_quantity(quantity,X)
#
for quantity in conserved_quantities:
    try:
        NaN = fid.variables[quantity].missing_value
    except:
        pass #quantity has no missing_value number
    X = interpolated_quantities[quantity]
    if very_verbose:
        print '       ',quantity
        print '        NaN =',NaN
        print '        max(X)'
        print '       ',max(X)
        print '        max(X)==NaN'
        print '       ',max(X)==NaN
        print ''
    if (max(X)==NaN) or (min(X)==NaN):
        if fail_if_NaN:
            msg = 'quantity "%s" contains no_data entry'%quantity
            raise msg
        else:
            data = (X<>NaN)
            X = (X*data)+(data==0)*NaN_filler
    if unique:
        X = resize(X,(X.shape[0]/3,3))
    domain.set_quantity(quantity,X)


"""
The following initialises visualiser and uses .sww file to update stage
at each timestep
"""

domain.initialise_visualiser()
#domain.visualiser.setup_all()
domain.visualiser.scale_z['stage'] = 20.0
domain.visualiser.coloring['stage'] = True
domain.visualiser.scale_z['elevation'] = 20.0
domain.visualiser.coloring['elevation'] = False
domain.visualiser.updating['elevation'] = True
#domain.visualiser.setup['stage'] = False
#domain.visualiser.updating['stage'] = False


domain.visualiser.start()
# Things go haywire if we start evolving before the vis is ready
domain.visualiser.idle.wait()
domain.visualiser.idle.clear()

index = 0 
ratio = 0.0
num_subt = 1 # use 1 if only want to display at yieldstep timeinterval

if len(time) > 1:
    yieldstep = time[1]-time[0]
else:
    yieldstep = 0.0

for t in time[:]:
    if t < time[-1]:
        if index%10 == 0:
            for subt in range(0,num_subt):
                ratio = 1-float(num_subt-subt)/float(num_subt)
                #print ratio
                time_interp = index, ratio
                update_viewer(time_interp, yieldstep, t)
    else:
        time_interp = index, 0.0
        update_viewer(time_interp, yieldstep, t)
            
    index = index + 1
    
print 'Finished'   
fid.close()