MODULE compute_pression_mod
  USE compute_diagnostics_mod
  USE icosa
  USE omp_para
  USE disvert_mod, ONLY : ap, bp, ap_bp_present
  IMPLICIT NONE
  PRIVATE

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

  PUBLIC :: pression, compute_pression_hex, compute_pression_unst, &
       pression_mid, compute_pression_mid_hex, compute_pression_mid_unst

CONTAINS

#ifdef BEGIN_DYSL

{%- macro compute_pression(llmax)%} 
{%- set inner_loop=caller() %}
{%- set llmax="'%s'"%llmax %}
IF(ap_bp_present) THEN
   IF(offset>0) THEN
     FORALL_CELLS_EXT('1',{{ llmax }})
       ON_PRIMAL
         {{ inner_loop }}
       END_BLOCK
     END_BLOCK
  ELSE
     FORALL_CELLS('1',{{ llmax }})
       ON_PRIMAL
         {{ inner_loop }}
       END_BLOCK
     END_BLOCK
  END IF
END IF
{%- endmacro %}

KERNEL(compute_pression)
{% call compute_pression('llm+1') %}
p(CELL) = AP(CELL) + BP(CELL) * ps(HIDX(CELL))  
{% endcall %}
END_BLOCK

KERNEL(compute_pmid)
{% call compute_pression('llm') %}
pmid(CELL) = .5*(AP(CELL)+AP(UP(CELL)) + (BP(CELL)+BP(UP(CELL))) * ps(HIDX(CELL)) )  
{% endcall %}
END_BLOCK

#endif END_DYSL

  SUBROUTINE pression(f_ps,f_p)
    TYPE(t_field), POINTER :: f_ps(:)
    TYPE(t_field), POINTER :: f_p(:)
  
    REAL(rstd), POINTER :: ps(:)
    REAL(rstd), POINTER :: p(:,:)
    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)
      p=f_p(ind)
      CALL compute_pression(ps, p,0)
    ENDDO
!$OMP BARRIER
  
  END SUBROUTINE pression

  SUBROUTINE pression_mid(f_ps,f_pmid)
    TYPE(t_field), POINTER :: f_ps(:)
    TYPE(t_field), POINTER :: f_pmid(:)
  
    REAL(rstd), POINTER :: ps(:)
    REAL(rstd), POINTER :: pmid(:,:)
    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)
      pmid=f_pmid(ind)
      CALL compute_pression_mid(ps, pmid,0)
    ENDDO
!$OMP BARRIER
  
  END SUBROUTINE pression_mid

!------------- hexagonal-mesh compute kernels --------

#define AP(ij,l) ap(l)
#define BP(ij,l) bp(l)

  SUBROUTINE compute_pression_hex(ps,p,offset)
    REAL(rstd),INTENT(IN) :: ps(iim*jjm)
    REAL(rstd),INTENT(OUT) :: p(iim*jjm,llm+1)
    INTEGER,INTENT(IN) :: offset
    INTEGER :: ij,l
#include "../kernels_hex/compute_pression.k90"
  END SUBROUTINE compute_pression_hex
  
  SUBROUTINE compute_pression_mid_hex(ps,pmid,offset)
    REAL(rstd),INTENT(IN) :: ps(iim*jjm)
    REAL(rstd),INTENT(OUT) :: pmid(iim*jjm,llm)
    INTEGER,INTENT(IN) :: offset
    INTEGER :: ij,l
#include "../kernels_hex/compute_pmid.k90"
  END SUBROUTINE compute_pression_mid_hex

#undef AP
#undef BP

!----------- unstructured-mesh compute kernels --------

#define AP(l,ij) ap(l)
#define BP(l,ij) bp(l)
  
  SUBROUTINE compute_pression_unst(ps, p, offset)
    FIELD_PS,     INTENT(IN)  :: ps
    FIELD_GEOPOT, INTENT(OUT) :: p
    INTEGER,      INTENT(IN)  :: offset
    DECLARE_INDICES
#include "../kernels_unst/compute_pression.k90"
  END SUBROUTINE compute_pression_unst

  SUBROUTINE compute_pression_mid_unst(ps, pmid, offset)
    FIELD_PS,   INTENT(IN)  :: ps
    FIELD_MASS, INTENT(OUT) :: pmid
    INTEGER,    INTENT(IN)  :: offset
    DECLARE_INDICES
#include "../kernels_unst/compute_pmid.k90"
  END SUBROUTINE compute_pression_mid_unst

#undef AP
#undef BP

END MODULE compute_pression_mod