MODULE compute_caldyn_Coriolis_mod
  USE prec, ONLY : rstd
  USE caldyn_vars_mod, ONLY : caldyn_conserv, conserv_energy, conserv_enstrophy, conserv_gassmann
  USE grid_param
  USE earth_const
  USE disvert_mod
  USE omp_para
  USE trace
  IMPLICIT NONE
  PRIVATE

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

  PUBLIC :: compute_caldyn_Coriolis_manual, &
       compute_caldyn_Coriolis_unst, compute_caldyn_Coriolis_hex

CONTAINS

#ifdef BEGIN_DYSL

KERNEL(coriolis)
!
  DO iq=1,nqdyn
    FORALL_CELLS_EXT()
      ON_EDGES
        Ftheta(EDGE) = .5*(theta(CELL1,iq)+theta(CELL2,iq))*hflux(EDGE)
      END_BLOCK
    END_BLOCK
    FORALL_CELLS()
      ON_PRIMAL
        divF=0.
        FORALL_EDGES
          divF = divF + Ftheta(EDGE)*SIGN
        END_BLOCK
        dtheta_rhodz(CELL,iq) = -divF / AI
      END_BLOCK
    END_BLOCK
  END DO ! iq
!
  FORALL_CELLS()
    ON_PRIMAL
      divF=0.
      FORALL_EDGES
        divF = divF + hflux(EDGE)*SIGN
      END_BLOCK
      convm(CELL) = -divF / AI
    END_BLOCK
  END_BLOCK
  !
  SELECT CASE(caldyn_conserv)
  CASE(conserv_energy) ! energy-conserving TRiSK
     FORALL_CELLS()
       ON_EDGES
         du_trisk=0.
         FORALL_TRISK
           du_trisk = du_trisk + WEE*hflux(EDGE_TRISK)*(qu(EDGE)+qu(EDGE_TRISK))
         END_BLOCK
         du(EDGE) = du(EDGE) + .5*du_trisk
       END_BLOCK
     END_BLOCK
  CASE(conserv_enstrophy) ! enstrophy-conserving TRiSK
     FORALL_CELLS()
       ON_EDGES
         du_trisk=0.
         FORALL_TRISK
           du_trisk = du_trisk + WEE*hflux(EDGE_TRISK)
         END_BLOCK
         du(EDGE) = du(EDGE) + du_trisk*qu(EDGE)
       END_BLOCK
     END_BLOCK
  END SELECT
END_BLOCK

#endif END_DYSL

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

    SUBROUTINE compute_caldyn_coriolis_unst(hflux,theta,qu, Ftheta, convm,dtheta_rhodz,du)
      REAL(rstd), INTENT(IN)    :: hflux(:,:), theta(:,:,:), qu(:,:)
      REAL(rstd), INTENT(OUT)   :: Ftheta(:,:), convm(:,:), dtheta_rhodz(:,:,:)
      REAL(rstd), INTENT(INOUT) :: du(:,:)
      CALL compute_caldyn_coriolis_unst_(hflux,theta,qu, Ftheta, convm,dtheta_rhodz,du)
    END SUBROUTINE compute_caldyn_coriolis_unst

    SUBROUTINE compute_caldyn_coriolis_hex(hflux,theta,qu, Ftheta, convm,dtheta_rhodz,du)
      REAL(rstd), INTENT(IN)    :: hflux(:,:), theta(:,:,:), qu(:,:)
      REAL(rstd), INTENT(OUT)   :: Ftheta(:,:), convm(:,:), dtheta_rhodz(:,:,:)
      REAL(rstd), INTENT(INOUT) :: du(:,:)
      CALL compute_caldyn_coriolis_hex_(hflux,theta,qu, Ftheta, convm,dtheta_rhodz,du)
    END SUBROUTINE compute_caldyn_coriolis_hex

!--------------------------------------------------------------

  SUBROUTINE compute_caldyn_coriolis_unst_(hflux,theta,qu, Ftheta, convm,dtheta_rhodz,du)
    USE ISO_C_BINDING, only : C_DOUBLE, C_FLOAT
    USE data_unstructured_mod, ONLY : enter_trace, exit_trace, &
         id_coriolis, left, right, primal_deg, primal_edge, primal_ne, &
         trisk_deg, trisk
         
    FIELD_U     :: hflux, Ftheta, qu, du ! IN, BUF, IN, INOUT
    FIELD_MASS  :: convm                 ! BUF
    FIELD_THETA :: theta, dtheta_rhodz   ! IN, OUT
    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 compute_caldyn_coriolis_unst_

  SUBROUTINE compute_caldyn_Coriolis_hex_(hflux,theta,qu, Ftheta, convm,dtheta_rhodz,du)
    USE icosa
    REAL(rstd),INTENT(IN)    :: hflux(3*iim*jjm,llm)  ! hflux in kg/s
    REAL(rstd),INTENT(IN)    :: theta(iim*jjm,llm,nqdyn) ! active scalars
    REAL(rstd),INTENT(IN)    :: qu(3*iim*jjm,llm)
    REAL(rstd), INTENT(OUT)  :: Ftheta(3*iim*jjm,llm)  ! potential temperature flux
    REAL(rstd),INTENT(OUT)   :: convm(iim*jjm,llm)  ! mass flux convergence
    REAL(rstd),INTENT(OUT)   :: dtheta_rhodz(iim*jjm,llm,nqdyn)
    REAL(rstd),INTENT(INOUT) :: du(3*iim*jjm,llm)
    REAL(rstd) :: uu_right, uu_lup, uu_ldown, du_trisk, divF
    INTEGER :: ij,iq,l,kdown
    CALL trace_start("compute_caldyn_Coriolis")

#include "../kernels_hex/coriolis.k90"

    CALL trace_end("compute_caldyn_Coriolis")
  END SUBROUTINE compute_caldyn_Coriolis_hex_

  SUBROUTINE compute_caldyn_Coriolis_manual(hflux,theta,qu, Ftheta, convm,dtheta_rhodz,du)
    USE icosa
    USE caldyn_vars_mod
    REAL(rstd),INTENT(IN)    :: hflux(3*iim*jjm,llm)  ! hflux in kg/s
    REAL(rstd),INTENT(IN)    :: theta(iim*jjm,llm,nqdyn) ! active scalars
    REAL(rstd),INTENT(IN)    :: qu(3*iim*jjm,llm)
    REAL(rstd),INTENT(IN)    :: Ftheta(3*iim*jjm,llm)  ! ignored in favor of local buffer
    REAL(rstd),INTENT(OUT)   :: convm(iim*jjm,llm)  ! mass flux convergence
    REAL(rstd),INTENT(OUT)   :: dtheta_rhodz(iim*jjm,llm,nqdyn)
    REAL(rstd),INTENT(INOUT) :: du(3*iim*jjm,llm)
    REAL(rstd)  :: buf_Ftheta(3*iim*jjm)  ! local buffer for potential temperature flux
    REAL(rstd) :: uu_right, uu_lup, uu_ldown, du_trisk, divF
    INTEGER :: ij,iq,l,kdown

    CALL trace_start("compute_caldyn_Coriolis")

#define FTHETA(ij) buf_Ftheta(ij)

    DO l=ll_begin, ll_end
       ! compute theta flux
       DO iq=1,nqdyn
       !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext      
             FTHETA(ij+u_right) = 0.5*(theta(ij,l,iq)+theta(ij+t_right,l,iq)) &
                                  * hflux(ij+u_right,l)
             FTHETA(ij+u_lup)   = 0.5*(theta(ij,l,iq)+theta(ij+t_lup,l,iq)) &
                  * hflux(ij+u_lup,l)
             FTHETA(ij+u_ldown) = 0.5*(theta(ij,l,iq)+theta(ij+t_ldown,l,iq)) &
                  * hflux(ij+u_ldown,l)
          END DO
          ! horizontal divergence of fluxes
       !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             ! dtheta_rhodz =  -div(flux.theta)
             dtheta_rhodz(ij,l,iq)= &
                  -1./Ai(ij)*(ne_right*FTHETA(ij+u_right)    +  &
                  ne_rup*FTHETA(ij+u_rup)        +  &  
                  ne_lup*FTHETA(ij+u_lup)        +  &  
                  ne_left*FTHETA(ij+u_left)      +  &  
                  ne_ldown*FTHETA(ij+u_ldown)    +  &
                  ne_rdown*FTHETA(ij+u_rdown) )
          END DO
       END DO

       !DIR$ SIMD
       DO ij=ij_begin,ij_end         
          ! convm = -div(mass flux), sign convention as in Ringler et al. 2012, eq. 21
          convm(ij,l)= -1./Ai(ij)*(ne_right*hflux(ij+u_right,l)   +  &
               ne_rup*hflux(ij+u_rup,l)       +  &  
               ne_lup*hflux(ij+u_lup,l)       +  &  
               ne_left*hflux(ij+u_left,l)     +  &  
               ne_ldown*hflux(ij+u_ldown,l)   +  &  
               ne_rdown*hflux(ij+u_rdown,l))
       END DO ! ij
    END DO ! llm

!!! Compute potential vorticity (Coriolis) contribution to du
    SELECT CASE(caldyn_conserv)

    CASE(conserv_energy) ! energy-conserving TRiSK

       DO l=ll_begin,ll_end
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             uu_right = &
                  wee(ij+u_right,1,1)*hflux(ij+u_rup,l)*(qu(ij+u_right,l)+qu(ij+u_rup,l))+                             &
                  wee(ij+u_right,2,1)*hflux(ij+u_lup,l)*(qu(ij+u_right,l)+qu(ij+u_lup,l))+                             &
                  wee(ij+u_right,3,1)*hflux(ij+u_left,l)*(qu(ij+u_right,l)+qu(ij+u_left,l))+                           &
                  wee(ij+u_right,4,1)*hflux(ij+u_ldown,l)*(qu(ij+u_right,l)+qu(ij+u_ldown,l))+                         &
                  wee(ij+u_right,5,1)*hflux(ij+u_rdown,l)*(qu(ij+u_right,l)+qu(ij+u_rdown,l))+                         & 
                  wee(ij+u_right,1,2)*hflux(ij+t_right+u_ldown,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_ldown,l))+         &
                  wee(ij+u_right,2,2)*hflux(ij+t_right+u_rdown,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_rdown,l))+         &
                  wee(ij+u_right,3,2)*hflux(ij+t_right+u_right,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_right,l))+         &
                  wee(ij+u_right,4,2)*hflux(ij+t_right+u_rup,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_rup,l))+             &
                  wee(ij+u_right,5,2)*hflux(ij+t_right+u_lup,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_lup,l))
             uu_lup = &
                  wee(ij+u_lup,1,1)*hflux(ij+u_left,l)*(qu(ij+u_lup,l)+qu(ij+u_left,l)) +                        &
                  wee(ij+u_lup,2,1)*hflux(ij+u_ldown,l)*(qu(ij+u_lup,l)+qu(ij+u_ldown,l)) +                      &
                  wee(ij+u_lup,3,1)*hflux(ij+u_rdown,l)*(qu(ij+u_lup,l)+qu(ij+u_rdown,l)) +                      &
                  wee(ij+u_lup,4,1)*hflux(ij+u_right,l)*(qu(ij+u_lup,l)+qu(ij+u_right,l)) +                      &
                  wee(ij+u_lup,5,1)*hflux(ij+u_rup,l)*(qu(ij+u_lup,l)+qu(ij+u_rup,l)) +                          & 
                  wee(ij+u_lup,1,2)*hflux(ij+t_lup+u_right,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_right,l)) +          &
                  wee(ij+u_lup,2,2)*hflux(ij+t_lup+u_rup,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_rup,l)) +              &
                  wee(ij+u_lup,3,2)*hflux(ij+t_lup+u_lup,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_lup,l)) +              &
                  wee(ij+u_lup,4,2)*hflux(ij+t_lup+u_left,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_left,l)) +            &
                  wee(ij+u_lup,5,2)*hflux(ij+t_lup+u_ldown,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_ldown,l))
             uu_ldown = &
                  wee(ij+u_ldown,1,1)*hflux(ij+u_rdown,l)*(qu(ij+u_ldown,l)+qu(ij+u_rdown,l)) +                      &
                  wee(ij+u_ldown,2,1)*hflux(ij+u_right,l)*(qu(ij+u_ldown,l)+qu(ij+u_right,l)) +                      &
                  wee(ij+u_ldown,3,1)*hflux(ij+u_rup,l)*(qu(ij+u_ldown,l)+qu(ij+u_rup,l)) +                          &
                  wee(ij+u_ldown,4,1)*hflux(ij+u_lup,l)*(qu(ij+u_ldown,l)+qu(ij+u_lup,l)) +                          &
                  wee(ij+u_ldown,5,1)*hflux(ij+u_left,l)*(qu(ij+u_ldown,l)+qu(ij+u_left,l)) +                        & 
                  wee(ij+u_ldown,1,2)*hflux(ij+t_ldown+u_lup,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_lup,l)) +          &
                  wee(ij+u_ldown,2,2)*hflux(ij+t_ldown+u_left,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_left,l)) +        &
                  wee(ij+u_ldown,3,2)*hflux(ij+t_ldown+u_ldown,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_ldown,l)) +      &
                  wee(ij+u_ldown,4,2)*hflux(ij+t_ldown+u_rdown,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_rdown,l)) +      &
                  wee(ij+u_ldown,5,2)*hflux(ij+t_ldown+u_right,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_right,l))
             du(ij+u_right,l) = du(ij+u_right,l) + .5*uu_right
             du(ij+u_lup,l)   = du(ij+u_lup,l)   + .5*uu_lup
             du(ij+u_ldown,l) = du(ij+u_ldown,l)   + .5*uu_ldown
          ENDDO
       ENDDO

    CASE(conserv_gassmann) ! energy-conserving TRiSK modified by Gassmann (2018)

       DO l=ll_begin,ll_end
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             uu_right = &
                  wee(ij+u_right,1,1)*hflux(ij+u_rup,l)  *qu(ij+t_right+u_lup,l)+                 &
                  wee(ij+u_right,2,1)*hflux(ij+u_lup,l)  *qu(ij+u_rup,l)+                         &
                .5*wee(ij+u_right,3,1)*hflux(ij+u_left,l)*(qu(ij+u_right,l)+qu(ij+u_left,l))+     &
                  wee(ij+u_right,4,1)*hflux(ij+u_ldown,l)*qu(ij+u_rdown,l)+                       &
                  wee(ij+u_right,5,1)*hflux(ij+u_rdown,l)*qu(ij+t_right+u_ldown,l)+               &
                  wee(ij+u_right,1,2)*hflux(ij+t_right+u_ldown,l)*qu(ij+u_rdown,l)+               &
                  wee(ij+u_right,2,2)*hflux(ij+t_right+u_rdown,l)*qu(ij+t_right+u_ldown,l)+       &
               .5*wee(ij+u_right,3,2)*hflux(ij+t_right+u_right,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_right,l))+         &
                  wee(ij+u_right,4,2)*hflux(ij+t_right+u_rup,l)*qu(ij+t_right+u_lup,l)+           &
                  wee(ij+u_right,5,2)*hflux(ij+t_right+u_lup,l)*qu(ij+u_rup,l)
             uu_lup = &
                  wee(ij+u_lup,1,1)*hflux(ij+u_left,l)*qu(ij+t_lup+u_ldown,l) +                   &
                  wee(ij+u_lup,2,1)*hflux(ij+u_ldown,l)*qu(ij+u_left,l) +                         &
               .5*wee(ij+u_lup,3,1)*hflux(ij+u_rdown,l)*(qu(ij+u_lup,l)+qu(ij+u_rdown,l)) +       &
                  wee(ij+u_lup,4,1)*hflux(ij+u_right,l)*qu(ij+u_rup,l) +                          &
                  wee(ij+u_lup,5,1)*hflux(ij+u_rup,l)*qu(ij+t_lup+u_right,l)+                     & 
                  wee(ij+u_lup,1,2)*hflux(ij+t_lup+u_right,l)*qu(ij+u_rup,l) +                   &
                  wee(ij+u_lup,2,2)*hflux(ij+t_lup+u_rup,l)*qu(ij+t_lup+u_right,l) +              &
               .5*wee(ij+u_lup,3,2)*hflux(ij+t_lup+u_lup,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_lup,l)) + &
                  wee(ij+u_lup,4,2)*hflux(ij+t_lup+u_left,l)*qu(ij+t_lup+u_ldown,l) +            &
                  wee(ij+u_lup,5,2)*hflux(ij+t_lup+u_ldown,l)*qu(ij+u_left,l)
             uu_ldown = &
                  wee(ij+u_ldown,1,1)*hflux(ij+u_rdown,l)*qu(ij+t_ldown,l+u_right) +               &
                  wee(ij+u_ldown,2,1)*hflux(ij+u_right,l)*qu(ij+u_rdown,l) +                       &
               .5*wee(ij+u_ldown,3,1)*hflux(ij+u_rup,l)*(qu(ij+u_ldown,l)+qu(ij+u_rup,l)) +        &
                  wee(ij+u_ldown,4,1)*hflux(ij+u_lup,l)*qu(ij+u_left,l) +                          &
                  wee(ij+u_ldown,5,1)*hflux(ij+u_left,l)*qu(ij+t_ldown+u_lup,l) +                  & 
                  wee(ij+u_ldown,1,2)*hflux(ij+t_ldown+u_lup,l)*qu(ij+u_left,l) +                    &
                  wee(ij+u_ldown,2,2)*hflux(ij+t_ldown+u_left,l)*qu(ij+t_ldown+u_lup,l) +          &
               .5*wee(ij+u_ldown,3,2)*hflux(ij+t_ldown+u_ldown,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_ldown,l)) +      &
                  wee(ij+u_ldown,4,2)*hflux(ij+t_ldown+u_rdown,l)*qu(ij+t_ldown+u_right,l) +      &
                  wee(ij+u_ldown,5,2)*hflux(ij+t_ldown+u_right,l)*qu(ij+u_rdown,l)
             du(ij+u_right,l) = du(ij+u_right,l) + uu_right
             du(ij+u_lup,l)   = du(ij+u_lup,l)   + uu_lup
             du(ij+u_ldown,l) = du(ij+u_ldown,l)   + uu_ldown
          ENDDO
       ENDDO

    CASE(conserv_enstrophy) ! enstrophy-conserving TRiSK

       DO l=ll_begin,ll_end
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             uu_right = &
                  wee(ij+u_right,1,1)*hflux(ij+u_rup,l)+                           &
                  wee(ij+u_right,2,1)*hflux(ij+u_lup,l)+                           &
                  wee(ij+u_right,3,1)*hflux(ij+u_left,l)+                          &
                  wee(ij+u_right,4,1)*hflux(ij+u_ldown,l)+                         &
                  wee(ij+u_right,5,1)*hflux(ij+u_rdown,l)+                         & 
                  wee(ij+u_right,1,2)*hflux(ij+t_right+u_ldown,l)+                 &
                  wee(ij+u_right,2,2)*hflux(ij+t_right+u_rdown,l)+                 &
                  wee(ij+u_right,3,2)*hflux(ij+t_right+u_right,l)+                 &
                  wee(ij+u_right,4,2)*hflux(ij+t_right+u_rup,l)+                   &
                  wee(ij+u_right,5,2)*hflux(ij+t_right+u_lup,l)
             uu_lup = &
                  wee(ij+u_lup,1,1)*hflux(ij+u_left,l)+                        &
                  wee(ij+u_lup,2,1)*hflux(ij+u_ldown,l)+                       &
                  wee(ij+u_lup,3,1)*hflux(ij+u_rdown,l)+                       &
                  wee(ij+u_lup,4,1)*hflux(ij+u_right,l)+                       &
                  wee(ij+u_lup,5,1)*hflux(ij+u_rup,l)+                         & 
                  wee(ij+u_lup,1,2)*hflux(ij+t_lup+u_right,l)+                 &
                  wee(ij+u_lup,2,2)*hflux(ij+t_lup+u_rup,l)+                   &
                  wee(ij+u_lup,3,2)*hflux(ij+t_lup+u_lup,l)+                   &
                  wee(ij+u_lup,4,2)*hflux(ij+t_lup+u_left,l)+                  &
                  wee(ij+u_lup,5,2)*hflux(ij+t_lup+u_ldown,l)
             uu_ldown = &
                  wee(ij+u_ldown,1,1)*hflux(ij+u_rdown,l)+                      &
                  wee(ij+u_ldown,2,1)*hflux(ij+u_right,l)+                      &
                  wee(ij+u_ldown,3,1)*hflux(ij+u_rup,l)+                        &
                  wee(ij+u_ldown,4,1)*hflux(ij+u_lup,l)+                        &
                  wee(ij+u_ldown,5,1)*hflux(ij+u_left,l)+                       & 
                  wee(ij+u_ldown,1,2)*hflux(ij+t_ldown+u_lup,l)+                &
                  wee(ij+u_ldown,2,2)*hflux(ij+t_ldown+u_left,l)+               &
                  wee(ij+u_ldown,3,2)*hflux(ij+t_ldown+u_ldown,l)+              &
                  wee(ij+u_ldown,4,2)*hflux(ij+t_ldown+u_rdown,l)+              &
                  wee(ij+u_ldown,5,2)*hflux(ij+t_ldown+u_right,l)

             du(ij+u_right,l) = du(ij+u_right,l) + uu_right*qu(ij+u_right,l)
             du(ij+u_lup,l)   = du(ij+u_lup,l)   + uu_lup*qu(ij+u_lup,l)     
             du(ij+u_ldown,l) = du(ij+u_ldown,l) + uu_ldown*qu(ij+u_ldown,l) 
          END DO
       END DO
    CASE DEFAULT
       STOP
    END SELECT
#undef FTHETA

    CALL trace_end("compute_caldyn_Coriolis")

  END SUBROUTINE compute_caldyn_Coriolis_manual
  
END MODULE compute_caldyn_Coriolis_mod