MODULE compute_pvort_only_mod
  USE compute_mod, ONLY : comp_pvort_only
  USE grid_param, ONLY : llm
  IMPLICIT NONE
  PRIVATE

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

  PUBLIC :: compute_pvort_only_hex, compute_pvort_only_unst

CONTAINS

  SUBROUTINE check_interface
    PROCEDURE(comp_pvort_only), POINTER :: ptr
    ptr => compute_pvort_only_unst
    ptr => compute_pvort_only_hex
  END SUBROUTINE check_interface
  
  SUBROUTINE compute_pvort_only_unst(u,rhodz,qu,qv, hv_)
    USE ISO_C_BINDING, only : C_DOUBLE, C_FLOAT
    USE data_unstructured_mod, ONLY : enter_trace, exit_trace, &
         id_pvort_only, primal_num, dual_num, edge_num, &
         dual_deg, dual_edge, dual_ne, dual_vertex, up, down, Av, fv, Riv2
    FIELD_MASS :: rhodz
    FIELD_U    :: u,qu
    FIELD_Z    :: qv, hv_
    DECLARE_INDICES
    DECLARE_EDGES
    DECLARE_VERTICES
    NUM :: etav, hv
    START_TRACE(id_pvort_only, 1,1,2) ! primal, dual, edge
#include "../kernels_unst/pvort_only.k90"
    STOP_TRACE
  END SUBROUTINE compute_pvort_only_unst

  SUBROUTINE compute_pvort_only_hex(u,rhodz,qu,qv,hv_)
    USE icosa
    USE trace, ONLY : trace_start, trace_end
    USE caldyn_vars_mod, ONLY : dysl_pvort_only
    USE omp_para, ONLY : ll_begin, ll_end
    REAL(rstd),INTENT(IN)  :: u(iim*3*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: qu(iim*3*jjm,llm)
    REAL(rstd),INTENT(OUT) :: qv(iim*2*jjm,llm)
    REAL(rstd),INTENT(OUT) :: hv_(iim*2*jjm,llm)

    INTEGER :: ij,l
    REAL(rstd) :: etav,hv,radius_m2

    CALL trace_start("compute_pvort_only")  
!!! Compute shallow-water potential vorticity
    IF(dysl_pvort_only) THEN
#include "../kernels_hex/pvort_only.k90"
    ELSE

    radius_m2=radius**(-2)
    DO l = ll_begin,ll_end
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext
          etav= 1./Av(ij+z_up)*(  ne_rup  * u(ij+u_rup,l)   &
                  + ne_left * u(ij+t_rup+u_left,l)          &
                  - ne_lup  * u(ij+u_lup,l) )                               
          hv =   Riv2(ij,vup)          * rhodz(ij,l)        &
               + Riv2(ij+t_rup,vldown) * rhodz(ij+t_rup,l)  &
               + Riv2(ij+t_lup,vrdown) * rhodz(ij+t_lup,l)
          qv(ij+z_up,l) = ( etav+fv(ij+z_up) )/hv
          hv_(ij+z_up,l) = hv
          
          etav = 1./Av(ij+z_down)*(  ne_ldown * u(ij+u_ldown,l)   &
               + ne_right * u(ij+t_ldown+u_right,l)               &
               - ne_rdown * u(ij+u_rdown,l) )
          hv =   Riv2(ij,vdown)        * rhodz(ij,l)              &
               + Riv2(ij+t_ldown,vrup) * rhodz(ij+t_ldown,l)      &
               + Riv2(ij+t_rdown,vlup) * rhodz(ij+t_rdown,l)
          qv(ij+z_down,l) =( etav+fv(ij+z_down) )/hv
          hv_(ij+z_down,l) = hv
       ENDDO

       !DIR$ SIMD
       DO ij=ij_begin,ij_end
          qu(ij+u_right,l) = 0.5*(qv(ij+z_rdown,l)+qv(ij+z_rup,l))  
          qu(ij+u_lup,l) = 0.5*(qv(ij+z_up,l)+qv(ij+z_lup,l))  
          qu(ij+u_ldown,l) = 0.5*(qv(ij+z_ldown,l)+qv(ij+z_down,l))  
       END DO

    ENDDO
   
    END IF ! dysl
    CALL trace_end("compute_pvort_only")

  END SUBROUTINE compute_pvort_only_hex

END MODULE compute_pvort_only_mod