MODULE compute_rhodz_mod
   USE icosa, ONLY: rstd
   USE earth_const, ONLY: g
   USE disvert_mod, ONLY: ap, bp
   USE grid_param, ONLY: llm
   USE ISO_C_BINDING, only: C_DOUBLE, C_FLOAT
   IMPLICIT NONE
   PRIVATE
   SAVE

   PUBLIC :: compute_rhodz_hex, compute_rhodz_unst

CONTAINS

!-------------- Wrappers for F2008 conformity -----------------
!--------------------------------------------------------------

   SUBROUTINE compute_rhodz_hex(comp, ps, rhodz)
      LOGICAL, INTENT(IN) :: comp ! .TRUE. to compute, .FALSE. to check
      REAL(rstd), INTENT(IN) :: ps(:)
      REAL(rstd), INTENT(INOUT) :: rhodz(:, :)
      CALL compute_rhodz_hex_(comp, ps, rhodz)
   END SUBROUTINE compute_rhodz_hex

   SUBROUTINE compute_rhodz_unst(comp, ps, rhodz)
      LOGICAL, INTENT(IN) :: comp ! .TRUE. to compute, .FALSE. to check
      REAL(rstd), INTENT(IN) :: ps(:)
      REAL(rstd), INTENT(INOUT) :: rhodz(:, :)
      CALL compute_rhodz_unst_(comp, ps, rhodz)
   END SUBROUTINE compute_rhodz_unst

!--------------------------------------------------------------

   SUBROUTINE compute_rhodz_unst_(comp, ps, rhodz)
      USE data_unstructured_mod, ONLY: primal_num
      LOGICAL, INTENT(IN)  :: comp
      REAL(C_DOUBLE), DIMENSION(primal_num), INTENT(IN) :: ps
      REAL(C_DOUBLE), DIMENSION(llm, primal_num), INTENT(INOUT) :: 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, err

      !--------------------------------------------------------------------------
      !---------------------------- compute_rhodz ----------------------------------
      IF (comp) THEN
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, primal_num
            !DIR$ SIMD
            DO l = 1, llm
               m = (ap(l) - ap(l + 1) + (bp(l) - bp(l + 1))*ps(ij))/g
               rhodz(l, ij) = m
            END DO
         END DO
         !$OMP END DO
      ELSE
         err = 0.
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, primal_num
            !DIR$ SIMD
            DO l = 1, llm
               m = (ap(l) - ap(l + 1) + (bp(l) - bp(l + 1))*ps(ij))/g
               err = MAX(err, ABS(m - rhodz(l, ij)))
            END DO
         END DO
         !$OMP END DO
         IF (err > 1e-10) THEN
            PRINT *, 'Discrepancy between ps and rhodz detected', err
            STOP
         END IF
      END IF
      !---------------------------- compute_rhodz ----------------------------------
      !--------------------------------------------------------------------------

   END SUBROUTINE compute_rhodz_unst_

   SUBROUTINE compute_rhodz_hex_(comp, ps, rhodz)
      USE icosa
      USE omp_para
      LOGICAL, INTENT(IN) :: comp ! .TRUE. to compute, .FALSE. to check
      REAL(rstd), INTENT(IN) :: ps(iim*jjm)
      REAL(rstd), INTENT(INOUT) :: rhodz(iim*jjm, llm)
      REAL(rstd) :: m, err
      INTEGER :: ij, l

      !--------------------------------------------------------------------------
      !---------------------------- compute_rhodz ----------------------------------
      IF (comp) THEN
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin, ij_end
               m = (ap(l) - ap(l + 1) + (bp(l) - bp(l + 1))*ps(ij))/g
               rhodz(ij, l) = m
            END DO
         END DO
      ELSE
         err = 0.
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               m = (ap(l) - ap(l + 1) + (bp(l) - bp(l + 1))*ps(ij))/g
               err = MAX(err, ABS(m - rhodz(ij, l)))
            END DO
         END DO
         IF (err > 1e-10) THEN
            PRINT *, 'Discrepancy between ps and rhodz detected', err
            STOP
         END IF
      END IF
      !---------------------------- compute_rhodz ----------------------------------
      !--------------------------------------------------------------------------

   END SUBROUTINE compute_rhodz_hex_

   SUBROUTINE compute_rhodz_handmade(comp, ps, rhodz)
      USE icosa
      USE omp_para
      LOGICAL, INTENT(IN) :: comp ! .TRUE. to compute, .FALSE. to check
      REAL(rstd), INTENT(IN) :: ps(iim*jjm)
      REAL(rstd), INTENT(INOUT) :: rhodz(iim*jjm, llm)
      REAL(rstd) :: m, err
      INTEGER :: l, i, j, ij, dd
      err = 0.

      IF (comp) THEN
         dd = 1
      ELSE
         dd = 0
      END IF

      DO l = ll_begin, ll_end
         DO j = jj_begin - dd, jj_end + dd
            DO i = ii_begin - dd, ii_end + dd
               ij = (j - 1)*iim + i
               m = (ap(l) - ap(l + 1) + (bp(l) - bp(l + 1))*ps(ij))/g
               IF (comp) THEN
                  rhodz(ij, l) = m
               ELSE
                  err = MAX(err, abs(m - rhodz(ij, l)))
               END IF
            END DO
         END DO
      END DO

      IF (.NOT. comp) THEN
         IF (err > 1e-10) THEN
            PRINT *, 'Discrepancy between ps and rhodz detected', err
            STOP
         ELSE
!          PRINT *, 'No discrepancy between ps and rhodz detected'
         END IF
      END IF

   END SUBROUTINE compute_rhodz_handmade

END MODULE compute_rhodz_mod
