## Python script to check model outputs within a given set of values from an external ASII file
# L. Fita, LMD. April 2015
## e.g. # checking_output.py -f histins.nc -v t2m -l 'limitsVariables.inf'
## e.g. # python checking_output.py -f limit.nc -v SST,ALB -l 'limitsVariables.inf'
# NOTE:
# `nc_var_tools.py' set of python scripts is required
# can be found in subversion `http://svn.lmd.jussieu.fr/LMDZ_WRF/trunk/tools/nc_var_tools.py'
import numpy as np
from netCDF4 import Dataset as NetCDFFile
import os
import re
from optparse import OptionParser
import nc_var_tools as ncvar

main = 'checking_output.py'
version = 0.1
author = 'L. Fita'
institution = 'Laboratoire de Meteorologie Dynamique'
city = ' Paris'
country = 'France'

errormsg = 'ERROR -- error -- ERROR -- error'
warnmsg = 'WARNING -- warning -- WARNING -- warning'

fillValue = 1.e20

def compute_stddev(vals):
    """ Function to compute standard deviation
      vals= values
    """
    fname = 'compute_stddev'

    mvals = np.mean(vals)
    m2vals = np.mean(vals*vals)

    stddev = np.sqrt(m2vals - mvals*mvals)

    return stddev

def checking(cn,valc,vern,lval,Vn):
    """ Function to check a variable with
      cn= checking name
      valc= value to check
      vern= verification value
      lval= value to use for checking
      Vn= variable name
    """
    fname = 'checking'

    VerN = vern[0:1]

    good = True
    if VerN == '<' and valc < lval:
        print warnmsg
        print '  ' + fname + "  variable '" + vn + "' has a '" +cn+ "'",valc,vern,lval
        good = False
    elif VerN == '>' and valc > lval:
        print warnmsg
        print '  ' + fname + "  variable '" + vn + "' has a '" +cn+ "'",valc,vern,lval
        good = False
    elif VerN == '%' and np.abs((valc-lval)*100./lval)>np.float(vern[1:len(vern)+1]):
        print warnmsg
        print '  ' + fname + "  variable '" + vn + "' has a '" +cn+ "'", valc, '>',  \
          VerN, np.float(vern[1:len(vern)+1]), 'than',lval
        good = False

    return good


####### ###### ##### #### ### ## #

parser = OptionParser()
parser.add_option("-f", "--file", dest="ncfile",
  help="file to check", metavar="FILE")
parser.add_option("-v", "--variable", dest="varns",
  help="var to check (',' list of variables; 'all', for all variables)", metavar="VALUE")
parser.add_option("-l", "--ASCIIvarlimits", dest="lfile",
  help="ASCII file with the limits for the variables (long explanation with -l h)", metavar="FILE")

(opts, args) = parser.parse_args()

#######    #######
## MAIN
    #######

# Available statistics
Availablestats = ['min', 'max', 'mean', 'stddev']
LAvailablestats = {'min':'minimum', 'max':'maximum', 'mean':'mean',                  \
  'stddev':'standard deviation'}

Availableverifs = ['<', '>', '%']
SAvailableverifs = {'<':'lt', '>':'gt', '%':'percen'}
LAvailableverifs = {'<':'smaller', '>':'bigger', '%':'percentage'}

####### 
ofile = 'wrongvariables.nc'

if not os.path.isfile(opts.ncfile):
    print errormsg
    print '   ' + main + ": file '" + opts.ncfile + "' does not exist !!"
    quit(-1)

if opts.varns is None:
    print warnmsg
    print '   ' + main + ": no variables given!!"
    print '    checking all'
    varns = 'all'
else:
    varns = opts.varns

if opts.lfile is None:
    print warnmsg
    print '   ' + main + ": no ASCII file with limits for the variables is given!!"
    quit(-1)
elif opts.lfile == 'h':
    print 'HELP: ASCII limits of the variables_____________'
    print 'First line: varname [stat1] [stat2] .... [statn]'
    print 'Second line: - [verif1] [verif2] ... [verifn]'
    print 'afterwards: [varname] [value1] [value2] ... [valuen]'
    print '  [stat]: statistics/parameter to compute. Has to exist wihin the script'
    print '    directly with an if or as function compute[stat]. Available ones are:'
    print '    ',Availablestats
    print '  [verif]: verification to make with the stat'
    print '    <: Error if variable has values smaller than the limit'
    print '    >: Error if variable has values bigger than the limit'
    print '    %N: Error if variable has values +/- N % than the limit'
    quit()

ncobj = NetCDFFile(opts.ncfile, 'r')
filevars = ncobj.variables.keys()

# Getting variables names
if varns == 'all':
    vns = filevars
else:
    vns = opts.varns.split(',')

Nvars = len(vns)
print "From file '" + opts.ncfile + "' checking",Nvars,'variables...'

# Getting ASCII file and its values
##
# Assuming first line with something like varname min max mean stddev
# Continuted with a line for each variable and value like:
# t2m -70. 50. 20. 40.
# (...)
# It is required than this script some function called compute_[value] must be placed
# '#' for comment
olval = open(opts.lfile, 'r')
ili = 0
statsverif = {}
limitvals = {}

for line in olval:
    linevals = ncvar.reduce_spaces(line)
    if len(linevals) != 0 and linevals[0][0:0] != '#':
#        print ili,linevals
        if ili == 0:
            NVals = len(linevals) - 1
            stats = linevals[1:NVals + 1]
        elif ili == 1:
            for ist in range(NVals):
                statsverif[stats[ist]] = linevals[1+ist]
        else:
            limitvals[linevals[0]] = np.zeros((NVals), dtype=np.float)
            for ist in range(NVals):
                limitvals[linevals[0]][ist] = np.float(linevals[1+ist])

    ili = ili + 1

# Checking stats to compute
for st in stats:
    if not ncvar.searchInlist(Availablestats,st):
        print errormsg
        print '  ' + main + ": statistics/value '" + st + "' not ready!!"
        print '    available parameters:', Availablestats
        quit(-1)

# Checking verifiaction to make
for st in stats:
    verin = statsverif[st]
    if not ncvar.searchInlist(Availableverifs,verin[0:1]):
        print errormsg
        print '  ' + main + ": verification '" + verin + "' not ready!!"
        print '    available verifications:', Availableverifs
        quit(-1)

print 'Checking',NVals,'parameters:',stats,'verifyed as:',statsverif.values()

# Creating output file
oobj = NetCDFFile(ofile, 'w')

# Dimensinos
newdim = oobj.createDimension('stats', NVals)
newdim = oobj.createDimension('verif', NVals)
newdim = oobj.createDimension('checks', 2)
newdim = oobj.createDimension('Lstring', 50)

# Variables with Dimensions values
newvar = oobj.createVariable('stats', 'c', ('stats', 'Lstring'))
newattr = ncvar.basicvardef(newvar, 'stats', 'statistics/parameters used to check',  \
  '-')
newvals = ncvar.writing_str_nc(newvar, stats, 50)

newvar = oobj.createVariable('verifs', 'c', ('stats', 'Lstring'))
newattr = ncvar.basicvardef(newvar, 'verifs', 'verifications used to check', '-')
newvals = ncvar.writing_str_nc(newvar, stats, 50)

newvar = oobj.createVariable('checks', 'c', ('checks', 'Lstring'))
newattr = ncvar.basicvardef(newvar, 'checks', 'values used to check', '-')
newvals = ncvar.writing_str_nc(newvar, ['computed', 'limits'], 50)

# Global attributes
newattr = ncvar.set_attribute(oobj, 'script', main)
newattr = ncvar.set_attribute(oobj, 'version', version)
newattr = ncvar.set_attribute(oobj, 'author', author)
newattr = ncvar.set_attribute(oobj, 'institution', institution)
newattr = ncvar.set_attribute(oobj, 'city', city)
newattr = ncvar.set_attribute(oobj, 'country', country)

newattr = ncvar.set_attribute(oobj, 'description', 'Variables from ' + opts.ncfile + \
  ' with a wrong value according to ' + opts.lfile)

oobj.sync()

# Getting netCDF file and its values
##
valsstats = np.zeros((NVals,Nvars), dtype = np.float)
wrongVars = {}
iv = 0
for iv in range(Nvars):
    vn = vns[iv]
    print vn + '...'
    if not ncvar.searchInlist(filevars, vn):
        print errormsg
        print '  ' + main + ": file '" + opts.ncfile + "' does not have variable '" +\
          vn + "' !!"
        quit(-1)

    vals = ncobj.variables[vn][:]
    wrongstats = []
    for ist in range(NVals):
        stn = stats[ist]
        Lstn = LAvailablestats[stn]
        statpos = stats.index(stn)
        limitval = limitvals[vn][statpos]
        verif = statsverif[stn]
#        print '  ', ist, stn, ':',limitval,verif
        if stn == 'max':
            valsstats[ist,iv] = np.max(vals)
        elif stn == 'mean':
            valsstats[ist,iv] = np.mean(vals)
        elif stn == 'min':
            valsstats[ist,iv] = np.min(vals)
        elif stn == 'stddev':
            valsstats[ist,iv] = compute_stddev(vals)

# checks
        if not checking(stn,valsstats[ist,iv],verif,limitval,vn):
            vval = valsstats[ist,iv]
            Sverif = SAvailableverifs[verif[0:1]]
            Lverif = LAvailableverifs[verif[0:1]]
            if verif[0:1] != '<' and verif[0:1] != '>':
                Vverif = np.float(verif[1:len(verif)])
                VverifS = str(Vverif)
            else:
                VverifS = ''

            wrongstats.append(stn)
            if not ncvar.searchInlist(oobj.variables.keys(), vn):
                oobj.close()
                ncvar.fvaradd(opts.ncfile + ':' + vn, ofile)

            oobj = NetCDFFile(ofile ,'a')
            vobj = oobj.variables[vn]
            if ncvar.searchInlist(vobj.ncattrs(),('units')):
                vunits = vobj.getncattr('units')
            else:
                vunits = '-'

            dims = vobj.dimensions
# 1/0 Matrix with wrong values
            if stn == 'max':
                wrongvals = np.where(vals > limitval, 1, 0)
                newvar = oobj.createVariable(vn+'Wrong'+stn+Sverif, 'i', dims)
                newattr = ncvar.basicvardef(newvar, vn+'_wrong_'+stn+'_'+Sverif,     \
                  'wrong ' + vn + ' ' + str(vval) + ' ' + Lstn + ' ' + VverifS + ' '+\
                  Lverif + ' ' + str(limitval), vunits)
                newvar[:] = wrongvals
            elif stn == 'mean':
                newattr = ncvar.set_attribute(vobj, 'wrong_'+stn+'_'+Sverif+VverifS, \
                  'wrong ' + vn + ' ' + str(vval) + ' ' + Lstn + ' ' + VverifS + ' '+\
                  Lverif + ' ' + str(limitval))
            elif stn == 'min':
                wrongvals = np.where(vals < limitval, 1, 0)
                newvar = oobj.createVariable(vn+'Wrong'+stn+Sverif, 'i', dims)
                newattr = ncvar.basicvardef(newvar, vn+'_wrong_'+stn+'_'+Sverif,     \
                  'wrong ' + vn + ' ' + str(vval) + ' ' + Lstn + ' ' + VverifS + ' '+\
                  Lverif + ' ' + str(limitval), vunits)
                newvar[:] = wrongvals
            elif stn == 'stddev':
                newattr = ncvar.set_attribute(vobj, 'wrong_'+stn+'_'+Sverif+VverifS, \
                  'wrong ' + vn + ' ' + str(vval) + ' ' + Lstn + ' ' + VverifS + ' '+\
                  Lverif + ' ' + str(limitval))
        
        if len(wrongstats) != 0: wrongVars[vn] = wrongstats
    newattr = ncvar.set_attribute(vobj, 'wrong',                                     \
      ncvar.numVector_String(wrongVars[vn], ' '))

    newvar = oobj.createVariable(vn + 'checks', 'f4', ('checks', 'stats'))
    newattr = ncvar.basicvardef(newvar, vn + '_checks', vn +                         \
      ' computed/checked', vunits)
    for ist in range(NVals):
        newvar[:,ist] = [valsstats[ist,iv], limitvals[vn][ist]]

oobj.sync()
oobj.close()

# Results
## 
print ' Variables with Wrong results _______________________________________________'
print wrongVars

print "succesfull writting of checking output file '" + ofile + "' !!"