#!/usr/bin/env python2.5
#
#  txt2netcdf.py
#  
#
#  Created by huik on 2/4/11.
#  Copyright (c) 2011 __MyCompanyName__. All rights reserved.
#

# import all the necessary modules
import sys, os, time
from numpy import *
from netCDF4 import Dataset
from scipy import loadtxt

# module specification
__version__ = '1.0'
__author__ = 'Ko van Huissteden'
__all__ = ['createnew', 'addvariable', 'text2cdf', 'listcdf', 'timeseries2cdf']

# functions: 
# createnew:  creates a new netcdf file with user specified coordinate dimensions
# addvariable: adds a variable to an existing netcdf file
# listcdf: lists the structure of an existing file
# text2cdf: appends the contents of a text file to an existing file
# timeseries2cdf: creates a simple time series net cdf file for one point in space

# FUNCTION DEFINITIONS

def createnew(outfile = 'x.nc', dimensions = ['lat', 'lon', 'time'], sizes = [1, 1, 0]):
# create a new netcdf file with name outfile
# BEWARE: an existing file of the same name is overwritten!
# and dimensions as specified in list dimensions
# and sizes as specified in list sizes
# 0 stands for an unlimited dimension

# open the file in overwrite mode
	nfile=Dataset(outfile, 'w')
	errors = 0;
	try:
	# add a basic comment attribute 
		nfile.history = 'Created from Python NetCDF4' + time.ctime(time.time())
	except:
		print 'File ', infile, ' not opened properly, cannot add attributes.\n'
		return 1
	for i in range(len(dimensions)):
	# add the coordinate dimensions
		try:
			nfile.createDimension(dimensions[i], sizes[i])
		except:
			print 'Cannot add dimensions to file ', infile
			errors +=1
	nfile.close()
	return errors

def addvariable(outfile = 'x.nc', varname = 'data', vartype= ['f8'], dimensions = ['lat', 'lon', 'time']):
# adds a varible to a netcdf file with name outfile
# BEWARE: an existing file of the same name is overwritten!
# BEWARE: no checking is done!
# BEWARE: data is not added with this function
# varname: string
# vartype: variable type
# dimensions: the list of dimensions which are the coordinate axes of the variable
# Valid datatype specifiers include: 'f4' (32-bit floating point), 'f8' (64-bit floating point), 'i4' (32-bit
# signed integer), 'i2' (16-bit signed integer), 'i8' (64-bit singed integer), 'i1' (8-bit signed integer), 'u1' (8-bit
# unsigned integer), 'u2' (16-bit unsigned integer), 'u4' (32-bit unsigned integer), 'u8' (64-bit unsigned integer), or
# 'S1' (single-character string).as specified in list dimensions

# open the file in append mode
	try:
		ncfile=Dataset(outfile, 'a') # open the file
	except:
		print 'File ', infile, ' does not exist or is not a valid NetCDF file.\n'
		return
	errors = 0;
	# add the variable
	try:
		varadded = ncfile.createVariable(varname, vartype, tuple(dimensions))
	except:
		print 'Addition of variable ', varname, ' failed\n'
		errors = 1;
	ncfile.close()
	return errors

def text2cdf(infile, outfile, varname, coordinates):
# adds data from text file infile to netcdf file outfile to variable varname
# infile: text file with numeric data
# outfile: netcdf file to append to
# dimensions: dimension names of the data; if in append mode, these should be defined in the file
# datatype: nemeric datatype, int, float or double

# load the input data
	try:
		inputdata = loadtxt(infile)
		l = len(inputdata) - 1
	except:
		print 'Cannot read text data input: ', infile
		return 1
# open the file in append mode
	try:
		nfile=Dataset(outfile, 'a') # open the netcdf file for appending
	except:
		print 'File ', infile, ' does not exist or is not a valid NetCDF file.\n'
		return 1
	datashape = nfile.variables[varname].shape # determine shape, not used now, but may be used in error checking later
	try:
		nfile.variables[varname][coordinates[0],coordinates[1],coordinates[2]:l]=inputdata[0:l]
		errors = 0
	except:
		print 'Addition of data to variable ', varname, 'at coordinates', coordinates, ' failed\n'
		errors = 1
	nfile.close()
	return errors

def listcdf(infile):
# lists the structure of a netcdf file
# needs to be extended later
	try:
		nfile=Dataset(infile, 'r') # open the file
	except:
		print 'File ', infile, ' does not exist or is not a valid NetCDF file.\n'
		return
	print 'Global attributes:\n' # file attributes
	for name in nfile.ncattrs():
		print 'Global attr', name, '=', getattr(nfile,name)
	print '\n'
	print 'Dimensions:\n' # file dimensions
	print 'Name            Length   Unlimited'
	for dimname, dimobj in nfile.dimensions.iteritems():
		print "%-15s %-8d %s" % (dimname, len(dimobj), dimobj.isunlimited())
	print '\n'
	print 'Variables:\n' # file variables
	for varname, varobj in nfile.variables.iteritems():
		print varname
	print '\n'
	nfile.close()
	
def timeseries2cdf(infile, netcdffile, varname = 'data', vartype= 'f8'):
# timeseries2cdf: creates a simple time series net cdf file for one point in space
# example: timeseries2cdf('watertable.txt', 'watertable.nc', 'watertable') or 
# timeseries2cdf('watertable.txt', 'watertable.nc', 'watertable', 'f4')
# the coordinates will be called 'lat', 'lon', 'time' with only one point for lat and lon
# the variable name can be specified, the variable type will be 64 bit floating point by default but can be specified
# if the input text file is a one column file, a time axes will be added by simply numbering the entries
# if the input text file is a two column file, the first column is assumed to be the time coordinate
	
	dimensions = ['lat', 'lon', 'time']
	error = createnew(netcdffile, dimensions, [1, 1, 0]) # create the file
	if not error:
		error = addvariable(netcdffile, varname, vartype, ['lat', 'lon', 'time']) # add the variable
		if not error:
			# add the data
			try: 
				inputdata = loadtxt(infile)
				
			except:
				print 'Cannot read text data input: ', infile, '\n'
				return 1
			size = inputdata.shape  # first determine size of input data
			l = size[0] - 1
			w = 1
			if len(size) > 1:
				w = 2
			ncfile=Dataset(netcdffile, 'a') # open the netcdf file for appending
			try:
				times = ncfile.createVariable('time',vartype,('time',))  # create the time variable
			except:
				ncfile.close()
				print 'Adding time axis failed\n'
				return 1
			try:
				if w == 2: # add the data
					ncfile.variables['time'][:] = inputdata[0:l, 0] # this is the case of a 2 column input file
					ncfile.variables[varname][0, 0, 0:l] = inputdata[0:l, 1]
				else:
					ncfile.variables['time'][:] = arange(1, l + 2, 1) # this is a one column input file
					ncfile.variables[varname][0, 0, 0:l] = inputdata[0:l]
				ncfile.close()
			except: 
				error = 0
				print 'Addition of data failed\n'
				error = 1
	if error:
		print 'An error occured while creating file ', netcdffile, 'from ', infile, '\n'
	return error