#!/usr/bin/env python

########################################################################
#
# py_interp alpha by Markel Garcia Diez February 2014
#
# py_interp 1.0.0 by Carlos Blanco September 2015 
# 
# See the README for info and/or type python py_interp.py -h for help.
#
########################################################################

__version__  = '1.0.0'
__author__   = 'Carlos Blanco'
__revision__ = "$Id: py_interp.py 2819 2015-09-23 15:24:22Z carlos $"

import numpy   as np
import netCDF4 as ncdf
import os, sys, logging
import py_interp.diags
from py_interp.fun import ( copy_n_filter_wrfout, BasicFields, 
                            add_pressure_axis, interp2plevs )
from optparse      import OptionParser
#
# Parse options
#
parser = OptionParser()
parser.set_defaults(quiet=False,singlerec=False)
parser.add_option("-i", dest="ifile", help="Input wrfout file")
parser.add_option("-v", dest="varlist", default="", help="Comma sepparated variable list")
parser.add_option("-p", "--plevs", dest="plevs", help="Comma sepparated pressure level list in hPa")
parser.add_option("--verbose", dest="verbose", action="store_true", default=False, help="Get a more detailed log")
(opt, args) = parser.parse_args()
if opt.verbose :
    logging.basicConfig(level=logging.DEBUG)
#
# Some global variables
#
ifile = opt.ifile
ofile = ifile + "_PLEV"
varlist = opt.varlist.split(",")[:]
plevs = np.array(opt.plevs.split(",")[:], dtype=np.float32)
logging.info( "Starting py_interp" )
logging.debug( "Input file: %s" % ifile )
logging.debug( "Output file: %s" % ofile )
logging.debug( "Variables to be processed: %s" % varlist )
logging.debug( "Levels to interpolate: %s" % plevs )
if not os.path.exists(ifile):
	logging.error( "ERROR: Input file %s does not exist" % ifile )
	sys.exit(1)
try :
	#
	# Reading input wrfout file
	#
	logging.info( "Reading file %s" % ifile )
	inc = ncdf.Dataset(ifile, "r")
	#
	# Separate the input variables in three lists: 2D variables (or soil variables)
	# to copy, 3D variables to interpolate and diagnostics that need specific
	# functions to be computed.
	#
	dims2D=("Time", "south_north", "west_east")
	diagnostics = ["VIQC", "VIQI", "VIM", "MSLP", "CLT", "CLT_OLD", "CLH", "CLM", "CLL", "GHT", "PRES", "RH", "TT"]
	copyvars = []
	interpvars = []
	diags = []
	for var in varlist:
		if var in diagnostics:
			diags.append(var)
		else:
			if var not in inc.variables.keys():
				logging.info( "Error: Variable %s not found in file, and is not a diagnostic" % var )
				sys.exit(1)
			vdims = inc.variables[var].dimensions
			if ("bottom_top" in vdims) or ("bottom_top_stag" in vdims):
				interpvars.append(var)
			else:
				copyvars.append(var)

	logging.debug( "Variables to copy %s" % copyvars )
	logging.debug( "3D variables to interpolate %s" % interpvars )
	logging.debug( "Diagnostics to compute %s" % diags )
	#
	# Pressure levels from hPa to Pa.
	#
	plevs = plevs*100
	#
	# Get some basic fields: psfc, pres_field, ght, hgt, temp
	#
	bf = BasicFields(inc)
	#
	# Copy wrfnc structure with the 2D vars,  which are left untouched
	#
	logging.debug( "Copying wrfout file structure and 2D fields" )
	onc = copy_n_filter_wrfout(inc, ofile, copyvars)
	#
	# Add the pressure axis and variable
	#
	logging.debug( "Adding pressure axis to output file" )
	onc = add_pressure_axis(onc, plevs)
	#
	# Interpolate to pressure levels 3D variables
	#
	for var in interpvars:
		logging.debug( "Interpolating %s to pressure levels" % var )
		onc = interp2plevs(var, inc, onc, bf, plevs)
	#
	# Compute diagnostics
	#
	for var in diags:
		logging.debug( "Computing diagnostic %s" % var )
		#
		# Use getattr to get the function from py_interp_diags.py and then call it
		#
		compute_diag = getattr(py_interp.diags, "compute_%s" % var)
		onc = compute_diag(var, inc, onc, bf, plevs)
	onc.sync()
except Exception as err :
	logging.error( "ERROR: processing %s file: %s" %(ifile,err) )
else :	
	logging.info( "SUCCESS: p_interp finished without errors" )
