MODULE compute_caldyn_fast_mod
  USE prec, ONLY : rstd
  USE grid_param
  USE earth_const
  USE disvert_mod
  USE omp_para
  USE trace
  IMPLICIT NONE
  PRIVATE

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

  PUBLIC :: compute_caldyn_fast_unst, compute_caldyn_fast_hex, compute_caldyn_fast_manual

CONTAINS

#ifdef BEGIN_DYSL

{% macro case_caldyn_fast(name) %}
  CASE({{name}})
    FORALL_CELLS()
      ON_PRIMAL
        Phi_ik = .5*(geopot(CELL)+geopot(UP(CELL)))
        {{ caller() }}
      END_BLOCK
    END_BLOCK
{%- endmacro %}

KERNEL(caldyn_fast) 
  !
  SELECT CASE(caldyn_thermo)
  CASE(thermo_boussinesq)
     FORALL_CELLS()
       ON_PRIMAL
         berni(CELL) = pk(CELL)
         ! from now on pk contains the vertically-averaged geopotential
         pk(CELL) = .5*(geopot(CELL)+geopot(UP(CELL)))
       END_BLOCK
     END_BLOCK
  {% call case_caldyn_fast('thermo_theta') %}
     berni(CELL) = Phi_ik
  {%- endcall %}
  {% call case_caldyn_fast('thermo_entropy') %}
     berni(CELL) = Phi_ik + pk(CELL)*(cpp-theta(CELL,1)) ! Gibbs = Cp.T-Ts = T(Cp-s)
  {%- endcall %}
  {% call case_caldyn_fast('thermo_variable_Cp') %}
    ! thermodynamics with variable Cp
    !           Cp(T) = Cp0 * (T/T0)^nu
    ! =>            h = Cp(T).T/(nu+1)
     cp_ik = cpp*(pk(ij,l)/Treff)**nu
     berni(CELL) = Phi_ik + pk(CELL)*(cp_ik/(nu+1.)-theta(CELL,1)) ! Gibbs = h-Ts = T(Cp/(nu+1)-s)
  {%- endcall %}
  CASE DEFAULT
    PRINT *, 'Unsupported value of caldyn_thermo : ',caldyn_thermo  ! FIXME
    STOP
  END SELECT
!
  FORALL_CELLS()
    ON_EDGES
      due = .5*(theta(CELL1,1)+theta(CELL2,1))*(pk(CELL2)-pk(CELL1)) + berni(CELL2)-berni(CELL1)
      du(EDGE) = du(EDGE) - SIGN*due
      u(EDGE) = u(EDGE) + tau*du(EDGE)
    END_BLOCK
  END_BLOCK  
!
END_BLOCK

#endif END_DYSL

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

  SUBROUTINE compute_caldyn_fast_hex(tau,theta,geopot, pk,berni,du,u)
    REAL(rstd),INTENT(IN)    :: tau, theta(:,:,:), geopot(:,:)
    REAL(rstd),INTENT(INOUT) :: pk(:,:), berni(:,:), du(:,:), u(:,:)
    CALL compute_caldyn_fast_hex_(tau,theta,geopot, pk,berni,du,u)
  END SUBROUTINE compute_caldyn_fast_hex

  SUBROUTINE compute_caldyn_fast_unst(tau,theta,geopot, pk,berni,du,u)
    REAL(rstd),INTENT(IN)    :: tau, theta(:,:,:), geopot(:,:)
    REAL(rstd),INTENT(INOUT) :: pk(:,:), berni(:,:), du(:,:), u(:,:)
    CALL compute_caldyn_fast_unst_(tau,theta,geopot, pk,berni,du,u)
  END SUBROUTINE compute_caldyn_fast_unst

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

  SUBROUTINE compute_caldyn_fast_unst_(tau,theta,geopot, pk,berni,du,u)
    USE ISO_C_BINDING, only : C_DOUBLE, C_FLOAT
    USE data_unstructured_mod, ONLY : enter_trace, exit_trace, &
         id_fast, dual_deg, dual_edge, dual_ne, dual_vertex, &
         up, down, left, right
    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, cp_ik, Phi_ik
    START_TRACE(id_fast, 4,0,2) ! primal, dual, edge
#include "../kernels_unst/caldyn_fast.k90"
    STOP_TRACE
  END SUBROUTINE compute_caldyn_fast_unst_

  SUBROUTINE compute_caldyn_fast_hex_(tau,theta,geopot, pk,berni,du,u)
    USE icosa
    REAL(rstd),INTENT(IN)    :: tau                ! "solve" u-tau*du/dt = rhs
    REAL(rstd),INTENT(IN)    :: theta(iim*jjm,llm,nqdyn)
    REAL(rstd),INTENT(IN)    :: geopot(iim*jjm,llm+1)
    REAL(rstd),INTENT(INOUT) :: pk(iim*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: berni(iim*jjm,llm)  ! partial Bernoulli function
    REAL(rstd),INTENT(INOUT) :: du(iim*3*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: u(iim*3*jjm,llm)   ! INOUT if tau>0
    REAL(rstd) :: berniv(iim*jjm,llm)  ! moist Bernoulli function

    INTEGER :: ij,l
    REAL(rstd) :: due, cp_ik, Phi_ik

    CALL trace_start("compute_caldyn_fast")
#include "../kernels_hex/caldyn_fast.k90"
    CALL trace_end("compute_caldyn_fast")
    
  END SUBROUTINE compute_caldyn_fast_hex_

  SUBROUTINE compute_caldyn_fast_manual(tau,theta,geopot, pk,berni,du,u)
    USE icosa
    REAL(rstd),INTENT(IN)    :: tau                ! "solve" u-tau*du/dt = rhs
    REAL(rstd),INTENT(IN)    :: theta(iim*jjm,llm,nqdyn)
    REAL(rstd),INTENT(IN)    :: geopot(iim*jjm,llm+1)
    REAL(rstd),INTENT(INOUT) :: pk(iim*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: berni(iim*jjm,llm)  ! partial Bernoulli function
    REAL(rstd),INTENT(INOUT) :: du(iim*3*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: u(iim*3*jjm,llm)   ! INOUT if tau>0
    REAL(rstd) :: berniv(iim*jjm,llm)  ! moist Bernoulli function

    INTEGER :: ij,l
    REAL(rstd) :: cp_ik, qv, temp, chi, log_p_preff, due, due_right, due_lup, due_ldown

    CALL trace_start("compute_caldyn_fast")

    ! Compute Bernoulli term
    IF(boussinesq) THEN
       DO l=ll_begin,ll_end
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             berni(ij,l) = pk(ij,l)
             ! from now on pk contains the vertically-averaged geopotential
             pk(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1))
          END DO
       END DO
    ELSE ! compressible

       DO l=ll_begin,ll_end
          SELECT CASE(caldyn_thermo)
          CASE(thermo_theta) ! vdp = theta.dpi => B = Phi
             !DIR$ SIMD
             DO ij=ij_begin,ij_end         
                berni(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1))
             END DO
          CASE(thermo_entropy) ! vdp = dG + sdT => B = Phi + G, G=h-Ts=T*(cpp-s) 
             !DIR$ SIMD
             DO ij=ij_begin,ij_end         
                berni(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1)) &
                     + pk(ij,l)*(cpp-theta(ij,l,1)) ! pk=temperature, theta=entropy
             END DO
          CASE(thermo_moist) 
             !DIR$ SIMD
             DO ij=ij_begin,ij_end         
                ! du/dt = grad(Bd)+rv.grad(Bv)+s.grad(T)
                ! Bd = Phi + gibbs_d
                ! Bv = Phi + gibbs_v
                ! pk=temperature, theta=entropy
                qv = theta(ij,l,2)
                temp = pk(ij,l)
                chi = log(temp/Treff)
                log_p_preff = (chi*(cpp+qv*cppv)-theta(ij,l,1))/(Rd+qv*Rv) ! log(p/preff)
                berni(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1)) &
                     + temp*(cpp*(1.-chi)+Rd*log_p_preff)
                berniv(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1)) &
                     + temp*(cppv*(1.-chi)+Rv*log_p_preff)
            END DO
          END SELECT
       END DO

    END IF ! Boussinesq/compressible

!!! u:=u+tau*du, du = -grad(B)-theta.grad(pi)
    DO l=ll_begin,ll_end
       IF(caldyn_thermo == thermo_moist) THEN
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             due_right =  berni(ij+t_right,l)-berni(ij,l)       &
                  + 0.5*(theta(ij,l,1)+theta(ij+t_right,l,1))   &
                       *(pk(ij+t_right,l)-pk(ij,l))             &
                  + 0.5*(theta(ij,l,2)+theta(ij+t_right,l,2))   &
                       *(berniv(ij+t_right,l)-berniv(ij,l))
             
             due_lup = berni(ij+t_lup,l)-berni(ij,l)            &
                  + 0.5*(theta(ij,l,1)+theta(ij+t_lup,l,1))     &
                       *(pk(ij+t_lup,l)-pk(ij,l))               &
                  + 0.5*(theta(ij,l,2)+theta(ij+t_lup,l,2))     &
                       *(berniv(ij+t_lup,l)-berniv(ij,l))
             
             due_ldown = berni(ij+t_ldown,l)-berni(ij,l)        &
                  + 0.5*(theta(ij,l,1)+theta(ij+t_ldown,l,1))   &
                       *(pk(ij+t_ldown,l)-pk(ij,l))             &
                  + 0.5*(theta(ij,l,2)+theta(ij+t_ldown,l,2))   &
                       *(berniv(ij+t_ldown,l)-berniv(ij,l))
             
             du(ij+u_right,l) = du(ij+u_right,l) - ne_right*due_right
             du(ij+u_lup,l)   = du(ij+u_lup,l)   - ne_lup*due_lup
             du(ij+u_ldown,l) = du(ij+u_ldown,l) - ne_ldown*due_ldown
             u(ij+u_right,l) = u(ij+u_right,l) + tau*du(ij+u_right,l)
             u(ij+u_lup,l)   = u(ij+u_lup,l)   + tau*du(ij+u_lup,l)
             u(ij+u_ldown,l) = u(ij+u_ldown,l) + tau*du(ij+u_ldown,l)
          END DO
       ELSE
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             due_right =  0.5*(theta(ij,l,1)+theta(ij+t_right,l,1))  &
                  *(pk(ij+t_right,l)-pk(ij,l))        &
                  +  berni(ij+t_right,l)-berni(ij,l)
             due_lup = 0.5*(theta(ij,l,1)+theta(ij+t_lup,l,1))    &
                  *(pk(ij+t_lup,l)-pk(ij,l))          &
                  +  berni(ij+t_lup,l)-berni(ij,l)
             due_ldown = 0.5*(theta(ij,l,1)+theta(ij+t_ldown,l,1)) &
                  *(pk(ij+t_ldown,l)-pk(ij,l))      &
                  +   berni(ij+t_ldown,l)-berni(ij,l)
             du(ij+u_right,l) = du(ij+u_right,l) - ne_right*due_right
             du(ij+u_lup,l)   = du(ij+u_lup,l)   - ne_lup*due_lup
             du(ij+u_ldown,l) = du(ij+u_ldown,l) - ne_ldown*due_ldown
             u(ij+u_right,l) = u(ij+u_right,l) + tau*du(ij+u_right,l)
             u(ij+u_lup,l)   = u(ij+u_lup,l)   + tau*du(ij+u_lup,l)
             u(ij+u_ldown,l) = u(ij+u_ldown,l) + tau*du(ij+u_ldown,l)
          END DO
       END IF
    END DO

    CALL trace_end("compute_caldyn_fast")

  END SUBROUTINE compute_caldyn_fast_manual

END MODULE compute_caldyn_fast_mod