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) :: 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
    REAL(rstd),POINTER :: temp(:,:)
    REAL(rstd) :: friction_time
    INTEGER :: ind

    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

    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
    ENDDO
  END SUBROUTINE init_physics_old

  SUBROUTINE physics(f_ps,f_theta_rhodz,f_u) 
    USE theta2theta_rhodz_mod
    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)
    ENDDO
    CALL temperature2theta_rhodz(f_ps,f_temp,f_theta_rhodz)
  END SUBROUTINE physics

  SUBROUTINE compute_physics(temp_eq, temp, u)
    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)
          ENDDO
       ENDDO
    ENDDO
    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)
       ENDDO
    ENDDO
    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
       ENDDO
    ENDDO
    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)
       ENDDO
    ENDDO

    ! 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 theta2theta_rhodz_mod
    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
    ENDDO
    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 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