"""Create several realisations of bathymetries for use with run_model.

This script will create a number of random samples of the input bathymetry and
then fit all of them to arrays that can be used in set_quantity inside run_model.

The reason for this is that doing all fits in one sweep saves computation time
"""

from RandomArray import normal
from Numeric import concatenate, NewAxis
import cPickle
from load_mesh.loadASCII import import_points_file, import_mesh_file
from pyvolution.least_squares import fit_to_mesh
import project

bathymetry_filename = project.bathymetry_filename
mesh_filename = project.mesh_filename


# Read mesh
mesh_dict = import_mesh_file(mesh_filename)
vertex_coordinates = mesh_dict['vertices']
triangles = mesh_dict['triangles']

# Read basic bathymetry from NetCDF .pts file
print 'Reading bathymetry'
base_bathymetry = import_points_file(bathymetry_filename)


# Extract relevant Numeric arrays
print 'Extract bathymetry'
point_coordinates = base_bathymetry['pointlist']
elevation = base_bathymetry['attributelist']['elevation']
N = len(elevation)
print 'read', N, 'points'

# Create different bathymetries based on sampling (Suresh to fill this in)
print 'Create different realisations'
Z = None
block_number = 0
for n in range(project.number_of_samples):

    print 'Creating sample #%d' %n
    # nth realisation of bathymetry
    z = elevation + normal(project.mean,
                           project.std_dev,
                           N)


    #Concatenate into array with each realisation a column
    if Z is None:
        Z = z[:, NewAxis] # The first realisation of this block
    else:    
        Z = concatenate((Z, z[:, NewAxis]), axis=1)

    if (n+1)%project.blocksize == 0 or n == project.number_of_samples:
        #Fit all of them to computational mesh
        print 'fit block %d to the mesh' %block_number
        V = fit_to_mesh(vertex_coordinates,
                        triangles,
                        point_coordinates,
                        Z,
                        alpha = 0.0001,
                        verbose = True,
                        expand_search = True,
                        precrop = True,
                        use_cache = True)


        Z = None

        #Store this block for use with model
        print 'Store fitted quantities in block %d for use with model'\
              %block_number

        filename = project.basename + '_block%d.pck'\
                   %block_number

        block_number += 1    
        
        
        fid = open(filename, 'wb')
        cPickle.dump(V, fid)
        fid.close()