MODULE thermodynamics_mod
   USE, INTRINSIC :: ISO_C_BINDING, ONLY: C_DOUBLE
   IMPLICIT NONE
   PRIVATE
   SAVE

   INTEGER, PARAMETER :: thermo_none = -99      ! undefined
   INTEGER, PARAMETER :: thermo_theta = 1       ! dry air, predicts potential temperature
   INTEGER, PARAMETER :: thermo_entropy = 2     ! dry air, predicts specific entropy
   INTEGER, PARAMETER :: thermo_variable_Cp = 3 ! dry air with variable Cp, predicts entropy
   INTEGER, PARAMETER :: thermo_binarygas = 4   ! mixture of two perfect gases, predicts potential enthalpy

   TYPE thermodynamics_t
      INTEGER :: which = thermo_none       ! selector for thermodynamics type
      INTEGER :: nqdyn = 1                 ! number of components (ex. 1 for dry air, 2 for air+vapor)
      REAL(C_DOUBLE) :: kappa = 0.2857143  ! ratio R/Cp for dry air
      REAL(C_DOUBLE) :: Treff = 273.       ! reference temperature, used in definition of entropy
      REAL(C_DOUBLE) :: preff = 101325.    ! reference pressure, used in definition of entropy and theta
      REAL(C_DOUBLE) :: cpp = 1004.70885   ! specific heat at constant pressure at t=Treff for dry air
      REAL(C_DOUBLE) :: cppv = 1860.       ! specific heat at constant pressure at t=Treff for water vapor
      REAL(C_DOUBLE) :: Rv = 461.5         ! specific perfect gas constant for water vapor
      REAL(C_DOUBLE) :: nu = 0.35          ! exponent in variable-Cp law Cp=cpp*(T/T0)^nu
      REAL(C_DOUBLE) :: T0 = 273.          ! temperature at which cp=cp0 for variable_Cp
   END TYPE

   ! ==== nomenclature for arguments to thermodynamic functions ====
   ! short        long
   ! p, pres      pressure
   ! T, temp      temperature
   ! V, vol       specific volume
   ! cons, theta  conservative variable

   PUBLIC :: thermodynamics_t, init_thermodynamics

   ! used at initialization (always on CPU, not performance-critical)
   PUBLIC :: temperature_from_pVq, theta_from_pTq, volume_from_pCq

   ! used for diagnostics (always on CPU, not performance-critical)
   PUBLIC :: temperature_from_pCq

   ! used for dynamics, on CPU or GPU : performance-critical !
   PUBLIC :: enthalpy_from_pq ! computes enthalpy and its derivatives wrt p, dynamical tracers

CONTAINS

   !==============================================================================================

   elemental real(C_DOUBLE) FUNCTION LOG1P(X) RESULT(Y)
      !$acc routine seq
      ! Returns log(1+x) without loss of accuracy for small x
      real(C_DOUBLE), INTENT(IN) :: X
      real(C_DOUBLE) :: Z
      real(C_DOUBLE), parameter :: ONE = 1.0

      Z = ONE + X
      Y = LOG(Z) - ((Z - ONE) - X)/Z                ! cancels errors with IEEE arithmetic

   END FUNCTION LOG1P

   !==============================================================================================

   SUBROUTINE init_thermodynamics(thermo, getin_c, getin_r, errstop)
      TYPE(thermodynamics_t), INTENT(OUT) :: thermo
      INTERFACE
         SUBROUTINE getin_c(name, val)
            CHARACTER(LEN=*), INTENT(IN)    :: name
            CHARACTER(LEN=*), INTENT(INOUT) :: val
         END SUBROUTINE
         SUBROUTINE getin_r(name, val)
            IMPORT
            CHARACTER(LEN=*), INTENT(IN)  :: name
            REAL(C_DOUBLE), INTENT(INOUT) :: val
         END SUBROUTINE
         SUBROUTINE errstop(message)
            CHARACTER(LEN=*), INTENT(IN), OPTIONAL :: message
         END SUBROUTINE
      END INTERFACE

      CHARACTER(len=255)  :: def
      CHARACTER(len=1024) :: msgbuf

      def = 'theta'
      CALL getin_c('thermo', def)
      SELECT CASE (TRIM(def))
      CASE ('theta')
         CALL getin_r("kappa", thermo%kappa)
         CALL getin_r("cpp", thermo%cpp)
         CALL getin_r("preff", thermo%preff)
         thermo%which = thermo_theta
      CASE ('entropy')
         CALL getin_r("kappa", thermo%kappa)
         CALL getin_r("cpp", thermo%cpp)
         CALL getin_r("preff", thermo%preff)
         CALL getin_r("Treff", thermo%Treff)
         thermo%which = thermo_entropy
      CASE ('binarygas')
         CALL getin_r("kappa", thermo%kappa)
         CALL getin_r("cpp", thermo%cpp)
         CALL getin_r("cppv", thermo%cppv)
         CALL getin_r("Rv", thermo%Rv)
         CALL getin_r("preff", thermo%preff)
         thermo%which = thermo_binarygas
         thermo%nqdyn = 2
      CASE ('variable_Cp')
         CALL getin_r("kappa", thermo%kappa)
         CALL getin_r("cpp", thermo%cpp)
         CALL getin_r("preff", thermo%preff)
         CALL getin_r("Treff", thermo%Treff)
         CALL getin_r("variable_Cp_nu", thermo%nu)
         CALL getin_r("variable_Cp_T0", thermo%T0)
         thermo%which = thermo_variable_Cp
      CASE DEFAULT
         WRITE (msgbuf, *) 'Bad selector for variable thermo : <', TRIM(def), &
            '> options are <theta>, <entropy>, <binarygas>, <variable_Cp>'
         CALL errstop(TRIM(msgbuf))
      END SELECT

   END SUBROUTINE

   SUBROUTINE temperature_from_pVq(thermo, npt, nq, temp, p, vol, q)
      TYPE(thermodynamics_t), INTENT(IN) :: thermo
      INTEGER, INTENT(IN) :: npt, nq
      REAL(C_DOUBLE), INTENT(OUT) :: temp(npt)
      REAL(C_DOUBLE), INTENT(IN)  :: p(npt), vol(npt), q(npt, nq)

      REAL(C_DOUBLE) :: Rd, inv_Rd, qv, Rmix
      INTEGER :: ij

      SELECT CASE (thermo%which)
      CASE (thermo_theta, thermo_entropy, thermo_variable_Cp)
         inv_Rd = 1.0/(thermo%kappa*thermo%cpp)
         DO ij = 1, npt
            ! pV=RT
            temp(ij) = inv_Rd*p(ij)*vol(ij)
         END DO

      CASE (thermo_binarygas)
         Rd = thermo%kappa*thermo%cpp
         DO ij = 1, npt
            ! pV=RT
            qv = q(ij, 1)
            Rmix = qv*thermo%Rv + (1.-qv)*Rd
            temp(ij) = (1./Rmix)*p(ij)*vol(ij)
         END DO

      END SELECT

   END SUBROUTINE

   SUBROUTINE theta_from_pTq(thermo, npt, nq, theta, p, temp, q)
      TYPE(thermodynamics_t), INTENT(IN) :: thermo
      INTEGER, INTENT(IN) :: npt, nq
      REAL(C_DOUBLE), INTENT(OUT) :: theta(npt) ! conservative variable, definition depends on thermo%which
      REAL(C_DOUBLE), INTENT(IN)  :: p(npt), temp(npt), q(npt, nq)

      REAL(C_DOUBLE) :: preff, Treff, kappa, Cppd, Rd, Cppv, Rv
      REAL(C_DOUBLE) :: log_p_preff, chi, qd, qv, R, cpp
      REAL(C_DOUBLE) :: nu
      INTEGER :: ij
      preff = thermo%preff
      Treff = thermo%Treff
      kappa = thermo%kappa
      cppd = thermo%cpp
      Rv = thermo%Rv
      cppv = thermo%cppv
      nu = thermo%nu
      Rd = kappa*cppd

      SELECT CASE (thermo%which)
      CASE (thermo_theta)
         ! theta = potential temperature
         DO ij = 1, npt
            theta(ij) = temp(ij)*(p(ij)/preff)**(-kappa)
         END DO
      CASE (thermo_entropy)
         ! theta = specific entropy
         DO ij = 1, npt
            log_p_preff = log(p(ij)/preff)
            chi = log(temp(ij)/Treff)
            theta(ij) = cppd*chi - Rd*log_p_preff
         END DO
      CASE (thermo_variable_Cp)
         ! theta = specific entropy
         DO ij = 1, npt
            log_p_preff = log(p(ij)/preff)
            chi = (temp(ij)/Treff)**nu - 1
            theta(ij) = (cppd/nu)*chi - Rd*log_p_preff
         END DO
      CASE (thermo_binarygas)
         ! theta = potential enthalpy = Cp*potential temperature
         DO ij = 1, npt
            qv = q(ij, 1)    ! mass fraction of 'secondary' gas
            qd = 1.-qv       ! mass fraction of 'main' gas
            R = qv*Rv + qd*Rd
            cpp = qv*cppv + qd*cppd
            kappa = R*(1./cpp)
            theta(ij) = cpp*temp(ij)*(p(ij)/preff)**(-kappa)
         END DO
      CASE DEFAULT
         STOP
      END SELECT
   END SUBROUTINE

   SUBROUTINE volume_from_pCq(thermo, npt, nq, vol, p, theta, q)
      TYPE(thermodynamics_t), INTENT(IN) :: thermo
      INTEGER, INTENT(IN) :: npt, nq
      REAL(C_DOUBLE), INTENT(OUT) :: vol(npt)
      REAL(C_DOUBLE), INTENT(IN)  :: p(npt), theta(npt), q(npt, nq)

      REAL(C_DOUBLE) :: preff, Treff, kappa, Cppd, Rd, Cppv, Rv
      REAL(C_DOUBLE) :: qd, qv, R, cpp, temp
      REAL(C_DOUBLE) :: nu
      INTEGER :: ij
      preff = thermo%preff
      Treff = thermo%Treff
      kappa = thermo%kappa
      cppd = thermo%cpp
      Rv = thermo%Rv
      cppv = thermo%cppv
      nu = thermo%nu
      Rd = kappa*cppd

      SELECT CASE (thermo%which)
      CASE (thermo_theta)
         ! theta = potential temperature
         DO ij = 1, npt
            temp = theta(ij)*(p(ij)/preff)**(kappa)
            vol(ij) = Rd*temp/p(ij)
         END DO
      CASE (thermo_entropy)
         ! theta = specific entropy
         DO ij = 1, npt
            temp = Treff*exp((theta(ij) + Rd*log(p(ij)/preff))/cppd)
            vol(ij) = Rd*temp/p(ij)
         END DO
      CASE (thermo_variable_Cp)
         ! theta = specific entropy
         DO ij = 1, npt
            temp = Treff*exp(LOG1P(nu*(theta(ij) + Rd*log(p(ij)/preff))/cppd)/nu)
            vol(ij) = Rd*temp/p(ij)
         END DO
      CASE (thermo_binarygas)
         ! theta = potential enthalpy = Cp*potential temperature
         DO ij = 1, npt
            qv = q(ij, 1)    ! mass fraction of 'secondary' gas
            qd = 1.-qv       ! mass fraction of 'main' gas
            R = qv*Rv + qd*Rd
            cpp = qv*cppv + qd*cppd
            kappa = R*(1./cpp)
            ! theta = potential enthalpy = Cp*potential temperature
            temp = (theta(ij)/cpp)*(p(ij)/preff)**(kappa)
            vol(ij) = R*temp/p(ij)
         END DO
      CASE DEFAULT
         STOP
      END SELECT

   END SUBROUTINE

   SUBROUTINE temperature_from_pCq(thermo, npt, nptq, nq, temp, p, theta, q)
      TYPE(thermodynamics_t), INTENT(IN) :: thermo
      INTEGER, INTENT(IN) :: npt, nptq, nq
      REAL(C_DOUBLE), INTENT(OUT) :: temp(npt)
      REAL(C_DOUBLE), INTENT(IN)  :: p(npt), theta(npt), q(nptq, nq)

      REAL(C_DOUBLE) :: preff, Treff, kappa, Cppd, Rd, Cppv, Rv
      REAL(C_DOUBLE) :: qd, qv, R, cpp, log_p_preff
      REAL(C_DOUBLE) :: nu
      INTEGER :: ij
      preff = thermo%preff
      Treff = thermo%Treff
      kappa = thermo%kappa
      cppd = thermo%cpp
      Rv = thermo%Rv
      cppv = thermo%cppv
      nu = thermo%nu
      Rd = kappa*cppd

      SELECT CASE (thermo%which)
      CASE (thermo_theta)
         ! theta = potential temperature
         DO ij = 1, npt
            temp(ij) = theta(ij)*(p(ij)/preff)**(kappa)
         END DO
      CASE (thermo_entropy)
         ! theta = specific entropy
         DO ij = 1, npt
            temp(ij) = Treff*exp((theta(ij) + Rd*log(p(ij)/preff))/cppd)
         END DO
      CASE (thermo_variable_Cp) ! modifié
         ! theta = specific entropy
         DO ij = 1, npt
            temp(ij) = Treff*exp(LOG1P(nu*((theta(ij) + Rd*log(p(ij)/preff)))/cppd)/nu)
         END DO
      CASE (thermo_binarygas)
         ! theta = potential enthalpy = Cp*potential temperature
         DO ij = 1, npt
            qv = q(ij, 1)    ! mass fraction of 'secondary' gas
            qd = 1.-qv       ! mass fraction of 'main' gas
            R = qv*Rv + qd*Rd
            cpp = qv*cppv + qd*cppd
            kappa = R*(1./cpp)
            ! theta = potential enthalpy = Cp*potential temperature
            temp(ij) = (theta(ij)/cpp)*(p(ij)/preff)**(kappa)
         END DO
      CASE DEFAULT
         STOP
      END SELECT

   END SUBROUTINE

   SUBROUTINE enthalpy_from_pq(thermo, on_gpu, ni, istart, istop, nj, jstart, jstop, nq, enthalpy, vol, exner, p, q)
      TYPE(thermodynamics_t), INTENT(IN) :: thermo
      LOGICAL, INTENT(IN) :: on_gpu
      INTEGER, INTENT(IN), VALUE :: ni, istart, istop, nj, jstart, jstop, nq
      REAL(C_DOUBLE), INTENT(OUT) :: enthalpy(ni, nj), vol(ni, nj), exner(ni, nj, nq)
      REAL(C_DOUBLE), INTENT(IN)  :: p(ni, nj), q(ni, nj, nq)

      REAL(C_DOUBLE) :: inv_preff, Treff, kappa, Cppd, Rd, Cppv, Rv
      REAL(C_DOUBLE) :: qd, qv, R, cpp, inv_cpp, hpot, pi0, log_pi0, temp
      REAL(C_DOUBLE) :: nu
      REAL(C_DOUBLE) :: XXX, log_p_preff
      INTEGER :: i, j
      inv_preff = 1./thermo%preff
      Treff = thermo%Treff
      kappa = thermo%kappa
      cppd = thermo%cpp
      Rv = thermo%Rv
      cppv = thermo%cppv
      nu = thermo%nu
      Rd = kappa*cppd

      !$acc kernels async if(on_gpu) present(enthalpy(:,:), vol(:,:), exner(:,:,:), p(:,:), q(:,:,:))

      SELECT CASE (thermo%which)
      CASE (thermo_theta)
         ! theta = potential temperature
         DO j = jstart, jstop
            DO i = istart, istop
               pi0 = (p(i, j)*inv_preff)**(kappa) ! exner factor (non-dimensional)
               temp = q(i, j, 1)*pi0
               enthalpy(i, j) = cppd*temp
               vol(i, j) = Rd*temp/p(i, j)
               exner(i, j, 1) = cppd*pi0
            END DO
         END DO

      CASE (thermo_entropy)
         ! theta = specific entropy
         DO j = jstart, jstop
            DO i = istart, istop
               pi0 = (p(i, j)*inv_preff)**(kappa) ! exner factor (non-dimensional)
               temp = Treff*exp((q(i, j, 1) + Rd*log(p(i, j)*inv_preff))/cppd)
               enthalpy(i, j) = cppd*temp
               vol(i, j) = Rd*temp/p(i, j)
               exner(i, j, 1) = temp ! variable conjugate to entropy is temperature
            END DO
         END DO

      CASE (thermo_variable_Cp)
         ! theta = specific entropy
         DO j = jstart, jstop
            DO i = istart, istop
               pi0 = (p(i, j)*inv_preff)**(kappa) ! exner factor (non-dimensional)
               ! ! log_p_preff = log(p(i,j)*inv_preff)
               ! ! XXX = (q(i,j,1) + Rd*log_p_preff)/cppd
               ! ! temp = Treff*exp( LOG1P(nu* XXX)/nu )
               ! ! temp = Treff*exp(  LOG1P(nu*(   q(i, j, 1) + Rd*log(  p(i, j)*inv_preff  )    )/cppd)/nu    )
               ! ! temp = Treff*( 1. + nu* (q(i, j, 1) + Rd*log(p(i, j)*inv_preff))/cppd   )**(1./nu) ! syntaxe qui devrait marché car équivalente au commit
               ! temp = Treff*( (nu/cppd) * ( q(i, j, 1) + Rd*log(p(i, j)*inv_preff) ) + 1 )**(1./nu) ! syntaxe qui marche dans le dernier commit
               temp = Treff*exp(LOG1P((nu/cppd)*(q(i, j, 1) + Rd*log(p(i, j)*inv_preff)))/nu)
               enthalpy(i, j) = (cppd/(nu + 1))*((temp**(nu + 1))/(Treff**nu))
               vol(i, j) = Rd*temp/p(i, j)
               exner(i, j, 1) = temp ! variable conjugate to entropy is temperature
            END DO
         END DO

      CASE (thermo_binarygas)
         ! theta = potential enthalpy = Cp*potential temperature
         DO j = jstart, jstop
            DO i = istart, istop
               hpot = q(i, j, 1) ! potential enthalpy = Cp*potential temperature
               qv = q(i, j, 2)   ! mass fraction of 'secondary' gas
               qd = 1.-qv        ! mass fraction of 'main' gas
               R = qv*Rv + qd*Rd
               cpp = qv*cppv + qd*cppd
               inv_cpp = 1./cpp
               kappa = R*inv_cpp
               log_pi0 = kappa*log(inv_preff*p(i, j))
               pi0 = EXP(log_pi0) ! exner factor (non-dimensional)
               enthalpy(i, j) = hpot*pi0
               temp = inv_cpp*enthalpy(i, j)
               vol(i, j) = R*temp*(1./p(i, j))
               exner(i, j, 1) = pi0 ! temp / potential temp
               exner(i, j, 2) = log_pi0*temp*((Rv - Rd)*(1./kappa) + cppd - cppv)
            END DO
         END DO
      CASE DEFAULT
         STOP
      END SELECT

      !$acc end kernels

   END SUBROUTINE

END MODULE
