"""Functions to store and retrieve data for the Shallow Water Wave equation.
There are two kinds of data

  1: Constant data: Vertex coordinates and field values. Stored once
  2: Variable data: Conserved quantities. Stored once per timestep.

All data is assumed to reside at vertex locations.  

"""

from Numeric import concatenate        


def make_filename(s):
    """Transform argument string into a suitable filename
    """

    s = s.strip()
    s = s.replace(' ', '_')
    s = s.replace('(', '')
    s = s.replace(')', '')
    s = s.replace('__', '_')                
        
    return s


def check_dir(path, verbose=None):
    """Check that specified path exists.
    If path does not exist it will be created if possible    

    USAGE:
       checkdir(path, verbose):

    ARGUMENTS:
        path -- Directory
        verbose -- Flag verbose output (default: None)

    RETURN VALUE:
        Verified path including trailing separator
        
    """

    import os, sys
    import os.path

    if sys.platform in ['nt', 'dos', 'win32', 'what else?']:
        unix = 0
    else:
        unix = 1


    if path[-1] != os.sep: 
        path = path + os.sep  # Add separator for directories

    path = os.path.expanduser(path) # Expand ~ or ~user in pathname
    if not (os.access(path,os.R_OK and os.W_OK) or path == ''):
        try:
            exitcode=os.mkdir(path)

            # Change access rights if possible
            #
            if unix:
                exitcode=os.system('chmod 775 '+path)
            else:
                pass  # FIXME: What about acces rights under Windows?
            
            if verbose: print 'MESSAGE: Directory', path, 'created.'
          
        except:
            print 'WARNING: Directory', path, 'could not be created.'
            if unix:
                path = '/tmp/'
            else:
                path = 'C:'
                
            print 'Using directory %s instead' %path

    return(path)

    

def del_dir(path):
    """Recursively delete directory path and all its contents
    """

    import os

    if os.path.isdir(path):
        for file in os.listdir(path):
            X = os.path.join(path, file)


            if os.path.isdir(X) and not os.path.islink(X):
                del_dir(X) 
            else:
                try:
                    os.remove(X)
                except:
                    print "Could not remove file %s" %X

        os.rmdir(path)


	
def create_filename(filename, format, size, time=None):

    import os
    from config import data_dir

           
    FN = check_dir(data_dir) + filename + '_size%d' %size
	
    if time is not None:
        FN += '_time%.2f' %time
        
    FN += '.' + format
        
    return FN
    

def get_files(filename, format, size):
    """Get all file (names) with gven name, size and format
    """

    import glob

    import os
    from config import data_dir

           
    dir = check_dir(data_dir)

    pattern = dir + os.sep + filename + '_size=%d*.%s' %(size, format)
    return glob.glob(pattern)



#Generic class for storing output to e.g. visualisation or checkpointing
class Data_format:
    """Generic interface to data formats
    """

        
    def __init__(self, domain, extension, mode = 'w'):
        assert mode in ['r', 'w', 'a'], '''Mode %s must be either:''' %mode +\
                                        '''   'w' (write)'''+\
                                        '''   'r' (read)''' +\
                                        '''   'a' (append)'''    

        #Create filename
        self.filename = create_filename(domain.get_name(), extension, len(domain))
        self.timestep = 0
        self.number_of_volumes = len(domain)
        self.domain = domain
    
    
    #FIXME: Should we have a general set_precision function?
    


#Class for storing output to e.g. visualisation
class Data_format_sww(Data_format):
    """Interface to native NetCDF format (.sww)
    """

        
    def __init__(self, domain, mode = 'w'):
        from Scientific.IO.NetCDF import NetCDFFile
        from Numeric import Int, Float, Float32 

        self.precision = Float32 #Use single precision
        
        Data_format.__init__(self, domain, 'sww', mode)


        # NetCDF file definition
        fid = NetCDFFile(self.filename, mode)
        
        if mode == 'w':
            
            #Create new file
            fid.institution = 'Geoscience Australia'
            fid.description = 'Output from pyvolution suitable for plotting'
            
            if domain.smooth:
                fid.smoothing = 'Yes'
            else:    
                fid.smoothing = 'No'

            fid.order = domain.default_order    

            # dimension definitions
            fid.createDimension('number_of_volumes', self.number_of_volumes)  
            fid.createDimension('number_of_vertices', 3)

            if domain.smooth is True:
                fid.createDimension('number_of_points', len(domain.vertexlist))
            else:    
                fid.createDimension('number_of_points', 3*self.number_of_volumes)
                
            fid.createDimension('number_of_timesteps', None) #extensible

            # variable definitions
            fid.createVariable('x', self.precision, ('number_of_points',))
            fid.createVariable('y', self.precision, ('number_of_points',))
            fid.createVariable('z', self.precision, ('number_of_points',))
        
            fid.createVariable('volumes', Int, ('number_of_volumes',
                                                'number_of_vertices'))

            fid.createVariable('time', self.precision,
                               ('number_of_timesteps',))
        
            fid.createVariable('stage', self.precision,
                               ('number_of_timesteps',
                                'number_of_points'))


        #Close
        fid.close()                


    def store_connectivity(self):
        """Specialisation of store_connectivity for net CDF format

        Writes x,y,z coordinates of triangles constituting
        the bed elevation.
        """

        from Scientific.IO.NetCDF import NetCDFFile

        from Numeric import concatenate

        domain = self.domain

        #Get NetCDF
        fid = NetCDFFile(self.filename, 'a')  #Open existing file for append
        
        # Get the variables
        x = fid.variables['x']
        y = fid.variables['y']
        z = fid.variables['z']
        
        volumes = fid.variables['volumes']

        # Get X, Y and bed elevation Z
	Q = domain.quantities['elevation']
        X,Y,Z,V = Q.get_vertex_values(xy=True,
				      precision = self.precision)
				      


        x[:] = X.astype(self.precision)
        y[:] = Y.astype(self.precision)
        z[:] = Z.astype(self.precision)

        volumes[:] = V
        
        #Close
        fid.close()


    def store_timestep(self, name):
        """Store time and named quantity to file
        """
        from Scientific.IO.NetCDF import NetCDFFile
	from time import sleep
		
        #Get NetCDF
	retries = 0
	file_open = False
	while not file_open and retries < 10:
  	    try:	
                fid = NetCDFFile(self.filename, 'a') #Open existing file
	    except IOError:
	        #This could happen if someone was reading the file.
		#In that case, wait a while and try again
		msg = 'Warning (store_timestep): File %s could not be opened'\
		%self.filename
		msg += ' - trying again'
		print msg
		retries += 1
		sleep(1)
	    else:
	       file_open = True
	       
	if not file_open:
	    msg = 'File %s could not be opened for append' %self.filename
	    raise msg
	    	
	    
        domain = self.domain
        
        # Get the variables
        time = fid.variables['time']
        stage = fid.variables['stage']        
        i = len(time)

        #Store stage
        time[i] = self.domain.time

        # Get quantity
	Q = domain.quantities[name]
        A,V = Q.get_vertex_values(xy=False,
				  precision = self.precision)
        
        stage[i,:] = A.astype(self.precision)

        #Flush and close
        fid.sync()
        fid.close()


#Class for handling checkpoints data
class Data_format_cpt(Data_format):
    """Interface to native NetCDF format (.cpt)
    """

        
    def __init__(self, domain, mode = 'w'):
        from Scientific.IO.NetCDF import NetCDFFile
        from Numeric import Int, Float, Float 

        self.precision = Float #Use full precision
        
        Data_format.__init__(self, domain, 'sww', mode)


        # NetCDF file definition
        fid = NetCDFFile(self.filename, mode)
        
        if mode == 'w':
            #Create new file
            fid.institution = 'Geoscience Australia'
            fid.description = 'Checkpoint data'
            #fid.smooth = domain.smooth                
            fid.order = domain.default_order    

            # dimension definitions
            fid.createDimension('number_of_volumes', self.number_of_volumes)  
            fid.createDimension('number_of_vertices', 3)

            #Store info at all vertices (no smoothing)
            fid.createDimension('number_of_points', 3*self.number_of_volumes)
            fid.createDimension('number_of_timesteps', None) #extensible

            # variable definitions

            #Mesh
            fid.createVariable('x', self.precision, ('number_of_points',))
            fid.createVariable('y', self.precision, ('number_of_points',))
            #fid.createVariable('z', self.precision, ('number_of_points',))
        
            fid.createVariable('volumes', Int, ('number_of_volumes',
                                                'number_of_vertices'))

            fid.createVariable('time', self.precision,
                               ('number_of_timesteps',))

            #Allocate space for all quantities
            for name in domain.quantities.keys():
                fid.createVariable(name, self.precision,
                                   ('number_of_timesteps',
                                    'number_of_points'))

        #Close
        fid.close()                


    def store_checkpoint(self):
        """
        Write x,y coordinates of triangles.
        Write connectivity (
        constituting
        the bed elevation.
        """

        from Scientific.IO.NetCDF import NetCDFFile

        from Numeric import concatenate

        domain = self.domain

        #Get NetCDF
        fid = NetCDFFile(self.filename, 'a')  #Open existing file for append
        
        # Get the variables
        x = fid.variables['x']
        y = fid.variables['y']
        
        volumes = fid.variables['volumes']

        # Get X, Y and bed elevation Z
	Q = domain.quantities['elevation']
        X,Y,Z,V = Q.get_vertex_values(xy=True,
				      precision = self.precision)
				      


        x[:] = X.astype(self.precision)
        y[:] = Y.astype(self.precision)
        z[:] = Z.astype(self.precision)

        volumes[:] = V
        
        #Close
        fid.close()


    def store_timestep(self, name):
        """Store time and named quantity to file
        """
        from Scientific.IO.NetCDF import NetCDFFile
	from time import sleep
		
        #Get NetCDF
	retries = 0
	file_open = False
	while not file_open and retries < 10:
  	    try:	
                fid = NetCDFFile(self.filename, 'a') #Open existing file
	    except IOError:
	        #This could happen if someone was reading the file.
		#In that case, wait a while and try again
		msg = 'Warning (store_timestep): File %s could not be opened'\
		%self.filename
		msg += ' - trying again'
		print msg
		retries += 1
		sleep(1)
	    else:
	       file_open = True
	       
	if not file_open:
	    msg = 'File %s could not be opened for append' %self.filename
	    raise msg
	    	
	    
        domain = self.domain
        
        # Get the variables
        time = fid.variables['time']
        stage = fid.variables['stage']        
        i = len(time)

        #Store stage
        time[i] = self.domain.time

        # Get quantity
	Q = domain.quantities[name]
        A,V = Q.get_vertex_values(xy=False,
				  precision = self.precision)
        
        stage[i,:] = A.astype(self.precision)

        #Flush and close
        fid.sync()
        fid.close()







#Function for storing xya output
#FIXME Not done yet for this version
class Data_format_xya(Data_format):
    """Generic interface to data formats
    """


    def __init__(self, domain, mode = 'w'):
        from Scientific.IO.NetCDF import NetCDFFile
        from Numeric import Int, Float, Float32 

        self.precision = Float32 #Use single precision
        
        Data_format.__init__(self, domain, 'xya', mode)
        
    
    

    def store_all(self):
        """Specialisation of store all for xya format

        Writes x,y,z coordinates of triangles constituting
        the bed elevation.
        """

        from Numeric import concatenate

        domain = self.domain
        
        fd = open(self.filename, 'w')


        if domain.smooth is True:
            number_of_points =  len(domain.vertexlist)
        else:    
            number_of_points = 3*self.number_of_volumes

        numVertAttrib = 3 #Three attributes is what is assumed by the xya format

        fd.write(str(number_of_points) + " " + str(numVertAttrib) +\
                 " # <vertex #> <x> <y> [attributes]" + "\n")


        # Get X, Y, bed elevation and friction (index=0,1)
        X,Y,A,V = domain.get_vertex_values(xy=True, value_array='field_values',
                                           indices = (0,1), precision = self.precision)

        bed_eles = A[:,0]
        fricts = A[:,1]

        # Get stage (index=0)
        B,V = domain.get_vertex_values(xy=False, value_array='conserved_quantities',
                                       indices = (0,), precision = self.precision)

        stages = B[:,0]

        #<vertex #> <x> <y> [attributes]
        for x, y, bed_ele, stage, frict in map(None, X, Y, bed_eles,
                                               stages, fricts):

            s = '%.6f %.6f %.6f %.6f %.6f\n' %(x, y, bed_ele, stage, frict)
            fd.write(s)
            
        #close
        fd.close()


    def store_timestep(self, t, V0, V1, V2):
        """Store time, water heights (and momentums) to file
        """
        pass
    

#Auxiliary
def write_obj(filename,x,y,z):
    """Store x,y,z vectors into filename (obj format)
       Vectors are assumed to have dimension (M,3) where
       M corresponds to the number elements.
       triangles are assumed to be disconnected

       The three numbers in each vector correspond to three vertices,

       e.g. the x coordinate of vertex 1 of element i is in x[i,1]
       
    """
    #print 'Writing obj to %s' % filename

    import os.path

    root, ext = os.path.splitext(filename)
    if ext == '.obj':
        FN = filename
    else:
        FN = filename + '.obj'
        

    outfile = open(FN, 'wb')
    outfile.write("# Triangulation as an obj file\n")

    M, N = x.shape
    assert N==3  #Assuming three vertices per element

    for i in range(M):
        for j in range(N):
            outfile.write("v %f %f %f\n" % (x[i,j],y[i,j],z[i,j]))

    for i in range(M):
        base = i*N
        outfile.write("f %d %d %d\n" % (base+1,base+2,base+3))

    outfile.close()



#Conversion routines
def sww2obj(basefilename, size):
    """Convert netcdf based data output to obj
    """
    from Scientific.IO.NetCDF import NetCDFFile        

    from Numeric import Float, zeros

    #Get NetCDF
    FN = create_filename(basefilename, 'sww', size)
    print 'Reading from ', FN
    fid = NetCDFFile(FN, 'r')  #Open existing file for read

        
    # Get the variables
    x = fid.variables['x']
    y = fid.variables['y']
    z = fid.variables['z']        
    time = fid.variables['time']
    stage = fid.variables['stage']
        
    M = size  #Number of lines
    xx = zeros((M,3), Float)
    yy = zeros((M,3), Float)
    zz = zeros((M,3), Float)

    for i in range(M):
        for j in range(3):
            xx[i,j] = x[i+j*M]
            yy[i,j] = y[i+j*M]
            zz[i,j] = z[i+j*M]
        
    #Write obj for bathymetry
    FN = create_filename(basefilename, 'obj', size)        
    write_obj(FN,xx,yy,zz)


    #Now read all the data with variable information, combine with
    #x,y info and store as obj

    for k in range(len(time)):
        t = time[k]
        print 'Processing timestep %f' %t

        for i in range(M):
            for j in range(3):
                zz[i,j] = stage[k,i+j*M]


        #Write obj for variable data
        #FN = create_filename(basefilename, 'obj', size, time=t)
        FN = create_filename(basefilename[:5], 'obj', size, time=t)   
        write_obj(FN,xx,yy,zz)


def dat2obj(basefilename):
    """Convert line based data output to obj
    """

    import glob, os
    from config import data_dir
   
    
    #Get bathymetry and x,y's
    lines = open(data_dir+os.sep+basefilename+'_geometry.dat', 'r').readlines()
    
    from Numeric import zeros, Float

    M = len(lines)  #Number of lines
    x = zeros((M,3), Float)
    y = zeros((M,3), Float)
    z = zeros((M,3), Float)

    ##i = 0
    for i, line in enumerate(lines):
        tokens = line.split()
        values = map(float,tokens)        

        for j in range(3):
            x[i,j] = values[j*3]
            y[i,j] = values[j*3+1]
            z[i,j] = values[j*3+2]        
        
        ##i += 1


    #Write obj for bathymetry
    write_obj(data_dir+os.sep+basefilename+'_geometry',x,y,z)


    #Now read all the data files with variable information, combine with
    #x,y info
    #and store as obj

    files = glob.glob(data_dir+os.sep+basefilename+'*.dat')

    for filename in files:
        print 'Processing %s' % filename

        lines = open(data_dir+os.sep+filename,'r').readlines()
        assert len(lines) == M
        root, ext = os.path.splitext(filename)

        #Get time from filename
        i0 = filename.find('_time=')
        if i0 == -1:
            #Skip bathymetry file
            continue
        
        i0 += 6  #Position where time starts           
        i1 = filename.find('.dat')	    

        if i1 > i0:
            t = float(filename[i0:i1])
        else:
            raise 'Hmmmm'
                
        
        
        ##i = 0
        for i, line in enumerate(lines):
            tokens = line.split()
            values = map(float,tokens)        

            for j in range(3):
                z[i,j] = values[j]        
        
            ##i += 1

        #Write obj for variable data
        write_obj(data_dir+os.sep+basefilename+'_time=%.4f' %t,x,y,z)


def filter_netcdf(filename1, filename2, first=0, last=None, step = 1):
    """Read netcdf filename1, pick timesteps first:step:last and save to
    nettcdf file filename2
    """
    from Scientific.IO.NetCDF import NetCDFFile        
    
    #Get NetCDF
    infile = NetCDFFile(filename1, 'r')  #Open existing file for read
    outfile = NetCDFFile(filename2, 'w')  #Open new file
    

    #Copy dimensions
    for d in infile.dimensions:
        outfile.createDimension(d, infile.dimensions[d])

    for name in infile.variables:
        var = infile.variables[name]
        outfile.createVariable(name, var.typecode(), var.dimensions)
        
    
    #Copy the static variables
    for name in infile.variables:
        if name == 'time' or name == 'stage':
            pass
        else:
            #Copy
            outfile.variables[name][:] = infile.variables[name][:]            

    #Copy selected timesteps 
    time = infile.variables['time']
    stage = infile.variables['stage']

    newtime = outfile.variables['time']
    newstage = outfile.variables['stage']    

    if last is None:
        last = len(time)

    selection = range(first, last, step)
    for i, j in enumerate(selection):
        print 'Copying timestep %d of %d (%f)' %(j, last-first, time[j])
        newtime[i] = time[j]
        newstage[i,:] = stage[j,:]        
        
    #Close
    infile.close()
    outfile.close()


#Get data objects
def get_dataobject(domain, mode='w'):
    """Return instance of class of given format using filename
    """

    cls = eval('Data_format_%s' %domain.format)
    return cls(domain, mode)




#OBSOLETE STUFF
#Native checkpoint format.
#Information needed to recreate a state is preserved
#FIXME: Rethink and maybe use netcdf format
def cpt_variable_writer(filename, t, v0, v1, v2):
    """Store all conserved quantities to file
    """

    M, N = v0.shape
	
    FN = create_filename(filename, 'cpt', M, t)
    #print 'Writing to %s' %FN
    
    fid = open(FN, 'w')
    for i in range(M):
        for j in range(N):        
            fid.write('%.16e ' %v0[i,j])            
        for j in range(N):
            fid.write('%.16e ' %v1[i,j])                        
        for j in range(N):
            fid.write('%.16e ' %v2[i,j])                        
            
        fid.write('\n')
    fid.close()


def cpt_variable_reader(filename, t, v0, v1, v2):
    """Store all conserved quantities to file
    """

    M, N = v0.shape
	
    FN = create_filename(filename, 'cpt', M, t)
    #print 'Reading from %s' %FN
    
    fid = open(FN)


    for i in range(M):
        values = fid.readline().split() #Get one line

        for j in range(N):
            v0[i,j] = float(values[j])
            v1[i,j] = float(values[3+j])
            v2[i,j] = float(values[6+j])
    
    fid.close()

def cpt_constant_writer(filename, X0, X1, X2, v0, v1, v2):
    """Writes x,y,z,z,z coordinates of triangles constituting the bed
    elevation.
    Not in use pt
    """

    M, N = v0.shape

    print X0
    import sys; sys.exit() 
    FN = create_filename(filename, 'cpt', M)
    print 'Writing to %s' %FN
    
    fid = open(FN, 'w')
    for i in range(M):
        for j in range(2):        
            fid.write('%.16e ' %X0[i,j])   #x, y
        for j in range(N):                    
            fid.write('%.16e ' %v0[i,j])       #z,z,z,
            
        for j in range(2):                    
            fid.write('%.16e ' %X1[i,j])   #x, y
        for j in range(N):                                
            fid.write('%.16e ' %v1[i,j])  
        
        for j in range(2):                                
            fid.write('%.16e ' %X2[i,j])   #x, y
        for j in range(N):                                            
            fid.write('%.16e ' %v2[i,j]) 
            
        fid.write('\n')
    fid.close()



#Function for storing out to e.g. visualisation
#FIXME: Do we want this?
#FIXME: Not done yet for this version
def dat_constant_writer(filename, X0, X1, X2, v0, v1, v2):
    """Writes x,y,z coordinates of triangles constituting the bed elevation.
    """

    M, N = v0.shape
	
    FN = create_filename(filename, 'dat', M)
    #print 'Writing to %s' %FN
    
    fid = open(FN, 'w')
    for i in range(M):
        for j in range(2):        
            fid.write('%f ' %X0[i,j])   #x, y
        fid.write('%f ' %v0[i,0])       #z
            
        for j in range(2):                    
            fid.write('%f ' %X1[i,j])   #x, y
        fid.write('%f ' %v1[i,0])       #z
        
        for j in range(2):                                
            fid.write('%f ' %X2[i,j])   #x, y
        fid.write('%f ' %v2[i,0])       #z
            
        fid.write('\n')
    fid.close()



def dat_variable_writer(filename, t, v0, v1, v2):
    """Store water height to file
    """

    M, N = v0.shape
	
    FN = create_filename(filename, 'dat', M, t)
    #print 'Writing to %s' %FN
    
    fid = open(FN, 'w')
    for i in range(M):
        fid.write('%.4f ' %v0[i,0])
        fid.write('%.4f ' %v1[i,0])
        fid.write('%.4f ' %v2[i,0])        

        fid.write('\n')
    fid.close()