MODULE compute_rhodz_mod
  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

#include "../unstructured/unstructured.h90"

  PUBLIC :: compute_rhodz_hex, compute_rhodz_unst

CONTAINS

#if BEGIN_DYSL

KERNEL(compute_rhodz)
IF(comp) THEN
   FORALL_CELLS()
     ON_PRIMAL
       m = ( AP(CELL)-AP(UP(CELL)) + (BP(CELL)-BP(UP(CELL)))*ps(HIDX(CELL)) )/g
       rhodz(CELL)=m
     END_BLOCK
   END_BLOCK
ELSE
   err=0.
   FORALL_CELLS_EXT()
     ON_PRIMAL
       m = ( AP(CELL)-AP(UP(CELL)) + (BP(CELL)-BP(UP(CELL)))*ps(HIDX(CELL)) )/g
       err = MAX(err, ABS(m-rhodz(CELL)))
     END_BLOCK
   END_BLOCK
   IF(err>1e-10) THEN
      PRINT *, 'Discrepancy between ps and rhodz detected', err
      STOP
   END IF
END IF
END_BLOCK

#endif END_DYSL

  SUBROUTINE compute_rhodz_unst(comp, ps, rhodz)
    USE data_unstructured_mod, ONLY : primal_num
    LOGICAL, INTENT(IN)  :: comp
    FIELD_PS, INTENT(IN) :: ps
    FIELD_MASS, INTENT(INOUT) :: rhodz
    DECLARE_INDICES
    NUM :: m, err
#define AP(l,ij) ap(l)
#define BP(l,ij) bp(l)
#include "../kernels_unst/compute_rhodz.k90"
#undef AP
#undef BP
  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
#define AP(ij,l) ap(l)
#define BP(ij,l) bp(l)
#include "../kernels_hex/compute_rhodz.k90"
#undef AP
#undef BP
  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
          ENDDO
       ENDDO
    ENDDO

    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