SUBMODULE(compute_caldyn_mod) compute_geopot_mod
   USE prec, ONLY: rstd
   USE earth_const
   USE disvert_mod
   USE omp_para
   USE trace
   IMPLICIT NONE
   SAVE

CONTAINS

   MODULE SUBROUTINE compute_geopot(rhodz, theta, ps, enthalpy, pk, geopot)
      USE grid_param, ONLY: grid_type, grid_ico
      REAL(rstd), INTENT(IN)    :: rhodz(:, :), theta(:, :, :) ! active scalars : theta/entropy, moisture, ...
      REAL(rstd), INTENT(INOUT) :: ps(:), geopot(:, :)       ! geopotential
      REAL(rstd), INTENT(OUT)   :: enthalpy(:, :), pk(:, :, :)   ! enthalpy, Exner function (compressible) /Lagrange multiplier (Boussinesq)
      IF (grid_type == grid_ico) THEN
         IF (boussinesq) THEN
            CALL compute_geopot_boussinesq_hex(rhodz, theta, ps, pk(:, :, 1), geopot)
         ELSE
            CALL compute_geopot_compressible_hex(rhodz, theta, ps, enthalpy, pk, geopot)
         END IF
      ELSE
         IF (boussinesq) THEN
            CALL compute_geopot_boussinesq_unst(rhodz, theta, ps, pk(:, :, 1), geopot)
         ELSE
            CALL compute_geopot_compressible_unst(rhodz, theta, ps, enthalpy, pk, geopot)
         END IF
      END IF
   END SUBROUTINE compute_geopot

   !**************************** Boussinesq *****************************

   SUBROUTINE compute_geopot_boussinesq_hex(rhodz, theta, ps, pk, geopot)
      USE grid_param
      REAL(rstd), INTENT(IN)    :: rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(IN)    :: theta(iim*jjm, llm, nqdyn) ! active scalars : theta/entropy, moisture, ...
      REAL(rstd), INTENT(INOUT) :: ps(iim*jjm)
      REAL(rstd), INTENT(OUT)   :: pk(iim*jjm, llm)         ! Lagrange multiplier (Boussinesq)
      REAL(rstd), INTENT(INOUT) :: geopot(iim*jjm, llm + 1) ! geopotential

      INTEGER :: ij, l, ij_omp_begin_ext, ij_omp_end_ext

      CALL trace_start("compute_geopot_boussinesq")

      !$OMP BARRIER

      CALL distrib_level(ij_begin_ext, ij_end_ext, ij_omp_begin_ext, ij_omp_end_ext)

      ! specific volume 1 = dphi/g/rhodz
      DO l = 1, llm
         !DIR$ SIMD
         DO ij = ij_omp_begin_ext, ij_omp_end_ext
            geopot(ij, l + 1) = geopot(ij, l) + g*rhodz(ij, l)
         END DO
      END DO
      ! use hydrostatic balance with theta*rhodz to find pk (Lagrange multiplier=pressure)
      ! uppermost layer
      !DIR$ SIMD
      DO ij = ij_begin_ext, ij_end_ext
         pk(ij, llm) = ptop + (.5*g)*theta(ij, llm, 1)*rhodz(ij, llm)
      END DO
      ! other layers
      DO l = llm - 1, 1, -1
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            pk(ij, l) = pk(ij, l + 1) + (.5*g)*(theta(ij, l, 1)*rhodz(ij, l) + theta(ij, l + 1, 1)*rhodz(ij, l + 1))
         END DO
      END DO
      ! now pk contains the Lagrange multiplier (pressure)

      !ym flush geopot
      !$OMP BARRIER

      CALL trace_end("compute_geopot_boussinesq")

   END SUBROUTINE compute_geopot_boussinesq_hex

   SUBROUTINE compute_geopot_boussinesq_unst(rhodz, theta, ps, pk, geopot)
      USE grid_param
      REAL(C_DOUBLE), DIMENSION(llm, primal_num)        :: rhodz, pk ! IN, OUT
      REAL(C_DOUBLE), DIMENSION(llm, primal_num, nqdyn) :: theta     ! IN
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num)    :: geopot    ! IN(l=1)/OUT(l>1)
      REAL(C_DOUBLE), DIMENSION(primal_num)             :: ps        ! OUT

      INTEGER :: ij, l
      REAL(C_DOUBLE) :: chi, gv, exner_ik, temp_ik, p_ik, qv, Rmix, Cp_ik

      CALL trace_start("compute_geopot_boussinesq")

      DO ij = 1, primal_num
         DO l = 1, llm
            geopot(l + 1, ij) = geopot(l, ij) + g*rhodz(l, ij)
         END DO
         ! use hydrostatic balance with theta*rhodz to find pk (Lagrange multiplier=pressure)
         ! uppermost layer
         pk(llm, ij) = ptop + (.5*g)*theta(llm, ij, 1)*rhodz(llm, ij)
         ! other layers
         DO l = llm - 1, 1, -1
            !DIR$ SIMD
            pk(l, ij) = pk(l + 1, ij) + (.5*g)*(theta(l, ij, 1)*rhodz(l, ij) + theta(l + 1, ij, 1)*rhodz(l + 1, ij))
         END DO
         ! now pk contains the Lagrange multiplier (pressure)
      END DO

      CALL trace_end("compute_geopot_boussinesq")

   END SUBROUTINE compute_geopot_boussinesq_unst

   !**************************** Compressible ***************************

   SUBROUTINE compute_geopot_compressible_hex(rhodz, theta, ps, enthalpy, pk, geopot)
      USE grid_param
      USE earth_const, ONLY: thermo
      USE thermodynamics_mod, ONLY: enthalpy_from_pq
      REAL(rstd), INTENT(IN)    :: rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(IN)    :: theta(iim*jjm, llm, nqdyn) ! active scalars : theta/entropy, moisture, ...
      REAL(rstd), INTENT(INOUT) :: ps(iim*jjm)
      REAL(rstd), INTENT(OUT)   :: enthalpy(iim*jjm, llm)   ! specific enthalpy
      REAL(rstd), INTENT(OUT)   :: pk(iim*jjm, llm, nqdyn)  ! Exner functions (compressible)
      REAL(rstd), INTENT(INOUT) :: geopot(iim*jjm, llm + 1) ! geopotential

      INTEGER :: i, j, ij, l
      REAL(rstd) :: p_ik, exner_ik, Cp_ik, temp_ik, qv, chi, Rmix, gv, Rd
      INTEGER    :: ij_omp_begin_ext, ij_omp_end_ext

      CALL trace_start("compute_geopot_compressible")

      !$OMP BARRIER

      CALL distrib_level(ij_begin_ext, ij_end_ext, ij_omp_begin_ext, ij_omp_end_ext)

      ! Pressure is computed first top-down (temporarily stored in pk)
      ! Then Exner pressure and geopotential are computed bottom-up
      ! Works also when caldyn_eta=eta_mass

      !$acc data present(rhodz(:,:), theta(:,:,:), ps(:), enthalpy(:,:), pk(:,:,:), geopot(:,:)) async

      ! uppermost layer
      !$acc parallel loop async
      !DIR$ SIMD
      DO ij = ij_omp_begin_ext, ij_omp_end_ext
         pk(ij, llm, 1) = ptop + (.5*g)*rhodz(ij, llm)
      END DO

      ! other layers
      ! We use kernels here instead of "loop seq" + "loop gang vector"
      ! to be sure the code really abides the standard. In practice,
      ! it seems like the compiler interchanges the loops.
      !$acc kernels async
      DO l = llm - 1, 1, -1
         !DIR$ SIMD
         DO ij = ij_omp_begin_ext, ij_omp_end_ext
            pk(ij, l, 1) = pk(ij, l + 1, 1) + (.5*g)*(rhodz(ij, l) + rhodz(ij, l + 1))
         END DO
      END DO
      !$acc end kernels

      ! surface pressure (for diagnostics)
      IF (caldyn_eta == eta_lag) THEN
         !$acc parallel loop async
         DO ij = ij_omp_begin_ext, ij_omp_end_ext
            ps(ij) = pk(ij, 1, 1) + (.5*g)*rhodz(ij, 1)
         END DO
      END IF

      CALL enthalpy_from_pq(thermo, .TRUE., &
      &                     iim*jjm, ij_omp_begin_ext, ij_omp_end_ext, &
      &                     llm, 1, llm, thermo%nqdyn, &
      &                     enthalpy, geopot(1, 2), pk, pk, theta)

      ! We use kernels here instead of "loop seq" + "loop gang vector"
      ! to be sure the code really abides the standard. In practice,
      ! it seems like the compiler interchanges the loops.
      !$acc kernels async
      DO l = 1, llm
         !DIR$ SIMD
         DO ij = ij_omp_begin_ext, ij_omp_end_ext
            ! specific volume v = dphi/g/rhodz
            geopot(ij, l + 1) = geopot(ij, l) + (g*rhodz(ij, l))*geopot(ij, l + 1)
         END DO
      END DO
      !$acc end kernels

      !$acc end data

      !flush geopot
      !$OMP BARRIER

      CALL trace_end("compute_geopot_compressible")

   END SUBROUTINE compute_geopot_compressible_hex

   SUBROUTINE compute_geopot_compressible_unst(rhodz, theta, ps, enthalpy, pk, geopot)
      USE grid_param
      USE earth_const, ONLY: thermo
      USE thermodynamics_mod, ONLY: enthalpy_from_pq
      REAL(C_DOUBLE), DIMENSION(llm, primal_num)        :: rhodz, enthalpy ! IN, OUT
      REAL(C_DOUBLE), DIMENSION(llm, primal_num, nqdyn) :: theta, pk       ! IN, OUT
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num)    :: geopot    ! IN(l=1)/OUT(l>1)
      REAL(C_DOUBLE), DIMENSION(primal_num)             :: ps        ! OUT

      REAL(C_DOUBLE), DIMENSION(llm, primal_num)        :: vol
      INTEGER :: ij, l
      REAL(rstd) :: p_ik, exner_ik, Cp_ik, temp_ik, qv, chi, Rmix, gv, Rd

      CALL trace_start("compute_geopot_compressible")

      DO ij = 1, primal_num
         ! uppermost layer
         pk(llm, ij, 1) = ptop + (.5*g)*rhodz(llm, ij)
         ! other layers
         DO l = llm - 1, 1, -1
            pk(l, ij, 1) = pk(l + 1, ij, 1) + (.5*g)*(rhodz(l, ij) + rhodz(l + 1, ij))
         END DO
         ! surface pressure (for diagnostics)
         IF (caldyn_eta == eta_lag) THEN
            ps(ij) = pk(1, ij, 1) + (.5*g)*rhodz(1, ij)
         END IF
      END DO

      CALL enthalpy_from_pq(thermo, .FALSE., &
      &                     llm, 1, llm, &
      &                     primal_num, 1, primal_num, thermo%nqdyn, &
      &                     enthalpy, vol, pk, pk, theta)

      DO ij = 1, primal_num
         DO l = 1, llm
            ! specific volume v = dphi/g/rhodz
            geopot(l + 1, ij) = geopot(l, ij) + (g*rhodz(l, ij))*vol(l, ij)
         END DO
      END DO
      CALL trace_end("compute_geopot_compressible")

   END SUBROUTINE compute_geopot_compressible_unst

END SUBMODULE compute_geopot_mod
