MODULE compute_temperature_mod
  USE prec, ONLY : rstd
  USE earth_const, ONLY : cpp, cppv, kappa, Rd, Rv, preff, Treff, nu, &
       caldyn_thermo, physics_thermo, thermo_fake_moist, &
       thermo_theta, thermo_entropy, thermo_variable_Cp, thermo_moist
  USE grid_param
  IMPLICIT NONE
  PRIVATE
  SAVE

  PUBLIC :: temperature, compute_temperature_unst, compute_temperature_hex, compute_temperature_manual

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

CONTAINS
  
  SUBROUTINE temperature(f_pmid,f_q,f_temp)
    USE compute_diagnostics_mod, ONLY : compute_temperature
    USE icosa
    TYPE(t_field), POINTER :: f_pmid(:)         ! IN
    TYPE(t_field), POINTER :: f_q(:)            ! IN
    TYPE(t_field), POINTER :: f_temp(:)         ! INOUT
    
    REAL(rstd), POINTER :: pmid(:,:)
    REAL(rstd), POINTER :: q(:,:,:)
    REAL(rstd), POINTER :: temp(:,:)
    INTEGER :: ind
    
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       pmid=f_pmid(ind)
       q=f_q(ind)
       temp=f_temp(ind)
       CALL compute_temperature(pmid,q,temp)
    END DO
  END SUBROUTINE temperature
  
! Macros loop_compute_temperature_fake_moist() and loop_compute_temperature() 
! are here to inline the thermodynamical formula inside the innermost loop.
! Tests are made outside the outer loop.

#ifdef BEGIN_DYSL

{%- macro loop_compute_temperature() %}
{%- set comp_temp=caller() %}
  FORALL_CELLS()
    ON_PRIMAL
      p_ik = pmid(CELL)
      theta_ik = temp(CELL)
      qv = q(CELL,1) ! water vapor mixing ratio = mv/md
      {{ comp_temp }}
      temp(CELL) = temp_ik
    END_BLOCK
  END_BLOCK
{%- endmacro %}

{%- macro loop_compute_temperature_fake_moist() %}
{%- set comp_temp=caller() %}
IF(physics_thermo==thermo_fake_moist) THEN
   {% call loop_compute_temperature() %}
     {{ comp_temp }}
     temp_ik = temp_ik/(1+0.608*qv)
   {% endcall %}
ELSE
   {% call loop_compute_temperature() %}
     {{ comp_temp }}
   {% endcall %}
END IF
{%- endmacro %}

KERNEL(compute_temperature)
  SELECT CASE(caldyn_thermo)
  CASE(thermo_theta)
     {% call loop_compute_temperature_fake_moist() %}
        temp_ik = theta_ik*((p_ik/preff)**kappa)
     {% endcall %}
  CASE(thermo_entropy)
     {% call loop_compute_temperature_fake_moist() %}
        temp_ik = Treff*exp((theta_ik + Rd*log(p_ik/preff))/cpp)
     {% endcall %}
  CASE(thermo_variable_Cp)
     {% call loop_compute_temperature() %}
       Cp_ik = nu*( theta_ik + Rd*log(p_ik/preff) )
       temp_ik = Treff* (Cp_ik/cpp)**(1./nu)  
     {% endcall %}
  CASE(thermo_moist)
     {% call loop_compute_temperature() %}
       Rmix = Rd+qv*Rv
       chi = ( theta_ik + Rmix*log(p_ik/preff) ) / (cpp + qv*cppv) ! log(T/Treff)
       temp_ik = Treff*exp(chi)
     {% endcall %}
  END SELECT
END_BLOCK

#endif END_DYSL

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

  SUBROUTINE compute_temperature_unst(pmid, q, temp)
    REAL(rstd),INTENT(IN)    :: pmid(:,:), q(:,:,:)
    REAL(rstd),INTENT(INOUT) :: temp(:,:)
    CALL compute_temperature_unst_(pmid, q, temp)
  END SUBROUTINE compute_temperature_unst

  SUBROUTINE compute_temperature_hex(pmid, q, temp)
    REAL(rstd),INTENT(IN)    :: pmid(:,:), q(:,:,:)
    REAL(rstd),INTENT(INOUT) :: temp(:,:)
    CALL compute_temperature_hex_(pmid, q, temp)
  END SUBROUTINE compute_temperature_hex

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

  SUBROUTINE compute_temperature_unst_(pmid, q, temp)
    REAL(rstd),INTENT(IN)    :: pmid(llm, primal_num)
    REAL(rstd),INTENT(IN)    :: q(llm, primal_num, nqtot)
    REAL(rstd),INTENT(INOUT) :: temp(llm, primal_num)
    REAL(rstd) :: p_ik, theta_ik, temp_ik, qv, chi, Rmix, Cp_ik
    DECLARE_INDICES
#include "../kernels_unst/compute_temperature.k90"
  END SUBROUTINE compute_temperature_unst_

  SUBROUTINE compute_temperature_hex_(pmid,q,temp)
    USE icosa
    USE omp_para
    REAL(rstd),INTENT(IN)    :: pmid(iim*jjm,llm)
    REAL(rstd),INTENT(IN)    :: q(iim*jjm,llm,nqtot)
    REAL(rstd),INTENT(INOUT) :: temp(iim*jjm,llm)

    REAL(rstd) :: p_ik, theta_ik, temp_ik, qv, chi, Rmix, Cp_ik
    INTEGER :: ij,l
#include "../kernels_hex/compute_temperature.k90"
  END SUBROUTINE compute_temperature_hex_

  SUBROUTINE compute_temperature_manual(pmid,q,temp)
    USE icosa
    USE omp_para
    REAL(rstd),INTENT(IN)    :: pmid(iim*jjm,llm)
    REAL(rstd),INTENT(IN)    :: q(iim*jjm,llm,nqtot)
    REAL(rstd),INTENT(INOUT) :: temp(iim*jjm,llm)

    REAL(rstd) :: p_ik, theta_ik, temp_ik, qv, chi, Rmix
    INTEGER :: ij,l

    Rd = kappa*cpp
    DO l=ll_begin,ll_end
       DO ij=ij_begin,ij_end
          p_ik = pmid(ij,l)
          theta_ik = temp(ij,l)
          qv = q(ij,l,1) ! water vapor mixing ratio = mv/md
          SELECT CASE(caldyn_thermo)
          CASE(thermo_theta)
             temp_ik = theta_ik*((p_ik/preff)**kappa)
          CASE(thermo_entropy)
             temp_ik = Treff*exp((theta_ik + Rd*log(p_ik/preff))/cpp)
          CASE(thermo_moist)
             Rmix = Rd+qv*Rv
             chi = ( theta_ik + Rmix*log(p_ik/preff) ) / (cpp + qv*cppv) ! log(T/Treff)
             temp_ik = Treff*exp(chi)
          END SELECT
          IF(physics_thermo==thermo_fake_moist) temp_ik=temp_ik/(1+0.608*qv)
          temp(ij,l)=temp_ik
       END DO
    END DO
  END SUBROUTINE compute_temperature_manual

  SUBROUTINE Tv2T(f_Tv, f_q, f_T)
    USE icosa
    TYPE(t_field), POINTER :: f_TV(:)
    TYPE(t_field), POINTER :: f_q(:)
    TYPE(t_field), POINTER :: f_T(:)
    
    REAL(rstd),POINTER :: Tv(:,:), q(:,:,:), T(:,:)
    INTEGER :: ind
    
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       Tv=f_Tv(ind)
       T=f_T(ind)
       SELECT CASE(physics_thermo)
       CASE(thermo_dry)
          T=Tv
       CASE(thermo_fake_moist)
          q=f_q(ind)
          T=Tv/(1+0.608*q(:,:,1))
       END SELECT
    END DO
  END SUBROUTINE Tv2T

END MODULE compute_temperature_mod