MODULE compute_pression_mod
   USE compute_diagnostics_mod
   USE icosa
   USE omp_para
   USE disvert_mod, ONLY: ap, bp, ap_bp_present, ptop, caldyn_eta, eta_lag
   IMPLICIT NONE
   PRIVATE

   PUBLIC :: compute_pression_hex, compute_pression_unst, &
             compute_pression_mid_hex, compute_pression_mid_unst, &
             hydrostatic_pressure, compute_hydrostatic_pressure_unst, compute_hydrostatic_pressure_hex

CONTAINS

   SUBROUTINE hydrostatic_pressure(f_rhodz, f_theta_rhodz, f_ps, f_p)
      TYPE(t_field), POINTER :: f_rhodz(:), f_theta_rhodz(:), f_ps(:), f_p(:)
      REAL(rstd), POINTER :: ps(:), rhodz(:, :), p(:, :), theta_rhodz(:, :, :)
      INTEGER :: ind
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         rhodz = f_rhodz(ind)
         theta_rhodz = f_theta_rhodz(ind)
         ps = f_ps(ind)
         p = f_p(ind)
         CALL compute_hydrostatic_pressure(rhodz, theta_rhodz, ps, p)
      END DO
   END SUBROUTINE hydrostatic_pressure

!-------------- Wrappers for F2008 conformity -----------------

   SUBROUTINE compute_pression_hex(ps, p, offset)
      REAL(rstd), INTENT(IN) :: ps(:)
      REAL(rstd), INTENT(OUT) :: p(:, :)
      INTEGER, INTENT(IN) :: offset
      CALL compute_pression_hex_(ps, p, offset)
   END SUBROUTINE compute_pression_hex

   SUBROUTINE compute_pression_unst(ps, p, offset)
      REAL(rstd), INTENT(IN) :: ps(:)
      REAL(rstd), INTENT(OUT) :: p(:, :)
      INTEGER, INTENT(IN) :: offset
      CALL compute_pression_unst_(ps, p, offset)
   END SUBROUTINE compute_pression_unst

   SUBROUTINE compute_pression_mid_hex(ps, p, offset)
      REAL(rstd), INTENT(IN) :: ps(:)
      REAL(rstd), INTENT(OUT) :: p(:, :)
      INTEGER, INTENT(IN) :: offset
      CALL compute_pression_mid_hex_(ps, p, offset)
   END SUBROUTINE compute_pression_mid_hex

   SUBROUTINE compute_pression_mid_unst(ps, p, offset)
      REAL(rstd), INTENT(IN) :: ps(:)
      REAL(rstd), INTENT(OUT) :: p(:, :)
      INTEGER, INTENT(IN) :: offset
      CALL compute_pression_mid_unst_(ps, p, offset)
   END SUBROUTINE compute_pression_mid_unst

   SUBROUTINE compute_hydrostatic_pressure_hex(rhodz, theta_rhodz, ps, p)
      REAL(rstd), INTENT(IN)  :: rhodz(:, :), theta_rhodz(:, :, :)
      REAL(rstd), INTENT(OUT) :: ps(:), p(:, :)
      CALL compute_hydrostatic_pressure_hex_(rhodz, theta_rhodz, ps, p)
   END SUBROUTINE compute_hydrostatic_pressure_hex

   SUBROUTINE compute_hydrostatic_pressure_unst(rhodz, theta_rhodz, ps, p)
      REAL(rstd), INTENT(IN)  :: rhodz(:, :), theta_rhodz(:, :, :)
      REAL(rstd), INTENT(OUT) :: ps(:), p(:, :)
      CALL compute_hydrostatic_pressure_unst_(rhodz, theta_rhodz, ps, p)
   END SUBROUTINE compute_hydrostatic_pressure_unst

!------------- hexagonal-mesh compute kernels --------

   SUBROUTINE compute_pression_hex_(ps, p, offset)
      REAL(rstd), INTENT(IN) :: ps(iim*jjm)
      REAL(rstd), INTENT(OUT) :: p(iim*jjm, llm + 1)
      INTEGER, INTENT(IN) :: offset
      INTEGER :: ij, l

      !--------------------------------------------------------------------------
      !---------------------------- compute_pression ----------------------------------
      IF (ap_bp_present) THEN
         IF (offset > 0) THEN
            DO l = ll_begin, ll_endp1
               !DIR$ SIMD
               DO ij = ij_begin_ext, ij_end_ext
                  p(ij, l) = ap(l) + bp(l)*ps(ij)
               END DO
            END DO
         ELSE
            DO l = ll_begin, ll_endp1
               !DIR$ SIMD
               DO ij = ij_begin, ij_end
                  p(ij, l) = ap(l) + bp(l)*ps(ij)
               END DO
            END DO
         END IF
      END IF
      !---------------------------- compute_pression ----------------------------------
      !--------------------------------------------------------------------------
   END SUBROUTINE compute_pression_hex_

   SUBROUTINE compute_pression_mid_hex_(ps, pmid, offset)
      REAL(rstd), INTENT(IN) :: ps(iim*jjm)
      REAL(rstd), INTENT(OUT) :: pmid(iim*jjm, llm)
      INTEGER, INTENT(IN) :: offset
      INTEGER :: ij, l

      !--------------------------------------------------------------------------
      !---------------------------- compute_pmid ----------------------------------
      IF (ap_bp_present) THEN
         IF (offset > 0) THEN
            DO l = ll_begin, ll_end
               !DIR$ SIMD
               DO ij = ij_begin_ext, ij_end_ext
                  pmid(ij, l) = .5*(ap(l) + ap(l + 1) + (bp(l) + bp(l + 1))*ps(ij))
               END DO
            END DO
         ELSE
            DO l = ll_begin, ll_end
               !DIR$ SIMD
               DO ij = ij_begin, ij_end
                  pmid(ij, l) = .5*(ap(l) + ap(l + 1) + (bp(l) + bp(l + 1))*ps(ij))
               END DO
            END DO
         END IF
      END IF
      !---------------------------- compute_pmid ----------------------------------
      !--------------------------------------------------------------------------
   END SUBROUTINE compute_pression_mid_hex_

   SUBROUTINE compute_hydrostatic_pressure_hex_(rhodz, theta_rhodz, ps, pk)
      REAL(rstd), INTENT(IN)  :: rhodz(iim*jjm, llm) ! mass per unit surface in each model level
      REAL(rstd), INTENT(IN)  :: theta_rhodz(iim*jjm, llm, nqdyn) ! dynamical tracers (theta/entropy)
      REAL(rstd), INTENT(OUT) :: ps(iim*jjm)        ! surface pressure, diagnosed if Lagrangian vertical coordinate
      REAL(rstd), INTENT(OUT) :: pk(iim*jjm, llm)    ! pressure at full levels
      INTEGER :: ij, l, ij_omp_begin_ext, ij_omp_end_ext
      !$OMP BARRIER
      CALL distrib_level(ij_begin_ext, ij_end_ext, ij_omp_begin_ext, ij_omp_end_ext)

      !--------------------------------------------------------------------------
      !---------------------------- compute_hydrostatic_pressure ----------------------------------
      SELECT CASE (caldyn_thermo)
      CASE (thermo_boussinesq)
         ! use hydrostatic balance with theta*rhodz to find pk (=Lagrange multiplier=pressure)
         !DIR$ SIMD
         DO ij = ij_omp_begin_ext, ij_omp_end_ext
            pk(ij, llm) = ptop + .5*g*theta_rhodz(ij, llm, 1)
         END DO
         DO l = llm - 1, 1, -1
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               pk(ij, l) = pk(ij, l + 1) + (.5*g)*(theta_rhodz(ij, l, 1) + theta_rhodz(ij, l + 1, 1))
            END DO
         END DO
         IF (caldyn_eta == eta_lag) THEN
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               ps(ij) = pk(ij, 1) + .5*g*theta_rhodz(ij, 1, 1)
            END DO
         END IF
      CASE (thermo_theta, thermo_entropy, thermo_variable_Cp)
         !DIR$ SIMD
         DO ij = ij_omp_begin_ext, ij_omp_end_ext
            pk(ij, llm) = ptop + .5*g*rhodz(ij, llm)
         END DO
         DO l = llm - 1, 1, -1
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               pk(ij, l) = pk(ij, l + 1) + (.5*g)*(rhodz(ij, l) + rhodz(ij, l + 1))
            END DO
         END DO
         IF (caldyn_eta == eta_lag) THEN
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               ps(ij) = pk(ij, 1) + .5*g*rhodz(ij, 1)
            END DO
         END IF
      CASE (thermo_moist)
         !DIR$ SIMD
         DO ij = ij_omp_begin_ext, ij_omp_end_ext
            pk(ij, llm) = ptop + .5*g*(rhodz(ij, llm) + theta_rhodz(ij, llm, 2))
         END DO
         DO l = llm - 1, 1, -1
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               pk(ij, l) = pk(ij, l + 1) + (.5*g)*((rhodz(ij, l) + theta_rhodz(ij, l, 2)) &
                                                   + (rhodz(ij, l + 1) + theta_rhodz(ij, l + 1, 2)))
            END DO
         END DO
         IF (caldyn_eta == eta_lag) THEN
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               ps(ij) = pk(ij, 1) + .5*g*(rhodz(ij, 1) + theta_rhodz(ij, 1, 2))
            END DO
         END IF
      END SELECT
      !---------------------------- compute_hydrostatic_pressure ----------------------------------
      !--------------------------------------------------------------------------
      !$OMP BARRIER
   END SUBROUTINE compute_hydrostatic_pressure_hex_

!----------- unstructured-mesh compute kernels --------

   SUBROUTINE compute_pression_unst_(ps, p, offset)
      REAL(C_DOUBLE), DIMENSION(primal_num), INTENT(IN)  :: ps
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num), INTENT(OUT) :: p
      INTEGER, INTENT(IN)  :: offset
      INTEGER ij, l, iq, iedge, edge, ivertex, vertex, ij_left, ij_right, ij_up, ij_down, itrisk, edge_trisk, kup, kdown

      !--------------------------------------------------------------------------
      !---------------------------- compute_pression ----------------------------------
      IF (ap_bp_present) THEN
         IF (offset > 0) THEN
            !$OMP DO SCHEDULE(STATIC)
            DO ij = 1, primal_num
               !DIR$ SIMD
               DO l = 1, llm + 1
                  p(l, ij) = ap(l) + bp(l)*ps(ij)
               END DO
            END DO
            !$OMP END DO
         ELSE
            !$OMP DO SCHEDULE(STATIC)
            DO ij = 1, primal_num
               !DIR$ SIMD
               DO l = 1, llm + 1
                  p(l, ij) = ap(l) + bp(l)*ps(ij)
               END DO
            END DO
            !$OMP END DO
         END IF
      END IF
      !---------------------------- compute_pression ----------------------------------
      !--------------------------------------------------------------------------
   END SUBROUTINE compute_pression_unst_

   SUBROUTINE compute_pression_mid_unst_(ps, pmid, offset)
      REAL(C_DOUBLE), DIMENSION(primal_num), INTENT(IN)  :: ps
      REAL(C_DOUBLE), DIMENSION(llm, primal_num), INTENT(OUT) :: pmid
      INTEGER, INTENT(IN)  :: offset
      INTEGER ij, l, iq, iedge, edge, ivertex, vertex, ij_left, ij_right, ij_up, ij_down, itrisk, edge_trisk, kup, kdown

      !--------------------------------------------------------------------------
      !---------------------------- compute_pmid ----------------------------------
      IF (ap_bp_present) THEN
         IF (offset > 0) THEN
            !$OMP DO SCHEDULE(STATIC)
            DO ij = 1, primal_num
               !DIR$ SIMD
               DO l = 1, llm
                  pmid(l, ij) = .5*(ap(l) + ap(l + 1) + (bp(l) + bp(l + 1))*ps(ij))
               END DO
            END DO
            !$OMP END DO
         ELSE
            !$OMP DO SCHEDULE(STATIC)
            DO ij = 1, primal_num
               !DIR$ SIMD
               DO l = 1, llm
                  pmid(l, ij) = .5*(ap(l) + ap(l + 1) + (bp(l) + bp(l + 1))*ps(ij))
               END DO
            END DO
            !$OMP END DO
         END IF
      END IF
      !---------------------------- compute_pmid ----------------------------------
      !--------------------------------------------------------------------------
   END SUBROUTINE compute_pression_mid_unst_

   SUBROUTINE compute_hydrostatic_pressure_unst_(rhodz, theta_rhodz, ps, pk)
      REAL(C_DOUBLE), DIMENSION(llm, primal_num), INTENT(IN)  :: rhodz
      REAL(C_DOUBLE), DIMENSION(llm, primal_num, nqdyn), INTENT(IN)  :: theta_rhodz
      REAL(C_DOUBLE), DIMENSION(primal_num), INTENT(OUT) :: ps
      REAL(C_DOUBLE), DIMENSION(llm, primal_num), INTENT(OUT) :: pk
      INTEGER ij, l, iq, iedge, edge, ivertex, vertex, ij_left, ij_right, ij_up, ij_down, itrisk, edge_trisk, kup, kdown

      !--------------------------------------------------------------------------
      !---------------------------- compute_hydrostatic_pressure ----------------------------------
      SELECT CASE (caldyn_thermo)
      CASE (thermo_boussinesq)
         ! use hydrostatic balance with theta*rhodz to find pk (=Lagrange multiplier=pressure)
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, primal_num
            pk(llm, ij) = ptop + .5*g*theta_rhodz(llm, ij, 1)
            DO l = llm - 1, 1, -1
               pk(l, ij) = pk(l + 1, ij) + (.5*g)*(theta_rhodz(l, ij, 1) + theta_rhodz(l + 1, ij, 1))
            END DO
            IF (caldyn_eta == eta_lag) THEN
               ps(ij) = pk(1, ij) + .5*g*theta_rhodz(1, ij, 1)
            END IF
         END DO
         !$OMP END DO
      CASE (thermo_theta, thermo_entropy, thermo_variable_Cp)
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, primal_num
            pk(llm, ij) = ptop + .5*g*rhodz(llm, ij)
            DO l = llm - 1, 1, -1
               pk(l, ij) = pk(l + 1, ij) + (.5*g)*(rhodz(l, ij) + rhodz(l + 1, ij))
            END DO
            IF (caldyn_eta == eta_lag) THEN
               ps(ij) = pk(1, ij) + .5*g*rhodz(1, ij)
            END IF
         END DO
         !$OMP END DO
      CASE (thermo_moist)
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, primal_num
            pk(llm, ij) = ptop + .5*g*(rhodz(llm, ij) + theta_rhodz(llm, ij, 2))
            DO l = llm - 1, 1, -1
               pk(l, ij) = pk(l + 1, ij) + (.5*g)*((rhodz(l, ij) + theta_rhodz(l, ij, 2)) &
                                                   + (rhodz(l + 1, ij) + theta_rhodz(l + 1, ij, 2)))
            END DO
            IF (caldyn_eta == eta_lag) THEN
               ps(ij) = pk(1, ij) + .5*g*(rhodz(1, ij) + theta_rhodz(1, ij, 2))
            END IF
         END DO
         !$OMP END DO
      END SELECT
      !---------------------------- compute_hydrostatic_pressure ----------------------------------
      !--------------------------------------------------------------------------
   END SUBROUTINE compute_hydrostatic_pressure_unst_

END MODULE compute_pression_mod
