#!/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.
#
########################################################################

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 not opt.ifile :
    parser.print_help()
    sys.exit(1)

if not os.path.exists(opt.ifile):
    logging.error( "ERROR: Input file %s does not exist" % opt.ifile  )
    sys.exit(1)

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 )

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" )
