{% macro compute_p_and_Aik() %}
{% set p_and_c2_from_theta=caller %}
    SEQUENCE_EXT
      BODY('1,llm')
          rho_ij = (g*m_ik(CELL))/(Phi_il(UP(CELL))-Phi_il(CELL))
          {{ p_and_c2_from_theta() }}
          A_ik(CELL) = c2_mik*(tau/g*rho_ij)**2
      END_BLOCK
    END_BLOCK
{%- endmacro %}

KERNEL(compute_NH_geopot)
  tau2_g=tau*tau/g
  g2=g*g
  gm2 = 1./g2
  gamma = 1./(1.-kappa)
    
  BARRIER

  ! compute Phi_star
  SEQUENCE_EXT
    BODY('1,llm+1')
      Phi_star_il(CELL) = Phi_il(CELL) + tau*g2*(W_il(CELL)/m_il(CELL)-tau)
    END_BLOCK
  END_BLOCK

  ! 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)
      {% call() compute_p_and_Aik() %}
        X_ij = (cpp/preff)*kappa*theta(CELL)*rho_ij
        p_ik(CELL) = preff*(X_ij**gamma)
        c2_mik = gamma*p_ik(CELL)/(rho_ij*m_ik(CELL)) ! c^2 = gamma*R*T = gamma*p/rho
      {% endcall %}
    CASE(thermo_entropy)
      {% call() compute_p_and_Aik() %}
        X_ij = Treff*exp(theta(CELL)/cpp) ! theta = Tref.exp(s/Cp)
        X_ij = (cpp/preff)*kappa*X_ij*rho_ij
        p_ik(CELL) = preff*(X_ij**gamma)
        c2_mik = gamma*p_ik(CELL)/(rho_ij*m_ik(CELL)) ! c^2 = gamma*R*T = gamma*p/rho
      {% endcall %}
    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

    SEQUENCE_EXT
      ! Compute residual R_il and B_il
      PROLOGUE(1)
        ! bottom interface l=1
        ml_g2 = gm2*m_il(CELL) 
        B_il(CELL) = A_ik(CELL) + ml_g2 + tau2_g*rho_bot
        R_il(CELL) = ml_g2*( Phi_il(CELL)-Phi_star_il(CELL))  &
                  + tau2_g*( p_ik(CELL)-pbot+rho_bot*(Phi_il(CELL)-PHI_BOT(HIDX(CELL))) )
      END_BLOCK
      BODY('2,llm')
        ! inner interfaces
        ml_g2 = gm2*m_il(CELL) 
        B_il(CELL) = A_ik(CELL)+A_ik(DOWN(CELL)) + ml_g2
        R_il(CELL) = ml_g2*( Phi_il(CELL)-Phi_star_il(CELL))  &
                  + tau2_g*(p_ik(CELL)-p_ik(DOWN(CELL)))
        ! consistency check : if Wil=0 and initial state is in hydrostatic balance
        ! then Phi_star_il(CELL) = Phi_il(CELL) - tau^2*g^2
        ! and residual = tau^2*(ml+(1/g)dl_pi)=0
      END_BLOCK
      EPILOGUE('llm+1')
        ! top interface l=llm+1
        ml_g2 = gm2*m_il(CELL) 
        B_il(CELL) = A_ik(DOWN(CELL)) + ml_g2
        R_il(CELL) = ml_g2*( Phi_il(CELL)-Phi_star_il(CELL))  &
                  + tau2_g*( ptop-p_ik(DOWN(CELL)) )
      END_BLOCK
      !
      ! 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))
      PROLOGUE(1)
        X_ij = 1./B_il(CELL)
        C_ik(CELL) = -A_ik(CELL) * X_ij 
        D_il(CELL) =  R_il(CELL) * X_ij
      END_BLOCK
      BODY('2,llm')
        X_ij = 1./( B_il(CELL) + A_ik(DOWN(CELL))*C_ik(DOWN(CELL)) )
        C_ik(CELL) = -A_ik(CELL) * X_ij 
        D_il(CELL) =  (R_il(CELL)+A_ik(DOWN(CELL))*D_il(DOWN(CELL))) * X_ij
      END_BLOCK
      EPILOGUE('llm+1')
        X_ij = 1./( B_il(CELL) + A_ik(DOWN(CELL))*C_ik(DOWN(CELL)) )
        D_il(CELL) =  (R_il(CELL)+A_ik(DOWN(CELL))*D_il(DOWN(CELL))) * 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(CELL) = D_il(CELL)
        Phi_il(CELL) = Phi_il(CELL) - x_il(CELL)
      END_BLOCK
      BODY('llm,1,-1')
        ! Back substitution at lower interfaces
        x_il(CELL) = D_il(CELL) - C_ik(CELL)*x_il(UP(CELL))
        Phi_il(CELL) = Phi_il(CELL) - x_il(CELL)
      END_BLOCK
    END_BLOCK

    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

  BARRIER

  debug_hevi_solver=.FALSE.
END_BLOCK

KERNEL(caldyn_solver)
  FORALL_CELLS_EXT('1', 'llm+1')
    ON_PRIMAL
       CST_IF(IS_BOTTOM_LEVEL,    m_il(CELL) = .5*rhodz(CELL)                     )
       CST_IF(IS_INNER_INTERFACE, m_il(CELL) = .5*(rhodz(CELL)+rhodz(DOWN(CELL))) )
       CST_IF(IS_TOP_INTERFACE,   m_il(CELL) = .5*rhodz(DOWN(CELL))               )
    END_BLOCK
  END_BLOCK
  !
  IF(tau>0) THEN ! solve implicit problem for geopotential
    CALL compute_NH_geopot(tau,PHI_BOT_VAR, rhodz, m_il, theta, W, geopot)
  END IF
  ! 
  ! 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
  FORALL_CELLS_EXT()
    SELECT CASE(caldyn_thermo)
    CASE(thermo_theta)
      ON_PRIMAL
        rho_ij = g*rhodz(CELL)/(geopot(UP(CELL))-geopot(CELL))
        X_ij = (cpp/preff)*kappa*theta(CELL,1)*rho_ij
        ! kappa.theta.rho = p/exner
        ! => X = (p/p0)/(exner/Cp)
        !      = (p/p0)^(1-kappa)
        pres(CELL) = preff*(X_ij**gamma)  ! pressure
        ! Compute Exner function (needed by compute_caldyn_fast)
        ! other formulae possible if exponentiation is slow
        pk(CELL)   = cpp*((pres(CELL)/preff)**kappa) ! Exner
      END_BLOCK
    CASE(thermo_entropy)
      ON_PRIMAL
        rho_ij = g*rhodz(CELL)/(geopot(UP(CELL))-geopot(CELL))
        T_ij = Treff*exp( (theta(CELL,1)+Rd*log(vreff*rho_ij))/Cvd )
        pres(CELL) = rho_ij*Rd*T_ij
        pk(CELL)   = T_ij
      END_BLOCK
    CASE DEFAULT
      STOP
    END SELECT
  END_BLOCK

  ! We need a barrier here because we compute pres above and do a vertical difference below
  BARRIER

  FORALL_CELLS_EXT('1', 'llm+1')
    ON_PRIMAL
      CST_IFTHEN(IS_BOTTOM_LEVEL)
        ! Lower BC
        dW(CELL)   = (1./g)*(pbot-rho_bot*(geopot(CELL)-PHI_BOT(HIDX(CELL)))-pres(CELL)) - m_il(CELL)
      CST_ELSEIF(IS_TOP_INTERFACE)
        ! Top BC
        dW(CELL)   = (1./g)*(pres(DOWN(CELL))-ptop) - m_il(CELL)
      CST_ELSE
        dW(CELL)   = (1./g)*(pres(DOWN(CELL))-pres(CELL)) - m_il(CELL)
      CST_ENDIF
      W(CELL)    = W(CELL)+tau*dW(CELL) ! update W
      dPhi(CELL) = g*g*W(CELL)/m_il(CELL)
    END_BLOCK
  END_BLOCK

  ! We need a barrier here because we update W above and do a vertical average below
  BARRIER

  FORALL_CELLS_EXT()
    ON_PRIMAL
      ! compute du = -0.5*g^2.grad(w^2)
      berni(CELL) = (-.25*g*g)*((W(CELL)/m_il(CELL))**2 + (W(UP(CELL))/m_il(UP(CELL)))**2 )
    END_BLOCK
  END_BLOCK
  FORALL_CELLS_EXT()
    ON_EDGES
      du(EDGE) = SIGN*(berni(CELL1)-berni(CELL2))
    END_BLOCK
  END_BLOCK
END_BLOCK

KERNEL(caldyn_vert_NH)
  FORALL_CELLS_EXT()
    ON_PRIMAL
      CST_IF(IS_BOTTOM_LEVEL,     w_ij = ( W(CELL)+.5*W(UP(CELL)) )/mass(CELL) )
      CST_IF(IS_INNER_LAYER , w_ij = .5*( W(CELL)+W(UP(CELL)) )/mass(CELL) )
      CST_IF(IS_TOP_LAYER,        w_ij = ( .5*W(CELL)+W(UP(CELL)) )/mass(CELL) )
      wflux_ij = .5*(wflux(CELL)+wflux(UP(CELL)))
      W_etadot(CELL) = wflux_ij*w_ij
      eta_dot(CELL) = wflux_ij / mass(CELL)
      wcov(CELL) = w_ij*(geopot(UP(CELL))-geopot(CELL))
    END_BLOCK
  END_BLOCK
  ! add NH term to du
  FORALL_CELLS()
    ON_EDGES
      du(EDGE) = du(EDGE) - .5*(wcov(CELL1)+wcov(CELL2))*SIGN*(eta_dot(CELL2)-eta_dot(CELL1))
    END_BLOCK
  END_BLOCK
  ! add NH terms to dW, dPhi
  ! FIXME : TODO top and bottom
  FORALL_CELLS('2','llm')
    ON_PRIMAL
      dPhi(CELL)=dPhi(CELL)-wflux(CELL)*(geopot(UP(CELL))-geopot(DOWN(CELL)))/(mass(DOWN(CELL))+mass(CELL))
    END_BLOCK
  END_BLOCK

  ! We need a barrier here because we compute W_etadot above and do a vertical difference below
  BARRIER

  FORALL_CELLS('1','llm+1')
    ON_PRIMAL
      CST_IFTHEN(IS_BOTTOM_LEVEL)
         dW(CELL) = dW(CELL) - W_etadot(CELL)
      CST_ELSEIF(IS_TOP_INTERFACE)
         dW(CELL) = dW(CELL) + W_etadot(DOWN(CELL))
      CST_ELSE
         dW(CELL) = dW(CELL) + W_etadot(DOWN(CELL)) - W_etadot(CELL)
      CST_ENDIF
    END_BLOCK
  END_BLOCK
END_BLOCK

KERNEL(caldyn_slow_NH)
  FORALL_CELLS_EXT('1', 'llm+1')
     ON_PRIMAL
        CST_IF(IS_INNER_INTERFACE, w_il(CELL) = 2.*W(CELL)/(rhodz(KDOWN(CELL))+rhodz(KUP(CELL))) )
        CST_IF(IS_BOTTOM_LEVEL,    w_il(CELL) = 2.*W(CELL)/rhodz(KUP(CELL)) )
        CST_IF(IS_TOP_INTERFACE,   w_il(CELL) = 2.*W(CELL)/rhodz(KDOWN(CELL)) )
     END_BLOCK
  END_BLOCK
  FORALL_CELLS_EXT('1', 'llm+1')
     ON_EDGES
       ! compute DePhi, v_el, G_el, F_el
       ! v_el, W2_el and therefore G_el incorporate metric factor le_de
       ! while DePhil, W_el and F_el do not
       W_el         = .5*( W(CELL2)+W(CELL1) )
       DePhil(EDGE) = SIGN*(Phi(CELL2)-Phi(CELL1))
       F_el(EDGE)   = DePhil(EDGE)*W_el
       W2_el        = .5*LE_DE * ( W(CELL1)*w_il(CELL1) + W(CELL2)*w_il(CELL2) )
       v_el(EDGE)   = .5*LE_DE*(u(KUP(EDGE))+u(KDOWN(EDGE))) ! checked
       G_el(EDGE)   = v_el(EDGE)*W_el - DePhil(EDGE)*W2_el
     END_BLOCK
  END_BLOCK

  FORALL_CELLS_EXT('1', 'llm+1')
    ! compute GradPhi2, dPhi, dW
    ON_PRIMAL
       gPhi2=0.
       dP=0.
       divG=0
       FORALL_EDGES       
         gPhi2 = gPhi2 + LE_DE*DePhil(EDGE)**2
         dP = dP + LE_DE*DePhil(EDGE)*v_el(EDGE)
         divG = divG + SIGN*G_el(EDGE) ! -div(G_el), G_el already has le_de
       END_BLOCK
       gradPhi2(CELL) = 1./(2.*AI) * gPhi2
       dPhi(CELL) = gradPhi2(CELL)*w_il(CELL) - 1./(2.*AI)*dP
       dW(CELL) = (-1./AI)*divG
    END_BLOCK
  END_BLOCK

  ! We need a barrier here because we compute gradPhi2, F_el and w_il above and do a vertical average below
  BARRIER

  FORALL_CELLS_EXT()
    ! Compute berni at scalar points
    ON_PRIMAL
       u2=0.
       FORALL_EDGES       
         u2 = u2 + LE_DE*u(EDGE)**2
       END_BLOCK
       berni(CELL) = 1./(4.*AI) * u2 - .25*( gradPhi2(CELL)*w_il(CELL)**2 + gradPhi2(UP(CELL))*w_il(UP(CELL))**2 )
    END_BLOCK
  END_BLOCK

  FORALL_CELLS_EXT()
    ON_EDGES
    ! Compute mass flux and grad(berni)
      uu = .5*(rhodz(CELL1)+rhodz(CELL2))*u(EDGE) - .5*( F_el(EDGE)+F_el(UP(EDGE)) )
      hflux(EDGE) = LE_DE*uu
      du(EDGE) = SIGN*(berni(CELL1)-berni(CELL2))
    END_BLOCK
  END_BLOCK

END_BLOCK