"""Stochastic study of the ANUGA implementation of the
shallow water wave equation.

Very simple flat bed with one spike in the middle where different values are sampled.
Three sides are reflective and one Dirichlet condition

Output is measured at gauges around the spike.

The system will evolve dynamically towards a steady state.

Suresh Kumar and Ole Nielsen 2006
"""


#------------------------------------------------------------------------------
# Import necessary modules
#------------------------------------------------------------------------------

# Standard modules
import os
import time

# Related major packages
from Numeric import allclose, zeros, Float
from RandomArray import normal, seed

# Application specific imports
from anuga.pyvolution.mesh_factory import rectangular
from anuga.pyvolution.shallow_water import Domain
from anuga.pyvolution.shallow_water import Reflective_boundary
from anuga.pyvolution.shallow_water import Dirichlet_boundary
from anuga.pyvolution.util import file_function
from caching.caching import cache
import project                 # Definition of file names and polygons


#-----------------------------------------------------------------------------
# Read in processor information
#-----------------------------------------------------------------------------

try:
    import pypar
except:
    print 'Could not import pypar'
    myid = 0
    numprocs = 1
    processor_name = 'local host'
else:    
    myid = pypar.rank()
    numprocs = pypar.size()
    processor_name = pypar.Get_processor_name()

print 'I am process %d of %d running on %s' %(myid, numprocs, processor_name)


#-----------------------------------------------------------------------------
# Setup computational domain
#----------------------------------------------------------------------------- 

# Create basic triangular mesh
points, vertices, boundary = rectangular(10, 10, 10, 10)

#Create shallow water domain
domain = Domain(points, vertices, boundary)
domain.set_name(project.basename)
domain.set_datadir(project.working_dir) 
print domain.statistics()

#------------------------------------------------------------------------------
# Setup boundary conditions
#------------------------------------------------------------------------------

Br = Reflective_boundary(domain) #Wall
Bd = Dirichlet_boundary([1, 0, 0])

# Bind boundary objects to tags
domain.set_boundary({'left': Br, 'right': Bd, 'top': Br, 'bottom': Br})


#------------------------------------------------------------------------------
# Setup initial conditions
#------------------------------------------------------------------------------
domain.set_quantity('friction', 0.0)
domain.set_quantity('stage', 0.0)

# Initial conditions
def f(x,y):
    z = zeros( x.shape, Float )
    for i in range(len(x)):
        if allclose([x[i], y[i]], 5):
            print i
            z[i] = 1
    return z

domain.set_quantity('elevation', f)


z = domain.get_quantity('elevation').get_values()
print z
#print z != 0

import sys; sys.exit() 

# Run model for different realisations of spike 
print 'Create %d different realisations' %project.number_of_realisations

# Array of bed elevation values at vertices

#print z
N = z.shape[0]

timestep = 0.1
finaltime = 1

seed(13,17)
for realisation in range(project.number_of_realisations):

    print 'Creating realisation #%d' %realisation
    #z[N/2-1,:] = normal(project.mean, project.std_dev, 1)

    # Distribute work in round-robin fashion
    if realisation%numprocs == myid:
        name = project.basename + '_P%d' %myid    
        domain.set_name(name)                # Output name
        #domain.set_quantity('elevation', z)  # Assign bathymetry
        domain.set_time(0.0)                 # Reset time

        #---------------------------------------------------
        # Evolve system through time
        #---------------------------------------------------
        print 'P%d: Running realisation %d of %d'\
              %(myid, realisation, project.number_of_realisations)

        t0 = time.time()
        for t in domain.evolve(yieldstep = timestep,
                               finaltime = finaltime):
            pass
            domain.write_time()
                    
        
        print 'P%d: Simulation of realisation %d took %.2f seconds'\
              %(myid, realisation, time.time()-t0)



        
        #---------------------------------------------------
        # Now extract the 3 timeseries (Ch 5-7-9) and store them
        # in three files for this realisation

        print 'P%d: Extracting time series for realisation %d from file %s'\
              %(myid, realisation, project.working_dir +\
                domain.get_name() + '.sww')
        f = file_function(project.working_dir +\
                          domain.get_name() + '.sww',
                          quantities='stage',
                          interpolation_points=project.gauges,
                          verbose=False)


        simulation_name = project.working_dir + \
                          project.basename + '_realisation_%d' %realisation

        print 'P%d: Writing to file %s'\
              %(myid, simulation_name + '_' + name + '.txt')
        
        for k, name in enumerate(project.gauge_names):
            fid = open(simulation_name + '_' + name + '.txt', 'w')
            for t in f.get_time():
                #For all precomputed timesteps
                val = f(t, point_id = k)[0]
                fid.write('%f %f\n' %(t, val))
                
            fid.close()

            
pypar.finalize()