# Python script to get vertical levels from original 17 WRF levels
# From WRFV3.3/WRFV3/dyn_em/module_initialize_real.F#compute_eta
import numpy as np
import os
import re
from optparse import OptionParser

main = 'interpolate_WRF_vertical_lev'
errmsg = 'ERROR -- errror -- ERROR -- error'
warnmsg = 'WARNING -- warning -- WARNING -- warning'

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

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

parser = OptionParser()
parser.add_option("-n", "--Level_Numbers", dest="nlevels",
                  help="number of levels to interpolate", metavar="VALUE")
(opts, args) = parser.parse_args()

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

Nwrflev = 17
znw_prac = np.zeros((Nwrflev), dtype=np.float)
znw_prac[:] = [ 1.000, 0.993, 0.983, 0.970, 0.954, 0.934, 0.909, 0.88, 0.8, 0.7,     \
  0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0.0]

print '#Standard eta WRF values_________'
for iwl in range(Nwrflev):
    print '#',iwl, znw_prac[iwl]

#Desired levels
deslev = int(opts.nlevels)
if deslev < Nwrflev:
    print errmsg
    print '  ' + main + ': less desired levels: ',deslev,                            \
      ' than standards from WRF: ', Nwrflev, '!!!!!'
    quit(-1)

deslevels = np.zeros((deslev), dtype=np.float)
desnivs = np.zeros((deslev), dtype=np.float)

il = 0
deslevels[il] = 1.

# Computing levels
# Number of levels between original WRF levels
Nbetween = int((deslev-1)/(Nwrflev-1))

# Number of levels to fullfill desired number of levels (1 more added to each lowest 
#    WRF until is filled)
fullfilldes = np.mod(deslev-1,Nwrflev-1)
print 'Nbetween: ', Nbetween, 'N extra levels:',fullfilldes

for iwl in range(1,Nwrflev):
    if il < deslev - 1:
        Nbtw = Nbetween
        if iwl <= fullfilldes: 
            print '# ',iwl, 'level with an extra value!'
            Nbtw = Nbetween + 1
        else:
            Nbtw = Nbetween
        for ibetl in range(Nbtw):
            il = il + 1
#            print il, iwl, ibetl, znw_prac[iwl-1], (iwl-1.)*1. + (ibetl+1.)/(Nbtw), znw_prac[iwl]
            if il == deslev - 1 or il == deslev:
                break
            else:
                deslevels[il] = znw_prac[iwl-1] + (ibetl+1)*(znw_prac[iwl] -         \
                  znw_prac[iwl-1])/(Nbtw)
                desnivs[il] = (iwl-1.)*1. + (ibetl+1.)/(Nbtw)
#    if il == deslev - 2:
#        il = il + 1
#        deslevels[il] = znw_prac[iwl]
#        desnivs[il] = iwl*1.

#quit()

deslevels[deslev-1] = 0.
desnivs[deslev-1] = (Nwrflev-1)*1.

print '#il equiv_WRFlev new_levels diff_levels________________'
for il in range(0,deslev):
    if il > 0:
        print il, desnivs[il], deslevels[il], deslevels[il-1] - deslevels[il]
    else:
        print il, desnivs[il], deslevels[il], 0.

print 'eta_levels                          = ',numVector_String(deslevels, ', ')