MODULE compute_pression_mod
  USE compute_diagnostics_mod
  USE icosa
  USE omp_para
  USE disvert_mod, ONLY : ap, bp, ap_bp_present, ptop, caldyn_eta, eta_lag
  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, &
       hydrostatic_pressure, compute_hydrostatic_pressure_unst, compute_hydrostatic_pressure_hex

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

  SUBROUTINE hydrostatic_pressure(f_rhodz, f_theta_rhodz, f_ps, f_p)
    TYPE(t_field), POINTER :: f_rhodz(:), f_theta_rhodz(:), f_ps(:), f_p(:)
    REAL(rstd), POINTER :: ps(:), rhodz(:,:), p(:,:), theta_rhodz(:,:,:)
    INTEGER :: ind
    DO ind=1,ndomain
      IF (.NOT. assigned_domain(ind)) CYCLE
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      rhodz=f_rhodz(ind)
      theta_rhodz=f_theta_rhodz(ind)
      ps=f_ps(ind)
      p=f_p(ind)
      CALL compute_hydrostatic_pressure(rhodz, theta_rhodz, ps, p)
    ENDDO
  END SUBROUTINE hydrostatic_pressure

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

  SUBROUTINE compute_pression_hex(ps,p,offset)
    REAL(rstd),INTENT(IN) :: ps(:)
    REAL(rstd),INTENT(OUT) :: p(:,:)
    INTEGER,INTENT(IN) :: offset
    CALL compute_pression_hex_(ps,p,offset)
  END SUBROUTINE compute_pression_hex

  SUBROUTINE compute_pression_unst(ps,p,offset)
    REAL(rstd),INTENT(IN) :: ps(:)
    REAL(rstd),INTENT(OUT) :: p(:,:)
    INTEGER,INTENT(IN) :: offset
    CALL compute_pression_unst_(ps,p,offset)
  END SUBROUTINE compute_pression_unst

  SUBROUTINE compute_pression_mid_hex(ps,p,offset)
    REAL(rstd),INTENT(IN) :: ps(:)
    REAL(rstd),INTENT(OUT) :: p(:,:)
    INTEGER,INTENT(IN) :: offset
    CALL compute_pression_mid_hex_(ps,p,offset)
  END SUBROUTINE compute_pression_mid_hex

  SUBROUTINE compute_pression_mid_unst(ps,p,offset)
    REAL(rstd),INTENT(IN) :: ps(:)
    REAL(rstd),INTENT(OUT) :: p(:,:)
    INTEGER,INTENT(IN) :: offset
    CALL compute_pression_mid_unst_(ps,p,offset)
  END SUBROUTINE compute_pression_mid_unst

  SUBROUTINE compute_hydrostatic_pressure_hex(rhodz, theta_rhodz, ps, p)
    REAL(rstd),INTENT(IN)  :: rhodz(:,:), theta_rhodz(:,:,:)
    REAL(rstd),INTENT(OUT) :: ps(:), p(:,:)
    CALL compute_hydrostatic_pressure_hex_(rhodz, theta_rhodz, ps, p)
  END SUBROUTINE compute_hydrostatic_pressure_hex

  SUBROUTINE compute_hydrostatic_pressure_unst(rhodz, theta_rhodz, ps, p)
    REAL(rstd),INTENT(IN)  :: rhodz(:,:), theta_rhodz(:,:,:)
    REAL(rstd),INTENT(OUT) :: ps(:), p(:,:)
    CALL compute_hydrostatic_pressure_unst_(rhodz, theta_rhodz, ps, p)
  END SUBROUTINE compute_hydrostatic_pressure_unst
  
!------------- 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

  SUBROUTINE compute_hydrostatic_pressure_hex_(rhodz, theta_rhodz, ps, pk)
    REAL(rstd),INTENT(IN)  :: rhodz(iim*jjm,llm) ! mass per unit surface in each model level
    REAL(rstd),INTENT(IN)  :: theta_rhodz(iim*jjm,llm, nqdyn) ! dynamical tracers (theta/entropy)
    REAL(rstd),INTENT(OUT) :: ps(iim*jjm)        ! surface pressure, diagnosed if Lagrangian vertical coordinate 
    REAL(rstd),INTENT(OUT) :: pk(iim*jjm,llm)    ! pressure at full levels
    INTEGER :: ij,l, ij_omp_begin_ext, ij_omp_end_ext
    !$OMP BARRIER
    CALL distrib_level(ij_begin_ext,ij_end_ext, ij_omp_begin_ext,ij_omp_end_ext)
#include "../kernels_hex/compute_hydrostatic_pressure.k90"
    !$OMP BARRIER
  END SUBROUTINE compute_hydrostatic_pressure_hex_

!----------- 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

  SUBROUTINE compute_hydrostatic_pressure_unst_(rhodz, theta_rhodz, ps, pk)
    FIELD_MASS,  INTENT(IN)  :: rhodz
    FIELD_THETA, INTENT(IN)  :: theta_rhodz
    FIELD_PS,    INTENT(OUT) :: ps
    FIELD_MASS,  INTENT(OUT) :: pk
    DECLARE_INDICES
#include "../kernels_unst/compute_hydrostatic_pressure.k90"
  END SUBROUTINE compute_hydrostatic_pressure_unst_
  

END MODULE compute_pression_mod