#!/usr/bin/python3
from	numpy		      import	*
import  numpy                 as        np
import  matplotlib.pyplot     as        mpl
from matplotlib.cm import get_cmap
import pylab
from matplotlib import ticker
import matplotlib.colors as colors
import datetime
from matplotlib import pyplot


######################################################################
# Input Parameters
######################################################################
input_file = 'z2sig.def_47'
ps=1.1 # in Pa
pa=0.5
preff=2

# first and last level difference in m
l0=0.0025
ln=4.6

# New levels
mynbl=201
n=float(mynbl)

plot = True
# plot = False


######################################################################
# Fonction  read  file
######################################################################
def readfile(name):
# test reading lire texte txt
    myfile = open(name, 'r')
    mylines=myfile.readlines()
    nbline=size(mylines)
    nbcolumn=len(mylines[0].split())
    print('nbline, nbcol:',nbline, nbcolumn)
    data=np.zeros(nbline-1,dtype='f')
    i=0
    nbl = 0
    for line in mylines:
        s=str.split(line)
        if i==0:
          hh=float(s[0])
        else:
          data[i-1]=float(s[0])
          nbl = nbl+1
        i=i+1
    return data,hh,nbl

data,hh,nbl = readfile(input_file)

#####################
# Get sigma levels
#####################
sig=np.ones(nbl,dtype='f')
for i in range(nbl-1):
 sig[i+1]=0.5*(exp(-data[i+1]/hh)+exp(-data[i]/hh))

#####################
# Get non-hybrid pressure levels
#####################
ap=0
bp=sig
pp=ap+bp*1.1

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

  newsig = sig
  x1=0
  x2=1
  if (sig>=1):
     newsig = sig
  elif (sig*preff/pa>=0.25):
     for j in range(9999):  # nombre d'iteration max
       F=((1-pa/preff)*exp(1-1./newsig**2)+(pa/preff)*newsig)/sig
       if (F>1):
          x2 = newsig
          newsig=(x1+newsig)*0.5
       else:
          x1=newsig
          newsig=(x2+newsig)*0.5
       #Test : on arete lorsque on approxime sig a moins de 0.01 m pres
       #(en pseudo altiude) :
       #if (abs(10.*log(F))<1.e-5):
       # break

  else:  # if (sig*preff/pa.le.0.25)
      newsig= sig*preff/pa

  return newsig
#----------------------------------

# hybrid pressure:
nsig=np.zeros(nbl,dtype='f')
for i in range(nbl):
  nsig[i]=sig_hybrid(sig[i],pa,preff)

bps=exp(1-1/(nsig)**2)
aps=pa*(nsig-bps)

pps=aps+bps*ps

#####################
# Get Pseudo altitude
#####################
# non hybrid
ph=-hh*log(pp/ps)
dzph=-(ph[0:-1]-ph[1:])
# hybrid
phs=-hh*log(pps/ps)
dzphs=-(phs[0:-1]-phs[1:])

#####################
# New choice delta-zh
#####################
myclev=np.linspace(1,mynbl,mynbl)
xx=myclev[0:mynbl-1] # xaxis for layer thickness
xx1=myclev[0:mynbl]
# Coefficient for the exp function
c=0.11
a=1/n*(np.log(ln/c)-np.log(l0/c))
b=-1/a*np.log(l0/c)
# Dz difference altitude between 2 layers :
mydz=exp(a*(xx-b))*c

# New pseudo altitude myph
myph=np.zeros(mynbl,dtype='f')
myph[0]=mydz[0]
print(myph[0])
for i in range(mynbl-1):
 myph[i+1]=mydz[i]+myph[i]
 print(myph[i+1])

to_file = np.concatenate([[hh], myph])
output = f"z2sig.def_{mynbl}"
np.savetxt(output,to_file, fmt="%g")
print(f"Saved to {output}")

# New Ps levels
mypp=ps*exp(-myph/hh)


if plot:
    clev=np.linspace(1,nbl,nbl)
    ### Plot pressure levels
    mpl.figure(figsize=(50/2.54, 30/2.54),facecolor='w')
    ax=mpl.subplot(131)
    mpl.plot(clev,pp,'k-',marker='*',label='P ini')
    mpl.plot(clev,pps,'r-',marker='*',label='P ini hybrid')
    mpl.plot(myclev,mypp,'b-',marker='*',label='P new')
    mpl.gca().invert_yaxis()
    pyplot.yscale('log')
    mpl.grid()

    ### Plot difference altitude between 2 layers
    ax=mpl.subplot(132)
    mpl.plot(clev[0:nbl-1],dzph,'k-',marker='*',label='dz ini')
    mpl.plot(clev[0:nbl-1],dzphs,'r-',marker='*',label='dz ini hybrid')
    mpl.plot(xx,mydz,'b-',marker='*',label='dz new')
    pyplot.yscale('log')
    mpl.grid()

    ### pseudo altitude
    ax=mpl.subplot(133)
    mpl.plot(clev[0:nbl],ph,'k-',marker='*',label='sig ini')
    mpl.plot(clev[0:nbl],phs,'r-',marker='*',label='sig ini hybrid')
    mpl.plot(xx1,myph,'b-',marker='*',label='sig new')
    #pyplot.yscale('log')
    mpl.grid()

    mpl.show()