# Python script to check namelist.input consistency with a given WRF version
# L. Fita, LMD, December 2014
import numpy as np
import os
from optparse import OptionParser
import sys
import nc_var_tools as ncvar

## e.g. # WRF_namelist_check.py -w /home/lluis/DATA/WRF/WRFV3.3/serial -n /home/lluis/DATA/WRF/WRFV3.3/serial/WRFV3/run/namelist.input
#* WRF structure
# - module_config_rec definition on: 
#   .- 'inc/namelist_defines.inc' definition of variables from 'namelist.input' called from 'frame/module_configure.F'
#   .- 'include/config_assigns.inc' assigment to the grid structure  called from 'frame/module_configure.F'

# - grid definition on:
#   .- 'inc/state_struct.inc' types of the grid content
#   .- 'frame/module_domain_type.F' definition of the grid
#   .- 'dyn_em/start_em.F' filling the grid

# - namelist reading in include: config_assigns.inc

main = 'WRF_namelist_check.py'

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

def searchInlist(listname, nameFind):
    """ Function to search a value within a list
    listname = list
    nameFind = value to find
    >>> searInlist(['1', '2', '3', '5'], '5')
    True
    """
    for x in listname:
      if x == nameFind:
        return True
    return False

def reduce_spaces(string):
    """ Function to give words of a line of text removing any extra space
    """
    values = string.replace('\n','').split(' ')
    vals = []
    for val in values:
         if len(val) > 0:
             vals.append(val)

    return vals

def numVector_String(vec,char):
    """ Function to transform a vector of numbers to a single string [char] separated
    numVector_String(vec,char)
      vec= vector with the numerical values
      char= single character to split the values
    >>> print numVector_String(np.arange(10),' ')
    0 1 2 3 4 5 6 7 8 9
    """
    fname = 'numVector_String'

    if vec == 'h':
        print fname + '_____________________________________________________________'
        print numVector_String.__doc__
        quit()

    Nvals = len(vec)

    string=''
    for i in range(Nvals):
        if i == 0:
            string = str(vec[i])
        else:
            string = string + char + str(vec[i])

    return string

def get_namelist_vars(values, namelist):
    """ Function to get namelist-like  values ([varname] = [value]) 
    get_namelist_vars(namelist)
      values= [sectionname],[kout]
        [sectionname]: name of the section from which one want the values ('all', for all)
        [kout]: kind of output
          'tex3': printed output as LaTeX table of three columns \verb+namelist_name+ & value
          'column': printed output as namelist_name value
          'dict': as two python dictionary object (namelistname, value and namelistname, sectionname)
      namelist= namelist_like file to retrieve values
    >>> get_namelist_vars('geogrid,dic','/home/lluis/etudes/domains/medic950116/namelist.wps')
    {'e_sn': '32, 97,', 'stand_lon': '0.', 'opt_geogrid_tbl_path': "'./'", 'geog_data_path': "'/home/lluis/DATA/WRF/geog'", 'pole_lat': '90.0', 'ref_lat': '35.0,', 'map_proj': "'lat-lon',", 'parent_id': '1, 1,', 'geog_data_res': "'2m','2m',", 'e_we': '32, 112,', 'dx': '0.35,', 'dy': '0.35,', 'parent_grid_ratio': '1, 3,', 'pole_lon': '0.0', 'ref_lon': '20,', 'j_parent_start': '1, 17,', 'i_parent_start': '1, 31,'}
    """

    fname = 'get_namelist_vars'

# List of characters which split pairs name/value
    valuessep = ['=']
# List of characters which indicate a comment
    commentchars = ['#']

    if values == 'h':
        print fname + '_____________________________________________________________'
        print get_namelist_vars.__doc__
        quit()

    expectargs = '[sectionname],[kout]'
    check_arguments(fname,len(expectargs.split(',')),values,',',expectargs)

    secname = values.split(',')[0]
    kout = values.split(',')[1]

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

    ncml = open(namelist, 'r')

    sections = {}
    namelistvals = {}
    namelistsecs = {}
    sectionnames = []
    namessec = []
    allnames = []
    namelistvalssec = {}
    namelistsecssec = {}
    nmlname = ''
    sectionname = ''

    for line in ncml:
        linevals = reduce_spaces(line)
        Nvals = len(linevals)

        if Nvals >= 1 and linevals[0][0:1] == '&':
            if len(sectionnames) > 1:
                sections[sectionname] = namessec

            sectionname = linevals[0][1:len(linevals[0])+1]
#            print '    ' + fname + ": new section '" + sectionname + "' !!!"
            sectionnames.append(sectionname)
            namessec = []
            nmlname = ''
        elif Nvals >= 1 and not searchInlist(commentchars,linevals[0][0:1]):
            if Nvals >= 3 and searchInlist(valuessep,linevals[1]):
                nmlname = linevals[0]
                nmlval = numVector_String(linevals[2:Nvals],' ')
            elif Nvals == 1:
                for valsep in valuessep:
                    if linevals[0].find(valsep) != -1:
                        nmlname = linevals[0].split(valsep)[0]
                        nmlval = linevals[0].split(valsep)[1]
                        break
            elif Nvals == 2:
                for valsep in valuessep:
                    if linevals[0].find(valsep) != -1:
                        nmlname = linevals[0].split(valsep)[0]
                        nmlval = linevals[1]
                        break
                    elif linevals[1].find(valsep) != -1:
                        nmlname = linevals[0]
                        nmlval = linevals[1].split(valsep)[0]
                        break
            else:
                print warnmsg
                print '  ' + fname + ': wrong number of values', Nvals,              \
                  'in the namelist line!'
                print '    line: ',line
                print '    line values:',linevals
#                quit(-1)

            namelistvals[nmlname] = nmlval
            namelistsecs[nmlname] = sectionname

            namessec.append(nmlname)
            allnames.append(nmlname)

    if len(sectionname) > 1:
        sections[sectionname] = namessec

    if secname != 'all':
        if not searchInlist(sections.keys(),secname):
            print errormsg
            print '  ' + fname + ": section '" + values + "' does not exist !!"
            print '    only found:',sectionnames
            quit(-1)

        namestouse = []
        for nml in allnames:
            for nnml in sections[secname]:
                namelistvalssec[nnml] = namelistvals[nnml]
                namelistsecssec[nnml] = secname
                if nml == nnml: namestouse.append(nml)
    else:
        namestouse = allnames
        namelistvalssec = namelistvals
        namelistsecssec = namelistsecs

    if kout == 'tex3':
        ofile='get_namelist_vars_3col.tex'
        fobj = open(ofile, 'w')

        vals = namestouse
        Nvals = len(vals)
        Nvals3 = int(Nvals/3)
        latextab = '\\begin{center}\n\\begin{tabular}{lclclc}\n'
        for il in range(2):
            latextab = latextab + '{\\bfseries{name}} & {\\bfseries{value}} & '
        latextab= latextab+ '{\\bfseries{name}} & {\\bfseries{value}} \\\\ \\hline\n'

        if np.mod(Nvals,3) != 0:
            Nvals0 = Nvals - np.mod(Nvals,3)
        else:
            Nvals0 = Nvals

        for ival in range(0,Nvals0,3):
            line = ''
            print '  ival:',ival
            for il in range(2):
                line = line + '\\verb+' + vals[ival+il] + '+ & ' +                   \
                   namelistvalssec[vals[ival+il]].replace('_','\\_') +' & '
            line = line + '\\verb+' + vals[ival+2] + '+ & ' +                       \
               namelistvalssec[vals[ival+2]].replace('_','\\_') + ' \\\\\n'
            latextab = latextab + line

        latextab = latextab + '%not multiple of three!!!!\n'
        print 'mod:', np.mod(Nvals,3),Nvals0,Nvals
        if np.mod(Nvals,3) != 0:
            ival = Nvals0
            line = ''
            for il in range(np.mod(Nvals,3)):
                print 'ival:',ival + il
                line = line + '\\verb+' + vals[ival+il] + '+ & ' +                   \
                      namelistvalssec[vals[ival+il]].replace('_','\\_') + ' & '
            for il in range(2-np.mod(Nvals,3)):
                line = line + ' & & '
            latextab = latextab + line + ' & \\\\\n'
        latextab = latextab + '\\end{tabular}\n\\end{center}\n'

#        print latextab
        fobj.write(latextab)
 
        fobj.close()
        print fname + ": successful writen '" + ofile + "' LaTeX tale file !!"

        return 
    elif kout == 'column':
        for dictv in namestouse:
            print dictv + ' = ' + namelistvalssec[dictv]
        return
    elif kout == 'dict':
        return namelistvalssec, namelistsecssec
    else:
        print errormsg
        print '  ' + fname + ": kind of output '" + kout + "' not ready!!!"
        quit(-1)

    return

def check_arguments(funcname,Nargs,args,char,expectargs):
    """ Function to check the number of arguments if they are coincident
    check_arguments(funcname,Nargs,args,char)
      funcname= name of the function/program to check
      Nargs= theoretical number of arguments
      args= passed arguments
      char= character used to split the arguments
    """

    fname = 'check_arguments'

    Nvals = len(args.split(char))
    if Nvals != Nargs:
        print errormsg
        print '  ' + fname + ': wrong number of arguments:',Nvals," passed to  '",   \
          funcname, "' which requires:",Nargs,'!!'
        print '    given arguments:',args.split(char)
        print '    expected arguments:',expectargs
        quit(-1)

    return

def check_conversion(value, trykind):
    """ Function to check a variable conversion to a python type
      from: http://stackoverflow.com/questions/4690600/python-exception-message-capturing
      value= value to convert
      trykind: kind to try to convert to
        integer: 'INTEGER', 'integer', 'int'
        float: 'REAL', 'real', 'float'
        intl: 'LOGICAL', 'logical', 'bool'
      >>> check_conversion('1.3','integer')
      ERROR -- error -- ERROR -- error
        check_conversion: Failed to convert '1.3' to 'int': invalid literal for int() with base 10: '1.3'
    """
    import sys, traceback
    fname = 'check_conversion'

    intk = ['INTEGER', 'integer', 'int']
    intr = ['REAL', 'real', 'float']
    intl = ['LOGICAL', 'logical', 'boolean', 'bool']
# Accepted values as boolean values:
    intlv = ['y', 'n', 'yes', 'no', '.TRUE.', '.FALSE.', '.true.', '.false.',        \
      '.T.', '.F.', '.t.', '.f.']

# All kinds
    allkinds = intk + intr + intl
    if not searchInlist(allkinds, trykind) and trykind[0:9] != 'character' and       \
      trykind[0:9] != 'CHARACTER':
        print errormsg
        print '  ' + fname + ": kind '" + trykind + "' not ready !!"
        quit(-1)

    try:
        if searchInlist(intk,trykind):
            conv = int(value)
        elif searchInlist(intr,trykind):
            conv = np.float(value)
        elif searchInlist(intl,trykind):
            if not searchInlist(intlv, value):
                print errormsg
                print '  ' + fname + ": Failed to convert '" + value + "' to '" +    \
                  trykind + "' !!"
                return -1
    except:
        exc_type, exc_value, exc_traceback = sys.exc_info()
        print errormsg
        print '  ' + fname + ": Failed to convert '" + value + "' to '" + trykind +  \
          "': " + str(exc_value)
        return -1

    return 0

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

parser = OptionParser()
parser.add_option("-w", "--WRF_folder", dest="wfold", 
                  help="compiled WRF folder to use", metavar="FOLDER")
parser.add_option("-n", "--WRF_namelist", dest="namelist", 
                  help="'namelist.input' to use", metavar="FILE")

(opts, args) = parser.parse_args()

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

if not os.path.isfile(opts.wfold + '/WRFV3/main/wrf.exe'):
    print errormsg
    print '  ' + main + ": compiled WRF '" + opts.wfold + "/WRFV3/main/wrf.exe'" +   \
      'does not exist !!'
    quit(-1)

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

infold = opts.wfold + '/WRFV3'

# Reading namelist characteristics
##
nmldef = infold + '/inc/namelist_defines.inc'

if not os.path.isfile(nmldef):
    print errormsg
    print '  ' + main + ": included WRF '" + nmldef + "' does not exist !!"
    quit(-1)

onmldef = open(nmldef, 'r')

nmlpar = []
nmlpark = {}
nmlpard = {}

for line in onmldef:
    if line[0:1] != '!':
        values = line.split(' ')
        par = line.split('::')[1].replace('\n','').replace(' ','')
        nmlpar.append(par)
        nmlpark[par] = values[0].replace(' ','')
        if len(line.split(',')) > 1:
            dim = line.split(',')[1].split('::')[0].split('(')[1].replace(')','')
            nmlpard[par] = dim.replace(' ','')

#        if nmlpard.has_key(par):
#            print par,':',nmlpark[par],'=',nmlpard[par]
#        else:
#            print par,':',nmlpark[par]

onmldef.close()

# Reading namelist sections
##
nmlsec = infold + '/inc/namelist_statements.inc'

if not os.path.isfile(nmlsec):
    print errormsg
    print '  ' + main + ": included WRF '" + nmlsec + "' does not exist !!"
    quit(-1)

onmlsec = open(nmlsec, 'r')

nmlpars = {}
nmlsecs = []

for line in onmlsec:
#    print 'line:',line
    if line[0:1] != '!':
        values = line.split('/')
        par = values[2].replace('\n','').replace(' ','')
        sec = values[1].replace(' ','')

        nmlpars[par] = sec
        if not searchInlist(nmlsecs,sec):
            nmlsecs.append(sec)

onmlsec.close()

# Getting values from 'namelist.input
##
readnmlvals, readnmlsecs = get_namelist_vars('all,dict', opts.namelist)

# Checking values
##
readparam = set(readnmlvals.keys())
configparam = set(nmlpark.keys())

coinread = list(readparam & configparam)
Wrongread = list(readparam - configparam)

# Wrong pramater name
## 
Nwrong = len(Wrongread)

if Nwrong > 0:
  print errormsg
  print '  ' + main + ': some readed parameters are wrong!'
  print '    wrong parameters: ', Wrongread

# Wrong section
##
for readpar in readnmlsecs.keys():
    if nmlpars.has_key(readpar) and readnmlsecs[readpar] != nmlpars[readpar]:
        print errormsg
        print '  ' + main + ": readed parameter '" + readpar + "' in section '" +    \
            readnmlsecs[readpar] + "' should be in section '" + nmlpars[readpar] +   \
            "' !!"

# Wrong type of value
## 
for readpar in readnmlvals.keys():
    if nmlpark.has_key(readpar):
        pkind = nmlpark[readpar]
        readvals = readnmlvals[readpar].replace(' ','').split(',')
        Nreadvals = len(readvals)

        if Nreadvals > 1:
            for ival in range(Nreadvals):
                if len(readvals[ival]) > 0:
                    ier = check_conversion(readvals[ival], pkind)
                    if ier != 0:
                        print errormsg
                        print '  ' + main + ": wrong readed '" + readpar +           \
                          "' value: '" + readvals[ival] + "' is not of type '" +     \
                          pkind + "' !!"
                    else:
                        if readpar == 'max_dom':
                            Ndomains = int(readvals[0])

        else:
            ier = check_conversion(readnmlvals[readpar], pkind)
            if ier != 0:
                print errormsg
                print '  ' + main + ": wrong readed '" + readpar + "' value: '" +    \
                  readvals + "' is not of type '", pkind, "' !!"
  
# Checking number of values
##

for readpar in readnmlvals.keys():
    readvals = readnmlvals[readpar].replace(' ','').split(',')
    Nreadvals = len(readvals)
# Removing spurious value due to ',' at the end of the parameter values in namelist
    if len(readvals[Nreadvals - 1]) == 0: Nreadvals = Nreadvals - 1

    if nmlpard.has_key(readpar):
#        print 'readpar:',readpar,'Nvals:',Nreadvals,'WRF:',nmlpard[readpar],'Ndoms:',Ndomains

        if nmlpard[readpar] == 'max_domains':
            if Nreadvals < Ndomains:
                print errormsg
                print '  ' + main + ': wrong number of values: ', Nreadvals, '=',    \
                  readnmlvals[readpar]," for parameter '" + readpar + "' should be:",\
                  Ndomains,'!!'

    else:
        if Nreadvals > 1:
            print errormsg
            print '  ' + main + ': wrong number of values: ', Nreadvals,             \
              " for parameter '" + readpar + "' should be: 1 !!"