MODULE disvert_strato_mod
  USE icosa
  USE mpipara
  USE getin_mod
  USE earth_const
  IMPLICIT NONE
  PRIVATE

  REAL(rstd), SAVE, ALLOCATABLE,TARGET :: ap(:)
!$OMP THREADPRIVATE(ap)
  REAL(rstd), SAVE, ALLOCATABLE,TARGET :: bp(:)
!$OMP THREADPRIVATE(bp)
  REAL(rstd), SAVE, ALLOCATABLE,TARGET :: presnivs(:)
!$OMP THREADPRIVATE(presnivs)
  REAL(rstd), SAVE, ALLOCATABLE,TARGET :: presinter(:)
  !$OMP THREADPRIVATE(presinter)

  PUBLIC :: ap,bp,presnivs,presinter, init_disvert, init_disvert_strato_custom

CONTAINS
  
  SUBROUTINE init_disvert
    REAL(rstd) :: dsigmin, snorm, x, dsig(llm),  sig(llm+1)
    INTEGER :: l

    ALLOCATE(ap(llm+1))
    ALLOCATE(bp(llm+1))
    ALLOCATE(presnivs(llm))
    ALLOCATE(presinter(llm+1))
      
    dsigmin=1.0  ! Should be 0.3 for CMIP5
    CALL getin('disvert_dsigmin', dsigmin) 
    
    snorm  = 0.
    DO l = 1, llm
      x = Pi*(l-0.5)/(llm+1)
      dsig(l) = (dsigmin + 7.0 * SIN(x)**2) * (0.5* (1.-tanh( (x-Pi/2)/(Pi/2) ) ))**2
      snorm = snorm + dsig(l)
    ENDDO    
    
    DO l = 1, llm
      dsig(l) = dsig(l)/snorm
    ENDDO

    sig(llm+1) = 0.
    DO l = llm, 1, -1
      sig(l) = sig(l+1) + dsig(l)
    ENDDO

    bp(llm+1) =   0.
    DO l = 1, llm
      bp(l) = EXP( 1. -1./( sig(l)*sig(l)) )
      ap(l) = pa * ( sig(l) - bp(l) )
    ENDDO
    bp(1)=1.
    ap(1)=0.
    ap(llm+1) = pa * ( sig(llm+1) - bp(llm+1) )
    DO l = 1, llm
      presnivs(l) = 0.5 *( ap(l)+bp(l)*preff + ap(l+1)+bp(l+1)*preff )
    ENDDO
    DO l=1, llm+1
      presinter(l)= ap(l)+bp(l)*preff
    END DO
    
    ! tell the world about it
    IF (is_mpi_root) THEN
!$OMP MASTER
      WRITE(*,*) "ap()=",ap
      WRITE(*,*) "bp()=",bp
      WRITE(*,*) "Approximative mid-layer pressure, assuming a surface pressure preff=",preff," Pa"
      WRITE(*,*) "and approximative mid-layer height, assuming an atmospheric scale height of ",scale_height/1000," (km)"
      DO l=1,llm
        WRITE(*,*) 'PRESNIVS(',l,')=',presnivs(l),'  Z ~ ',log(preff/presnivs(l))*scale_height/1000,       &
                   ' DZ ~ ',scale_height/1000*log((ap(l)+bp(l)*preff)/ max(ap(l+1)+bp(l+1)*preff,1.e-10))
      ENDDO
!$OMP END MASTER
    ENDIF
  
  END SUBROUTINE init_disvert
  
  SUBROUTINE init_disvert_strato_custom
    REAL(rstd) :: vert_scale,vert_dzmin,vert_dzlow,vert_z0low,vert_dzmid,vert_z0mid,vert_h_mid,vert_dzhig,vert_z0hig,vert_h_hig
    REAL(rstd) :: z, sig(llm+1), sig0(llm+1), zz(llm+1)
    INTEGER :: l
    ALLOCATE(ap(llm+1))
    ALLOCATE(bp(llm+1))
    ALLOCATE(presnivs(llm))
    ALLOCATE(presinter(llm+1))
!===================================================================
! David Cugnet, François Lott, Lionel Guez, Ehouoarn Millour, Fredho
! 2014/05
! custumize strato distribution
! Al the parameter are given in km assuming a given scalehigh
    vert_scale=7.     ! scale hight
    vert_dzmin=0.02   ! width of first layer
    vert_dzlow=1.     ! dz in the low atmosphere
    vert_z0low=8.     ! height at which resolution recches dzlow
    vert_dzmid=3.     ! dz in the mid atmsophere
    vert_z0mid=70.    ! height at which resolution recches dzmid
    vert_h_mid=20.    ! width of the transition
    vert_dzhig=11.    ! dz in the high atmsophere
    vert_z0hig=80.    ! height at which resolution recches dz
    vert_h_hig=20.    ! width of the transition
!===================================================================

    call getin('vert_scale',vert_scale)
    call getin('vert_dzmin',vert_dzmin)
    call getin('vert_dzlow',vert_dzlow)
    call getin('vert_z0low',vert_z0low)
    CALL getin('vert_dzmid',vert_dzmid)
    CALL getin('vert_z0mid',vert_z0mid)
    call getin('vert_h_mid',vert_h_mid)
    call getin('vert_dzhig',vert_dzhig)
    call getin('vert_z0hig',vert_z0hig)
    call getin('vert_h_hig',vert_h_hig)
  
!!!!!    scaleheight=vert_scale ! for consistency with further computations
    ! Build geometrical distribution
    zz(1)=0.
    DO l=1,llm
       z=zz(l)
       zz(l+1)=zz(l)+vert_dzmin+vert_dzlow*TANH(z/vert_z0low)+                &
                  (vert_dzmid-vert_dzlow)* &
                       (TANH((z-vert_z0mid)/vert_h_mid)-TANH((-vert_z0mid)/vert_h_mid)) &
                 +(vert_dzhig-vert_dzmid-vert_dzlow)*                                  &
                       (TANH((z-vert_z0hig)/vert_h_hig)-TANH((-vert_z0hig)/vert_h_hig))
    ENDDO


!===================================================================
! Comment added Fredho 2014/05/18, Saint-Louis, Senegal
! From approximate z to ap, bp, so that p=ap+bp*p0 and p/p0=exp(-z/H)
! sig0 is p/p0
! sig is an intermediate distribution introduce to estimate bp
! 1.   sig0=exp(-z/H)
! 2.   inversion of sig0=(1-pa/p0)*sig+(1-pa/p0)*exp(1-1/sig**2)
! 3.   bp=exp(1-1/sig**2)
! 4.   ap deduced from  the combination of 2 and 3 so that sig0=ap/p0+bp
!===================================================================

    sig0(llm+1)=0.
    DO l = 1, llm
       sig0(l)=EXP(-zz(l)/vert_scale)
!!!     sig(l)=racinesig(sig0(l))
    ENDDO
    sig=racinesig(sig0)
    
    bp(llm+1) =   0.
    DO l = 2, llm
       bp(l) = EXP( 1. -1./( sig(l)*sig(l)) )
       ap(l) = pa * ( sig(l) - bp(l) )
    ENDDO
    bp(1)=1.
    !!    ap=pa*(sig-bp)
    ap(1)=0.
    ap(llm+1) = pa * ( sig(llm+1) - bp(llm+1) )
    DO l = 1, llm
       presnivs(l) = 0.5 *( ap(l)+bp(l)*preff + ap(l+1)+bp(l+1)*preff )
    ENDDO
    DO l=1, llm+1
      presinter(l)= ap(l)+bp(l)*preff
    END DO
      
    ! tell the world about it
    IF (is_mpi_root) THEN
       !$OMP MASTER
       WRITE(*,*) "ap()=",ap
       WRITE(*,*) "bp()=",bp
       WRITE(*,*) "Approximative mid-layer pressure, assuming a surface pressure preff=",preff," Pa"
       WRITE(*,*) "and approximative mid-layer height, assuming an atmospheric scale height of ",scale_height/1000," (km)"
       DO l=1,llm
          WRITE(*,*) 'PRESNIVS(',l,')=',presnivs(l),'  Z ~ ',log(preff/presnivs(l))*scale_height/1000,       &
               ' DZ ~ ',scale_height/1000*log((ap(l)+bp(l)*preff)/ max(ap(l+1)+bp(l+1)*preff,1.e-10))
       ENDDO
       !$OMP END MASTER
    ENDIF
    
  CONTAINS
    
    FUNCTION racinesig(sig) RESULT(sg)
      !
      !-------------------------------------------------------------------------------
      ! Fredho 2014/05/18
      ! Purpose: Search for s solving (Pa/Preff)*sg+(1-Pa/Preff)*EXP(1-1./sg**2)=s
      ! Notes:   Uses Newton Raphson search
      !-------------------------------------------------------------------------------
      ! Arguments:
      REAL, INTENT(IN)  :: sig(:)
      REAL(rstd)        :: sg(SIZE(sig))
      !-------------------------------------------------------------------------------
      ! Local variables:
      INTEGER      :: it, ns, maxit
      REAL(rstd)   :: c1, c2, f1
      !-------------------------------------------------------------------------------
      ns=SIZE(sig); maxit=100
      c1=Pa/Preff; c2=1.-c1
      DO l=2,ns-1
         sg(l)=sig(l)
         DO it=1,maxit
            f1=exp(1-1./sg(l)**2)*(1.-c1)
            sg(l)=sg(l)-(c1*sg(l)+f1-sig(l))/(c1+2*f1*sg(l)**(-3))
         ENDDO
         !   print*,'SSSSIG ',sig(l),sg(l),c1*sg(l)+exp(1-1./sg(l)**2)*(1.-c1)
      ENDDO
      sg(1)=1.; sg(ns)=0.
      
    END FUNCTION racinesig
    
  END SUBROUTINE init_disvert_strato_custom
  
END MODULE disvert_strato_mod
