MODULE compute_theta_mod
  USE grid_param, ONLY : llm
  IMPLICIT NONE
  PRIVATE

  PUBLIC :: compute_theta

CONTAINS

  SUBROUTINE compute_theta(ps,theta_rhodz, rhodz,theta)
    USE icosa
    USE trace, ONLY : trace_start, trace_end
    USE omp_para, ONLY : ll_begin, ll_end
    USE disvert_mod, ONLY : mass_dak, mass_dbk, caldyn_eta, eta_mass, ptop
    REAL(rstd),INTENT(IN)    :: ps(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)
    INTEGER :: ij,l,iq
    REAL(rstd) :: m
    CALL trace_start("compute_theta")

    IF(caldyn_eta==eta_mass) THEN ! Compute mass
       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
          END DO
       END DO
    END IF

    DO l = ll_begin,ll_end
       DO iq=1,nqdyn
          !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

    CALL trace_end("compute_theta")
  END SUBROUTINE compute_theta
  
END MODULE compute_theta_mod