#!/opt/env/minihycom/bin/python

"""
Drill North Atlantic HYCOM THREDDS catalog for for input to Ichthyop.

Catalog is sliced by time, bounding box, and area of interest.
Depth dimension is eliminated by using only top layer.
Dimensions are transformed by specialized functions to simulate ROMS input
(TBD with Trond Kristiansen).

Delete this module from the importing namespace after use
to free up memory consumed by the input catalog.
"""

import datetime

__authors__ = ('Chris Calloway',
               'Trond Kristiansen',
               'Nathan F. Putman',
               'Thomas Shay',
             )
__copyright__ = 'Copyright 2008 UNC-Chapel Hill Department of Marine Sciences'
__license__ = 'GPL2'
__version__ = '1.0.0'
__created__ = datetime.date(2008, 6, 11)
__modified__ = datetime.date(2008, 6, 11)

import os
import dap.client
import Nio

# Move this out to a config file
dataSource = 'http://hycom.coaps.fsu.edu/thredds/dodsC/atl_ops'
dataDestination = '/seacoos/data/minihycom/minihycom.nc'

# Nio cannot create new file when old exists
try:
    os.remove(dataDestination)
except:
    # FIX THIS: Narrow down source of exception
    print dataDestination, 'does not already exist.'

# Note: is it not necessary (or possible) to close the input catalog,
# as backing files are only opened during data access functions.
dataIn = dap.client.open(dataSource, verbose=1)

# Map numpy types to Nio types.
# Expand as necessary to cover more types.
# Expand as necessary to cover Numeric and Numarray.
# Among other things, check this if you see a glibc error.
datatypes = {'Float32' : 'f',
             'Float64' : 'd',
            }

# Dimensional constraints.
timeName = dataIn.MT.name
timeType = datatypes[dataIn.MT.type]
minTime = 0
maxTime = 3
timeLen = maxTime - minTime

# Note: depth is not referenced as a dimension
# by any other variable in the output except depth.
depthName = dataIn.Depth.name
depthType = datatypes[dataIn.Depth.type]
minDepth = 0
maxDepth = dataIn.Depth.shape[0]
depthLen = maxDepth - minDepth

latitudeName = dataIn.Latitude.name
latitudeType = datatypes[dataIn.Latitude.type]
minLatitude = 365
maxLatitude = 1214
latitudeLen = maxLatitude - minLatitude

longitudeName = dataIn.Longitude.name
longitudeType = datatypes[dataIn.Longitude.type]
minLongitude = 0
maxLongitude = dataIn.Longitude.shape[0]
longitudeLen = maxLongitude - minLongitude

# Variables of interest.
dateName = dataIn.Date.name
dateType = datatypes[dataIn.Date.type]
uName = dataIn.u.name
uType = datatypes[dataIn.u.type]
vName = dataIn.v.name
vType = datatypes[dataIn.v.type]
temperatureName = dataIn.temperature.name
temperatureType = datatypes[dataIn.temperature.type]
salinityName = dataIn.salinity.name
salinityType = datatypes[dataIn.salinity.type]
densityName = dataIn.density.name
densityType = datatypes[dataIn.density.type]
sshName = dataIn.ssh.name
sshType = datatypes[dataIn.ssh.type]
thicknessName = dataIn.surface_boundary_layer_thickness.name
thicknessType = datatypes[dataIn.surface_boundary_layer_thickness.type]

def createVariableAndAttributes(dataOut, varName, varType, dims, dimSlices):
    """
    Create variable and all attributes at once.
    
    createVariableAndAttributes(dataOut, varName, varType, dims, dimSlices) -> 0
    
    varName - string name of variable
    varType - string for Nio type code
    dims - tuple of dimension names
    dimSlices - tuple of slice objects constraining input dataset
    """
    dummy = dataOut.create_variable(varName, varType, dims)
    dummy[:] = dataIn[varName][dimSlices]
    for attributeName, attributeValue in dataIn[varName].attributes.items():
        setattr(dataOut.variables[varName], attributeName, attributeValue)
    return 0

def _main():
    """
    Make the output NetCDF dataset.
    
    Run module as script to invoke.
    """
    dataOut = Nio.open_file(dataDestination, mode="c")
    
    # Copy global attributes.
    for attributeName, attributeValue in dataIn.attributes['NC_GLOBAL'].items():
        setattr(dataOut, attributeName, attributeValue)
    
    # Create dimensions
    dataOut.create_dimension(timeName, timeLen)
    # Note: depth is not referenced as a dimension
    # by any other variable in the output except depth.
    dataOut.create_dimension(depthName, depthLen)
    dataOut.create_dimension(latitudeName, latitudeLen)
    dataOut.create_dimension(longitudeName, longitudeLen)
    
    # Create variables and their attributes.
    createVariableAndAttributes(dataOut,
                                timeName,
                                timeType,
                                (timeName,),
                                (slice(minTime, maxTime),))
    
    # Note: depth is not referenced as a dimension
    # by any other variable in the output except depth.
    createVariableAndAttributes(dataOut,
                                depthName,
                                depthType,
                                (depthName,),
                                (slice(minDepth, maxDepth),))
    
    createVariableAndAttributes(dataOut,
                                latitudeName,
                                latitudeType,
                                (latitudeName,),
                                (slice(minLatitude, maxLatitude),))
    
    createVariableAndAttributes(dataOut,
                                longitudeName,
                                longitudeType,
                                (longitudeName,),
                                (slice(minLongitude, maxLongitude),))
    
    # The Date variable is another expression of the time dimension.
    createVariableAndAttributes(dataOut,
                                dateName,
                                dateType,
                                (timeName,),
                                (slice(minTime, maxTime),))
    
    # Eliminate depth as a dimension by slicing a scalar depth.
    createVariableAndAttributes(dataOut,
                                uName,
                                uType,
                                (timeName, latitudeName, longitudeName,),
                                (slice(minTime, maxTime),
                                 0,
                                 slice(minLatitude, maxLatitude),
                                 slice(minLongitude, maxLongitude),))
    
    createVariableAndAttributes(dataOut,
                                vName,
                                vType,
                                (timeName, latitudeName, longitudeName,),
                                (slice(minTime, maxTime),
                                 0,
                                 slice(minLatitude, maxLatitude),
                                 slice(minLongitude, maxLongitude),))
    
    createVariableAndAttributes(dataOut,
                                temperatureName,
                                temperatureType,
                                (timeName, latitudeName, longitudeName,),
                                (slice(minTime, maxTime),
                                 0,
                                 slice(minLatitude, maxLatitude),
                                 slice(minLongitude, maxLongitude),))
    
    createVariableAndAttributes(dataOut,
                                salinityName,
                                salinityType,
                                (timeName, latitudeName, longitudeName,),
                                (slice(minTime, maxTime),
                                 0,
                                 slice(minLatitude, maxLatitude),
                                 slice(minLongitude, maxLongitude),))
    
    createVariableAndAttributes(dataOut,
                                densityName,
                                densityType,
                                (timeName, latitudeName, longitudeName,),
                                (slice(minTime, maxTime),
                                 0,
                                 slice(minLatitude, maxLatitude),
                                 slice(minLongitude, maxLongitude),))
    
    # Depth is not a dimension in the input for the remaining variables.
    createVariableAndAttributes(dataOut,
                                sshName,
                                sshType,
                                (timeName, latitudeName, longitudeName,),
                                (slice(minTime, maxTime),
                                 slice(minLatitude, maxLatitude),
                                 slice(minLongitude, maxLongitude),))
    
    createVariableAndAttributes(dataOut,
                                thicknessName,
                                thicknessType,
                                (timeName, latitudeName, longitudeName,),
                                (slice(minTime, maxTime),
                                 slice(minLatitude, maxLatitude),
                                 slice(minLongitude, maxLongitude),))
    
    # Write the output dataset
    dataOut.close()
    return 0

# Run as script.
if __name__ == "__main__":
    _main()