!> @file solver_state_2d.f90
!> @brief 2D Euler solver state: neq2d, the state type, init and allocation.
!!
!! The local block may be a sub-block of the global grid under MPI
!! decomposition (a `decomp_2d` component + per-rank `nx_local`/`ny_local`);
!! for a single rank the local block equals the global grid.
module solver_state_2d
  use precision, only: wp
  use config_2d, only: config_2d_t
  use solver_interfaces, only: reconstructor_iface
  use domain_decomposition_2d, only: decomp_2d_t
  use mesh_2d, only: mesh_2d_t, build_mesh_2d_global, build_mesh_2d_uniform_nodal, &
                     build_mesh_2d_fdm_curvilinear
  use block_mod, only: block_t
  use option_registry, only: is_fds_flux_scheme, method_fdm
  implicit none
  private
  public :: solver_state_2d_t, neq2d
  public :: init_from_config_2d, allocate_work_arrays_2d

  !> Number of conserved variables for 2D Euler: rho, rho*u, rho*v, E.
  integer, parameter :: neq2d = 4

  type :: solver_state_2d_t
    type(config_2d_t) :: cfg
    integer :: nx = 0          !< GLOBAL cells in x (mesh % nx goes rank-local after build_mesh_2d_local)
    integer :: ny = 0          !< GLOBAL cells in y (mesh % ny goes rank-local after build_mesh_2d_local)
    integer :: nx_local = 0   !< per-rank local cells in x (= nx at np=1)
    integer :: ny_local = 0   !< per-rank local cells in y (= ny at np=1)
    real(wp) :: dx = 0.0_wp     !< cell size in x
    real(wp) :: dy = 0.0_wp     !< cell size in y
    !> Curvilinear grid metrics (uniform fast-path keeps dx/dy above).
    type(mesh_2d_t) :: mesh
    integer :: halo_width = 3  !< ghost-cell width (matches default weno5 stencil)
    !> 2D block decomposition for this rank (whole-domain at np=1).
    type(decomp_2d_t) :: decomp_2d
    !> Per-block discretization descriptor (one block; phases beyond 1 may extend).
    type(block_t) :: blocks(1)
    !> .true. when the active flux scheme is flux-difference splitting (FDS);
    !! .false. for flux-vector splitting (FVS).  Set by init_from_config_2d.
    logical :: use_fds = .true.
    !> Conserved variables, halo-padded: (neq2d, 1-h : nx_local+h, 1-h : ny_local+h).
    real(wp), allocatable :: ub(:, :, :)
    !> Spatial residual R(Q), halo-padded like ub.
    real(wp), allocatable :: resid(:, :, :)
    !> RK stage-save scratch array, halo-padded like ub.
    real(wp), allocatable :: scratch1(:, :, :)
    !> second RK stage-save scratch (rk4/ssprk54), halo-padded like ub.
    real(wp), allocatable :: scratch2(:, :, :)
    !> Third RK stage-save scratch (ssprk54: holds u3; L(u3) stays in
    !! resid), halo-padded like ub.
    real(wp), allocatable :: scratch3(:, :, :)
    !> Current time-step size Δt [s].  Set by the time-integration layer.
    real(wp) :: dt = 0.0_wp
    !> Reconstruction scheme dispatch + stencil geometry (bound by schemes_2d).
    procedure(reconstructor_iface), pointer, nopass :: reconstruct => null()
    integer :: stencil_width = 0
    integer :: stencil_start_offset = 0
    !> Persistent residual-assembly scratch (allocated once in compute_resid_2d,
    !! reused every RK substage; sizes fixed by nx_local/ny_local/stencil_width).
    real(wp), allocatable :: rs_fx(:, :)       !< x-face fluxes (neq2d, nx_local+1)
    real(wp), allocatable :: rs_gy(:, :)       !< y-face fluxes (neq2d, ny_local+1)
    real(wp), allocatable :: rs_stencil(:, :)  !< reconstruction stencil (neq2d, stencil_width)
    !> FVS split-flux scratch, halo-padded like ub (neq2d, 1-h:nx+h, 1-h:ny+h).
    !! Used ONLY by the FDM+FVS residual (compute_resid_fdm_fvs_2d); allocated
    !! lazily on first FVS call (ensure_fvs_scratch_2d) so FDS/FVM runs never pay
    !! for them.  fp_all holds F^+ (positive split flux), fm_all holds F^-; both
    !! are re-filled per sweep (x then y) since the split is direction-dependent.
    real(wp), allocatable :: fp_all(:, :, :)
    real(wp), allocatable :: fm_all(:, :, :)
  end type solver_state_2d_t

contains

  !> Build mesh from config and derive grid scalars.
  subroutine init_from_config_2d(state, cfg)
    type(solver_state_2d_t), intent(inout) :: state
    type(config_2d_t), intent(in) :: cfg
    integer :: nx_io, ny_io
    logical :: ok
    character(len=256) :: msg

    state % cfg = cfg
    if (len_trim(cfg % method) > 8) &
      error stop 'init_from_config_2d: method token exceeds block_t field width (8)'
    state % blocks(1) % method = cfg % method(1:8)
    state % use_fds = is_fds_flux_scheme(trim(cfg % flux_scheme))
    nx_io = cfg % nx; ny_io = cfg % ny
    ! The work-array shapes (allocate_work_arrays_2d) are halo-padded and shared
    ! by both methods; only the global counts and mesh layout differ. For FDM the
    ! counts are NODE counts (nx+1, ny+1) on a uniform nodal mesh; for FVM they
    ! are CELL counts on the (curvilinear-capable) FV mesh. Branch on the block's
    ! method token (set above), not a raw config string.
    select case (trim(state % blocks(1) % method))
    case (method_fdm)
      select case (trim(cfg % grid_type))
      case ('plot3d')
        call build_mesh_2d_fdm_curvilinear(state % mesh, cfg % grid_file, cfg % grid_binary, ok, msg)
        if (.not. ok) error stop 'init_from_config_2d: '//trim(msg)
      case default   ! 'uniform'
        call build_mesh_2d_uniform_nodal(state % mesh, cfg % nx + 1, cfg % ny + 1, &
                                         cfg % x_left, cfg % x_right, cfg % y_left, cfg % y_right)
      end select
    case default   ! fvm
      call build_mesh_2d_global(state % mesh, cfg % grid_type, cfg % grid_file, &
                                cfg % grid_binary, nx_io, ny_io, &
                                cfg % x_left, cfg % x_right, cfg % y_left, cfg % y_right, &
                                ok, msg)
      if (.not. ok) error stop 'init_from_config_2d: '//trim(msg)
    end select

    state % nx = state % mesh % nx
    state % ny = state % mesh % ny
    state % dx = state % mesh % dx_uniform   ! 0 for curvilinear (unused there)
    state % dy = state % mesh % dy_uniform

    state % nx_local = state % nx
    state % ny_local = state % ny
    ! serial whole-domain decomposition (overridden by setup_decomp_2d for MPI runs)
    state % decomp_2d % nx_global = state % nx
    state % decomp_2d % ny_global = state % ny
    state % decomp_2d % nx_local = state % nx
    state % decomp_2d % ny_local = state % ny
    state % decomp_2d % ix_first_global = 1
    state % decomp_2d % iy_first_global = 1
    state % decomp_2d % ix_last_global = state % nx
    state % decomp_2d % iy_last_global = state % ny
    state % decomp_2d % coord_x = 0; state % decomp_2d % coord_y = 0
    state % decomp_2d % dim_x = 1; state % decomp_2d % dim_y = 1
  end subroutine init_from_config_2d

  !> Allocate the conserved-variable array with halo layers on both axes.
  subroutine allocate_work_arrays_2d(state)
    type(solver_state_2d_t), intent(inout) :: state
    integer :: h, stat
    character(len=*), parameter :: ctx = 'solver_state_2d: allocate ub'

    h = state % halo_width
    if (allocated(state % ub)) then
      deallocate (state % ub, stat=stat)
      if (stat /= 0) error stop 'solver_state_2d: ub deallocate failed'
    end if
    allocate (state % ub(neq2d, 1 - h:state % nx_local + h, 1 - h:state % ny_local + h), stat=stat)
    if (stat /= 0) error stop ctx
    state % ub = 0.0_wp
    if (allocated(state % resid)) then
      deallocate (state % resid, stat=stat)
      if (stat /= 0) error stop 'solver_state_2d: resid deallocate failed'
    end if
    if (allocated(state % scratch1)) then
      deallocate (state % scratch1, stat=stat)
      if (stat /= 0) error stop 'solver_state_2d: scratch1 deallocate failed'
    end if
    if (allocated(state % scratch2)) then
      deallocate (state % scratch2, stat=stat)
      if (stat /= 0) error stop 'solver_state_2d: scratch2 deallocate failed'
    end if
    if (allocated(state % scratch3)) then
      deallocate (state % scratch3, stat=stat)
      if (stat /= 0) error stop 'solver_state_2d: scratch3 deallocate failed'
    end if
    allocate (state % resid(neq2d, 1 - h:state % nx_local + h, 1 - h:state % ny_local + h), stat=stat)
    if (stat /= 0) error stop ctx
    allocate (state % scratch1(neq2d, 1 - h:state % nx_local + h, 1 - h:state % ny_local + h), stat=stat)
    if (stat /= 0) error stop ctx
    allocate (state % scratch2(neq2d, 1 - h:state % nx_local + h, 1 - h:state % ny_local + h), stat=stat)
    if (stat /= 0) error stop ctx
    allocate (state % scratch3(neq2d, 1 - h:state % nx_local + h, 1 - h:state % ny_local + h), stat=stat)
    if (stat /= 0) error stop ctx
    state % resid = 0.0_wp
    state % scratch1 = 0.0_wp
    state % scratch2 = 0.0_wp
    state % scratch3 = 0.0_wp
  end subroutine allocate_work_arrays_2d

end module solver_state_2d
