MODULE theta2theta_rhodz_mod
   USE field_mod
   PRIVATE

   TYPE(t_field), POINTER, SAVE  :: f_p(:)

   PUBLIC :: init_theta2theta_rhodz, theta_rhodz2theta, &
             theta_rhodz2temperature, temperature2theta_rhodz, &
             theta2theta_rhodz, &
             compute_theta2theta_rhodz, compute_theta_rhodz2theta

CONTAINS

   SUBROUTINE init_theta2theta_rhodz
      USE icosa
      USE field_mod
      IMPLICIT NONE
      CALL allocate_field(f_p, field_t, type_real, llm + 1, name='p (theta2theta_rhodz_mod)')
   END SUBROUTINE init_theta2theta_rhodz

   SUBROUTINE theta_rhodz2theta(f_ps, f_theta_rhodz, f_theta)
      USE icosa
      IMPLICIT NONE
      TYPE(t_field), POINTER :: f_ps(:)
      TYPE(t_field), POINTER :: f_theta_rhodz(:)
      TYPE(t_field), POINTER :: f_theta(:)

      REAL(rstd), POINTER :: ps(:)
      REAL(rstd), POINTER :: theta_rhodz(:, :)
      REAL(rstd), POINTER :: theta(:, :)
      INTEGER :: ind

!$OMP BARRIER
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ps = f_ps(ind)
         theta_rhodz = f_theta_rhodz(ind)
         theta = f_theta(ind)
         CALL compute_theta_rhodz2theta(ps, theta_rhodz, theta, 0, ondevice=.false.)
      END DO
!$OMP BARRIER

   END SUBROUTINE theta_rhodz2theta

   SUBROUTINE theta_rhodz2temperature(f_ps, f_theta_rhodz, f_temp)
      USE icosa
      USE pression_mod, ONLY: compute_pression
      IMPLICIT NONE
      TYPE(t_field), POINTER :: f_ps(:)
      TYPE(t_field), POINTER :: f_theta_rhodz(:)
      TYPE(t_field), POINTER :: f_temp(:)

      REAL(rstd), POINTER :: ps(:)
      REAL(rstd), POINTER :: theta_rhodz(:, :, :)
      REAL(rstd), POINTER :: temp(:, :)
      REAL(rstd), POINTER :: p(:, :)
      INTEGER :: ind

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ps = f_ps(ind)
         p = f_p(ind)
         theta_rhodz = f_theta_rhodz(ind)
         temp = f_temp(ind)

!$OMP BARRIER
         CALL compute_pression(ps, p, 0)
!$OMP BARRIER
         CALL compute_theta_rhodz2temperature(p, theta_rhodz(:, :, 1), temp, 0)
      END DO
!$OMP BARRIER

   END SUBROUTINE theta_rhodz2temperature

   SUBROUTINE temperature2theta_rhodz(f_ps, f_temp, f_theta_rhodz)
      USE icosa
      USE pression_mod, ONLY: compute_pression
      IMPLICIT NONE
      TYPE(t_field), POINTER :: f_ps(:)
      TYPE(t_field), POINTER :: f_theta_rhodz(:)
      TYPE(t_field), POINTER :: f_temp(:)

      REAL(rstd), POINTER :: ps(:)
      REAL(rstd), POINTER :: theta_rhodz(:, :, :)
      REAL(rstd), POINTER :: temp(:, :)
      REAL(rstd), POINTER :: p(:, :)
      INTEGER :: ind

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ps = f_ps(ind)
         p = f_p(ind)
         theta_rhodz = f_theta_rhodz(ind)
         temp = f_temp(ind)

!$OMP BARRIER
         CALL compute_pression(ps, p, 0)
!$OMP BARRIER
         CALL compute_temperature2theta_rhodz(p, temp, theta_rhodz(:, :, 1), 0)
      END DO
!$OMP BARRIER

   END SUBROUTINE temperature2theta_rhodz

   SUBROUTINE theta2theta_rhodz(f_ps, f_theta, f_theta_rhodz)
      USE icosa
      IMPLICIT NONE
      TYPE(t_field), POINTER :: f_ps(:)
      TYPE(t_field), POINTER :: f_theta(:)
      TYPE(t_field), POINTER :: f_theta_rhodz(:)

      REAL(rstd), POINTER :: ps(:)
      REAL(rstd), POINTER :: theta(:, :)
      REAL(rstd), POINTER :: theta_rhodz(:, :)
      INTEGER :: ind

!$OMP BARRIER
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ps = f_ps(ind)
         theta = f_theta(ind)
         theta_rhodz = f_theta_rhodz(ind)
         CALL compute_theta2theta_rhodz(ps, theta, theta_rhodz, 0, ondevice=.false.)
      END DO
!$OMP BARRIER

   END SUBROUTINE theta2theta_rhodz

   SUBROUTINE compute_theta2theta_rhodz(ps, theta, theta_rhodz, offset, ondevice)
      USE icosa
      USE disvert_mod
      USE omp_para
      IMPLICIT NONE
      REAL(rstd), INTENT(IN) :: ps(iim*jjm)
      REAL(rstd), INTENT(IN) :: theta(iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: theta_rhodz(iim*jjm, llm)
      INTEGER, INTENT(IN) :: offset
      LOGICAL, INTENT(IN) :: ondevice
      REAL(rstd) :: rhodz
      INTEGER :: i, j, ij, l

!$OMP BARRIER
      !$acc parallel loop collapse(3) default(present) async if(ondevice)
      DO l = ll_begin, ll_end
         DO j = jj_begin - offset, jj_end + offset
            DO i = ii_begin - offset, ii_end + offset
               ij = (j - 1)*iim + i
               rhodz = (ap(l) - ap(l + 1) + (bp(l) - bp(l + 1))*ps(ij))/g
               theta_rhodz(ij, l) = theta(ij, l)*rhodz
            END DO
         END DO
      END DO
!$OMP BARRIER

   END SUBROUTINE compute_theta2theta_rhodz

   SUBROUTINE compute_theta_rhodz2theta(ps, theta_rhodz, theta, offset, ondevice)
      USE icosa
      USE disvert_mod
      USE omp_para
      IMPLICIT NONE
      REAL(rstd), INTENT(IN) :: ps(iim*jjm)
      REAL(rstd), INTENT(IN) :: theta_rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: theta(iim*jjm, llm)
      INTEGER, INTENT(IN) :: offset
      LOGICAL, INTENT(IN) :: ondevice
      REAL(rstd) :: rhodz
      INTEGER :: i, j, ij, l

!$OMP BARRIER
      !$acc parallel loop collapse(3) default(present) async if(ondevice)
      DO l = ll_begin, ll_end
         DO j = jj_begin - offset, jj_end + offset
            DO i = ii_begin - offset, ii_end + offset
               ij = (j - 1)*iim + i
               rhodz = (ap(l) - ap(l + 1) + (bp(l) - bp(l + 1))*ps(ij))/g
               theta(ij, l) = theta_rhodz(ij, l)/rhodz
            END DO
         END DO
      END DO
!$OMP BARRIER

   END SUBROUTINE compute_theta_rhodz2theta

   SUBROUTINE compute_theta_rhodz2temperature(p, theta_rhodz, temp, offset)
      USE icosa
      USE omp_para
      IMPLICIT NONE
      REAL(rstd), INTENT(IN) :: p(iim*jjm, llm + 1)
      REAL(rstd), INTENT(IN) :: theta_rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: temp(iim*jjm, llm)
      INTEGER, INTENT(IN) :: offset
      REAL(rstd) :: pk_ij
      INTEGER :: i, j, ij, l

! flush p
!$OMP BARRIER
      DO l = ll_begin, ll_end
         DO j = jj_begin - offset, jj_end + offset
            DO i = ii_begin - offset, ii_end + offset
               ij = (j - 1)*iim + i
               pk_ij = ((.5/preff)*(p(ij, l) + p(ij, l + 1)))**kappa
               temp(ij, l) = (theta_rhodz(ij, l)/((p(ij, l) - p(ij, l + 1))/g))*pk_ij
            END DO
         END DO
      END DO
!$OMP BARRIER

   END SUBROUTINE compute_theta_rhodz2temperature

   SUBROUTINE compute_temperature2theta_rhodz(p, temp, theta_rhodz, offset)
      USE icosa
      USE exner_mod
      USE omp_para
      IMPLICIT NONE
      REAL(rstd), INTENT(IN)  :: p(iim*jjm, llm + 1)
      REAL(rstd), INTENT(OUT) :: theta_rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(IN)  :: temp(iim*jjm, llm)
      INTEGER, INTENT(IN) :: offset
      REAL(rstd) :: pk_ij
      INTEGER :: i, j, ij, l

! flush p
!$OMP BARRIER

      DO l = ll_begin, ll_end
         DO j = jj_begin - offset, jj_end + offset
            DO i = ii_begin - offset, ii_end + offset
               ij = (j - 1)*iim + i
               pk_ij = ((.5/preff)*(p(ij, l) + p(ij, l + 1)))**kappa
               theta_rhodz(ij, l) = temp(ij, l)*((p(ij, l) - p(ij, l + 1))/g)/pk_ij
            END DO
         END DO
      END DO
!$OMP BARRIER

   END SUBROUTINE compute_temperature2theta_rhodz

END MODULE theta2theta_rhodz_mod
