SUBMODULE(compute_caldyn_mod) compute_caldyn_solver_mod
   USE prec, only: rstd
   IMPLICIT NONE

CONTAINS

   MODULE SUBROUTINE compute_caldyn_solver(tau, phis, rhodz, theta, pk, geopot, W, m_il, pres, dPhi, dW, du)
      REAL(rstd), INTENT(IN) :: tau ! "solve" Phi-tau*dPhi/dt = Phi_rhs
      REAL(rstd), INTENT(IN)    :: phis(:)
      REAL(rstd), INTENT(IN)    :: rhodz(:, :)
      REAL(rstd), INTENT(IN)    :: theta(:, :, :)
      REAL(rstd), INTENT(OUT)   :: pk(:, :)
      REAL(rstd), INTENT(INOUT) :: geopot(:, :)
      REAL(rstd), INTENT(INOUT) :: W(:, :) ! OUT if tau>0
      REAL(rstd), INTENT(OUT)   :: m_il(:, :)        ! rhodz averaged to interfaces
      REAL(rstd), INTENT(OUT)   :: pres(:, :)          ! pressure
      REAL(rstd), INTENT(OUT)   :: dW(:, :)
      REAL(rstd), INTENT(OUT)   :: dPhi(:, :)
      REAL(rstd), INTENT(OUT)   :: du(:, :)
      CALL compute_caldyn_solver_hex(tau, phis, rhodz, theta, pk, geopot, W, m_il, pres, dPhi, dW, du)
   END SUBROUTINE

   SUBROUTINE compute_caldyn_solver_hex(tau, phis, rhodz, theta, pk, geopot, W, m_il, pres, dPhi, dW, du)
      USE dimensions ! iim, jjm, offsets ...
      USE geometry ! signs
      USE grid_param, ONLY: llm, nqdyn
      USE trace, ONLY: trace_start, trace_end
      USE omp_para, ONLY: ll_begin, ll_beginp1, ll_end, ll_endp1
      USE disvert_mod, ONLY: ptop
      USE caldyn_vars_mod, ONLY: dysl
      USE earth_const, ONLY: thermo_theta, thermo_entropy, caldyn_thermo, cpp, kappa, g, preff, Treff
      REAL(rstd), INTENT(IN) :: tau ! "solve" Phi-tau*dPhi/dt = Phi_rhs
      REAL(rstd), INTENT(IN)    :: phis(iim*jjm)
      REAL(rstd), INTENT(IN)    :: rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(IN)    :: theta(iim*jjm, llm, nqdyn)
      REAL(rstd), INTENT(OUT)   :: pk(iim*jjm, llm)
      REAL(rstd), INTENT(INOUT) :: geopot(iim*jjm, llm + 1)
      REAL(rstd), INTENT(INOUT) :: W(iim*jjm, llm + 1) ! OUT if tau>0
      REAL(rstd), INTENT(OUT)   :: m_il(iim*jjm, llm + 1)        ! rhodz averaged to interfaces
      REAL(rstd), INTENT(OUT)   :: pres(iim*jjm, llm)          ! pressure
      REAL(rstd), INTENT(OUT)   :: dW(iim*jjm, llm + 1)
      REAL(rstd), INTENT(OUT)   :: dPhi(iim*jjm, llm + 1)
      REAL(rstd), INTENT(OUT)   :: du(3*iim*jjm, llm)

      REAL(rstd) :: berni(iim*jjm, llm)         ! (W/m_il)^2
      REAL(rstd) :: berni1(iim*jjm)         ! (W/m_il)^2
      REAL(rstd) :: gamma, rho_ij, T_ij, X_ij, vreff, Rd, Cvd
      INTEGER    :: ij, l

      CALL trace_start("compute_caldyn_solver")

      Rd = cpp*kappa

      IF (dysl) THEN

!$OMP BARRIER

         !--------------------------------------------------------------------------
         !---------------------------- caldyn_mil ----------------------------------
         IF (ll_begin == 1) THEN
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               m_il(ij, 1) = .5*rhodz(ij, 1)
            END DO
         END IF
         DO l = ll_beginp1, ll_end
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               m_il(ij, l) = .5*(rhodz(ij, l) + rhodz(ij, l - 1))
            END DO
         END DO
         IF (ll_endp1 == llm + 1) THEN
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               m_il(ij, llm + 1) = .5*rhodz(ij, llm + 1 - 1)
            END DO
         END IF
         !---------------------------- caldyn_mil ----------------------------------
         !--------------------------------------------------------------------------
         IF (tau > 0) THEN ! solve implicit problem for geopotential
            CALL compute_NH_geopot_hex(tau, phis, rhodz, m_il, theta, W, geopot)
         END IF

         !--------------------------------------------------------------------------
         !---------------------------- caldyn_solver ----------------------------------
         !
         ! Compute pressure (pres) and Exner function (pk)
         ! kappa = R/Cp
         ! 1-kappa = Cv/Cp
         ! Cp/Cv = 1/(1-kappa)
         gamma = 1./(1.-kappa)
         vreff = Rd*Treff/preff ! reference specific volume
         Cvd = cpp - Rd
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               SELECT CASE (caldyn_thermo)
               CASE (thermo_theta)
                  rho_ij = g*rhodz(ij, l)/(geopot(ij, l + 1) - geopot(ij, l))
                  X_ij = (cpp/preff)*kappa*theta(ij, l, 1)*rho_ij
                  ! kappa.theta.rho = p/exner
                  ! => X = (p/p0)/(exner/Cp)
                  ! = (p/p0)^(1-kappa)
                  pres(ij, l) = preff*(X_ij**gamma) ! pressure
                  ! Compute Exner function (needed by compute_caldyn_fast)
                  ! other formulae possible if exponentiation is slow
                  pk(ij, l) = cpp*((pres(ij, l)/preff)**kappa) ! Exner
               CASE (thermo_entropy)
                  rho_ij = g*rhodz(ij, l)/(geopot(ij, l + 1) - geopot(ij, l))
                  T_ij = Treff*exp((theta(ij, l, 1) + Rd*log(vreff*rho_ij))/Cvd)
                  pres(ij, l) = rho_ij*Rd*T_ij
                  pk(ij, l) = T_ij
               CASE DEFAULT
                  STOP
               END SELECT
            END DO
         END DO
         ! We need a barrier here because we compute pres above and do a vertical difference below
         !$OMP BARRIER
         IF (ll_begin == 1) THEN
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               ! Lower BC
               dW(ij, 1) = (1./g)*(pbot - rho_bot*(geopot(ij, 1) - phis(ij)) - pres(ij, 1)) - m_il(ij, 1)
               W(ij, 1) = W(ij, 1) + tau*dW(ij, 1) ! update W
               dPhi(ij, 1) = g*g*W(ij, 1)/m_il(ij, 1)
            END DO
         END IF
         DO l = ll_beginp1, ll_end
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               dW(ij, l) = (1./g)*(pres(ij, l - 1) - pres(ij, l)) - m_il(ij, l)
               W(ij, l) = W(ij, l) + tau*dW(ij, l) ! update W
               dPhi(ij, l) = g*g*W(ij, l)/m_il(ij, l)
            END DO
         END DO
         IF (ll_endp1 == llm + 1) THEN
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               ! Top BC
               dW(ij, llm + 1) = (1./g)*(pres(ij, llm + 1 - 1) - ptop) - m_il(ij, llm + 1)
               W(ij, llm + 1) = W(ij, llm + 1) + tau*dW(ij, llm + 1) ! update W
               dPhi(ij, llm + 1) = g*g*W(ij, llm + 1)/m_il(ij, llm + 1)
            END DO
         END IF
         ! We need a barrier here because we update W above and do a vertical average below
         !$OMP BARRIER
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               ! compute du = -0.5*g^2.grad(w^2)
               berni(ij, l) = (-.25*g*g)*((W(ij, l)/m_il(ij, l))**2 + (W(ij, l + 1)/m_il(ij, l + 1))**2)
            END DO
         END DO
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               du(ij + u_right, l) = ne_right*(berni(ij, l) - berni(ij + t_right, l))
               du(ij + u_lup, l) = ne_lup*(berni(ij, l) - berni(ij + t_lup, l))
               du(ij + u_ldown, l) = ne_ldown*(berni(ij, l) - berni(ij + t_ldown, l))
            END DO
         END DO
         !---------------------------- caldyn_solver ----------------------------------
         !--------------------------------------------------------------------------

!$OMP BARRIER

      ELSE

         ! FIXME : vertical OpenMP parallelism will not work

         ! average m_ik to interfaces => m_il
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            m_il(ij, 1) = .5*rhodz(ij, 1)
         END DO
         DO l = 2, llm
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               m_il(ij, l) = .5*(rhodz(ij, l - 1) + rhodz(ij, l))
            END DO
         END DO
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            m_il(ij, llm + 1) = .5*rhodz(ij, llm)
         END DO

         IF (tau > 0) THEN ! solve implicit problem for geopotential
            CALL compute_NH_geopot_hex(tau, phis, rhodz, m_il, theta, W, geopot)
         END IF

         ! Compute pressure, stored temporarily in pk
         ! kappa = R/Cp
         ! 1-kappa = Cv/Cp
         ! Cp/Cv = 1/(1-kappa)
         gamma = 1./(1.-kappa)
         DO l = 1, llm
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               rho_ij = (g*rhodz(ij, l))/(geopot(ij, l + 1) - geopot(ij, l))
               X_ij = (cpp/preff)*kappa*theta(ij, l, 1)*rho_ij
               ! kappa.theta.rho = p/exner
               ! => X = (p/p0)/(exner/Cp)
               !      = (p/p0)^(1-kappa)
               pk(ij, l) = preff*(X_ij**gamma)
            END DO
         END DO

         ! Update W, compute tendencies
         DO l = 2, llm
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               dW(ij, l) = (1./g)*(pk(ij, l - 1) - pk(ij, l)) - m_il(ij, l)
               W(ij, l) = W(ij, l) + tau*dW(ij, l) ! update W
               dPhi(ij, l) = g*g*W(ij, l)/m_il(ij, l)
            END DO
            !       PRINT *,'Max dPhi', l,ij_begin,ij_end, MAXVAL(abs(dPhi(ij_begin:ij_end,l)))
            !       PRINT *,'Max dW', l,ij_begin,ij_end, MAXVAL(abs(dW(ij_begin:ij_end,l)))
         END DO
         ! Lower BC (FIXME : no orography yet !)
         DO ij = ij_begin, ij_end
            dPhi(ij, 1) = 0
            W(ij, 1) = 0
            dW(ij, 1) = 0
            dPhi(ij, llm + 1) = 0 ! rigid lid
            W(ij, llm + 1) = 0
            dW(ij, llm + 1) = 0
         END DO
         ! Upper BC p=ptop
         !    DO ij=ij_omp_begin_ext,ij_omp_end_ext
         !       dPhi(ij,llm+1) = W(ij,llm+1)/rhodz(ij,llm)
         !       dW(ij,llm+1) = (1./g)*(pk(ij,llm)-ptop) - .5*rhodz(ij,llm)
         !    ENDDO

         ! Compute Exner function (needed by compute_caldyn_fast) and du=-g^2.grad(w^2)
         DO l = 1, llm
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               pk(ij, l) = cpp*((pk(ij, l)/preff)**kappa) ! other formulae possible if exponentiation is slow
               berni1(ij) = (-.25*g*g)*( &
                            (W(ij, l)/m_il(ij, l))**2 &
                            + (W(ij, l + 1)/m_il(ij, l + 1))**2)
            END DO
            DO ij = ij_begin, ij_end
               du(ij + u_right, l) = ne_right*(berni1(ij) - berni1(ij + t_right))
               du(ij + u_lup, l) = ne_lup*(berni1(ij) - berni1(ij + t_lup))
               du(ij + u_ldown, l) = ne_ldown*(berni1(ij) - berni1(ij + t_ldown))
            END DO
         END DO

      END IF ! dysl

      CALL trace_end("compute_caldyn_solver")

   END SUBROUTINE compute_caldyn_solver_hex

   SUBROUTINE compute_NH_geopot_hex(tau, phis, m_ik, m_il, theta, W_il, Phi_il)
      USE disvert_mod
      USE caldyn_vars_mod
      USE omp_para
      USE grid_param
      REAL(rstd), INTENT(IN)    :: tau ! solve Phi-tau*dPhi/dt = Phi_rhs
      REAL(rstd), INTENT(IN)    :: phis(iim*jjm)
      REAL(rstd), INTENT(IN)    :: m_ik(iim*jjm, llm)
      REAL(rstd), INTENT(IN)    :: m_il(iim*jjm, llm + 1)
      REAL(rstd), INTENT(IN)    :: theta(iim*jjm, llm)
      REAL(rstd), INTENT(IN)    :: W_il(iim*jjm, llm + 1)   ! vertical momentum
      REAL(rstd), INTENT(INOUT) :: Phi_il(iim*jjm, llm + 1) ! geopotential

      REAL(rstd) :: Phi_star_il(iim*jjm, llm + 1)
      REAL(rstd) :: p_ik(iim*jjm, llm)      ! pressure
      REAL(rstd) :: R_il(iim*jjm, llm + 1)    ! rhs of tridiag problem
      REAL(rstd) :: x_il(iim*jjm, llm + 1)    ! solution of tridiag problem
      REAL(rstd) :: A_ik(iim*jjm, llm)      ! off-diagonal coefficients of tridiag problem
      REAL(rstd) :: B_il(iim*jjm, llm + 1)    ! diagonal coefficients of tridiag problem
      REAL(rstd) :: C_ik(iim*jjm, llm)      ! Thomas algorithm
      REAL(rstd) :: D_il(iim*jjm, llm + 1)    ! Thomas algorithm
      REAL(rstd) :: gamma, rho_ij, X_ij
      REAL(rstd) :: tau2_g, g2, gm2, ml_g2, c2_mik, vreff

      INTEGER    :: iter, ij, l, ij_omp_begin_ext, ij_omp_end_ext

      CALL distrib_level(ij_begin_ext, ij_end_ext, ij_omp_begin_ext, ij_omp_end_ext)

      IF (dysl) THEN

         !--------------------------------------------------------------------------
         !---------------------------- compute_NH_geopot ----------------------------------
         tau2_g = tau*tau/g
         g2 = g*g
         gm2 = 1./g2
         ! vreff = Treff*cpp/preff*kappa
         gamma = 1./(1.-kappa)
         !$OMP BARRIER
         ! compute Phi_star
         DO l = 1, llm + 1
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               Phi_star_il(ij, l) = Phi_il(ij, l) + tau*g2*(W_il(ij, l)/m_il(ij, l) - tau)
            END DO
         END DO
         ! Newton-Raphson iteration : Phi_il contains current guess value
         DO iter = 1, 2 ! 2 iterations should be enough
            ! Compute pressure, A_ik
            SELECT CASE (caldyn_thermo)
            CASE (thermo_theta)
               DO l = 1, llm
                  !DIR$ SIMD
                  DO ij = ij_omp_begin_ext, ij_omp_end_ext
                     rho_ij = (g*m_ik(ij, l))/(Phi_il(ij, l + 1) - Phi_il(ij, l))
                     X_ij = (cpp/preff)*kappa*theta(ij, l)*rho_ij
                     p_ik(ij, l) = preff*(X_ij**gamma)
                     c2_mik = gamma*p_ik(ij, l)/(rho_ij*m_ik(ij, l)) ! c^2 = gamma*R*T = gamma*p/rho
                     A_ik(ij, l) = c2_mik*(tau/g*rho_ij)**2
                  END DO
               END DO
            CASE (thermo_entropy)
               DO l = 1, llm
                  !DIR$ SIMD
                  DO ij = ij_omp_begin_ext, ij_omp_end_ext
                     rho_ij = (g*m_ik(ij, l))/(Phi_il(ij, l + 1) - Phi_il(ij, l))
!                     X_ij = log(vreff*rho_ij) + theta(ij, l)/cpp
!                     p_ik(ij, l) = preff*exp(X_ij*gamma)
                     X_ij = Treff*exp(theta(ij, l)/cpp) ! theta = Tref.exp(s/Cp)
                     X_ij = (cpp/preff)*kappa*X_ij*rho_ij
                     p_ik(ij, l) = preff*(X_ij**gamma)
                     c2_mik = gamma*p_ik(ij, l)/(rho_ij*m_ik(ij, l)) ! c^2 = gamma*R*T = gamma*p/rho
                     A_ik(ij, l) = c2_mik*(tau/g*rho_ij)**2
                  END DO
               END DO
            CASE DEFAULT
               PRINT *, 'caldyn_thermo not supported by compute_NH_geopot', caldyn_thermo
               STOP
            END SELECT
            ! NB : A(1), A(llm), R(1), R(llm+1) = 0 => x(l)=0 at l=1,llm+1 => flat, rigid top and bottom
            ! Solve -A(l-1)x(l-1) + B(l)x(l) - A(l)x(l+1) = R(l) using Thomas algorithm
            ! Compute residual R_il and B_il
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               ! bottom interface l=1
               ml_g2 = gm2*m_il(ij, 1)
               B_il(ij, 1) = A_ik(ij, 1) + ml_g2 + tau2_g*rho_bot
               R_il(ij, 1) = ml_g2*(Phi_il(ij, 1) - Phi_star_il(ij, 1)) &
                             + tau2_g*(p_ik(ij, 1) - pbot + rho_bot*(Phi_il(ij, 1) - phis(ij)))
            END DO
            DO l = 2, llm
               !DIR$ SIMD
               DO ij = ij_omp_begin_ext, ij_omp_end_ext
                  ! inner interfaces
                  ml_g2 = gm2*m_il(ij, l)
                  B_il(ij, l) = A_ik(ij, l) + A_ik(ij, l - 1) + ml_g2
                  R_il(ij, l) = ml_g2*(Phi_il(ij, l) - Phi_star_il(ij, l)) &
                                + tau2_g*(p_ik(ij, l) - p_ik(ij, l - 1))
                  ! consistency check : if Wil=0 and initial state is in hydrostatic balance
                  ! then Phi_star_il(ij,l) = Phi_il(ij,l) - tau^2*g^2
                  ! and residual = tau^2*(ml+(1/g)dl_pi)=0
               END DO
            END DO
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               ! top interface l=llm+1
               ml_g2 = gm2*m_il(ij, llm + 1)
               B_il(ij, llm + 1) = A_ik(ij, llm + 1 - 1) + ml_g2
               R_il(ij, llm + 1) = ml_g2*(Phi_il(ij, llm + 1) - Phi_star_il(ij, llm + 1)) &
                                   + tau2_g*(ptop - p_ik(ij, llm + 1 - 1))
            END DO
            !
            ! Forward sweep :
            ! C(0)=0, C(l) = -A(l) / (B(l)+A(l-1)C(l-1)),
            ! D(0)=0, D(l) = (R(l)+A(l-1)D(l-1)) / (B(l)+A(l-1)C(l-1))
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               X_ij = 1./B_il(ij, 1)
               C_ik(ij, 1) = -A_ik(ij, 1)*X_ij
               D_il(ij, 1) = R_il(ij, 1)*X_ij
            END DO
            DO l = 2, llm
               !DIR$ SIMD
               DO ij = ij_omp_begin_ext, ij_omp_end_ext
                  X_ij = 1./(B_il(ij, l) + A_ik(ij, l - 1)*C_ik(ij, l - 1))
                  C_ik(ij, l) = -A_ik(ij, l)*X_ij
                  D_il(ij, l) = (R_il(ij, l) + A_ik(ij, l - 1)*D_il(ij, l - 1))*X_ij
               END DO
            END DO
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               X_ij = 1./(B_il(ij, llm + 1) + A_ik(ij, llm + 1 - 1)*C_ik(ij, llm + 1 - 1))
               D_il(ij, llm + 1) = (R_il(ij, llm + 1) + A_ik(ij, llm + 1 - 1)*D_il(ij, llm + 1 - 1))*X_ij
               ! Back substitution :
               ! x(i) = D(i)-C(i)x(i+1), x(llm+1)=0
               ! + Newton-Raphson update
               ! top interface l=llm+1
               x_il(ij, llm + 1) = D_il(ij, llm + 1)
               Phi_il(ij, llm + 1) = Phi_il(ij, llm + 1) - x_il(ij, llm + 1)
            END DO
            DO l = llm, 1, -1
               !DIR$ SIMD
               DO ij = ij_omp_begin_ext, ij_omp_end_ext
                  ! Back substitution at lower interfaces
                  x_il(ij, l) = D_il(ij, l) - C_ik(ij, l)*x_il(ij, l + 1)
                  Phi_il(ij, l) = Phi_il(ij, l) - x_il(ij, l)
               END DO
            END DO
            IF (debug_hevi_solver) THEN
               PRINT *, '[hevi_solver] A,B', iter, MAXVAL(ABS(A_ik)), MAXVAL(ABS(B_il))
               PRINT *, '[hevi_solver] C,D', iter, MAXVAL(ABS(C_ik)), MAXVAL(ABS(D_il))
               DO l = 1, llm + 1
                  WRITE (*, '(A,I2.1,I3.2,E9.2)') '[hevi_solver] x_il', iter, l, MAXVAL(ABS(x_il(l, :)))
               END DO
               DO l = 1, llm + 1
                  WRITE (*, '(A,I2.1,I3.2,E9.2)') '[hevi_solver] R_il', iter, l, MAXVAL(ABS(R_il(l, :)))
               END DO
            END IF
         END DO ! Newton-Raphson
         !$OMP BARRIER
         !---------------------------- compute_NH_geopot ----------------------------------
         !--------------------------------------------------------------------------

      ELSE
!    FIXME : vertical OpenMP parallelism will not work

         tau2_g = tau*tau/g
         g2 = g*g
         gm2 = g**(-2)
         gamma = 1./(1.-kappa)

         ! compute Phi_star
         DO l = 1, llm + 1
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               Phi_star_il(ij, l) = Phi_il(ij, l) + tau*g2*(W_il(ij, l)/m_il(ij, l) - tau)
            END DO
         END DO

         ! Newton-Raphson iteration : Phi_il contains current guess value
         DO iter = 1, 5 ! 2 iterations should be enough

            ! Compute pressure, A_ik
            DO l = 1, llm
               !DIR$ SIMD
               DO ij = ij_begin_ext, ij_end_ext
                  rho_ij = (g*m_ik(ij, l))/(Phi_il(ij, l + 1) - Phi_il(ij, l))
                  X_ij = (cpp/preff)*kappa*theta(ij, l)*rho_ij
                  p_ik(ij, l) = preff*(X_ij**gamma)
                  c2_mik = gamma*p_ik(ij, l)/(rho_ij*m_ik(ij, l)) ! c^2 = gamma*R*T = gamma*p/rho
                  A_ik(ij, l) = c2_mik*(tau/g*rho_ij)**2
               END DO
            END DO

            ! Compute residual, B_il
            ! bottom interface l=1
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               ml_g2 = gm2*m_il(ij, 1)
               B_il(ij, 1) = A_ik(ij, 1) + ml_g2 + tau2_g*rho_bot
               R_il(ij, 1) = ml_g2*(Phi_il(ij, 1) - Phi_star_il(ij, 1)) &
                             + tau2_g*(p_ik(ij, 1) - pbot + rho_bot*(Phi_il(ij, 1) - phis(ij)))
            END DO
            ! inner interfaces
            DO l = 2, llm
               !DIR$ SIMD
               DO ij = ij_begin_ext, ij_end_ext
                  ml_g2 = gm2*m_il(ij, l)
                  B_il(ij, l) = A_ik(ij, l) + A_ik(ij, l - 1) + ml_g2
                  R_il(ij, l) = ml_g2*(Phi_il(ij, l) - Phi_star_il(ij, l)) &
                                + tau2_g*(p_ik(ij, l) - p_ik(ij, l - 1))
                  ! consistency check : if Wil=0 and initial state is in hydrostatic balance
                  ! then Phi_star_il(ij,l) = Phi_il(ij,l) - tau^2*g^2
                  ! and residual = tau^2*(ml+(1/g)dl_pi)=0
               END DO
            END DO
            ! top interface l=llm+1
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               ml_g2 = gm2*m_il(ij, llm + 1)
               B_il(ij, llm + 1) = A_ik(ij, llm) + ml_g2
               R_il(ij, llm + 1) = ml_g2*(Phi_il(ij, llm + 1) - Phi_star_il(ij, llm + 1)) &
                                   + tau2_g*(ptop - p_ik(ij, llm))
            END DO

            ! FIXME later
            ! the lines below modify the tridiag problem
            ! for flat, rigid boundary conditions at top and bottom :
            ! zero out A(1), A(llm), R(1), R(llm+1)
            ! => x(l)=0 at l=1,llm+1
            DO ij = ij_begin_ext, ij_end_ext
               A_ik(ij, 1) = 0.
               A_ik(ij, llm) = 0.
               R_il(ij, 1) = 0.
               R_il(ij, llm + 1) = 0.
            END DO

            IF (debug_hevi_solver) THEN ! print Linf(residual)
               PRINT *, '[hevi_solver] R,p', iter, MAXVAL(ABS(R_il)), MAXVAL(p_ik)
            END IF

            ! Solve -A(l-1)x(l-1) + B(l)x(l) - A(l)x(l+1) = R(l) using Thomas algorithm
            ! Forward sweep :
            ! C(0)=0, C(l) = -A(l) / (B(l)+A(l-1)C(l-1)),
            ! D(0)=0, D(l) = (R(l)+A(l-1)D(l-1)) / (B(l)+A(l-1)C(l-1))
            ! bottom interface l=1
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               X_ij = 1./B_il(ij, 1)
               C_ik(ij, 1) = -A_ik(ij, 1)*X_ij
               D_il(ij, 1) = R_il(ij, 1)*X_ij
            END DO
            ! inner interfaces/layers
            DO l = 2, llm
               !DIR$ SIMD
               DO ij = ij_begin_ext, ij_end_ext
                  X_ij = 1./(B_il(ij, l) + A_ik(ij, l - 1)*C_ik(ij, l - 1))
                  C_ik(ij, l) = -A_ik(ij, l)*X_ij
                  D_il(ij, l) = (R_il(ij, l) + A_ik(ij, l - 1)*D_il(ij, l - 1))*X_ij
               END DO
            END DO
            ! top interface l=llm+1
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               X_ij = 1./(B_il(ij, llm + 1) + A_ik(ij, llm)*C_ik(ij, llm))
               D_il(ij, llm + 1) = (R_il(ij, llm + 1) + A_ik(ij, llm)*D_il(ij, llm))*X_ij
            END DO

            ! Back substitution :
            ! x(i) = D(i)-C(i)x(i+1), x(N+1)=0
            ! + Newton-Raphson update
            x_il = 0. ! FIXME
            ! top interface l=llm+1
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               x_il(ij, llm + 1) = D_il(ij, llm + 1)
               Phi_il(ij, llm + 1) = Phi_il(ij, llm + 1) - x_il(ij, llm + 1)
            END DO
            ! lower interfaces
            DO l = llm, 1, -1
               !DIR$ SIMD
               DO ij = ij_begin_ext, ij_end_ext
                  x_il(ij, l) = D_il(ij, l) - C_ik(ij, l)*x_il(ij, l + 1)
                  Phi_il(ij, l) = Phi_il(ij, l) - x_il(ij, l)
               END DO
            END DO

            IF (debug_hevi_solver) THEN
               PRINT *, '[hevi_solver] A,B', iter, MAXVAL(ABS(A_ik)), MAXVAL(ABS(B_il))
               PRINT *, '[hevi_solver] C,D', iter, MAXVAL(ABS(C_ik)), MAXVAL(ABS(D_il))
               DO l = 1, llm + 1
                  WRITE (*, '(A,I2.1,I3.2,E9.2)') '[hevi_solver] x', iter, l, MAXVAL(ABS(x_il(:, l)))
               END DO
            END IF

         END DO ! Newton-Raphson

      END IF ! dysl

   END SUBROUTINE compute_NH_geopot_hex

   SUBROUTINE compute_caldyn_solver_unst(tau, rhodz, theta, berni, pres, m_il, pk, geopot, W, dPhi, dW, du)
      USE ISO_C_BINDING, only: C_DOUBLE, C_FLOAT
      USE earth_const
      USE trace
      USE grid_param
      USE disvert_mod, ONLY: ptop
      USE data_unstructured_mod, ONLY: enter_trace, exit_trace, &
                                       id_solver, left, right, PHI_BOT
      REAL(C_DOUBLE), INTENT(IN) :: tau
      REAL(C_DOUBLE), DIMENSION(llm, primal_num)   :: rhodz, pk, berni, pres    ! IN, OUT, BUF*2
      REAL(C_DOUBLE), DIMENSION(llm, primal_num, nqdyn)  :: theta                  ! IN
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num) :: geopot, W, dPhi, dW, m_il ! INOUT,INOUT, OUT,OUT, BUF
      REAL(C_DOUBLE), DIMENSION(llm, edge_num)      :: du                     ! OUT
      INTEGER ij, l, iq, iedge, edge, ivertex, vertex, ij_left, ij_right, ij_up, ij_down, itrisk, edge_trisk, kup, kdown
      REAL(C_DOUBLE) :: X_ij, rho_ij, T_ij, gamma, Cvd, vreff, Rd_preff

      !--------------------------------------------------------------------------
      !---------------------------- caldyn_mil ----------------------------------
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         l = 1
         m_il(l, ij) = .5*rhodz(l, ij)
         !DIR$ SIMD
         DO l = 2, llm
            m_il(l, ij) = .5*(rhodz(l, ij) + rhodz(l - 1, ij))
         END DO
         l = llm + 1
         m_il(l, ij) = .5*rhodz(l - 1, ij)
      END DO
      !$OMP END DO
      !---------------------------- caldyn_mil ----------------------------------
      !--------------------------------------------------------------------------
      IF (tau > 0) THEN ! solve implicit problem for geopotential
         CALL compute_NH_geopot_unst(tau, rhodz, m_il, theta, W, geopot)
      END IF
      CALL enter_trace(id_solver, 8*llm*((7)*primal_num + (0)*dual_num + (1)*edge_num))

      !--------------------------------------------------------------------------
      !---------------------------- caldyn_solver ----------------------------------
      !
      ! Compute pressure (pres) and Exner function (pk)
      ! kappa = R/Cp
      ! 1-kappa = Cv/Cp
      ! Cp/Cv = 1/(1-kappa)
      gamma = 1./(1.-kappa)
      vreff = Rd*Treff/preff ! reference specific volume
      Cvd = 1./(cpp - Rd)
      Rd_preff = kappa*cpp/preff
      SELECT CASE (caldyn_thermo)
      CASE (thermo_theta)
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, primal_num
            !DIR$ SIMD
            DO l = 1, llm
               rho_ij = 1./(geopot(l + 1, ij) - geopot(l, ij))
               rho_ij = rho_ij*g*rhodz(l, ij)
               X_ij = Rd_preff*theta(l, ij, 1)*rho_ij
               ! kappa.theta.rho = p/exner
               ! => X = (p/p0)/(exner/Cp)
               ! = (p/p0)^(1-kappa)
               pres(l, ij) = preff*(X_ij**gamma) ! pressure
               ! Compute Exner function (needed by compute_caldyn_fast)
               ! other formulae possible if exponentiation is slow
               pk(l, ij) = cpp*((pres(l, ij)/preff)**kappa) ! Exner
            END DO
         END DO
         !$OMP END DO
      CASE (thermo_entropy)
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, primal_num
            !DIR$ SIMD
            DO l = 1, llm
               rho_ij = 1./(geopot(l + 1, ij) - geopot(l, ij))
               rho_ij = rho_ij*g*rhodz(l, ij)
               T_ij = Treff*exp((theta(l, ij, 1) + Rd*log(vreff*rho_ij))*Cvd)
               pres(l, ij) = rho_ij*Rd*T_ij
               pk(l, ij) = T_ij
            END DO
         END DO
         !$OMP END DO
      CASE DEFAULT
         STOP
      END SELECT
      ! We need a barrier here because we compute pres above and do a vertical difference below
      !$OMP BARRIER
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         l = 1
         ! Lower BC
         dW(l, ij) = (1./g)*(pbot - rho_bot*(geopot(l, ij) - Phi_bot) - pres(l, ij)) - m_il(l, ij)
         W(l, ij) = W(l, ij) + tau*dW(l, ij) ! update W
         dPhi(l, ij) = g*g*W(l, ij)/m_il(l, ij)
         !DIR$ SIMD
         DO l = 2, llm
            dW(l, ij) = (1./g)*(pres(l - 1, ij) - pres(l, ij)) - m_il(l, ij)
            W(l, ij) = W(l, ij) + tau*dW(l, ij) ! update W
            dPhi(l, ij) = g*g*W(l, ij)/m_il(l, ij)
         END DO
         l = llm + 1
         ! Top BC
         dW(l, ij) = (1./g)*(pres(l - 1, ij) - ptop) - m_il(l, ij)
         W(l, ij) = W(l, ij) + tau*dW(l, ij) ! update W
         dPhi(l, ij) = g*g*W(l, ij)/m_il(l, ij)
      END DO
      !$OMP END DO
      ! We need a barrier here because we update W above and do a vertical average below
      !$OMP BARRIER
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         !DIR$ SIMD
         DO l = 1, llm
            ! compute du = -0.5*g^2.grad(w^2)
            berni(l, ij) = (-.25*g*g)*((W(l, ij)/m_il(l, ij))**2 + (W(l + 1, ij)/m_il(l + 1, ij))**2)
         END DO
      END DO
      !$OMP END DO
      !$OMP DO SCHEDULE(STATIC)
      DO edge = 1, edge_num
         ij_left = left(edge)
         ij_right = right(edge)
         !DIR$ SIMD
         DO l = 1, llm
            du(l, edge) = 1.*(berni(l, ij_left) - berni(l, ij_right))
         END DO
      END DO
      !$OMP END DO
      !---------------------------- caldyn_solver ----------------------------------
      !--------------------------------------------------------------------------
      CALL exit_trace()

   END SUBROUTINE compute_caldyn_solver_unst

   SUBROUTINE compute_NH_geopot_unst(tau, m_ik, m_il, theta, W_il, Phi_il)
      USE grid_param
      USE ISO_C_BINDING, only: C_DOUBLE, C_FLOAT
      USE disvert_mod, only: g, Treff, cpp, preff, kappa, caldyn_thermo, thermo_theta, &
                             thermo_entropy
      USE disvert_mod, ONLY: ptop
      USE data_unstructured_mod, ONLY: enter_trace, exit_trace, &
                                       id_NH_geopot, &
                                       PHI_BOT, pbot, rho_bot
      REAL(C_DOUBLE), DIMENSION(llm, primal_num)   :: m_ik, theta   ! IN*2
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num) :: m_il, W_il, Phi_il, Phi_star_il  ! IN,INOUT*2, LOCAL
      REAL(C_DOUBLE) :: tau, gamma, tau2_g, tau2_g2, g2, gm2, vreff, Rd_preff
      INTEGER :: iter
      INTEGER ij, l, iq, iedge, edge, ivertex, vertex, ij_left, ij_right, ij_up, ij_down, itrisk, edge_trisk, kup, kdown
      REAL(C_DOUBLE) :: rho_ij, X_ij, Y_ij, wil, rho_c2_mik, c2_mik, ml_g2
      REAL(C_DOUBLE), DIMENSION(llm, primal_num) :: p_ik, A_ik, C_ik
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num) :: R_il, B_il, D_il, x_il

      CALL enter_trace(id_NH_geopot, 8*llm*((7)*primal_num + (0)*dual_num + (0)*edge_num))

      !--------------------------------------------------------------------------
      !---------------------------- compute_NH_geopot ----------------------------------
      tau2_g = tau*tau/g
      g2 = g*g
      gm2 = 1./g2
      vreff = Treff*cpp/preff*kappa
      gamma = 1./(1.-kappa)
      !$OMP BARRIER
      ! compute Phi_star
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         DO l = 1, llm + 1
            Phi_star_il(l, ij) = Phi_il(l, ij) + tau*g2*(W_il(l, ij)/m_il(l, ij) - tau)
         END DO
      END DO
      !$OMP END DO
      ! Newton-Raphson iteration : Phi_il contains current guess value
      DO iter = 1, 2 ! 2 iterations should be enough
         ! Compute pressure, A_ik
         SELECT CASE (caldyn_thermo)
         CASE (thermo_theta)
            !$OMP DO SCHEDULE(STATIC)
            DO ij = 1, primal_num
               DO l = 1, llm
                  rho_ij = (g*m_ik(l, ij))/(Phi_il(l + 1, ij) - Phi_il(l, ij))
                  X_ij = (cpp/preff)*kappa*theta(l, ij)*rho_ij
                  p_ik(l, ij) = preff*(X_ij**gamma)
                  c2_mik = gamma*p_ik(l, ij)/(rho_ij*m_ik(l, ij)) ! c^2 = gamma*R*T = gamma*p/rho
                  A_ik(l, ij) = c2_mik*(tau/g*rho_ij)**2
               END DO
            END DO
            !$OMP END DO
         CASE (thermo_entropy)
            !$OMP DO SCHEDULE(STATIC)
            DO ij = 1, primal_num
               DO l = 1, llm
                  rho_ij = (g*m_ik(l, ij))/(Phi_il(l + 1, ij) - Phi_il(l, ij))
                  X_ij = log(vreff*rho_ij) + theta(l, ij)/cpp
                  p_ik(l, ij) = preff*exp(X_ij*gamma)
                  c2_mik = gamma*p_ik(l, ij)/(rho_ij*m_ik(l, ij)) ! c^2 = gamma*R*T = gamma*p/rho
                  A_ik(l, ij) = c2_mik*(tau/g*rho_ij)**2
               END DO
            END DO
            !$OMP END DO
         CASE DEFAULT
            PRINT *, 'caldyn_thermo not supported by compute_NH_geopot', caldyn_thermo
            STOP
         END SELECT
         ! NB : A(1), A(llm), R(1), R(llm+1) = 0 => x(l)=0 at l=1,llm+1 => flat, rigid top and bottom
         ! Solve -A(l-1)x(l-1) + B(l)x(l) - A(l)x(l+1) = R(l) using Thomas algorithm
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, primal_num
            ! Compute residual R_il and B_il
            ! bottom interface l=1
            ml_g2 = gm2*m_il(1, ij)
            B_il(1, ij) = A_ik(1, ij) + ml_g2 + tau2_g*rho_bot
            R_il(1, ij) = ml_g2*(Phi_il(1, ij) - Phi_star_il(1, ij)) &
                          + tau2_g*(p_ik(1, ij) - pbot + rho_bot*(Phi_il(1, ij) - Phi_bot))
            DO l = 2, llm
               ! inner interfaces
               ml_g2 = gm2*m_il(l, ij)
               B_il(l, ij) = A_ik(l, ij) + A_ik(l - 1, ij) + ml_g2
               R_il(l, ij) = ml_g2*(Phi_il(l, ij) - Phi_star_il(l, ij)) &
                             + tau2_g*(p_ik(l, ij) - p_ik(l - 1, ij))
               ! consistency check : if Wil=0 and initial state is in hydrostatic balance
               ! then Phi_star_il(l,ij) = Phi_il(l,ij) - tau^2*g^2
               ! and residual = tau^2*(ml+(1/g)dl_pi)=0
            END DO
            ! top interface l=llm+1
            ml_g2 = gm2*m_il(llm + 1, ij)
            B_il(llm + 1, ij) = A_ik(llm + 1 - 1, ij) + ml_g2
            R_il(llm + 1, ij) = ml_g2*(Phi_il(llm + 1, ij) - Phi_star_il(llm + 1, ij)) &
                                + tau2_g*(ptop - p_ik(llm + 1 - 1, ij))
            !
            ! Forward sweep :
            ! C(0)=0, C(l) = -A(l) / (B(l)+A(l-1)C(l-1)),
            ! D(0)=0, D(l) = (R(l)+A(l-1)D(l-1)) / (B(l)+A(l-1)C(l-1))
            X_ij = 1./B_il(1, ij)
            C_ik(1, ij) = -A_ik(1, ij)*X_ij
            D_il(1, ij) = R_il(1, ij)*X_ij
            DO l = 2, llm
               X_ij = 1./(B_il(l, ij) + A_ik(l - 1, ij)*C_ik(l - 1, ij))
               C_ik(l, ij) = -A_ik(l, ij)*X_ij
               D_il(l, ij) = (R_il(l, ij) + A_ik(l - 1, ij)*D_il(l - 1, ij))*X_ij
            END DO
            X_ij = 1./(B_il(llm + 1, ij) + A_ik(llm + 1 - 1, ij)*C_ik(llm + 1 - 1, ij))
            D_il(llm + 1, ij) = (R_il(llm + 1, ij) + A_ik(llm + 1 - 1, ij)*D_il(llm + 1 - 1, ij))*X_ij
            ! Back substitution :
            ! x(i) = D(i)-C(i)x(i+1), x(llm+1)=0
            ! + Newton-Raphson update
            ! top interface l=llm+1
            x_il(llm + 1, ij) = D_il(llm + 1, ij)
            Phi_il(llm + 1, ij) = Phi_il(llm + 1, ij) - x_il(llm + 1, ij)
            DO l = llm, 1, -1
               ! Back substitution at lower interfaces
               x_il(l, ij) = D_il(l, ij) - C_ik(l, ij)*x_il(l + 1, ij)
               Phi_il(l, ij) = Phi_il(l, ij) - x_il(l, ij)
            END DO
         END DO
         !$OMP END DO
         IF (debug_hevi_solver) THEN
            PRINT *, '[hevi_solver] A,B', iter, MAXVAL(ABS(A_ik)), MAXVAL(ABS(B_il))
            PRINT *, '[hevi_solver] C,D', iter, MAXVAL(ABS(C_ik)), MAXVAL(ABS(D_il))
            DO l = 1, llm + 1
               WRITE (*, '(A,I2.1,I3.2,E9.2)') '[hevi_solver] x_il', iter, l, MAXVAL(ABS(x_il(l, :)))
            END DO
            DO l = 1, llm + 1
               WRITE (*, '(A,I2.1,I3.2,E9.2)') '[hevi_solver] R_il', iter, l, MAXVAL(ABS(R_il(l, :)))
            END DO
         END IF
      END DO ! Newton-Raphson
      !$OMP BARRIER
      !---------------------------- compute_NH_geopot ----------------------------------
      !--------------------------------------------------------------------------
      CALL exit_trace()

   END SUBROUTINE compute_NH_geopot_unst

END SUBMODULE compute_caldyn_solver_mod
