MODULE compute_caldyn_mod
   USE prec, ONLY: rstd
   IMPLICIT NONE

   REAL(rstd), PARAMETER :: pbot = 1e5, rho_bot = 1e6
   LOGICAL, PARAMETER :: debug_hevi_solver = .FALSE.

   INTERFACE

      MODULE SUBROUTINE compute_planetvel(planetvel)
         REAL(rstd), INTENT(OUT)  :: planetvel(:)
      END SUBROUTINE

      MODULE SUBROUTINE compute_pvort(ps, u, theta_rhodz, rhodz, theta, qu, qv)
         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(:, :)
      END SUBROUTINE

      MODULE SUBROUTINE compute_caldyn_horiz(u, rhodz, qu, theta, pk, geopot, hflux, convm, dtheta_rhodz, du)
         REAL(rstd), INTENT(IN)  :: u(:, :)    ! prognostic "velocity"
         REAL(rstd), INTENT(IN)  :: rhodz(:, :)
         REAL(rstd), INTENT(IN)  :: qu(:, :)
         REAL(rstd), INTENT(IN)  :: theta(:, :)  ! potential temperature
         REAL(rstd), INTENT(INOUT) :: pk(:, :) ! Exner function
         REAL(rstd), INTENT(IN)  :: geopot(:, :)    ! geopotential

         REAL(rstd), INTENT(OUT) :: hflux(:, :) ! hflux in kg/s
         REAL(rstd), INTENT(OUT) :: convm(:, :)  ! mass flux convergence
         REAL(rstd), INTENT(OUT) :: dtheta_rhodz(:, :)
         REAL(rstd), INTENT(OUT) :: du(:, :)
      END SUBROUTINE

      MODULE SUBROUTINE compute_pvort_only(u, rhodz, qu, qv)
         REAL(rstd), INTENT(IN)    :: u(:, :)
         REAL(rstd), INTENT(INOUT) :: rhodz(:, :)
         REAL(rstd), INTENT(OUT)   :: qu(:, :)
         REAL(rstd), INTENT(OUT)   :: qv(:, :)
      END SUBROUTINE compute_pvort_only

      MODULE SUBROUTINE compute_theta(ps, theta_rhodz, rhodz, theta)
         REAL(rstd), INTENT(IN)    :: ps(:)
         REAL(rstd), INTENT(IN)    :: theta_rhodz(:, :, :)
         REAL(rstd), INTENT(INOUT) :: rhodz(:, :)
         REAL(rstd), INTENT(OUT)   :: theta(:, :, :)
      END SUBROUTINE compute_theta

      MODULE SUBROUTINE compute_geopot(rhodz, theta, ps, enthalpy, pk, geopot)
         REAL(rstd), INTENT(IN)    :: rhodz(:, :)     ! rho*dz = mass per unit surface in each full model level
         REAL(rstd), INTENT(IN)    :: theta(:, :, :)  ! active scalars : theta/entropy, moisture, ...
         REAL(rstd), INTENT(INOUT) :: ps(:)           ! surface pressure
         REAL(rstd), INTENT(OUT)   :: enthalpy(:, :)  ! Bernoulli function
         REAL(rstd), INTENT(OUT)   :: pk(:, :, :)     ! Exner functions (compressible) / Lagrange multiplier (Boussinesq)
         REAL(rstd), INTENT(INOUT) :: geopot(:, :)    ! geopotential
      END SUBROUTINE compute_geopot

      MODULE SUBROUTINE compute_caldyn_fast(tau, theta, geopot, pk, berni, du, u)
         REAL(rstd), INTENT(IN)    :: tau           ! "solve" u-tau*du/dt = rhs
         REAL(rstd), INTENT(IN)    :: theta(:, :, :)
         REAL(rstd), INTENT(IN)    :: geopot(:, :)
         REAL(rstd), INTENT(INOUT) :: pk(:, :, :)
         REAL(rstd), INTENT(INOUT) :: berni(:, :)    ! partial Bernoulli function
         REAL(rstd), INTENT(INOUT) :: du(:, :)
         REAL(rstd), INTENT(INOUT) :: u(:, :)        ! INOUT if tau>0
      END SUBROUTINE compute_caldyn_fast

      MODULE SUBROUTINE compute_caldyn_slow_hydro(zero, u, rhodz, berni, hflux, du)
         LOGICAL, INTENT(IN)    :: zero
         REAL(rstd), INTENT(IN)    :: u(:, :)      ! prognostic "velocity"
         REAL(rstd), INTENT(IN)    :: rhodz(:, :)
         REAL(rstd), INTENT(OUT)   :: berni(:, :)  ! Bernoulli function
         REAL(rstd), INTENT(OUT)   :: hflux(:, :)  ! hflux in kg/s
         REAL(rstd), INTENT(INOUT) :: du(:, :)
      END SUBROUTINE compute_caldyn_slow_hydro

      MODULE SUBROUTINE compute_caldyn_coriolis(hflux, theta, qu, Ftheta, convm, dtheta_rhodz, du)
         REAL(rstd), INTENT(IN)    :: hflux(:, :)   ! hflux in kg/s
         REAL(rstd), INTENT(IN)    :: theta(:, :, :) ! active scalars
         REAL(rstd), INTENT(IN)    :: qu(:, :)
         REAL(rstd), INTENT(OUT)   :: Ftheta(:, :)  ! potential temperature flux
         REAL(rstd), INTENT(OUT)   :: convm(:, :)   ! mass flux convergence
         REAL(rstd), INTENT(OUT)   :: dtheta_rhodz(:, :, :)
         REAL(rstd), INTENT(INOUT) :: du(:, :)
      END SUBROUTINE compute_caldyn_coriolis

      MODULE SUBROUTINE compute_caldyn_vert(u, theta, rhodz, convm, wflux, wwuu, dps, dtheta_rhodz, du)
         REAL(rstd), INTENT(IN)    :: u(:, :)              ! momentum (covariant)
         REAL(rstd), INTENT(IN)    :: theta(:, :, :)        ! dynamical tracers (scalars)
         REAL(rstd), INTENT(IN)    :: rhodz(:, :)          ! mass (density)
         REAL(rstd), INTENT(INOUT) :: convm(:, :)          ! horizontal mass flux convergence (density)
         REAL(rstd), INTENT(INOUT) :: wflux(:, :)          ! vertical mass flux
         REAL(rstd), INTENT(OUT)   :: wwuu(:, :)           ! wflux * u
         REAL(rstd), INTENT(INOUT)  :: du(:, :)             ! momentum tendency
         REAL(rstd), INTENT(INOUT)  :: dtheta_rhodz(:, :, :) ! dynamical tracers tendencies (densities)
         REAL(rstd), INTENT(OUT)    :: dps(:)              ! column mass tendency
      END SUBROUTINE compute_caldyn_vert

      ! Unused / untested / incomplete

      MODULE SUBROUTINE compute_caldyn_kv(ue, Kv)
         REAL(rstd), INTENT(IN) :: ue(:, :)
         REAL(rstd), INTENT(OUT) :: Kv(:, :)
      END SUBROUTINE

      ! Non-hydrostatic

      MODULE SUBROUTINE compute_caldyn_vert_NH(mass, geopot, W, wflux, W_etadot, du, dPhi, dW)
         REAL(rstd), INTENT(IN) :: mass(:, :)
         REAL(rstd), INTENT(IN) :: geopot(:, :)
         REAL(rstd), INTENT(IN) :: W(:, :)
         REAL(rstd), INTENT(IN) :: wflux(:, :)
         REAL(rstd), INTENT(INOUT) :: du(:, :)
         REAL(rstd), INTENT(INOUT) :: dPhi(:, :)
         REAL(rstd), INTENT(INOUT) :: dW(:, :)
         REAL(rstd) :: W_etadot(:, :) ! vertical flux of vertical momentum
      END SUBROUTINE

      MODULE SUBROUTINE compute_caldyn_solver(tau, phis, rhodz, theta, pk, geopot, W, m_il, pres, dPhi, dW, du)
         REAL(rstd), INTENT(IN) :: tau ! "solve" Phi-tau*dPhi/dt = Phi_rhs
         REAL(rstd), INTENT(IN)    :: phis(:)
         REAL(rstd), INTENT(IN)    :: rhodz(:, :)
         REAL(rstd), INTENT(IN)    :: theta(:, :, :)
         REAL(rstd), INTENT(OUT)   :: pk(:, :)
         REAL(rstd), INTENT(INOUT) :: geopot(:, :)
         REAL(rstd), INTENT(INOUT) :: W(:, :) ! OUT if tau>0
         REAL(rstd), INTENT(OUT)   :: m_il(:, :)        ! rhodz averaged to interfaces
         REAL(rstd), INTENT(OUT)   :: pres(:, :)          ! pressure
         REAL(rstd), INTENT(OUT)   :: dW(:, :)
         REAL(rstd), INTENT(OUT)   :: dPhi(:, :)
         REAL(rstd), INTENT(OUT)   :: du(:, :)
      END SUBROUTINE

      MODULE SUBROUTINE compute_caldyn_slow_NH(u, rhodz, Phi, W, F_el, gradPhi2, w_il, hflux, du, dPhi, dW)
         REAL(rstd), INTENT(IN)  :: u(:, :)   ! prognostic "velocity"
         REAL(rstd), INTENT(IN)  :: rhodz(:, :)  ! rho*dz
         REAL(rstd), INTENT(IN)  :: Phi(:, :)  ! prognostic geopotential
         REAL(rstd), INTENT(IN)  :: W(:, :)    ! prognostic vertical momentum

         REAL(rstd), INTENT(OUT) :: F_el(:, :) ! NH mass flux
         REAL(rstd), INTENT(OUT) :: gradPhi2(:, :) ! grad_Phi**2
         REAL(rstd), INTENT(OUT) :: w_il(:, :) ! Wil/mil

         REAL(rstd), INTENT(OUT) :: hflux(:, :) ! hflux in kg/s
         REAL(rstd), INTENT(OUT) :: du(:, :)
         REAL(rstd), INTENT(OUT) :: dW(:, :)
         REAL(rstd), INTENT(OUT) :: dPhi(:, :)
      END SUBROUTINE

   END INTERFACE

END MODULE compute_caldyn_mod
