## L. Fita, LMD October 2014. 
# Generation of initial conditions for an aqua-planet from the LMDZ model (r1818) 'iniacademic.F90' program
#   Author:    Frederic Hourdin      original: 15/01/93
# The forcing defined here is from Held and Suarez, 1994, Bulletin
# of the American Meteorological Society, 75, 1825.

from optparse import OptionParser
import numpy as np
from netCDF4 import Dataset as NetCDFFile
import os
import re
import nc_var_tools as ncvar
from lmdz_const import *

main = 'iniaqua.py'
errormsg='ERROR -- error -- ERROR -- error'
warnmsg='WARNING -- warning -- WARNING -- warning'

filekinds = ['CF', 'WRF']

## g.e. # iniaqua.py -d 32,32,39 -o WRF -t 19791201000000 -y true

def sig_hybrid(sig,pa,preff):
    """ Function utilisee pour calculer des valeurs de sigma modifie
     pour conserver les coordonnees verticales decrites dans
     esasig.def/z2sig.def lors du passage en coordonnees hybrides
     F. Forget 2002
       sig= sigma level
       pa=
       preff= reference pressure
     Connaissant sig (niveaux "sigma" ou on veut mettre les couches)
     L'objectif est de calculer newsig telle que
       (1 -pa/preff)*exp(1-1./newsig**2)+(pa/preff)*newsig = sig
     Cela ne se resoud pas analytiquement: 
     => on resoud par iterration bourrine 
     ----------------------------------------------
     Information  : where exp(1-1./x**2) become << x
           x      exp(1-1./x**2) /x
           1           1
           0.68       0.5
           0.5        1.E-1
           0.391      1.E-2
           0.333      1.E-3
           0.295      1.E-4
           0.269      1.E-5
           0.248      1.E-6
        => on peut utiliser newsig = sig*preff/pa si sig*preff/pa < 0.25
    """
    fname = 'sig_hybrid'

# maximum number of iterations
    maxiter = 9999

    nsig = sig
    x1=0
    x2=1
    if sig >= 1.:
        nsig = sig
    elif sig*preff/pa >= 0.25:
        for j in range(maxiter):
            F = ((1. -pa/preff)*np.exp(1.-1./nsig**2)+(pa/preff)*nsig)/sig
#            print J,'nsig=', newsig, 'F=', F
            if F > 1:
                X2 = newsig
                nsig = (X1+nsig)*0.5
            else:
                X1 = newsig
                nsig = (X2+nsig)*0.5

#       Test : on arete lorsque on approxime sig a moins de 0.01 m pres 
#         (en pseudo altitude) :
            if np.abs(10.*np.log(F)) < 1.e-5: break
    else:
        nsig= sig*preff/pa

    return nsig

def presnivs_calc(hybrid, dz):
    """ From dyn3d/disvert_noterre.F calculation of vertical pressure levels
      hybrid= whether hydbrid coordinates have to be used
      dz= numbef of vertical layers
    """

    fname = 'presnivs_calc'

#-----------------------------------------------------------------------
#    ....  Calculs  de ap(l) et de bp(l)  ....
#    .........................................
#
#   .....  pa et preff sont lus  sur les fichiers start par dynetat0 .....
#-----------------------------------------------------------------------
#
    llmp1 = dz + 1
    pnivs = np.zeros((dz), dtype=np.float)
    s = np.zeros((dz), dtype=np.float)
    sig = np.zeros((dz+1), dtype=np.float)
    ap = np.zeros((dz+1), dtype=np.float)
    bp = np.zeros((dz+1), dtype=np.float)

# Ouverture possible de fichiers typiquement E.T.

    if os.path.isfile('easig.def'):
#-----------------------------------------------------------------------
#   cas 1 on lit les options dans esasig.def:
#   ----------------------------------------

        ofet = open('esasig.def', 'r')
#        Lecture de esasig.def :
#        Systeme peu souple, mais qui respecte en theorie
#        La conservation de l'energie (conversion Energie potentielle
#        <-> energie cinetique, d'apres la note de Frederic Hourdin...

        print '*****************************'
        print "WARNING reading 'esasig.def'"
        print '*****************************'
        for line in ofet:
            linevalues = ncvar.reduce_spaces(line)
            scaleheight = np.float(linevalues[0])
            dz0 = np.float(linevalues[1])
            dz1 = np.float(linevalues[2])
            nhaut = np.float(linevalues[3])

        ofet.close()
        dz0 = dz0/scaleheight
        dz1 = dz1/scaleheight

        sig1=(1.-dz1)/tanh(.5*(dz-1)/nhaut)

        esig=1.
        for l in range(19):
            esig=-np.log((1./sig1-1.)*np.exp(-dz0)/esig)/(dz-1.)

        csig=(1./sig1-1.)/(np.exp(esig)-1.)

        for l in range(1,dz):
            zz=csig*(np.exp(esig*(l-1.))-1.)
            sig[l] = 1./(1.+zz)*np.tanh(.5*(dz+1-l)/nhaut)

        sig[0] = 1.
        sig[dz] = 0.
        quoi = 1. + 2.* kappa
        s[dz-1]  = 1.
        s[dz-2] = quoi
        if dz > 1:
            for ll in range(1, dz-2):
                l = dz+1 - ll
                quand = sig[l+1]/sig[l]
                s[l-1] = quoi*(1.-quand)*s[l] + quand*s[l+1]
#
        snorm=(1.-.5*sig[1]+kappa*(1.-sig[1]))*s[0]+.5*sig[1]*s[1]
        for l in range(dz):
            s[l] = s[l]/ snorm
    elif os.path.isfile('z2sig.def'):
        fet = open('z2sig.def', 'r')
#-----------------------------------------------------------------------
#   cas 2 on lit les options dans z2sig.def:
#   ----------------------------------------
        print '****************************'
        print 'Reading z2sig.def'
        print '****************************'

        for line in ofet:
            linevalues = ncvar.reduce_spaces(line)
            scaleheight = np.float(linevalues[0])
            for l in range(dz):
                zsig[l] = linevalues[l+1]

        ofet.close()

        sig[0] = 1.
        for l in range(1,dz):
            sig[l] = 0.5 * ( np.exp(-zsig[l]/scaleheight) +                          \
              np.exp(-zsig[l-1]/scaleheight) )

        sig[dz+1] = 0.

#-----------------------------------------------------------------------
    else:
        print errormsg
        print '  ' + fname + ": didn't you forget something ???"
        print "    We need file  'z2sig.def' ! (OR 'esasig.def')"
        quit(-1)
#-----------------------------------------------------------------------

    nivsigs = np.arange(dz)*1.
    nivsig = np.arange(llmp1)*1.

    if hybrid:
# use hybrid coordinates
        print "*********************************"
        print "Using hybrid vertical coordinates"
        print 
#        Coordonnees hybrides avec mod
        for l in range(dz):
            newsig = sig_hybrid(sig[l],pa,preff)
            bp[l] = np.exp(1.-1./(newsig**2))
            ap[l] = pa * (newsig - bp[l] )

        bp[llmp1-1] = 0.
        ap[llmp1-1] = 0.
    else:
# use sigma coordinates
        print "********************************"
        print "Using sigma vertical coordinates"
        print
#       Pour ne pas passer en coordonnees hybrides
        for l in range(dz):
            ap[l] = 0.
            bp[l] = sig[l]

        ap[llmp1-1] = 0.

    bp[llmp1-1] = 0.

    print 'BP________ ', bp
    print 'AP________ ', ap

#   Calcul au milieu des couches  (llm = dz):
#   WARNING : le choix de placer le milieu des couches au niveau de
#   pression intermediaire est arbitraire et pourrait etre modifie.
#   Le calcul du niveau pour la derniere couche 
#   (on met la meme distance (en log pression)  entre P(llm)
#   et P(llm -1) qu'entre P(llm-1) et P(llm-2) ) est
#   Specifique.  Ce choix est specifie ici ET dans exner_milieu.F

    for l in range(dz-1):
        aps[0:dz-1] = 0.5*( ap[0:dz-1]+ap[1:dz]) 
        bps[0:dz-1] = 0.5*( bp[0:dz-1]+bp[1:dz]) 

    if hybrid:
        aps[dz-1] = aps[dz-2]**2 / aps[dz-3] 
        bps[dz-1] = 0.5*(bp[dz-1] + bp[dz])
    else:
        bps[dz-1] = bps[dz-2]**2 / bps[dz-3] 
        aps[dz-1] = 0.

    print 'BPs_______ ', bps
    print 'APs_______ ', aps

    for l in range(dz):
        psnivs[l] = aps[l]+bps[l]*preff
        psalt[l] = -scaleheight*np.log(presnivs[l]/preff)

    return psnivs, psalt

def global_lonlat(dx,dy):
    """ Function to generate 2D matrices with the global longitude, latitudes
      dx, dy: dimension of the desired matrix
    >>> global_lonlat(5,5)
    array([[  36.,  108.,  180.,  252.,  324.],
          [  36.,  108.,  180.,  252.,  324.],
          [  36.,  108.,  180.,  252.,  324.],
          [  36.,  108.,  180.,  252.,  324.],
          [  36.,  108.,  180.,  252.,  324.]]), array([[-72., -72., -72., -72., -72.],
          [-36., -36., -36., -36., -36.],
          [  0.,   0.,   0.,   0.,   0.],
          [ 36.,  36.,  36.,  36.,  36.],
          [ 72.,  72.,  72.,  72.,  72.]]))
    """ 

    fname = 'global_lonlat'

    longitude = np.zeros((dy,dx), dtype=np.float)
    latitude = np.zeros((dy,dx), dtype=np.float)

    for ix in range(dx):
        longitude[:,ix] = 360.*(1./2. + ix )/(dx)

    for iy in range(dy):
        latitude[iy,:] =  180.*(1./2. + iy )/(dy) - 90.

    return longitude, latitude

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

filekindsnames = "'" + ncvar.numVector_String(filekinds, "', '") + "'"

parser = OptionParser()
parser.add_option("-o", "--outputkind", type='choice', dest="okind", 
  choices=filekinds, help="kind of output: " + filekindsnames, metavar="KIND")
parser.add_option("-d", "--dimensions", dest="dims", 
  help="dimensions of the initial conditions: dimx,dimy,dimz", metavar="VALUES")
parser.add_option("-t", "--time", dest="time", 
  help="time of the initial conditions ([YYYY][MM][DD][HH][MI][SS] format)", metavar="VALUE")
parser.add_option("-y", "--hybrid", dest="hybrid",
  help="whether vertical levels have to compute hydbid or not (true/false)", metavar="VAR")

(opts, args) = parser.parse_args()

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

#  dynamic variables
#  vcov, ucov: covariant winds
#  teta: potential temperature
#  q: advecting fields (humidity species)
#  ps: surface pressure
#  masse: air mass
#  phis: surface geopotential

#  Local:
#  p: pressure at the interface between layers (half-sigma)
#  pks: Exner function at the surface
#  pk: Exner functino at the half-sigma layers
#  pkf: Filtred Exner function at the half-sigma layers
#  phi: geopotential height
#  ddsin,zsig,tetapv,w_pv: auxiliar variables
#  tetastrat: potential temporeature in the stratosphere (K)
#  teta0,ttp,delt_y,delt_z,eps: constants for the T profile
#  k_f,k_c_a,k_c_s: calling constants
#  ok_geost: Initialisation geostrohic wind
#  ok_pv: Polar Vortex
#  phi_pv,dphi_pv,gam_pv: polar vortex constants

dimx = int(opts.dims.split(',')[0])
dimy = int(opts.dims.split(',')[1])
dimz = int(opts.dims.split(',')[2])

if opts.hybrid == 'true':
    hybrid_calc = True
else:
    hybrid_calc = False

ofile = 'iniaqua.nc'

print 'dimensions: ',dimx,', ',dimy,', ',dimz

if opts.okind == 'CF':
    varnames = ['zg', 'ta', 'ua', 'va', 'pres', 'r']
# Reference time from 1949-12-01 00:00:00 UTC
    timev = ncvar.datetimeStr_conversion(opts.time,'YmdHMS','cfTime')
    dimnames = ['time', 'z', 'y', 'x']
elif opts.okind == 'WRF':
    varnames = ['Geopot', 'T', 'U', 'V', 'PRES', 'QVAPOR']
    timev = ncvar.datetimeStr_conversion(opts.time,'YmdHMS','WRFdatetime')
    dimnames = ['Time', 'bottom_top', 'south_north', 'west_east']

# Constants
##
llm = dimz

ok_geost = True
# Constants for Newtonian relaxation and friction
k_f = 1.
k_f = 1./(daysec*k_f)
# cooling surface
k_c_s=4.  
k_c_s=1./(daysec*k_c_s)
# cooling free atm
k_c_a=40. 
k_c_a=1./(daysec*k_c_a)
# Constants for Teta equilibrium profile
# mean Teta (S.H. 315K)
teta0=315.
# Tropopause temperature (S.H. 200K)
ttp=200.
# Deviation to N-S symmetry(~0-20K)
eps=0.
# Merid Temp. Gradient (S.H. 60K)
delt_y=60.
# Vertical Gradient (S.H. 10K)
delt_z=10.
# Polar vortex
ok_pv = False
# Latitude of edge of vortex
phi_pv=-50.
phi_pv=phi_pv*pi/180.
# Width of the edge
dphi_pv=5.
dphi_pv=dphi_pv*pi/180.
# -dT/dz vortex (in K/km)
gam_pv=4.

presnivs, pseudoalt = presnivs_calc(hybrid_calc, dimz)
lon, lat = global_lonlat(dx,dy)
lonu, latu = global_lonlat(dx+1,dy)
lonv, latv = global_lonlat(dx,dy+1)

# 2. Initialize fields towards which to relax
##

knewt_t = np.zeros((dimz), dtype=np.float)
kfrict = np.zeros((dimz), dtype=np.float)
clat4 =  np.zeros((dimy+1, dimx), dtype=np.float)

# Friction
knewt_g = k_c_a
for l in range(dimz):
    zsig=presnivs[l]/preff
    knewt_t[l]=(k_c_s-k_c_a)*np.max([0.,(zsig-0.7)/0.3])
    kfrict[l]=k_f*np.max([0.,(zsig-0.7)/0.3])

    for j in 1,range(dimy+1):
        clat4[j,:]=np.cos(rlatu[j])**4

# Vertical profile
tetajl =  np.zeros((dimy+1, dimx), dtype=np.float)
teta = np.zeros((dimy+1, dimx+1), dtype=np.float)
tetarappel = np.zeros((dimy+1, dimx+1), dtype=np.float)

for l in range (dz):
    zsig = presnivs[l]/preff
    tetastrat = ttp*zsig**(-kappa)
    tetapv = tetastrat
    if ok_pv and zsig < 0.1:
        tetapv = tetastrat*(zsig*10.)**(kappa*cpp*gam_pv/1000./g)

    ddsin = np.sin(latu)
    tetajl[l,:,:] = teta0-delt_y*ddsin+eps*ddsin-delt_z*(1.-ddsin*ddsin)*np.log(zsig)
    if planet_type == 'giant':
        tetajl[l,:,:] = teta0+(delt_y*((np.sin(latu*3.14159*eps+0.0001))**2) /       \
          ((rlatu*3.14159*eps+0.0001)**2))-delt_z*np.log(zsig)

# Profile stratospheric isotherm (+vortex)
    w_pv=(1.-np.tanh((rlatu-phi_pv)/dphi_pv))/2.
    tetastrat=tetastrat*(1.-w_pv)+tetapv*w_pv             
    for iy in range(dy):
        for ix in range(dx):
            tetajl[l,iy,ix]=np.max([tetajl(l,ix,iy),tetastrat])

for iz in range(dimz):
    for iy in range(dimy+1):
        tetarappel[iz,iy,0:dimx+1] = tetajl[iz,iy,:]

    tetarappel[iz,iy,0:dimx] = tetajl[iz,iy,dimx-1]


ncf = NetCDFFile(ofile, 'w')

# Dimensions creation
ncf.createDimension(dimnames[0], None)
ncf.createDimension(dimnames[1], dimz)
ncf.createDimension(dimnames[2], dimy)
ncf.createDimension(dimnames[3], dimx)

ncf.sync()
ncf.close()

print main + ": successfull writing of file '" + ofile + "' !!"
