""" Interpolation of a sww file.
Used to interpolate height information from sww files.

When using as a stand alone application,
Input;
 - The sww file with stage and bed elevation information
 - An xya file specifying the points (x,y) where the height values
   (w.r.t. time) are needed.

Ouput;
An xya file with x,y and height values w.r.t. time.

The first row of the output xya file has the time value for each height
Each following row has the format x,y,[height,]

NOTE: stage = bed elevation + height

   Ole Nielsen, Stephen Roberts, Duncan Gray, Christopher Zoppou
   Geoscience Australia, 2004.   
"""

from Numeric import transpose
from least_squares import Interpolation 

DEFAULT_QUANTITY = "height"

class Interpolate_sww:
    def __init__(self,file_name, quantity_name):

        #It's bad to have the quantity_name passed in here.
        # But it works with how the program is currently used.
        # Refactor when necessary. - DSG
        
        x, y, volumes, time, quantity = self.read_sww(file_name, quantity_name)
        vertex_coordinates = transpose([x,y])
        
        if False:
            print "****************************"
            print "x "
            print x
            print "****************************"
            print "Y "
            print y
            print "****************************"
            print "V "
            print volumes
            print "****************************"
            print "Time "
            print time
            print "****************************"
            print "quantity "
            print quantity
            print "****************************"
            print "vertex_coordinates"
            print vertex_coordinates
            print "****************************"

        self.interp = Interpolation(vertex_coordinates, volumes )
        self.time = time

        self.quantity_name = quantity_name
        self.quantity = quantity 
        
    def interpolate_xya(self,file_name):
        """
        Given a point file, interpolate the height w.r.t. time at the points
        specified in the point file.

        Input;
        file_name - the xya file
        """
        self.point_coordinates = self.read_xya(file_name)
        self.interp.build_interpolation_matrix_A(self.point_coordinates)
        self.interpolated_quantity = self.interp.interpolate( transpose(self.quantity))
        
    def read_sww(self,file_name, quantity_name):
        """
        Read in an sww file.

        Input;
        file_name - the sww file

        Output;
        x - Vector of x values
        y - Vector of y values
        z - Vector of bed elevation
        volumes - Array.  Each row has 3 values, representing
                  the vertices that define the volume
        time - Vector of the times where there is stage information
        stage - array with respect to time and vertices (x,y)
        """
        
        #FIXME Have this reader as part of data_manager?

        from Scientific.IO.NetCDF import NetCDFFile        	
            
        #Check contents
        #Get NetCDF
        fid = NetCDFFile(file_name, 'r') 
      
        # Get the variables
        x = fid.variables['x'][:]
        y = fid.variables['y'][:]
        volumes = fid.variables['volumes'][:] 
        time = fid.variables['time'][:]
        try:
            if quantity_name == 'height':
                z = fid.variables['z'][:]
                stage = fid.variables['stage'][:,:]
                quantity = stage - z  # 2D, using broadcasting 
            else:
                quantity = fid.variables[quantity_name][:,:]   # 2D 
        except (KeyError, IndexError),e:
            fid.close()
            raise KeyError
            
        fid.close()
        return x, y, volumes, time, quantity

    def read_xya(self,point_file):
        """
        Read in an sww file.

        Input;
        point_file - the xya file

        Output;
        point_coordinates - the points (x,y) where the height values
                             (w.r.t. time) are needed
        """
        from load_mesh.loadASCII import  load_xya_file
        
        point_dict = load_xya_file(point_file)
        return point_dict['pointlist']


    def write_depth_xya(self,point_file, delimiter = ','):
        """
        pre condition:
        point attributes have been determined
        (interpolate_xya has been called)
        The time list is defined
        """
        from load_mesh.loadASCII import  export_xya_file

        #create the first line, giving the times 
        if len(self.time) != 0:
            for i in range(len(self.time)):
                if i == 0:
                    title = str(self.time[0])
                else:
                    title = title + delimiter + str(self.time[i])
        else:
            title  = ""
        
        #print ">" +  str(self.time) + "<"
        #print ">" + title + "<"
        
        xya_dict = {}
        xya_dict['pointlist'] = self.point_coordinates
        xya_dict['pointattributelist'] = self.interpolated_quantity
        
        export_xya_file(point_file, xya_dict, title, delimiter = delimiter)
        
#-------------------------------------------------------------
if __name__ == "__main__":
    """
    Load in an sww file and an xya file and return an xya file 
    """
    import os, sys
    usage = "usage: %s pyvolution_results.sww points.xya depth.xya [height|stage|(other quantities)]" %         os.path.basename(sys.argv[0])
    if len(sys.argv) < 4:
        print usage
    else:
        sww_file = sys.argv[1]
        point_file_in = sys.argv[2]
        point_file_out = sys.argv[3]
        if len(sys.argv) == 5:
            quantity_name = sys.argv[4]
        else:
            quantity_name = DEFAULT_QUANTITY   
        #print "quantity",quantity
        try:
            interp = Interpolate_sww(sww_file, quantity_name)
        except KeyError:
            print "Error: Unknown quantity"
            sys.exit(1)

        interp.interpolate_xya(point_file_in)
        interp.write_depth_xya(point_file_out)