MODULE etat0_venus_mod
   USE icosa
   IMPLICIT NONE
   PRIVATE
   SAVE

   TYPE(t_field), POINTER :: f_temp_eq(:)
   TYPE(t_field), POINTER :: f_temp(:) ! buffer used for physics
   REAL(rstd), ALLOCATABLE :: temp_eq_packed(:, :)

   REAL(rstd), SAVE :: kfrict, kv
!$OMP THREADPRIVATE(kfrict, kv)

   REAL(rstd), PARAMETER :: tauCLee = 86400*25 ! 25 Earth days, cf Lebonnois 2012

   PUBLIC :: etat0, init_physics, full_physics

CONTAINS

!-------------------------------- "Physics" ----------------------------------------

   SUBROUTINE init_physics_old
      USE getin_mod
      IMPLICIT NONE
      REAL(rstd), POINTER :: temp(:, :)
      REAL(rstd) :: friction_time
      INTEGER :: ind

      kv = 0.15  ! vertical turbulent viscosity
      CALL getin('venus_diffusion', kv)
      ! @getin_doc rayleigh_friction_type realm=dissipation type=real default=86400
      !  Friction time scale used when `etat0=="venus"`.
      friction_time = 86400.  !friction
      CALL getin('venus_friction_time', friction_time)
      kfrict = 1./friction_time

      CALL allocate_field(f_temp, field_t, type_real, llm) ! Buffer for later use by physics

      PRINT *, 'Initializing Temp_eq (venus)'
      CALL allocate_field(f_temp_eq, field_t, type_real, llm)
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         temp = f_temp_eq(ind)
         CALL compute_temp_ref(.TRUE., iim*jjm, lat_i, temp) ! With meridional gradient
      END DO
   END SUBROUTINE init_physics_old

   SUBROUTINE physics(f_ps, f_theta_rhodz, f_u)
      USE icosa
      USE theta2theta_rhodz_mod
      IMPLICIT NONE
      TYPE(t_field), POINTER :: f_theta_rhodz(:)
      TYPE(t_field), POINTER :: f_u(:)
      TYPE(t_field), POINTER :: f_ps(:)
      REAL(rstd), POINTER :: temp(:, :)
      REAL(rstd), POINTER :: temp_eq(:, :)
      REAL(rstd), POINTER :: u(:, :)
      INTEGER :: ind

      CALL theta_rhodz2temperature(f_ps, f_theta_rhodz, f_temp)
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         u = f_u(ind)
         temp_eq = f_temp_eq(ind)
         temp = f_temp(ind)
         CALL compute_physics(temp_eq, temp, u)
      END DO
      CALL temperature2theta_rhodz(f_ps, f_temp, f_theta_rhodz)
   END SUBROUTINE physics

   SUBROUTINE compute_physics(temp_eq, temp, u)
      USE icosa
      USE theta2theta_rhodz_mod
      IMPLICIT NONE
      REAL(rstd), INTENT(IN)    :: temp_eq(iim*jjm, llm)
      REAL(rstd), INTENT(INOUT) :: temp(iim*jjm, llm)
      REAL(rstd), INTENT(INOUT) :: u(3*iim*jjm, llm)
      INTEGER :: i, j, l, ij
      DO l = 1, llm
         DO j = jj_begin - 1, jj_end + 1
            DO i = ii_begin - 1, ii_end + 1
               ij = (j - 1)*iim + i
               temp(ij, l) = temp(ij, l) - (temp(ij, l) - temp_eq(ij, l))*(dt*itau_physics/tauCLee)
            END DO
         END DO
      END DO
      u(:, 1) = u(:, 1)*(1.-dt*itau_physics*kfrict)
   END SUBROUTINE compute_physics

!----------------------- Re-implementation using physics_interface_mod -----------------

   SUBROUTINE init_physics
      USE getin_mod
      USE physics_interface_mod
      REAL(rstd), POINTER :: temp(:, :)
      REAL(rstd) :: friction_time
      INTEGER :: ngrid

      kv = 0.15  ! vertical turbulent viscosity
      CALL getin('venus_diffusion', kv)
      friction_time = 86400.  !friction
      CALL getin('venus_friction_time', friction_time)
      kfrict = 1./friction_time

      ngrid = physics_inout%ngrid
      ALLOCATE (temp_eq_packed(ngrid, llm))
      CALL compute_temp_ref(.TRUE., ngrid, physics_inout%lat, temp_eq_packed) ! With meridional gradient
   END SUBROUTINE init_physics

   SUBROUTINE full_physics
      USE physics_interface_mod
      CALL compute_physics_column(physics_inout%ngrid, physics_inout%dt_phys, &
                                  physics_inout%p, physics_inout%geopot, &
                                  physics_inout%Temp, physics_inout%ulon, physics_inout%ulat, &
                                  physics_inout%dTemp, physics_inout%dulon, physics_inout%dulat)
   END SUBROUTINE full_physics

   SUBROUTINE compute_physics_column(ngrid, dt_phys, p, geopot, Temp, u, v, dTemp, du, dv)
      USE earth_const, only: g
      INTEGER, INTENT(IN)   :: ngrid
      REAL(rstd), INTENT(IN) :: dt_phys, &
                                p(ngrid, llm + 1), geopot(ngrid, llm + 1), &
                                Temp(ngrid, llm), u(ngrid, llm), v(ngrid, llm)
      REAL(rstd), INTENT(OUT) :: dTemp(ngrid, llm), du(ngrid, llm), dv(ngrid, llm)
      REAL(rstd) :: mass(ngrid, llm), &  ! rho.dz
                    A(ngrid, llm + 1), &  ! off-diagonal coefficients
                    B(ngrid, llm), &  ! diagonal coefficients
                    Ru(ngrid, llm), Rv(ngrid, llm), & ! right-hand-sides
                    C(ngrid, llm), &  ! LU factorization (Thomas algorithm)
                    xu(ngrid, llm), xv(ngrid, llm)      ! solution
      REAL(rstd) :: rho, X_ij
      INTEGER :: l, ij
      ! Vertical diffusion : rho.du/dt = d/dz (rho*kappa*du/dz)
      ! rho.dz = mass.deta
      ! => mass.du/dt = d/deta ((rho^2 kappa/m)du/deta)
      ! Backward Euler : (M-S)u_new = M.u_old
      ! with M.u = mass(l)*u(l)
      ! and  S.u = A(l+1)*(u(l+1)-u(l)) - A(l)*(u(l)-u(l-1))
      ! => solve -A(l)u(l-1) + B(l)u(l) - A(l+1)u(l+1) = Ru(l) using Thomas algorithm
      ! with A=tau*kappa*rho^2/m and B(l) = mass(l)+A(l)+A(l+1)
      DO l = 1, llm
         !DIR$ SIMD
         DO ij = 1, ngrid
            mass(ij, l) = (p(ij, l) - p(ij, l + 1))*(1./g)
            rho = (p(ij, l) - p(ij, l + 1))/(geopot(ij, l + 1) - geopot(ij, l))
            ! A = kappa.tau.rho^2/m
            A(ij, l) = (kv*dt_phys)*(rho**2)/mass(ij, l)
            Ru(ij, l) = mass(ij, l)*u(ij, l)
            Rv(ij, l) = mass(ij, l)*v(ij, l)
         END DO
      END DO
      A(:, llm + 1) = 0.
      DO l = llm, 2, -1
         !DIR$ SIMD
         DO ij = 1, ngrid
            A(ij, l) = .5*(A(ij, l) + A(ij, l - 1)) ! average A to interfaces
         END DO
      END DO
      A(:, 1) = 0.
      DO l = 1, llm
         !DIR$ SIMD
         DO ij = 1, ngrid
            B(ij, l) = mass(ij, l) + A(ij, l) + A(ij, l + 1)
         END DO
      END DO

      ! Solve -A(l)x(l-1) + B(l)x(l) - A(l+1)x(l+1) = R(l) using Thomas algorithm
      ! Forward sweep :
      ! C(0)=0, C(l) = -A(l+1) / (B(l)+A(l)C(l-1)),
      ! D(0)=0, D(l) = (R(l)+A(l)D(l-1)) / (B(l)+A(l)C(l-1))
      !DIR$ SIMD
      DO ij = 1, ngrid
         X_ij = 1./B(ij, 1)
         C(ij, 1) = -A(ij, 2)*X_ij
         xu(ij, 1) = Ru(ij, 1)*X_ij
         xv(ij, 1) = Rv(ij, 1)*X_ij
      END DO
      DO l = 2, llm
         !DIR$ SIMD
         DO ij = 1, ngrid
            X_ij = 1./(B(ij, l) + A(ij, l)*C(ij, l - 1))
            C(ij, l) = -A(ij, l + 1)*X_ij ! zero for l=llm
            xu(ij, l) = (Ru(ij, l) + A(ij, l)*xu(ij, l - 1))*X_ij
            xv(ij, l) = (Rv(ij, l) + A(ij, l)*xv(ij, l - 1))*X_ij
         END DO
      END DO
      ! Back substitution :
      ! x(i) = D(i)-C(i)x(i+1), x(llm)=D(llm)
      !DIR$ SIMD
      DO ij = 1, ngrid
         ! top layer l=llm
         du(ij, llm) = (xu(ij, llm) - u(ij, llm))/dt_phys
         dv(ij, llm) = (xv(ij, llm) - v(ij, llm))/dt_phys
      END DO
      ! Back substitution at lower layers
      DO l = llm - 1, 1, -1
         !DIR$ SIMD
         DO ij = 1, ngrid
            xu(ij, l) = xu(ij, l) - C(ij, l)*xu(ij, l + 1)
            xv(ij, l) = xv(ij, l) - C(ij, l)*xv(ij, l + 1)
            du(ij, l) = (xu(ij, l) - u(ij, l))/dt_phys
            dv(ij, l) = (xv(ij, l) - v(ij, l))/dt_phys
         END DO
      END DO

      IF (.FALSE.) THEN
         ! check correctness of solution of tridiagonal problem
         PRINT *, 'Max 1 residual of tridiagonal solver', MAXVAL(ABS(Ru))
         DO l = 2, llm - 1
            !DIR$ SIMD
            DO ij = 1, ngrid
               Ru(ij, l) = Ru(ij, l) - B(ij, l)*xu(ij, l) + A(ij, l)*xu(ij, l - 1) + A(ij, l + 1)*xu(ij, l + 1)
            END DO
         END DO
         DO ij = 1, ngrid
            Ru(ij, 1) = Ru(ij, 1) - B(ij, 1)*xu(ij, 1) + A(ij, 2)*xu(ij, 2)
            Ru(ij, llm) = Ru(ij, llm) - B(ij, llm)*xu(ij, llm) + A(ij, llm)*xu(ij, llm - 1)
         END DO
         PRINT *, 'Max 2 residual of tridiagonal solver', MAXVAL(ABS(Ru))
      END IF

      ! bottom friction + thermal relaxation
      du(:, 1) = du(:, 1) - kfrict*u(:, 1)
      dTemp(:, :) = (1./tauCLee)*(temp_eq_packed(:, :) - temp(:, :))
   END SUBROUTINE compute_physics_column

!----------------------------- Initialize to T_eq --------------------------------------

   SUBROUTINE etat0(f_ps, f_phis, f_theta_rhodz, f_u, f_q)
      USE icosa
      USE theta2theta_rhodz_mod
      IMPLICIT NONE
      TYPE(t_field), POINTER :: f_ps(:)
      TYPE(t_field), POINTER :: f_phis(:)
      TYPE(t_field), POINTER :: f_theta_rhodz(:)
      TYPE(t_field), POINTER :: f_u(:)
      TYPE(t_field), POINTER :: f_q(:)

      TYPE(t_field), POINTER :: f_temp(:)
      REAL(rstd), POINTER :: temp(:, :)
      REAL(rstd), POINTER :: ps(:)
      REAL(rstd), POINTER :: phis(:)
      REAL(rstd), POINTER :: u(:, :)
      REAL(rstd), POINTER :: q(:, :, :)
      INTEGER :: ind

      CALL allocate_field(f_temp, field_t, type_real, llm, name='temp_buf_venus')
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ps = f_ps(ind)
         ps(:) = preff
         phis = f_phis(ind)
         phis(:) = 0.
         u = f_u(ind)
         u(:, :) = 0
         q = f_q(ind)
         q(:, :, :) = 1e2
         temp = f_temp(ind)
         CALL compute_temp_ref(.FALSE., iim*jjm, lat_i, temp) ! Without meridional gradient
      END DO
      CALL temperature2theta_rhodz(f_ps, f_temp, f_theta_rhodz)
      CALL deallocate_field(f_temp)

   END SUBROUTINE etat0

!------------------------- Compute reference temperature field ------------------------

   SUBROUTINE compute_temp_ref(gradient, ngrid, lat, temp_eq)
      USE icosa
      USE disvert_mod
      USE omp_para
      USE math_const
      LOGICAL, INTENT(IN) :: gradient
      INTEGER, INTENT(IN) :: ngrid
      REAL(rstd), INTENT(IN) :: lat(ngrid) ! latitude
      REAL(rstd), INTENT(OUT) :: temp_eq(ngrid, llm)

      REAL(rstd)  :: clat(ngrid)
      INTEGER :: level
      INTEGER, PARAMETER :: nlev = 30
      REAL ::  pressCLee(nlev + 1), tempCLee(nlev + 1), dt_epCLee(nlev + 1), etaCLee(nlev + 1)

      REAL(rstd) ::  pplay, ztemp, zdt, fact
      INTEGER :: ij, l, ll

      data etaCLee/9.602e-1, 8.679e-1, 7.577e-1, 6.420e-1, 5.299e-1, &
         4.273e-1, 3.373e-1, 2.610e-1, 1.979e-1, 1.472e-1, &
         1.074e-1, 7.672e-2, 5.361e-2, 3.657e-2, 2.430e-2, &
         1.569e-2, 9.814e-3, 5.929e-3, 3.454e-3, 1.934e-3, &
         1.043e-3, 5.400e-4, 2.710e-4, 1.324e-4, 6.355e-5, &
         3.070e-5, 1.525e-5, 7.950e-6, 4.500e-6, 2.925e-6, &
         2.265e-6/
      data tempCLee/728.187, 715.129, 697.876, 677.284, 654.078, 628.885, &
         602.225, 574.542, 546.104, 517.339, 488.560, 459.932, &
         431.741, 404.202, 377.555, 352.042, 327.887, 305.313, &
         284.556, 265.697, 248.844, 233.771, 220.368, 208.247, &
         197.127, 187.104, 178.489, 171.800, 167.598, 165.899, &
         165.676/
      data dt_epCLee/6.101, 6.136, 6.176, 6.410, 6.634, 6.678, &
         6.719, 6.762, 7.167, 7.524, 9.840, 14.948, &
         21.370, 28.746, 36.373, 43.315, 48.534, 51.175, &
         50.757, 47.342, 41.536, 34.295, 26.758, 19.807, &
         14.001, 9.599, 6.504, 4.439, 3.126, 2.370, &
         2.000/

      pressCLee(:) = etaCLee(:)*9.2e6
      clat(:) = COS(lat(:))

      DO ij = 1, ngrid
            DO l = 1, llm

               pplay = .5*(ap(l) + ap(l + 1) + (bp(l) + bp(l + 1))*preff) ! ps=preff
               ! look for largest level such that pressCLee(level) > pplay(ij,l))
               ! => pressClee(level+1) < pplay(ij,l) < pressClee(level)
               level = 1
               DO ll = 1, nlev ! nlev data levels
                  IF (pressCLee(ll) > pplay) THEN
                     level = ll
                  END IF
               END DO

               ! interpolate between level and level+1
               ! interpolation is linear in log(pressure)
               fact = (log10(pplay) - log10(pressCLee(level))) &
                      /(log10(pressCLee(level + 1)) - log10(pressCLee(level)))
               ztemp = tempCLee(level)*(1 - fact) + tempCLee(level + 1)*fact
               zdt = dt_epCLee(level)*(1 - fact) + dt_epCLee(level + 1)*fact

               IF (gradient) THEN
                  temp_eq(ij, l) = ztemp + zdt*(clat(ij) - Pi/4.)
               ELSE
                  temp_eq(ij, l) = ztemp
               END IF
         END DO
      END DO
   END SUBROUTINE compute_temp_ref

END MODULE etat0_venus_mod
