#!/usr/bin/env python
# Last modified:  Time-stamp: <2008-01-08 16:10:22 haines>
"""
Create, update and load utilities for netcdf files
"""

from pycdf import *
import os
import numpy

def nc_create(ncFile, (global_atts, var_atts, dim_inits, var_inits, var_data)):
    """
    Create new netcdf file

    :Parameters:
        ncFile : string
          Path and name of file to create
    """
    try:
        # Open new netCDF file, overwrite if it exists, create if does not
        nc = CDF(ncFile, NC.WRITE|NC.CREATE|NC.TRUNC)
        # Automatically set define and data modes.
        nc.automode()
        #
        # GLOBALS
        for attrName in global_atts.keys():
            setattr(nc, attrName, global_atts[attrName])
        
        # DIMENSIONS
        for dim in dim_inits:
            dimName, dimValue = dim
            # print '%s = %d' % (dimName, dimValue)
            ncdim = nc.def_dim(dimName, dimValue)
        
        # VARIABLES
        for var in var_inits:
            varName, varType, varDim = var
            ncvar = nc.def_var(varName, varType, varDim)
            # add attributes
            for attrName in var_atts[varName].keys():
                setattr(ncvar, attrName, var_atts[varName][attrName])
            # setattr(ncvar, '_FillValue', numpy.nan)
            
        # add data
        nrecs = nc.inq_unlimlen()
        for var in var_data:
            varName, varData = var
            # print varName
            ncvar = nc.var(varName)
            # e.g. lat = array(var_data['lat'])
            # if an array
            if type(varData) == numpy.ndarray:
                if ncvar.isrecord():
                    # time, ens, u, v
                    ncvar[nrecs:nrecs+len(varData)] = varData.tolist()
                else:
                    ncvar[:] = varData.tolist() # z
            else:
                # if tuple, sequence or scalar
                ncvar[:] = varData
        
        nc.close()
    except CDFError, msg:
        print "CDFError:", msg
        # if nc:
        #     nc.close()
        #     del(nc)

def nc_update(ncFile, (global_atts, var_atts, var_data)):
    """
    Create new netcdf file

    :Parameters:
        ncFile : string
          Path and name of file to create
    """
    try:
        # Open netCDF in write mode
        nc = CDF(ncFile, NC.WRITE)
        # Automatically set define and data modes.
        nc.automode()
        #
        # GLOBALS
        for attrName in global_atts.keys():
            setattr(nc, attrName, global_atts[attrName])
        
        # VARIABLES
        # update attributes
        for var in var_atts:
            varName, atts = var
            ncvar = nc.var(varName)
            for attrName in atts.keys():
                setattr(ncvar, attrName, atts[attrName])
            
        # update data
        nrecs = nc.inq_unlimlen()
        for var in var_data:
            varName, varData = var
            ncvar = nc.var(varName)
            # e.g. lat = array(var_data['lat'])
            # if an array
            if type(varData) == numpy.ndarray:
                if ncvar.isrecord():
                    # time, ens, u, v (with unlimited dimension)
                    ncvar[nrecs:nrecs+len(varData)] = varData.tolist()
                else:
                    ncvar[:] = varData.tolist() # z (limited dimension)
            else:
                # if tuple, sequence or scalar
                ncvar[:] = varData

        nc.close()
    except CDFError, msg:
        print "CDFError:", msg
        # if nc:
        #     nc.close()
        #     del(nc)

def nc_get_time(ncFile):
    """get time array from file """
    try:
        nc = CDF(ncFile)
        ncvars = nc.variables()
        if 'time' in ncvars.keys():
            es = nc.var('time')[:]
            units = nc.var('time').units
        nc.close()
        return (es, units)
    except CDFError, msg:
        print "CDFError:", msg


    

def nc_load(ncFile, nameType='variable_name',
            varNames='all', ga_flag=True, va_flag=True):
    """
    Load netcdf file

    :Parameters:
        ncFile : string
            Path and name of file to load

    :Other Parameters:
        nameType : string 'variable_name' (default) or 'standard_name'
            Defines naming convention to use for variable names as data
            are loaded.  Variable name is the name used to store data
            in file.  'standard_name' means use variable name based on
            variable attribute called 'standard_name' of netcdf variable.
        varNames : string or tuple of strings
            specific variable names to be loaded into a sequence or scalar
            in python following specification set in nameType
            By default, all variables will be loaded.
        ga_flag : boolean flag
            By default, load the global file attributes
        va_flag : boolean flag
            By default, load the variable file attributes
            
    """
    try:
        nc = CDF(ncFile)
        attr = nc.attributes(full=1)
        dims = nc.dimensions(full=1)
        ncvars = nc.variables()
        for var in ncvars.keys():
            # load each variable by name??
            pass
        
    except CDFError, msg:
        print "CDFError:", msg