"""Example of shallow water wave equation.

This script sets up a 2D version of the 1D LWRU1 benchmark with initial condition stated in the file benchmark_1.txt.

See also

http://www.cee.cornell.edu/longwave/index.cfm?page=benchmark&problem=1


"""

######################
# Module imports

from pyvolution.shallow_water import Domain, Reflective_boundary,\
     Dirichlet_boundary,Transmissive_boundary, Constant_height, Constant_stage

from pyvolution.mesh_factory import rectangular_cross
from Numeric import array, zeros, Float, allclose


#######################
# Domain
#


print 'Creating domain'
#Create basic mesh
#
#The initial condition extends 50km off shore
#and 5,000m is allowed on shore for wetting
#(only about 200m is expected, though)

points, vertices, boundary = rectangular_cross(150, 15,
                                         len1=55000, len2=5000,
                                         origin = (-5000, 0.0))

#points, vertices, boundary = rectangular_cross(100, 10,
#                                         len1=55000, len2=5000,
#                                         origin = (-5000, 0.0))


#Create shallow water domain
domain = Domain(points, vertices, boundary)

domain.check_integrity()
domain.default_order = 2

#Output params
domain.smooth = True
domain.reduction = min  #Looks a lot better on top of steep slopes
print "Number of triangles = ", len(domain)

domain.visualise = False
domain.store = True    #Store for visualisation purposes
domain.format = 'sww'   #Native netcdf visualisation format

import sys, os
base = os.path.basename(sys.argv[0])
domain.filename, _ = os.path.splitext(base)


#Set initial values
def slope(x, y):
    return -x/10


class IC_x:
    """
    Read 1D initial condition and provide values at any x, y

    File is assumed to list x values in the first column and
    stage in the second.
    """

    def __init__(self, filename):

        self.x = []
        self.w = []
        fid = open(filename)
        for line in fid.readlines():
            fields = line.split()
            assert len(fields) == 2, '%s' %fields
            self.x.append( float(fields[0]) )
            self.w.append( float(fields[1]) )

        #print 'X', self.x, len(self.x)
        #print 'W', self.w, len(self.w)
        #from pylab import plot, show
        #plot(self.x, self.w)
        #show()
        #import sys; sys.exit()

    def __call__(self, x, y):

        w = zeros( len(x), Float )
        for i in range(len(x)):
            xi = x[i]


            #Find slot

            if xi < self.x[0]:
                w[i] = self.w[0]
            elif xi > self.x[-1]:
                w[i] = self.w[-1]
            else:
                index = 0
                while xi > self.x[index]: index += 1
                while xi < self.x[index]: index -= 1

                #print xi, index, self.x[index], self.w[index]

                if xi == self.x[index]:
                #if allclose(xi, self.x[index]):
                    #Protect against case where x is the last value
                    # - also works in general when x == self.x[i]
                    ratio = 0
                else:
                    #x is now between index and index+1
                    ratio = (xi - self.x[index])/\
                            (self.x[index+1] - self.x[index])

                #print xi, index, self.x[index], ratio

                #Compute interpolated value
                if ratio > 0:
                    w[i] = self.w[index] +\
                           ratio*(self.w[index+1] - self.w[index])
                else:
                    w[i] = self.w[index]

        #print x, w
        return w



print 'Field values'
domain.set_quantity('elevation', slope)
domain.set_quantity('friction', 0.0)
domain.set_quantity('stage', IC_x('lwru1_IC.txt'))

#import sys; sys.exit()

#print domain.quantities['stage'].centroid_values

######################
# Boundary conditions
#
print 'Boundaries'
Br = Reflective_boundary(domain)
Bt = Transmissive_boundary(domain)

#Constant inflow
Bd = Dirichlet_boundary([0.0, 0.0, 0.0])

#Set boundary conditions
domain.set_boundary({'left': Br, 'right': Br, 'bottom': Br, 'top': Br})


#Evolve
import time
t0 = time.time()




pt = []
xes = []
y = 2500
x0 = -500
step = 5
for i in range(1000):
    x = x0+i*step
    xes.append(x)
    pt.append( [x,y] )

from pylab import *
from pyvolution.least_squares import Interpolation


V = domain.get_vertex_coordinates(obj=True) #Why?
T = domain.get_triangles(obj=True)


I = Interpolation(V,
                  T,
                  point_coordinates = pt,
                  verbose = True)


f = domain.quantities['elevation'].vertex_values.flat
z = I.interpolate( f )

print 'xxxxx'


f = domain.quantities['stage'].vertex_values.flat
y = I.interpolate( f )

#ion()
#plot(xes, y, '-b', xes, z, '-k', [-500, 50000], [0.0, 0.0], '-k')
ion()
clf()
hold(True)
plot(xes, y, '-b')
plot(xes, z, '-k')
plot([-500, 50000], [0.0, 0.0], '-k')
set( gca(), Ylim=(-100,100) )
set( gca(), Xlim=(-500,2000) )
draw()
ioff()

#raw_input('go')
for t in domain.evolve(yieldstep = 10, finaltime = 300.0):
    domain.write_time()


    f = domain.quantities['stage'].vertex_values.flat
    y = I.interpolate( f )

    clf()
    hold(True)
    plot(xes, y, '-b')
    plot(xes, z, '-k')
    plot([-500, 50000], [0.0, 0.0], '-k')
    set( gca(), Ylim=(-100,100) )
    set( gca(), Xlim=(-500,2000) )
    draw()


    #raw_input('go')

    #print y[:], y.shape





print 'That took %.2f seconds' %(time.time()-t0)
show()