import RT
#
#-- This is a custom module! ask RT.py to fpaletou@irap.omp.eu in case
#
#  made python3 compliant Apr 2020, frederic.paletou@univ-tlse3.fr

import numpy as np
import matplotlib.pyplot as plt

#--- Generates semi-infite slab grid
dtau1=1.e-2
taumax=1.e4
npdec=5
#--- Collisional destruction probability, or (1-albedo)
eps=1.e-4
#--- Number of iterations for the iterative process
niter=150
#--- the above will be an INPUT

#--- The logarithmic opacity/spatial grid ---
fac = 10.**(1./np.float(npdec))
dtau=dtau1
t=0.
k=1
while (t < taumax):
    k=k+1
    t=t+dtau
    dtau=fac*dtau
npts=k
z=np.zeros((npts))
z[0]=0.
z[1]=dtau1
dtau=dtau1*fac
for k in np.arange(2,npts):
    z[k]=z[k-1]+dtau
    dtau=fac*dtau
print(z)

#--- Initial source function (Planck) ---
s=np.ones((npts))
b=np.ones((npts))

#--- Absorption coefficient set to unity throughout the slab ---
chi=np.ones((npts))

#--- Upper boundary condition (in general MU and FREQ dependent)
bcu=0.

#--- Lower boundary condition (in general MU and FREQ dependent)
bcd=1.

#--- Zero boundary conditions (for Lstar computation)
zerobc=0.

#--- Two-stream approximation/Eddington's approximation
mu=1./np.sqrt(3.)
wtdir=0.5
#--- In general, use your preferred quadrature (e.g., Gauss-Legendre)

#--- Analytical reference solution of Eddington
#sedd=1. - (1.-eps)*( exp(-z*sqrt(3.*eps)) )/(1. + sqrt(eps) )
sedd=1. - (1.-np.sqrt(eps))*( np.exp(-np.sqrt(3.*eps)*z) )

#--------------------------------
#--- Compute diagonal operator  |
#--------------------------------
lstar=RT.formal(z,npts,chi,zerobc,zerobc,s,mu,wtdir,True)

shist=np.zeros((npts,niter))
relerr=np.zeros((niter))
dsk=np.zeros((npts))
te=np.zeros((niter))

#---------------------------
#--- ALI-Jacobi iterations |
#---------------------------
for i in np.arange(niter):
    jbar=RT.formal(z,npts,chi,bcu,bcd,s,mu,wtdir,False)
    dsk=((1.-eps)*jbar+eps*b-s)/(1.-(1.-eps)*lstar)
    relerr[i]=max(abs(dsk)/s)
    te[i]=max(abs(s-sedd)/sedd)
    s=s + dsk
    shist[:,i]=s[:]
    print('Iter: ',i, 'Max.rel.err.S: ',relerr[i], 'True.err: ', te[i])

plt.figure(1)
plt.plot(np.log10(z[1:npts]),np.log10(shist[1:npts,:]))
plt.plot(np.log10(z[1:npts]),np.log10(sedd[1:npts]),'o')
plt.xlabel(r"log($\tau$)",fontsize=18)
plt.ylabel('log(S/B)',fontsize=18)

plt.figure(2)
plt.semilogy(relerr, '--k', linewidth=2)
plt.semilogy(te,'k',linewidth=2)
plt.xlabel('iteration number',fontsize=18)
plt.ylabel('Max. relative correction and true error',fontsize=18)

plt.show()