!> @file time_integration_2d.f90
!> @brief Explicit RK time integration for 2D Euler + CFL dt — all steppers route through the shared stepper_kernels (Q2a).
module time_integration_2d
  use precision, only: wp
  use solver_state_2d, only: solver_state_2d_t, neq2d
  use spatial_discretization_2d, only: compute_resid_2d
  use parallel_reductions, only: par_max_real
  use option_registry, only: time_euler, time_ssprk22, time_rk4, time_ssprk54
  use stepper_kernels, only: rhs_fn, kernel_euler, kernel_ssprk22, kernel_tvd_rk3, &
                             kernel_rk4, kernel_ssprk54
  implicit none
  private
  public :: euler_step_2d, ssprk22_step_2d, tvd_rk3_step_2d
  public :: rk4_step_2d, ssprk54_step_2d
  public :: resolve_time_scheme_2d, compute_dt_2d, stepper_2d_iface

  abstract interface
    subroutine stepper_2d_iface(state)
      import :: solver_state_2d_t
      type(solver_state_2d_t), intent(inout), target :: state
    end subroutine stepper_2d_iface
  end interface

contains

  function compute_dt_2d(state) result(dt)
    type(solver_state_2d_t), intent(in) :: state
    real(wp) :: dt, rho, u, v, p, c, sig, sig_max, gam
    integer :: i, j
    if (state % cfg % cfl <= 0.0_wp) then
      dt = state % cfg % dt          ! fixed-timestep mode
      return
    end if
    gam = state % cfg % gam
    sig_max = 0.0_wp
    if (state % mesh % fdm_curvilinear) then
      ! Curvilinear-FDM nodal metrics: same spectral-radius estimate as the FVM
      ! curvilinear branch, but with the nodal transformation metrics
      ! (sx_xi/sx_eta) and the nodal Jacobian (jac) in place of the FV face-area
      ! vectors and cell volume.  sx_xi(:,i,j)/sx_xi(:,i+1,j) are the two xi-faces
      ! bounding node (i,j); likewise sx_eta for eta.
      associate (mesh => state % mesh)
        do j = 1, state % ny_local
          do i = 1, state % nx_local
            rho = state % ub(1, i, j)
            u = state % ub(2, i, j) / rho
            v = state % ub(3, i, j) / rho
            p = (gam - 1.0_wp) * (state % ub(4, i, j) - 0.5_wp * rho * (u * u + v * v))
            c = sqrt(gam * p / rho)
            sig = curvilinear_sigma(u, v, c, &
                                    0.5_wp * (mesh % sx_xi(:, i, j) + mesh % sx_xi(:, i + 1, j)), &
                                    0.5_wp * (mesh % sx_eta(:, i, j) + mesh % sx_eta(:, i, j + 1)), &
                                    mesh % jac(i, j))
            sig_max = max(sig_max, sig)
          end do
        end do
      end associate
    else if (.not. state % mesh % uniform) then
      associate (mesh => state % mesh)
        do j = 1, state % ny_local
          do i = 1, state % nx_local
            rho = state % ub(1, i, j)
            u = state % ub(2, i, j) / rho
            v = state % ub(3, i, j) / rho
            p = (gam - 1.0_wp) * (state % ub(4, i, j) - 0.5_wp * rho * (u * u + v * v))
            c = sqrt(gam * p / rho)
            sig = curvilinear_sigma(u, v, c, &
                                    0.5_wp * (mesh % s_xi(:, i, j) + mesh % s_xi(:, i + 1, j)), &
                                    0.5_wp * (mesh % s_eta(:, i, j) + mesh % s_eta(:, i, j + 1)), &
                                    mesh % vol(i, j))
            sig_max = max(sig_max, sig)
          end do
        end do
      end associate
    else
      do j = 1, state % ny_local
        do i = 1, state % nx_local
          rho = state % ub(1, i, j)
          u = state % ub(2, i, j) / rho
          v = state % ub(3, i, j) / rho
          p = (gam - 1.0_wp) * (state % ub(4, i, j) - 0.5_wp * rho * (u * u + v * v))
          c = sqrt(gam * p / rho)
          sig = (abs(u) + c) / state % dx + (abs(v) + c) / state % dy
          sig_max = max(sig_max, sig)
        end do
      end do
    end if
    sig_max = par_max_real(sig_max)
    if (sig_max <= 0.0_wp) error stop 'compute_dt_2d: non-positive max wave speed'
    dt = state % cfg % cfl / sig_max
  end function compute_dt_2d

  !> stepper_kernels rhs callback: recompute the 2D residual (halo exchange +
  !! BCs + interior resid happen inside compute_resid_2d).
  subroutine rhs_2d(ctx)
    class(*), intent(inout), target :: ctx

    select type (ctx)
    type is (solver_state_2d_t)
      call compute_resid_2d(ctx)
    class default
      error stop 'time_integration_2d: rhs_2d received unexpected context type'
    end select
  end subroutine rhs_2d

  !> Classic RK4 via the shared kernel. Whole-array flat update (halos
  !! included) — legal under the halo-zero resid invariant; halos are
  !! refreshed inside compute_resid_2d before each use. New in R2; no
  !! bit-history to preserve.
  subroutine rk4_step_2d(state)
    type(solver_state_2d_t), intent(inout), target :: state
    real(wp), pointer :: u(:), r(:), s1(:), s2(:)

    u(1:size(state % ub)) => state % ub
    r(1:size(state % resid)) => state % resid
    s1(1:size(state % scratch1)) => state % scratch1
    s2(1:size(state % scratch2)) => state % scratch2
    call kernel_rk4(u, r, s1, s2, state % dt, rhs_2d, state)
  end subroutine rk4_step_2d

  !> SSPRK(5,4) via the shared kernel (see rk4_step_2d notes).
  subroutine ssprk54_step_2d(state)
    type(solver_state_2d_t), intent(inout), target :: state
    real(wp), pointer :: u(:), r(:), s1(:), s2(:), s3(:)

    u(1:size(state % ub)) => state % ub
    r(1:size(state % resid)) => state % resid
    s1(1:size(state % scratch1)) => state % scratch1
    s2(1:size(state % scratch2)) => state % scratch2
    s3(1:size(state % scratch3)) => state % scratch3
    call kernel_ssprk54(u, r, s1, s2, s3, state % dt, rhs_2d, state)
  end subroutine ssprk54_step_2d

  !> Explicit Euler via the shared kernel (whole-array flat update under the
  !! halo-zero resid invariant — see rk4_step_2d notes). Retired from the
  !! legacy interior-slice form in Q2a: `u + dt*resid` is the identical
  !! expression and halo resid is zero, so this one is bit-neutral.
  subroutine euler_step_2d(state)
    type(solver_state_2d_t), intent(inout), target :: state
    real(wp), pointer :: u(:), r(:)

    u(1:size(state % ub)) => state % ub
    r(1:size(state % resid)) => state % resid
    call kernel_euler(u, r, state % dt, rhs_2d, state)
  end subroutine euler_step_2d

  !> SSPRK(2,2) via the shared kernel (see rk4_step_2d notes). Q2a authorized
  !! last-bit change: the kernel's 1D-verbatim distributed stage groupings
  !! replace the legacy factored groupings (spec 2026-07-10 §5).
  subroutine ssprk22_step_2d(state)
    type(solver_state_2d_t), intent(inout), target :: state
    real(wp), pointer :: u(:), r(:), s1(:)

    u(1:size(state % ub)) => state % ub
    r(1:size(state % resid)) => state % resid
    s1(1:size(state % scratch1)) => state % scratch1
    call kernel_ssprk22(u, r, s1, state % dt, rhs_2d, state)
  end subroutine ssprk22_step_2d

  !> TVD-RK3 via the shared kernel (see rk4_step_2d notes). Q2a authorized
  !! last-bit change, as ssprk22 above. This is every 2D deck's default
  !! integrator, so all 2D golden outputs move at the last bit.
  subroutine tvd_rk3_step_2d(state)
    type(solver_state_2d_t), intent(inout), target :: state
    real(wp), pointer :: u(:), r(:), s1(:)

    u(1:size(state % ub)) => state % ub
    r(1:size(state % resid)) => state % resid
    s1(1:size(state % scratch1)) => state % scratch1
    call kernel_tvd_rk3(u, r, s1, state % dt, rhs_2d, state)
  end subroutine tvd_rk3_step_2d

  subroutine resolve_time_scheme_2d(stepper, scheme)
    procedure(stepper_2d_iface), pointer, intent(out) :: stepper
    character(len=*), intent(in) :: scheme
    select case (trim(scheme))
    case (time_euler); stepper => euler_step_2d
    case (time_ssprk22); stepper => ssprk22_step_2d
    case (time_rk4); stepper => rk4_step_2d
    case (time_ssprk54); stepper => ssprk54_step_2d
    case default; stepper => tvd_rk3_step_2d
    end select
  end subroutine resolve_time_scheme_2d

  !> Per-cell contravariant spectral radius for the curvilinear CFL limit:
  !! (|u.Sxi| + c|Sxi| + |u.Seta| + c|Seta|) / V, with cell-averaged face vectors.
  pure function curvilinear_sigma(u, v, c, sxi, seta, vol) result(sig)
    real(wp), intent(in) :: u, v, c, sxi(2), seta(2), vol
    real(wp) :: sig, uxi, ueta, mxi, meta
    uxi = u * sxi(1) + v * sxi(2)
    ueta = u * seta(1) + v * seta(2)
    mxi = sqrt(sxi(1)**2 + sxi(2)**2)
    meta = sqrt(seta(1)**2 + seta(2)**2)
    sig = (abs(uxi) + c * mxi + abs(ueta) + c * meta) / vol
  end function curvilinear_sigma

end module time_integration_2d
