MODULE caldyn_unstructured_mod
  USE ISO_C_BINDING
  USE OMP_LIB
  USE data_unstructured_mod
  IMPLICIT NONE
  SAVE

CONTAINS

#include "unstructured.h90"

#define INDICES1 ij,l,iq,iedge,edge,ivertex,vertex,ij_left,ij_right
#define INDICES2 ij_up,ij_down,itrisk,edge_trisk,kup,kdown
#define EDGES edge1,edge2,edge3,edge4,edge5,edge6
#define VERTICES vertex1,vertex2,vertex3,vertex4,vertex5,vertex6
#define SIGNS sign1,sign2,sign3,sign4,sign5,sign6
#define EDGE_ENDS ij_up1,ij_up2,ij_up3,ij_up4,ij_up5,ij_up6,ij_down1,ij_down2,ij_down3,ij_down4,ij_down5,ij_down6
#define LENGTHS le_de1,le_de2,le_de3,le_de4,le_de5,le_de6
#define DECLARE_INDICES INTEGER INDICES1,INDICES2
#define DECLARE_EDGES NUM SIGNS,LENGTHS ; INTEGER EDGES, EDGE_ENDS
#define DECLARE_VERTICES INTEGER VERTICES
#define PHI_BOT(ij) Phi_bot

#define HASNAN(field) (ANY(.NOT.ABS(field)<1e20))

#define START_TRACE(id,nprimal,ndual,nedge) CALL enter_trace(id, 8*llm*((nprimal)*primal_num+(ndual)*dual_num+(nedge)*edge_num) )
#define STOP_TRACE CALL exit_trace()

!----------------------------- Non-Hydrostatic -----------------------------

SUBROUTINE compute_NH_geopot(tau, m_ik, m_il, theta, W_il, Phi_il)
  FIELD_MASS   :: m_ik, theta   ! IN*2
  FIELD_GEOPOT :: m_il, W_il, Phi_il, Phi_star_il  ! IN,INOUT*2, LOCAL
  NUM :: tau, gamma, tau2_g, tau2_g2, g2, gm2, vreff, Rd_preff
  INTEGER :: iter
  LOGICAL :: debug_hevi_solver
  DECLARE_INDICES
  NUM :: rho_ij, X_ij, Y_ij, wil, rho_c2_mik, c2_mik, ml_g2
#define COLUMN 0
#if COLUMN 
  NUM1(llm)  :: pk, Ak, Ck
  NUM1(llm+1):: Rl, Bl, Dl, xl
#define p_ik(l,ij) pk(l)
#define A_ik(l,ij) Ak(l)
#define C_ik(l,ij) Ck(l)
#define R_il(l,ij) Rl(l)
#define B_il(l,ij) Bl(l)
#define D_il(l,ij) Dl(l)
#define x_il(l,ij) xl(l)
#else
  FIELD_MASS :: p_ik, A_ik, C_ik
  FIELD_GEOPOT :: R_il, B_il, D_il, x_il
#endif

  debug_hevi_solver=.FALSE.
!$OMP MASTER
  debug_hevi_solver = debug_hevi_solver_
!$OMP END MASTER

  START_TRACE(id_NH_geopot, 7,0,0)
#include "../kernels_unst/compute_NH_geopot.k90"
  STOP_TRACE  

!$OMP MASTER
  debug_hevi_solver_ = debug_hevi_solver
!$OMP END MASTER

#if COLUMN
#undef p_ik
#undef A_ik
#undef C_ik
#undef R_il
#undef B_il
#undef D_il
#undef x_il
#endif
#undef COLUMN
END SUBROUTINE compute_NH_geopot

SUBROUTINE compute_caldyn_slow_NH(u,rhodz,Phi,W, w_il,berni,gradPhi2,DePhil,v_el,G_el,F_el, hflux,du,dPhi,dW)
  FIELD_U      :: u, hflux, du   ! IN, OUT, OUT
  FIELD_MASS   :: rhodz, berni   ! IN, BUF
  FIELD_GEOPOT :: Phi,W,dPhi,dW, w_il, gradPhi2  ! IN,IN, OUT,OUT, BUF*2
  FIELD_UL     :: DePhil, v_el, G_el, F_el ! BUF*4
  DECLARE_INDICES
  DECLARE_EDGES
  NUM :: W_el, W2_el, gPhi2, dP, divG, u2, uu
  START_TRACE(id_slow_NH, 5,0,3)
#include "../kernels_unst/caldyn_slow_NH.k90"
  STOP_TRACE
END SUBROUTINE compute_caldyn_slow_NH

SUBROUTINE compute_caldyn_solver(tau,rhodz,theta, berni,pres,m_il, pk,geopot,W,dPhi,dW,du)
  NUM, INTENT(IN) :: tau
  FIELD_MASS   :: rhodz,pk,berni,pres    ! IN, OUT, BUF*2
  FIELD_THETA  :: theta                  ! IN
  FIELD_GEOPOT :: geopot,W,dPhi,dW, m_il ! INOUT,INOUT, OUT,OUT, BUF
  FIELD_U      :: du                     ! OUT
  DECLARE_INDICES
  NUM :: X_ij, rho_ij, T_ij, gamma, Cvd, vreff, Rd_preff
#include "../kernels_unst/caldyn_mil.k90"
  IF(tau>0) THEN ! solve implicit problem for geopotential
    CALL compute_NH_geopot(tau, rhodz, m_il, theta, W, geopot)
  END IF
  START_TRACE(id_solver, 7,0,1)
#include "../kernels_unst/caldyn_solver.k90"
  STOP_TRACE
END SUBROUTINE compute_caldyn_solver

SUBROUTINE compute_caldyn_vert_NH(mass,geopot,W,wflux, eta_dot,wcov,W_etadot, du,dPhi,dW)
  FIELD_MASS   :: mass, eta_dot, wcov, W_etadot ! IN, BUF*3
  FIELD_GEOPOT :: geopot,W,wflux,dPhi,dW ! IN*3, INOUT*2
  FIELD_U      :: du ! INOUT
  DECLARE_INDICES
  NUM :: w_ij, wflux_ij
  START_TRACE(id_vert_NH, 6,0,1)
#include "../kernels_unst/caldyn_vert_NH.k90"
  STOP_TRACE
END SUBROUTINE compute_caldyn_vert_NH

!----------------------------- Hydrostatic -------------------------------

SUBROUTINE compute_geopot(rhodz,theta,ps,pk,geopot)
  FIELD_MASS  :: rhodz,pk   ! IN, OUT
  FIELD_THETA :: theta      ! IN
  FIELD_GEOPOT :: geopot    ! IN(l=1)/OUT(l>1)
  FIELD_PS     :: ps        ! OUT
  DECLARE_INDICES
  NUM :: gdz, ke, uu, chi, gv, exner_ik, temp_ik, p_ik, qv, Rmix
  START_TRACE(id_geopot, 3,0,3)
#include "../kernels_unst/compute_geopot.k90"
  STOP_TRACE
END SUBROUTINE compute_geopot
!
SUBROUTINE compute_caldyn_slow_hydro(rhodz,theta,u, berni,hflux,du)
  FIELD_MASS  :: rhodz,berni! IN
  FIELD_THETA :: theta      ! IN
  FIELD_U     :: u,hflux,du ! IN, OUT, OUT
  DECLARE_INDICES
  DECLARE_EDGES
  LOGICAL, PARAMETER :: zero=.TRUE.
  NUM :: ke, uu
  START_TRACE(id_slow_hydro, 3,0,3)
#include "../kernels_unst/caldyn_slow_hydro.k90"
  STOP_TRACE
END SUBROUTINE compute_caldyn_slow_hydro

!---------------------------------- Generic ------------------------------

SUBROUTINE caldyn_vert(convm,rhodz,theta,u, dmass_col,wflux,dtheta_rhodz,du, wwuu)
  FIELD_PS    :: dmass_col
  FIELD_MASS  :: convm, rhodz
  FIELD_U     :: u,du
  FIELD_THETA :: dtheta_rhodz, theta
  FIELD_GEOPOT :: wflux
  FIELD_UL     :: wwuu
  DECLARE_INDICES
  NUM :: dF_deta, dFu_deta
  wwuu=0.
  !$OMP BARRIER
  START_TRACE(id_vert, 5,0,3)
#include "../kernels_unst/caldyn_wflux.k90"
#include "../kernels_unst/caldyn_dmass.k90"
#include "../kernels_unst/caldyn_vert.k90"
  STOP_TRACE
END SUBROUTINE caldyn_vert

SUBROUTINE compute_coriolis(hflux,theta,qu,Ftheta, convm,dtheta_rhodz,du)
  FIELD_U     :: hflux, Ftheta, qu, du
  FIELD_MASS  :: convm
  FIELD_THETA :: theta, dtheta_rhodz
  DECLARE_INDICES
  DECLARE_EDGES
  NUM :: divF, du_trisk
  START_TRACE(id_coriolis, 3,4,0) ! primal, dual, edge
#include "../kernels_unst/coriolis.k90"
  STOP_TRACE
END SUBROUTINE

SUBROUTINE compute_theta(mass_col,rhodz,theta_rhodz, theta)
  FIELD_PS :: mass_col
  FIELD_MASS :: rhodz
  FIELD_THETA :: theta, theta_rhodz
  DECLARE_INDICES
  NUM :: m
  START_TRACE(id_theta, 3,0,0) ! primal, dual, edge
#include "../kernels_unst/theta.k90"
  STOP_TRACE
END SUBROUTINE

SUBROUTINE compute_pvort_only(rhodz,u,qv,qu)
  FIELD_MASS :: rhodz
  FIELD_U    :: u,qu
  FIELD_Z    :: qv
  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

SUBROUTINE compute_caldyn_fast(tau, pk,berni,theta,geopot, du,u)
  NUM, INTENT(IN) :: tau
  FIELD_MASS   :: pk,berni  ! INOUT, OUT
  FIELD_THETA  :: theta     ! IN
  FIELD_GEOPOT :: geopot    ! IN
  FIELD_U      :: u,du      ! INOUT,INOUT
  DECLARE_INDICES
  DECLARE_EDGES
  NUM          :: due
  START_TRACE(id_fast, 4,0,2) ! primal, dual, edge
#include "../kernels_unst/caldyn_fast.k90"
  STOP_TRACE
END SUBROUTINE compute_caldyn_fast

!----------------------------- Unused -----------------------------

SUBROUTINE gradient(b,grad) BINDC(gradient)
  FIELD_MASS :: b
  FIELD_U    :: grad
  DECLARE_INDICES
#include "../kernels_unst/gradient.k90"
END SUBROUTINE
!
SUBROUTINE div(u,divu) BINDC(div)
  FIELD_MASS :: divu
  FIELD_U    :: u
  DECLARE_INDICES
  DECLARE_EDGES
  NUM :: div_ij
  !$OMP PARALLEL NUM_THREADS(nb_threads)
#include "../kernels_unst/div.k90"
  !$OMP END PARALLEL
END SUBROUTINE

SUBROUTINE compute_scalar_laplacian(b,grad,divu)
  FIELD_MASS   :: b, divu ! IN, OUT
  FIELD_U      :: grad ! BUF
  DECLARE_INDICES
  DECLARE_EDGES
  NUM :: div_ij
  START_TRACE(id_scalar_laplacian, 0,0,1)
#include "../kernels_unst/scalar_laplacian.k90"
  STOP_TRACE
END SUBROUTINE compute_scalar_laplacian

END MODULE caldyn_unstructured_mod