SUBMODULE(compute_caldyn_mod) compute_caldyn_pvort_mod

   IMPLICIT NONE

CONTAINS

   MODULE SUBROUTINE compute_pvort(ps, u, theta_rhodz, rhodz, theta, qu, qv)
      USE prec, ONLY: rstd
      REAL(rstd), INTENT(IN) :: u(:, :)
      REAL(rstd), INTENT(IN) :: ps(:)
      REAL(rstd), INTENT(IN) :: theta_rhodz(:, :)
      REAL(rstd), INTENT(INOUT) :: rhodz(:, :)
      REAL(rstd), INTENT(OUT) :: theta(:, :)
      REAL(rstd), INTENT(OUT) :: qu(:, :)
      REAL(rstd), INTENT(OUT) :: qv(:, :)
      CALL compute_pvort_hex(ps, u, theta_rhodz, rhodz, theta, qu, qv)
   END SUBROUTINE

   SUBROUTINE compute_pvort_hex(ps, u, theta_rhodz, rhodz, theta, qu, qv)
      USE icosa
      USE disvert_mod, ONLY: mass_dak, mass_dbk, caldyn_eta, eta_mass, ptop
      USE trace
      USE omp_para
      REAL(rstd), INTENT(IN) :: u(iim*3*jjm, llm)
      REAL(rstd), INTENT(IN) :: ps(iim*jjm)
      REAL(rstd), INTENT(IN) :: theta_rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(INOUT) :: rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: theta(iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: qu(iim*3*jjm, llm)
      REAL(rstd), INTENT(OUT) :: qv(iim*2*jjm, llm)

      INTEGER :: ij, l, i, j
      REAL(rstd) :: etav, hv, m
      CALL trace_start("compute_pvort")

      IF (caldyn_eta == eta_mass) THEN ! Compute mass & theta
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               m = mass_dak(l) + (ps(ij)*g + ptop)*mass_dbk(l) ! ps is actually Ms
               rhodz(ij, l) = m/g
               theta(ij, l) = theta_rhodz(ij, l)/rhodz(ij, l)
            END DO
         END DO
      ELSE ! Compute only theta
         DO l = ll_begin, ll_end
            !         CALL test_message(req_u)
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               theta(ij, l) = theta_rhodz(ij, l)/rhodz(ij, l)
            END DO
         END DO
      END IF

      !   CALL wait_message(req_u) ! COM02

    !!! Compute shallow-water potential vorticity
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            etav = 1./Av(ij + z_up)*(ne_rup*u(ij + u_rup, l) &
                                     + ne_left*u(ij + t_rup + u_left, l) &
                                     - ne_lup*u(ij + u_lup, l))
            hv = Riv2(ij, vup)*rhodz(ij, l) &
                 + Riv2(ij + t_rup, vldown)*rhodz(ij + t_rup, l) &
                 + Riv2(ij + t_lup, vrdown)*rhodz(ij + t_lup, l)
            qv(ij + z_up, l) = (etav + fv(ij + z_up))/hv

            etav = 1./Av(ij + z_down)*(ne_ldown*u(ij + u_ldown, l) &
                                       + ne_right*u(ij + t_ldown + u_right, l) &
                                       - ne_rdown*u(ij + u_rdown, l))
            hv = Riv2(ij, vdown)*rhodz(ij, l) &
                 + Riv2(ij + t_ldown, vrup)*rhodz(ij + t_ldown, l) &
                 + Riv2(ij + t_rdown, vlup)*rhodz(ij + t_rdown, l)
            qv(ij + z_down, l) = (etav + fv(ij + z_down))/hv
         END DO

         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            qu(ij + u_right, l) = 0.5*(qv(ij + z_rdown, l) + qv(ij + z_rup, l))
            qu(ij + u_lup, l) = 0.5*(qv(ij + z_up, l) + qv(ij + z_lup, l))
            qu(ij + u_ldown, l) = 0.5*(qv(ij + z_ldown, l) + qv(ij + z_down, l))
         END DO

      END DO

      CALL trace_end("compute_pvort")
   END SUBROUTINE compute_pvort_hex

END SUBMODULE
