#!/usr/bin/env python

import os, sys, string
from Scientific.IO.NetCDF import NetCDFFile
from optparse import OptionParser

##############################################################
# create input options
def CreateParser():
    usage = "%prog [options]\n\n"
    usage += "An arbitrary dimension slice extractor for NetCDF 3.x files.\n"
    usage += "Dimensionality is preserved for all sliced variables.\n\n"
    usage += "Version 0.0\n"
    usage += "Remik.Lastname@noaa.gov where Lastname = Ziemlinski"

    parser = OptionParser(usage=usage)
    parser.add_option("-f", "--file", dest="filename1",
                      help="input filename", metavar="FILE")
    parser.add_option("-o", "--out", dest="filename2",
                      help="output filename", metavar="FILE")
    parser.add_option("-v", "--var", dest="var",
                      help="comma-delimited list of variables to extract.  if not given, processes all variables", \
                      metavar="VAR[,VARn]")
    parser.add_option("-d", "--dim", dest="dim",
                      help="extracted dimension name", metavar="DIM")
    parser.add_option("-i", "--index", dest="index",
                      help="comma-delimited list of dimension indices to extract (0-based)", metavar="INDEX[,INDEXn]")
    parser.add_option("-V", "--verbose",
                      action="store_true", dest="verbose", default=False,
                      help="shows progress")
    return parser
   
##############################################################
def CopyGlobalAtts(f1,f2):
    # copy global attributes

    # ignore the standard members for NetCDFFile
    ignore = ['close', 'createDimension', 'createVariable', 'filename', 'flush', 'sync']

    dirf1 = dir(f1)
    for k in dirf1:
        # is the attribute a standard interface member we should ignore?
        if ignore.count( k ) == 0:
            # its a valid attribute, copy it
            exec('f2.%s = f1.%s' % (k,k))           

##############################################################
def CopyVarAtts(f1,f2,v):
    # ignore the standard members for NetCDFVariable
    ignore = ['assignValue', 'getValue', 'typecode']
   
    dirv = dir(f1.variables[v])
    for a in dirv:
        if ignore.count(a) == 0:
            exec('f2.variables[v].%s = f1.variables[v].%s' % (a, a))

##############################################################
def GetDims(f, variables):
    """
    Returns list of necessary dimensions.
    """
    dims = []

    for v in variables:
        vdims = f.variables[v].dimensions
        
        for vd in vdims:
            if dims.count(vd) == 0:
                dims.append(vd)
    
    return dims
    
##############################################################    
def CreateDimDefs(f1,f2,variables,index,dim):
    """
    Creates dimension definitions in f2 based on f1, variables, and selected indices.
    """
    # 
    # find which dimensions are needed
    #
    dims = GetDims(f1, variables)
    
    #
    # copy their sizes
    #
    for d in dims:
        if d <> dim:
            f2.createDimension(d, f1.dimensions[d])
        
    #
    # define new dimension size
    #
    n = len(index)
     
    f2.createDimension(dim, n)
    
##############################################################    
def CreateVarDefs(f1,f2,variables):
    vars_dims = []
    vars_dims.extend(variables)
    
    #
    # create coord var defitions; 
    # test if dimension name is also a defined variable
    #   
    for d in f2.dimensions.keys():
        if f1.variables.has_key(d):
            if (vars_dims.count(d) == 0):
                vars_dims.append(d)
                
    #
    # remove occurance of record dimension var name 
    # we assume defined outside this function
    #
    recdim = GetRecDimName(f2)    
    while vars_dims.count(recdim) <> 0:
        vars_dims.remove(recdim)
    
    #
    # copy name,type,dimname definitions from in file to out file
    #
    for v in vars_dims:
        vdims = f1.variables[v].dimensions
        typecode = f1.variables[v].typecode()
        f2.createVariable(v, typecode, vdims)                   
    
##############################################################
def CopyVarAttsFileToFile(f1,f2):
    for v in f2.variables.keys():
        CopyVarAtts(f1,f2,v)
    
##############################################################
def UsesDim(var, dim):
    vdim = var.dimensions        
    for d in vdim:
        if dim == d:
            return True
    
    return False

##############################################################
def GetRecDimName(f):
    for d in f.dimensions.keys():
        if f.dimensions[d] == None:
            return d
    
    return None
    
##############################################################
def ExtractVarSlices(f1,f2,v,index,dim):
    """
    Assumes that this variable DOES use the dimension being sliced.
    """
    
    # find which axis uses the dim, and make a bracketed indexing string
    vdim = f1.variables[v].dimensions
    slice = '['
    for i in range(len(vdim)):
        if vdim[i] == dim:
            slice += '%s'
            
            if (i+1) < len(vdim):
                slice += ',:'            
            
            slice += ']'
            break
        else:
            slice += ':,'
   
    #
    # extract the indexed slices
    #
    v1 = f1.variables[v]
    v2 = f2.variables[v]
    ctr = 0
    for i in index:
        s1 = slice % (i)
        s2 = slice % (ctr)
        cmd = 'v2%s = v1%s' % (s2, s1)
        exec(cmd)
        ctr += 1
    
##############################################################
def CopyRecVar(f1,f2):
    """
    Must be called to establish size of record dimension because
    it's initially None.  
    Also, must be called after the output file's dimensions were defined.
    After this function, any variable may be copied.
    """
    recdim = GetRecDimName(f1)
    recdimOut = GetRecDimName(f2)
    
    # ensure the output will use the record dimension
    if (recdimOut == None):
        return
    
    # create the definition
    recdims  = f1.variables[recdim].dimensions
    typecode = f1.variables[recdim].typecode()
    f2.createVariable(recdim, typecode, recdims) 

    # must copy one value at a time to make record dimension GROW despite defined None size
    n = len(f1.variables[recdim][:])
    for r in range(n):
        f2.variables[recdim][r] = f1.variables[recdim][r]
        
##############################################################
def Extract(f1,f2,index,dim,verbose):
    #
    # for each variable, check if 'dim' is a dimension
    # if yes, then find which axis, and copy appropriate indices
    # otherwise, just copy over all the variable's data
    #
    for v in f2.variables.keys():
        if verbose:
            print 'Processing %s...' % (v)
            
        # check for 'dim'
        if not UsesDim(f1.variables[v], dim):
            # doesn't use the slicing vertical dimension, copy everything over
            f2.variables[v][:] = f1.variables[v][:]            
        else:
            ExtractVarSlices(f1,f2,v,index,dim)
                    
##############################################################

def main():
    parser = CreateParser()
    (options, args) = parser.parse_args()
        
    filename1 = options.filename1
    filename2 = options.filename2
    variables = options.var
    index     = options.index
    dim       = options.dim
    verbose   = options.verbose
        
    if not filename1:
        print >>sys.stderr,'ERROR: [-f|--file=] input filename required'
        sys.exit()

    if not os.path.exists(filename1):
        print >>sys.stderr,'ERROR: file %s does not exist' % filename
        sys.exit()

    if not filename2:
        print >>sys.stderr,'ERROR: [-o|--out=] output filename required'
        sys.exit()
        
    if variables:      
        variables = string.split(variables, ',')
        
    if not index:
        print >>sys.stderr,'ERROR: [-i|--index] list of dimension indices required'
        sys.exit()
    else:
        index = string.split(index, ',')
        # convert strings to integers
        for i in range(len(index)):
            index[i] = int(index[i])
                
    if not dim:
        print >>sys.stderr,'ERROR: [-d|--dim=] extracted dimension name required'
        sys.exit()
   
    # open the input file
    f1 = NetCDFFile(filename1)

    # create output file
    f2 = NetCDFFile(filename2, 'w')

    if not variables:
        # user didn't specify any variables, so process all variables in input file
        variables = f1.variables.keys()
        

    if verbose:
        print 'Copying global attributes...'
        
    CopyGlobalAtts(f1,f2)

    if verbose:
        print 'Copying dimension definitions...'

    CreateDimDefs(f1,f2,variables,index,dim)

    if verbose:
        print 'Copying record definition...'

    CopyRecVar(f1,f2) # must do this before defining/copying any other variables to establish rec dim

    if verbose:
        print 'Copying variable definitions...'

    CreateVarDefs(f1,f2,variables)
    
    if verbose:
        print 'Copying variable attributes...'

    CopyVarAttsFileToFile(f1,f2)
    
    Extract(f1,f2,index,dim,verbose)    
    
    if verbose:
        print 'Closing files...'
        
    f1.close()
    f2.close()

##############################################################
    
if __name__=='__main__':
    main()













#======================================================================
#  U.S. Department of Commerce (DOC) Software License for script 
#  developed at the Geophysical Fluid Dynamics Laboratory/NOAA.
#
#  1. Scope of License
#
#   Subject to all the terms and conditions of this license, DOC grants
#   USER the royalty-free, nonexclusive, nontransferable, and worldwide
#   rights to reproduce, modify, and distribute this script
#   developed at the Geophysical Fluid Dynamics Laboratory/NOAA, herein
#   referred to as the Product.
#
#  2. Conditions and Limitations of Use
#
#   Warranties.  Neither the U.S. Government, nor any agency or
#     employee thereof, makes any warranties, expressed or implied,
#     with respect to the Product provided under this License,
#     including but not limited to the implied warranties or
#     merchantability and fitness for any particular purpose.
#
#   Liability.  In no event shall the U.S. Government, nor any agency
#     or employee thereof, be liable for any  direct, indirect, or
#     consequential damages flowing from the use of the Product
#     provided under this License.
#
#   Non-Assignment.  Neither this License nor any rights granted
#     hereunder are transferable or assignable without the explicit
#     prior written consent of DOC. 
#
#   Names and Logos.  USER shall not substitute its name or logo for the
#     name or logo of DOC, or any of its agencies, in identification of
#     the Product.
#
#   Export of technology.  USER shall comply with all U.S. laws and
#     regulations restricting the export of the Product to other
#     countries.  
#
#   Governing Law.  This License shall be governed by the laws of 
#     United States as interpreted and applied by the Federal courts 
#     in the District of Columbia.
#
#  3. Term of License This License shall remain in effect as long as USER
#     uses the Product in accordance with Paragraphs 1 and 2.
#======================================================================