SUBMODULE(compute_caldyn_mod) compute_theta_mod
   USE prec, ONLY: rstd
   IMPLICIT NONE

CONTAINS

   MODULE SUBROUTINE compute_theta(ps, theta_rhodz, rhodz, theta)
      USE grid_param, ONLY: grid_type, grid_ico
      USE disvert_mod, ONLY: mass_dak, mass_dbk
      USE trace
      REAL(rstd), INTENT(IN)    :: ps(:), theta_rhodz(:, :, :)
      REAL(rstd), INTENT(INOUT) :: rhodz(:, :)
      REAL(rstd), INTENT(OUT)   :: theta(:, :, :)
      CALL trace_start("compute_theta")
      IF (grid_type == grid_ico) THEN
         CALL compute_theta_hex_(ps, theta_rhodz, rhodz, theta, mass_dak, mass_dbk)
      ELSE
         CALL compute_theta_unst_(ps, theta_rhodz, rhodz, theta)
      END IF
      CALL trace_end("compute_theta")
   END SUBROUTINE compute_theta

   SUBROUTINE compute_theta_unst_(mass_col, theta_rhodz, rhodz, theta)
      USE ISO_C_BINDING, only: C_DOUBLE, C_FLOAT
      USE data_unstructured_mod, ONLY: enter_trace, exit_trace, id_theta
      USE earth_const, ONLY: g
      USE disvert_mod, ONLY: mass_dak, mass_dbk, caldyn_eta, eta_mass, ptop
      USE grid_param, ONLY: primal_num, dual_num, edge_num, llm, nqdyn
      REAL(C_DOUBLE), DIMENSION(primal_num) :: mass_col
      REAL(C_DOUBLE), DIMENSION(llm, primal_num) :: rhodz
      REAL(C_DOUBLE), DIMENSION(llm, primal_num, nqdyn) :: theta, theta_rhodz
      INTEGER ij, l, iq, iedge, edge, ivertex, vertex, ij_left, ij_right, ij_up, ij_down, itrisk, edge_trisk, kup, kdown
      REAL(C_DOUBLE) :: m
      CALL enter_trace(id_theta, 8*llm*((3)*primal_num + (0)*dual_num + (0)*edge_num)) ! primal, dual, edge
      !--------------------------------------------------------------------------
      !---------------------------- compute_theta ----------------------------------
      IF (caldyn_eta == eta_mass) THEN ! Compute mass
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, primal_num
            !DIR$ SIMD
            DO l = 1, llm
               m = mass_dak(l) + (mass_col(ij)*g + ptop)*mass_dbk(l)
               rhodz(l, ij) = m/g
            END DO
         END DO
         !$OMP END DO
      END IF
      DO iq = 1, nqdyn
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, primal_num
            !DIR$ SIMD
            DO l = 1, llm
               theta(l, ij, iq) = theta_rhodz(l, ij, iq)/rhodz(l, ij)
            END DO
         END DO
         !$OMP END DO
      END DO
      !---------------------------- compute_theta ----------------------------------
      !--------------------------------------------------------------------------
      CALL exit_trace()
   END SUBROUTINE compute_theta_unst_

   SUBROUTINE compute_theta_hex_(mass_col, theta_rhodz, rhodz, theta, mass_dak, mass_dbk)
      USE icosa
      USE trace, ONLY: trace_start, trace_end
      USE omp_para, ONLY: ll_begin, ll_end
      USE disvert_mod, ONLY: caldyn_eta, eta_mass, ptop
      REAL(rstd), INTENT(IN)    :: mass_col(iim*jjm)
      REAL(rstd), INTENT(IN)    :: theta_rhodz(iim*jjm, llm, nqdyn)
      REAL(rstd), INTENT(INOUT) :: rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(OUT)   :: theta(iim*jjm, llm, nqdyn)
      REAL(rstd), INTENT(IN)    :: mass_dak(llm)
      REAL(rstd), INTENT(IN)    :: mass_dbk(llm)
      INTEGER :: ij, l, iq
      REAL(rstd) :: m

      !$acc data async &
      !$acc present(mass_col(:), theta_rhodz(:,:,:), rhodz(:,:), theta(:,:,:), mass_dak(:), mass_dbk(:))

      IF (caldyn_eta == eta_mass) THEN ! Compute mass
         !$acc parallel loop collapse(2) async
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               m = mass_dak(l) + (mass_col(ij)*g + ptop)*mass_dbk(l)
               rhodz(ij, l) = m/g
            END DO
         END DO
      END IF

      !$acc parallel loop collapse(3) async
      DO iq = 1, nqdyn
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               theta(ij, l, iq) = theta_rhodz(ij, l, iq)/rhodz(ij, l)
            END DO
         END DO
      END DO
      !$acc end data
   END SUBROUTINE compute_theta_hex_

END SUBMODULE compute_theta_mod
