MODULE tropopause_m

  IMPLICIT NONE

  PRIVATE

  PUBLIC :: dyn_tropopause

  REAL,    PARAMETER :: DynPTrMin = 8.E+3                  !--- Dyn tropopause pressures < DynPTrMin are set to DynPTrMin  (Pa)
  REAL,    PARAMETER :: DynPTrMax = 4.E+4                  !--- Dyn tropopause pressures > DynPTrMax are set to DynPTrMax  (Pa)
  REAL,    PARAMETER :: theta0 = 380.                      !--- Default threshold for theta-defined tropopause              (K)
  REAL,    PARAMETER :: pVort0 = 2.0                       !--- Default threshold for PV-defined tropopause               (PVU)
  REAL,    PARAMETER :: sg0  = 0.75                        !--- Bottom->top PV=pv0e search loop starts at sigma=sg0 level
  INTEGER, PARAMETER :: nadj = 3                           !--- Threshold must be exceeded on nadj adjacent levels
  INTEGER, PARAMETER :: ns   = 2                           !--- Number of neighbours used each side for vertical smoothing

CONTAINS

!===============================================================================================================================
FUNCTION dyn_tropopause(t, ts, paprs, pplay, rot, itrop, thet0, potV0) RESULT(pTrop)
  USE assert_m,     ONLY: assert
  USE assert_eq_m,  ONLY: assert_eq
  USE dimphy,       ONLY: klon, klev
  USE geometry_mod, ONLY: latitude
  USE strings_mod,  ONLY: maxlen
  USE yomcst_mod_h, ONLY: ROMEGA, RKAPPA, RG
  USE vertical_layers_mod,    ONLY: aps, bps, preff
  USE lmdz_reprobus_wrappers, ONLY: itroprep
  USE lmdz_cppkeys_wrapper,   ONLY: CPPKEY_REPROBUS
  USE mod_phys_lmdz_para,     ONLY: is_master
  USE mod_phys_lmdz_transfert_para, ONLY : bcast
  IMPLICIT NONE
  REAL                           :: pTrop(klon)            !--- Pressure at dynamical tropopause                           (Pa)
  REAL,              INTENT(IN)  ::      t(:,:)            !--- Temperature at layers centers                               (K)
  REAL,              INTENT(IN)  ::     ts(:)              !--- Temperature on surface layer interface                      (K)
  REAL,              INTENT(IN)  ::  paprs(:,:)            !--- Pressure at layers interfaces                              (Pa)
  REAL,              INTENT(IN)  ::  pplay(:,:)            !--- Pressure at layers centers                                 (Pa)
  REAL,              INTENT(IN)  ::    rot(:,:)            !--- Relative vorticity at layers centers                      (s-1)
  INTEGER, OPTIONAL, INTENT(OUT) :: itrop(klon)            !--- Last tropospheric layer idx
  REAL,    OPTIONAL, INTENT(IN)  :: thet0                  !--- Potential temperature at the tropopause (tropical region)   (K)
  REAL,    OPTIONAL, INTENT(IN)  :: potV0                  !--- Potential vorticity   at the tropopause (rest of globe)   (PVU)
!------------------------------------------------------------------------------------------------------------------------------
  CHARACTER(LEN=maxlen) :: modname                         !--- Current routine name
  REAL                  ::    Temp_edg(klon,klev)          !--- Regular   temperature at layers interfaces (except last one)(K)
  REAL                  :: potTemp_edg(klon,klev)          !--- Potential temperature at layers interfaces (except last one)(K)
  REAL                  :: potTemp_cen(klon,klev)          !--- Potential temperature at layers centers                     (K)
  REAL                  :: potVort_cen(klon,klev)          !--- Potential vorticity   at layers centers                     (K)
  REAL                  :: p_th0(klon)                     !--- Pressures at theta=380K                                    (Pa)
  REAL                  :: p_pv0(klon)                     !--- Pressures at PV=2PVU                                       (Pa)
  REAL                  :: al, th0, pv0                    !--- Interpolation coefficient + potential temp. and PV thresholds
  INTEGER               :: i, k, kb, kt, kp, ib, ie, nw, n
  INTEGER               :: ith(klon)                       !--- Indices of first TH=380K layers (top -> bottom search)
  INTEGER               :: ipv(klon)                       !--- Indices of first PV=2PVU layers (top -> bottom search)
  INTEGER               :: ipv0(klon)                      !--- Indices of first PV=2PVU layers (bottom -> top search)
  INTEGER               :: ncons(klon)                     !--- Number of consecutive matching values found in vertical loops
  INTEGER               :: itr(klon)                       !--- Index of last layer with a center pressure lower than pTrop
  INTEGER               :: co(2*ns+1)                      !--- Binomial coefficients used compute smoothing weights "w(:,:)"
  INTEGER,           SAVE :: k0                            !--- Start index (sigma=sg0) for 2PVU bottom->top search loop
  REAL, ALLOCATABLE, SAVE :: fac(:)                        !--- Coriolis parameter: 2*ROMEGA*SIN(cells centers latitudes) (s-1)
  REAL, ALLOCATABLE, SAVE :: w(:,:)                        !--- Coefficients for vertical smoothing froutine "smooth"
  LOGICAL,           SAVE :: lFirst = .TRUE.
!$OMP THREADPRIVATE(k0, fac, w, lFirst)
!------------------------------------------------------------------------------------------------------------------------------
  modname = 'dyn_tropopause'
  CALL assert(SIZE(t,  DIM=1) == klon,        TRIM(modname)//" t klon")
  CALL assert(SIZE(t,  DIM=2) == klev,        TRIM(modname)//" t klev")
  CALL assert(SIZE(ts, DIM=1) == klon,        TRIM(modname)//" ts klon")
  CALL assert(SHAPE(paprs) == [klon, klev+1], TRIM(modname)//" paprs shape")
  CALL assert(SHAPE(pplay) == [klon, klev  ], TRIM(modname)//" pplay shape")
  CALL assert(SHAPE(rot)   == [klon, klev  ], TRIM(modname)//" rot shape")

  !--- MODIFY THE THRESHOLDS FOR THE DYNAMICAL TROPOPAUSE DEFINITION IN CASE THE CORRESPONDING OPTIONAL ARGUMENTS ARE USED
  th0 = theta0; IF(PRESENT(thet0)) th0 = thet0            !--- Potential temperature at the tropopause (tropical region)   (K)
  pv0 = pVort0; IF(PRESENT(potV0)) pv0 = potV0            !--- Potential vorticity   at the tropopause (rest of globe)   (PVU)

  IF(lFirst) THEN
     ALLOCATE(fac(klon), w(ns+1, ns+1))

     !--- COMPUTE THE CORIOLIS PARAMETER FOR PV ALCULATION ROUTINE "potentialVorticity"
     DO i = 1, klon
        fac(i) = 2. * ROMEGA * SIN(latitude(i))
     END DO
!$OMP BARRIER

     IF(is_master) THEN

       !--- GET THE INDEX "k0" OF THE FIRST LOWER INTERFACE LAYER WITH SIGMA COORDINATE LOWER THAN "sg0"
       !--- NOTE: "k0" DEPENDS ON VERTICAL DISCRETIZATION ONLY (VIA HYBRID COEFFS aps, bps) AND IS NOT SIMULATION-DEPENDENT
        DO k0 = 1, klev; IF( aps(k0) / preff + bps(k0) < sg0 ) EXIT; END DO     !--- START INDEX FOR BOTTOM->TOP PV SEARCH LOOP

        !--- COMPUTE THE WEIGHTS FOR THE VERTICAL SMOOTHING ROUTINE "smooth"
        co(:) = 0;  w(:, :) = 0.
        co(1) = 1;  w(1, 1) = 1.
        DO i = 1, ns
           co(2:2*ns+1) = co(2:2*ns+1) + co(1:2*ns)        !--- C(n+1,p+1) = C(n,p+1) + C(n,p)
           co(2:2*ns+1) = co(2:2*ns+1) + co(1:2*ns)        !--- C(n+1,p+1) = C(n,p+1) + C(n,p) AGAIN
           w(i+1, 1:i+1) = REAL(co(i+1:2*i+1))/REAL(SUM(co(i+1:2*i+1)))
        END DO

        lFirst=.FALSE.
     END IF
     CALL bcast(k0)
     CALL bcast(w)
     CALL bcast(lFirst)
  END IF

  !=== DETERMINE THE PRESSURE AT WHICH THETA = th0 ============================================================================
  CALL potentialTemperature(pplay, t, potTemp_cen)                             !--- POTENTIAL TEMPERATURE @ LAYERS CENTERS

  !--- INDEX OF FIRST LAYERS WITH THETA<380K @ CENTER ON "nadj" CONSECUTIVE LAYERS
  CALL getLayerIdx(potTemp_cen, th0, -1, nadj, ith)                            !--- FROM TOP TO BOTTOM

  CALL getPressure(potTemp_cen, th0, ith, pplay, paprs, p_th0)                 !--- PRESSURE @ THETA = th0 SURFACE

  !=== DETERMINE THE PRESSURE AT WHICH PV = pv0 ===============================================================================
  CALL cen2edg(t, ts, pplay, paprs(:,1:klev), temp_edg)                        !--- TEMP @ LAYERS INTERFACES (EXCEPT LAST ONE)

  CALL potentialTemperature (paprs(:,1:klev), temp_edg, potTemp_edg)           !--- TPOT @ LAYERS INTERFACES (EXCEPT LAST ONE)

  CALL potentialVorticity(rot, potTemp_edg, paprs(:,1:klev), potVort_cen)      !--- ERTEL POTENTIAL VORTICITY @ LAYERS CENTERS

  !--- INDEX OF FIRST LAYERS WITH PV<=2PVU @ CENTER ON "nadj" CONSECUTIVE LAYERS
  CALL getLayerIdx(potVort_cen, pv0, -1, nadj, ipv)                            !--- FROM TOP TO BOTTOM
  CALL getLayerIdx(potVort_cen, pv0, k0, nadj, ipv0)                           !--- FROM LAYER @ sig=sig0 TO TOP
  DO i = 1, klon; n = 0                                                        !--- CHOOSE BETWEEN BOTTOM AND TOP INDEX
     IF(ipv0(i) == k0-1 .OR. ipv0(i) > ipv(i)) CYCLE                           !--- ipv0 CAN'T BE USED
     DO k = ipv0(i), ipv(i); IF(potVort_cen(i, k) > pv0) n = n+1; END DO       !--- NUMBER OF POINTS WITH PV>2PVU
     IF(2 * n >= ipv(i)-ipv0(i)+1) ipv(i) = ipv0(i)                            !--- MORE THAN 50% > pv0 => LOWER POINT KEPT
  END DO

  CALL getPressure(potVort_cen, pv0, ipv, pplay, paprs, p_pv0)                  !--- PRESSURE @ PV = pv0 SURFACE

  !=== DETERMINE THE UNFILTERED DYNAMICAL TROPOPAUSE PRESSURE FIELD (LOWER POINT BETWEEN THETA=380K AND PV=2PVU) ==============
  DO i = 1, klon
     pTrop(i) = MAX(p_th0(i), p_pv0(i))
  END DO

  !=== FILTER THE PRESSURE FIELD: TOO HIGH AND TOO LOW VALUES ARE CLIPPED =====================================================
  DO i = 1, klon
     IF(pTrop(i) < DynPTrMin) pTrop(i) = DynPTrMin
     IF(pTrop(i) > DynPTrMax) pTrop(i) = DynPTrMax
  END DO

  !=== LAST VERTICAL INDEX WITH A PRESSURE HIGHER THAN TROPOPAUSE PRESSURE ====================================================
  IF(.NOT.(PRESENT(itrop) .OR. CPPKEY_REPROBUS)) RETURN
  DO i = 1, klon
     DO k = 1, klev
        IF(pplay(i,k+1) <= pTrop(i)) EXIT
     END DO
     IF(PRESENT(itrop )) itrop(i)    = k
     IF(CPPKEY_REPROBUS) itroprep(i) = k
  END DO

CONTAINS

!===============================================================================================================================
SUBROUTINE cen2edg(v_cen, v0_edg, p_cen, p_edg, v_edg)
  IMPLICIT NONE
  REAL, DIMENSION(klon, klev), INTENT(IN)  :: v_cen, p_cen, p_edg
  REAL, DIMENSION(klon),       INTENT(IN)  :: v0_edg
  REAL, DIMENSION(klon, klev), INTENT(OUT) :: v_edg
  INTEGER :: i, k
  DO i = 1, klon
     v_edg(i, 1) = v0_edg(i)
  END DO
  DO k = 2, klev
     DO i = 1, klon
        al = LOG(p_edg(i, k-1)/p_cen(i, k)) / LOG(p_cen(i, k-1)/p_cen(i, k))   !--- CENTER -> INTERFACE INTERPOLATION COEFF
        v_edg(i, k) = v_cen(i, k-1) + al * (v_cen(i, k) - v_cen(i, k-1))       !--- FIELD AT LAYER INTERFACE
     END DO
  END DO
END SUBROUTINE cen2edg
!===============================================================================================================================
SUBROUTINE getPressure(v_cen, v0, ix, p_cen, p_int, pre_v0)
  IMPLICIT NONE
  REAL,    INTENT(IN)  :: v_cen(klon, klev), v0
  INTEGER, INTENT(IN)  ::    ix(klon)
  REAL,    INTENT(IN)  :: p_cen(klon, klev), p_int(klon, klev+1)
  REAL,    INTENT(OUT) :: pre_v0(klon)
  REAL    :: al
  INTEGER :: i, k
  DO i = 1, klon; k = ix(i)
     IF(k == 0) THEN
        pre_v0(i) = p_int(i,1)
     ELSE IF(k == klev) THEN
        pre_v0(i) = p_int(i,klev+1)
     ELSE
        al =  (v0 - v_cen(i, k+1)) / (v_cen(i, k) - v_cen(i, k+1))
        pre_v0(i) = p_cen(i, k+1)  * (p_cen(i, k) / p_cen(i, k+1))**al
     END IF
  END DO
END SUBROUTINE getPressure
!===============================================================================================================================
SUBROUTINE getLayerIdx(v, v0, k0, nadj, ix)
! Purpose: Search for the index of the last layer ix(i) with a value v(i,k) lower than or equal to v0.
!          At least nadj adjacent layers must satisfy the criterium (less - as much as possible - near top or bottom).
!          The search is done from:    * top to bottom if k0 < 0 (from k=klev to k=|k0|)
!                                      * bottom to top if k0 > 0 (from k=k0   to k=klev)
!          - nominal case: k0 <= ix(i) < klev
!          - special case: ix(i) == klev:   ALL(v(i,k0:klev) <= v0)
!          - special case: ix(i) == |k0|-1: ALL(v(i,k0:klev) >  v0)
  IMPLICIT NONE
  REAL,    INTENT(IN)  ::  v(klon, klev), v0
  INTEGER, INTENT(IN)  :: k0, nadj
  INTEGER, INTENT(OUT) :: ix(klon)
  INTEGER :: i, k, nc(klon)
  nc(:) = 0
  ix(:) = 0
  IF(k0 < 0) THEN
     !=== SEARCH FROM TOP TO BOTTOM: klev -> -k0
     !--- ix(i) depends on nc(i), the number of adjacent layers with v(i,:) <= v0 (k is the index of the last tested layer)
     !---  *     nc(i) == nadj   nominal case: enough matching values   => ix(i) = k+nadj-1   (|k0|+nadj-1 <= k <= klev-nadj+1)
     !---                     particular case: all values are matching  => ix(i) = klev       (k = klev-nadj+1)
     !---  * 0 < nc(i) < nadj  bottom reached: nc<nadj matching values  => ix(i) = k+nc(i)-1  (k = |k0|)
     !---  *     nc(i) == 0    bottom reached:      no matching values  => ix(i) = k          (k = |k0|-1)
     !--- So ix(i) = MAX(k, k+nc(i)-1) fits for each case.
     DO k = klev, -1, -k0
        DO i = 1, klon
           IF(ix(i) /= 0) CYCLE                                                !--- ADEQUATE LAYER ALREADY FOUND
           nc(i) = nc(i) + 1
           IF(ABS(v(i, k)) > v0) nc(i) = 0
           IF(nc(i) /= nadj) CYCLE                                             !--- nc<nadj ADJACENT LAYERS WITH v<=v0 FOUND
           ix(i) = 1                                                           !--- FAKE /=0 VALUE TO SKIP FOLLOWING ITERATIONS
        END DO
     END DO
     DO i = 1, klon
        ix(i) = MAX(k, k+nc(i)-1)                                              !--- INDEX OF LOWEST LAYER WITH v<=v0
     END DO
  ELSE
     !=== SEARCH FROM BOTTOM TO TOP: k0 -> klev
     !--- ix(i) depends on nc(i), the number of adjacent layers with v(i,:) > v0 (k is the index of the last tested layer)
     !---  *     nc(i) == nadj   nominal case: enough matching values   => ix(i) = k-nadj     ( k0 +nadj-1 <= k <= klev-nadj+1)
     !---                     particular case: all values are matching  => ix(i) = k0-1       (k = k0+nadj-1)
     !---  * 0 < nc(i) < nadj     top reached: nc<nadj matching values  => ix(i) = k-nc(i)    (k = klev)
     !---  *     nc(i) == 0       top reached:      no matching values  => ix(i) = k          (k = klev)
     !--- So ix(i) = k-nc(i) fits for each case.
     DO k = k0, klev
        DO i = 1, klon
           IF(ix(i) /= 0) CYCLE                                                !--- ADEQUATE LAYER ALREADY FOUND
           nc(i) = nc(i) + 1
           IF(ABS(v(i, k)) <= v0) nc(i) = 0
           IF(nc(i) /= nadj) CYCLE                                             !--- nc<nadj ADJACENT LAYERS WITH v<=v0 FOUND
           ix(i) = 1                                                           !--- FAKE /=0 VALUE TO SKIP FOLLOWING ITERATIONS
        END DO
     END DO
     DO i = 1, klon
        ix(i) = k-nc(i)                                                        !--- INDEX OF LOWEST LAYER WITH v<=v0
     END DO
  END IF
END SUBROUTINE getLayerIdx
!===============================================================================================================================
SUBROUTINE potentialTemperature(pre, temp, tPot)
  IMPLICIT NONE
  REAL, DIMENSION(:, :),                       INTENT(IN)  :: pre, temp
  REAL, DIMENSION(SIZE(pre, 1), SIZE(pre, 2)), INTENT(OUT) :: tPot
  REAL, ALLOCATABLE :: tmp(:,:)
  CHARACTER(LEN=maxlen) :: modname
  INTEGER :: i, k, ni, nk
  modname = 'potentialTemperature'
  ni = SIZE(pre, 1)
  nk = SIZE(pre, 2)
  CALL assert(SIZE(temp, DIM=1) == ni, TRIM(modname)//" SIZE(temp,1) SIZE(pre,1)")
  CALL assert(SIZE(temp, DIM=2) == nk, TRIM(modname)//" SIZE(temp,2) SIZE(pre,2)")
  ALLOCATE(tmp(ni, nk))
  DO k = 1, nk                                                                 !--- COMPUTE RAW FIELD
     DO i = 1, ni
        tmp(i, k) = temp(i, k) * (100000. / pre(i, k))**RKAPPA
     END DO
  END DO
  DO k = 2, nk                                                                 !--- ENSURE GROWING FIELD WITH ALTITUDE
     DO i = 1, ni
        IF(tmp(i, k)< tmp(i, k-1)) tmp(i, k) = tmp(i, k-1) + 1.E-5
     END DO
  END DO
  CALL smooth(tmp, tPot)                                                       !--- FILTER THE FIELD
END SUBROUTINE potentialTemperature
!===============================================================================================================================
SUBROUTINE potentialVorticity(rot_cen, th_int, pint, pVor_cen)
  IMPLICIT NONE
  REAL, DIMENSION(klon, klev), INTENT(IN)  :: rot_cen, th_int, pint
  REAL, DIMENSION(klon, klev), INTENT(OUT) :: pVor_cen
  REAL ::     tmp(klon, klev)
  INTEGER :: i, k, kp
  DO k = 1, klev-1                                                             !--- COMPUTE RAW FIELD
     DO i = 1, klon
        tmp(i, k) = -1.E6 * RG * (rot_cen(i, k) + fac(i)) * (th_int(i, k+1)-th_int(i, k)) / (pint(i, k+1)-pint(i, k))
     END DO
  END DO
  DO i = 1, klon
     tmp(i, klev) = tmp(i, klev-1)
  END DO
  CALL smooth(tmp, pVor_cen)                                                   !--- FILTER THE FIELD
END SUBROUTINE potentialVorticity
!===============================================================================================================================
SUBROUTINE smooth(v, vs)
! Purpose: Vertical smoothing of each profile v(i,:) using 2*ns+1 centered binomial weights (+/- ns points).
! Note:    For levels near the bottom (k <= ns) or the top (k > klev-ns), a narrower set of weights (n<ns) is used.
!          => in particular, first and last levels are left untouched.
  IMPLICIT NONE
  REAL,    INTENT(IN)  :: v (klon, klev)
  REAL,    INTENT(OUT) :: vs(klon, klev)
  INTEGER :: i, j, k
  vs(:, :) = 0.
  DO k = 1, klev
     n = MIN(k-1, klev-k, ns)
     DO j = k-n, k+n
        DO i = 1, klon
           vs(i, k) = vs(i, k) + v(i, j) * w(n+1, 1+ABS(j-k))
        END DO
     END DO
  END DO
END SUBROUTINE smooth

END FUNCTION dyn_tropopause

END MODULE tropopause_m