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

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

  PUBLIC :: compute_geopot_unst, compute_geopot_hex, compute_geopot_manual

CONTAINS

#ifdef BEGIN_DYSL

{# ---------------- macro to generate code computing pressure top-down ---------------
   formula = formula to compute 'gravitational' mass
   = rhodz (dry) rhodz*theta (boussinesq) rhodz*(1+qv) (moist)         #}

#define BALANCE(formula) {% call(thecell) balance() %} formula {% endcall %}
{% macro balance() %} {% set formula=caller %}
SEQUENCE_C1
  PROLOGUE(llm)
    pk(CELL) = ptop + .5*g*{{ formula('CELL') }}
  END_BLOCK
  BODY('llm-1,1,-1') 
    pk(CELL) = pk(UP(CELL)) + (.5*g)*({{ formula('CELL') }}+{{ formula('UP(CELL)') }})
  END_BLOCK
  IF(caldyn_eta == eta_lag) THEN
    EPILOGUE(1)
      ps(HIDX(CELL)) = pk(CELL) + .5*g*{{ formula('CELL') }}
    END_BLOCK
  END IF
END_BLOCK
{%- endmacro %}

{# ------------ macro to generate code computing geopotential bottom-up --------------
   var = variable to be stored in pk(CELL)
   caller() computes gv = g*v where v = specific volume
   details depend on caldyn_thermo #}

#define GEOPOT(var) {% call geopot(var) %}
{% macro geopot(var) %} {% set formula=caller %}
  SEQUENCE_C1
    BODY('1,llm')
      p_ik = pk(CELL)
      {{ formula() }}
      pk(CELL) = {{ var }}
      geopot(UP(CELL)) = geopot(CELL) + gv*rhodz(CELL)
    END_BLOCK
  END_BLOCK
{%- endmacro %}

#define END_GEOPOT {% endcall %}

#define THECELL {{ thecell }}

KERNEL(compute_hydrostatic_pressure)
  SELECT CASE(caldyn_thermo)
  CASE(thermo_boussinesq)
    ! use hydrostatic balance with theta*rhodz to find pk (=Lagrange multiplier=pressure)
    BALANCE( theta_rhodz(THECELL,1) )
  CASE(thermo_theta, thermo_entropy, thermo_variable_Cp)
    BALANCE( rhodz(THECELL) )
  CASE(thermo_moist)
    BALANCE( (rhodz(THECELL)+theta_rhodz(THECELL,2)) )
  END SELECT
END_BLOCK

KERNEL(compute_geopot)
  SELECT CASE(caldyn_thermo)
  CASE(thermo_boussinesq)
    ! use hydrostatic balance with theta*rhodz to find pk (=Lagrange multiplier=pressure)
    BALANCE( theta(THECELL,1)*rhodz(THECELL) )
    ! now pk contains the Lagrange multiplier (pressure)
    ! specific volume 1 = dphi/g/rhodz
    SEQUENCE_C1
      BODY('1,llm')
        geopot(UP(CELL)) = geopot(CELL) + g*rhodz(CELL)
      END_BLOCK
    END_BLOCK
  CASE(thermo_theta)
    BALANCE( rhodz(THECELL) )
    GEOPOT('exner_ik')
      exner_ik = cpp * (p_ik/preff) ** kappa
      gv = (g*kappa)*theta(CELL,1)*exner_ik/p_ik
    END_GEOPOT
  CASE(thermo_entropy)
    BALANCE( rhodz(THECELL) )
    GEOPOT('temp_ik')
      temp_ik = Treff*exp((theta(CELL,1) + Rd*log(p_ik/preff))/cpp)
      gv = (g*Rd)*temp_ik/p_ik   ! specific volume v = Rd*T/p
    END_GEOPOT
  CASE(thermo_variable_Cp)
    ! thermodynamics with variable Cp
    !      Cp.dT = dh = Tds + vdp
    !              pv = RT
    ! =>           ds = (dh+v.dp)/T = Cp.dT/T - R dp/p
    ! 		Cp(T) = Cp0 * (T/T0)^nu
    ! =>       s(p,T) = Cp(T)/nu - R log(p/preff) 
    ! 	       	    h = Cp(T).T/(nu+1)
    BALANCE( rhodz(THECELL) )
    GEOPOT('temp_ik')
      Cp_ik = nu*( theta(CELL,1) + Rd*log(p_ik/preff) )
      temp_ik = Treff* (Cp_ik/cpp)**(1./nu)
      gv = (g*Rd)*temp_ik/p_ik   ! specific volume v = Rd*T/p
    END_GEOPOT
  CASE(thermo_moist)
    BALANCE( rhodz(THECELL)*(1.+theta(THECELL,2)) )
    GEOPOT('temp_ik')
      qv = theta(CELL,2) ! water vaper mixing ratio = mv/md
      Rmix = Rd+qv*Rv
      chi = ( theta(CELL,1) + Rmix*log(p_ik/preff) ) / (cpp + qv*cppv) ! log(T/Treff)
      temp_ik = Treff*exp(chi)
      ! specific volume v = R*T/p
      ! R = (Rd + qv.Rv)/(1+qv)
      gv = g*Rmix*temp_ik/(p_ik*(1+qv))
    END_GEOPOT
  END SELECT
END_BLOCK

#endif END_DYSL

  !**************************** Geopotential *****************************

  SUBROUTINE compute_geopot_unst(rhodz,theta,ps,pk,geopot)
        USE ISO_C_BINDING, only : C_DOUBLE, C_FLOAT
    USE data_unstructured_mod, ONLY : enter_trace, exit_trace, &
         id_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 :: chi, gv, exner_ik, temp_ik, p_ik, qv, Rmix, Cp_ik
    START_TRACE(id_geopot, 3,0,3)
#include "../kernels_unst/compute_geopot.k90"
    STOP_TRACE
  END SUBROUTINE compute_geopot_unst
  
  SUBROUTINE compute_geopot_hex(rhodz,theta, ps,pk,geopot) 
    REAL(rstd),INTENT(IN)    :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(IN)    :: theta(iim*jjm,llm,nqdyn) ! active scalars : theta/entropy, moisture, ...
    REAL(rstd),INTENT(INOUT) :: ps(iim*jjm)
    REAL(rstd),INTENT(OUT)   :: pk(iim*jjm,llm)       ! Exner function (compressible) /Lagrange multiplier (Boussinesq)
    REAL(rstd),INTENT(INOUT) :: geopot(iim*jjm,llm+1) ! geopotential

    INTEGER :: i,j,ij,l
    REAL(rstd) :: p_ik, exner_ik, Cp_ik, temp_ik, qv, chi, Rmix, gv
    INTEGER    :: ij_omp_begin_ext, ij_omp_end_ext

    CALL trace_start("compute_geopot")
    !$OMP BARRIER
    CALL distrib_level(ij_begin_ext,ij_end_ext, ij_omp_begin_ext,ij_omp_end_ext)
#include "../kernels_hex/compute_geopot.k90"
    !$OMP BARRIER

    CALL trace_end("compute_geopot")
  END SUBROUTINE compute_geopot_hex

  SUBROUTINE compute_geopot_manual(rhodz,theta, ps,pk,geopot) 
    REAL(rstd),INTENT(IN)    :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(IN)    :: theta(iim*jjm,llm,nqdyn) ! active scalars : theta/entropy, moisture, ...
    REAL(rstd),INTENT(INOUT) :: ps(iim*jjm)
    REAL(rstd),INTENT(OUT)   :: pk(iim*jjm,llm)       ! Exner function (compressible) /Lagrange multiplier (Boussinesq)
    REAL(rstd),INTENT(INOUT) :: geopot(iim*jjm,llm+1) ! geopotential

    INTEGER :: i,j,ij,l
    REAL(rstd) :: p_ik, exner_ik, Cp_ik, temp_ik, qv, chi, Rmix, gv
    INTEGER    :: ij_omp_begin_ext, ij_omp_end_ext

    CALL trace_start("compute_geopot")

!$OMP BARRIER

    CALL distrib_level(ij_begin_ext,ij_end_ext, ij_omp_begin_ext,ij_omp_end_ext)

    ! Pressure is computed first top-down (temporarily stored in pk)
    ! Then Exner pressure and geopotential are computed bottom-up
    ! Works also when caldyn_eta=eta_mass          

    IF(boussinesq) THEN ! compute geopotential and pk=Lagrange multiplier
       ! specific volume 1 = dphi/g/rhodz
       !         IF (is_omp_level_master) THEN ! no openMP on vertical due to dependency
       DO l = 1,llm
          !DIR$ SIMD
          DO ij=ij_omp_begin_ext,ij_omp_end_ext         
             geopot(ij,l+1) = geopot(ij,l) + g*rhodz(ij,l)
          ENDDO
       ENDDO
       ! use hydrostatic balance with theta*rhodz to find pk (Lagrange multiplier=pressure) 
       ! uppermost layer
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext         
          pk(ij,llm) = ptop + (.5*g)*theta(ij,llm,1)*rhodz(ij,llm)
       END DO
       ! other layers
       DO l = llm-1, 1, -1
          !          !$OMP DO SCHEDULE(STATIC) 
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext         
             pk(ij,l) = pk(ij,l+1) + (.5*g)*(theta(ij,l,1)*rhodz(ij,l)+theta(ij,l+1,1)*rhodz(ij,l+1))
          END DO
       END DO
       ! now pk contains the Lagrange multiplier (pressure)
    ELSE ! non-Boussinesq, compute pressure, Exner pressure or temperature, then geopotential
       ! uppermost layer
       
       SELECT CASE(caldyn_thermo)
          CASE(thermo_theta, thermo_entropy)
             !DIR$ SIMD
             DO ij=ij_omp_begin_ext,ij_omp_end_ext
                pk(ij,llm) = ptop + (.5*g)*rhodz(ij,llm)
             END DO
             ! other layers
             DO l = llm-1, 1, -1
                !DIR$ SIMD
                DO ij=ij_omp_begin_ext,ij_omp_end_ext         
                   pk(ij,l) = pk(ij,l+1) + (.5*g)*(rhodz(ij,l)+rhodz(ij,l+1))
                END DO
             END DO
             ! surface pressure (for diagnostics)
             IF(caldyn_eta==eta_lag) THEN
                DO ij=ij_omp_begin_ext,ij_omp_end_ext         
                   ps(ij) = pk(ij,1) + (.5*g)*rhodz(ij,1)
                END DO
             END IF
          CASE(thermo_moist) ! theta(ij,l,2) = qv = mv/md
             !DIR$ SIMD
             DO ij=ij_omp_begin_ext,ij_omp_end_ext
                pk(ij,llm) = ptop + (.5*g)*rhodz(ij,llm)*(1.+theta(ij,l,2))
             END DO
             ! other layers
             DO l = llm-1, 1, -1
                !DIR$ SIMD
                DO ij=ij_omp_begin_ext,ij_omp_end_ext         
                   pk(ij,l) = pk(ij,l+1) + (.5*g)*(          &
                        rhodz(ij,l)  *(1.+theta(ij,l,2)) +   &
                        rhodz(ij,l+1)*(1.+theta(ij,l+1,2)) )
                END DO
             END DO
             ! surface pressure (for diagnostics)
             IF(caldyn_eta==eta_lag) THEN
                DO ij=ij_omp_begin_ext,ij_omp_end_ext         
                   ps(ij) = pk(ij,1) + (.5*g)*rhodz(ij,1)*(1.+theta(ij,l,2))
                END DO
             END IF
          END SELECT

       DO l = 1,llm
          SELECT CASE(caldyn_thermo)
          CASE(thermo_theta)
             !DIR$ SIMD
             DO ij=ij_omp_begin_ext,ij_omp_end_ext
                p_ik = pk(ij,l)
                exner_ik = cpp * (p_ik/preff) ** kappa
                pk(ij,l) = exner_ik
                ! specific volume v = kappa*theta*pi/p = dphi/g/rhodz
                geopot(ij,l+1) = geopot(ij,l) + (g*kappa)*rhodz(ij,l)*theta(ij,l,1)*exner_ik/p_ik
             ENDDO
          CASE(thermo_entropy) ! theta is in fact entropy = cpp*log(theta/Treff) = cpp*log(T/Treff) - Rd*log(p/preff)
             !DIR$ SIMD
             DO ij=ij_omp_begin_ext,ij_omp_end_ext
                p_ik = pk(ij,l)
                temp_ik = Treff*exp((theta(ij,l,1) + Rd*log(p_ik/preff))/cpp)
                pk(ij,l) = temp_ik
                ! specific volume v = Rd*T/p = dphi/g/rhodz
                geopot(ij,l+1) = geopot(ij,l) + (g*Rd)*rhodz(ij,l)*temp_ik/p_ik
             ENDDO
          CASE(thermo_moist) ! theta is moist pseudo-entropy per dry air mass
             DO ij=ij_omp_begin_ext,ij_omp_end_ext
                p_ik = pk(ij,l)
                qv = theta(ij,l,2) ! water vaper mixing ratio = mv/md
                Rmix = Rd+qv*Rv
                chi = ( theta(ij,l,1) + Rmix*log(p_ik/preff) ) / (cpp + qv*cppv) ! log(T/Treff)
                temp_ik = Treff*exp(chi)
                pk(ij,l) = temp_ik
                ! specific volume v = R*T/p = dphi/g/rhodz
                ! R = (Rd + qv.Rv)/(1+qv)
                geopot(ij,l+1) = geopot(ij,l) + g*Rmix*rhodz(ij,l)*temp_ik/(p_ik*(1+qv))
             ENDDO
          CASE DEFAULT
             STOP
          END SELECT
       ENDDO
    END IF

    !ym flush geopot
    !$OMP BARRIER

    CALL trace_end("compute_geopot")

  END SUBROUTINE compute_geopot_manual

END MODULE compute_geopot_mod