!> @file initial_conditions_2d.f90
!> @brief 2D Euler initial conditions, selected by config % problem_type.
module initial_conditions_2d
  use precision, only: wp
  use config_2d, only: config_2d_t
  use solver_state_2d, only: solver_state_2d_t, neq2d
  use mpi_runtime, only: parallel_fatal
  use option_registry, only: problem_2d_uniform, problem_2d_isentropic_vortex, &
                             problem_2d_sedov, problem_2d_toro_explosion, &
                             problem_2d_riemann2d_config3, problem_2d_freestream, &
                             method_fdm
  implicit none
  private
  public :: apply_initial_condition_2d

contains

  !> Fill state % ub over interior cells according to cfg % problem_type.
  subroutine apply_initial_condition_2d(state, cfg)
    type(solver_state_2d_t), intent(inout) :: state
    type(config_2d_t), intent(in) :: cfg

    select case (trim(cfg % problem_type))
    case (problem_2d_uniform)
      call ic_uniform(state, cfg)
    case (problem_2d_isentropic_vortex)
      call ic_isentropic_vortex(state, cfg)
    case (problem_2d_sedov)
      call ic_sedov(state, cfg)
    case (problem_2d_toro_explosion)
      call ic_toro_explosion(state, cfg)
    case (problem_2d_riemann2d_config3)
      call ic_config3(state, cfg)
    case (problem_2d_freestream)
      call ic_freestream(state, cfg)
    case default
      error stop 'initial_conditions_2d: unknown problem_type "'// &
        trim(cfg % problem_type)//'"'
    end select
  end subroutine apply_initial_condition_2d

  !> Constant state rho=1, u=0, v=0, p=1.
  subroutine ic_uniform(state, cfg)
    type(solver_state_2d_t), intent(inout) :: state
    type(config_2d_t), intent(in) :: cfg
    integer :: i, j
    real(wp) :: e_tot

    e_tot = 1.0_wp / (cfg % gam - 1.0_wp)   ! p/(gam-1) with rho=1,u=v=0,p=1
    do j = 1, state % ny_local
      do i = 1, state % nx_local
        state % ub(:, i, j) = [1.0_wp, 0.0_wp, 0.0_wp, e_tot]
      end do
    end do
  end subroutine ic_uniform

  !> Classic isentropic vortex (Shu 1998). Mean flow rho=1,p=1,u=v=1; strength
  !! beta=5; centred at the domain midpoint. Exact, smooth, periodic.
  subroutine ic_isentropic_vortex(state, cfg)
    type(solver_state_2d_t), intent(inout) :: state
    type(config_2d_t), intent(in) :: cfg
    integer :: i, j
    real(wp) :: x, y, xc, yc, xb, yb, r2, du, dv, dtemp, rho, u, v, p
    real(wp) :: pi, beta, gm1, fac
    real(wp), parameter :: u_inf = 1.0_wp, v_inf = 1.0_wp

    pi = acos(-1.0_wp)
    beta = 5.0_wp
    gm1 = cfg % gam - 1.0_wp
    xc = 0.5_wp * (cfg % x_left + cfg % x_right)
    yc = 0.5_wp * (cfg % y_left + cfg % y_right)
    fac = beta / (2.0_wp * pi)

    do j = 1, state % ny_local
      do i = 1, state % nx_local
        if (trim(state % blocks(1) % method) == method_fdm .and. .not. state % mesh % uniform) then
          ! Curvilinear-FDM nodal grid: the node coordinates ARE the solution
          ! points and are already sliced to this rank's local block by
          ! build_mesh_2d_local, so index x_node/y_node directly with the local
          ! i/j (no global-index dx arithmetic — the mapping is non-uniform).
          x = state % mesh % x_node(i, j)
          y = state % mesh % y_node(i, j)
        else if (trim(state % blocks(1) % method) == method_fdm) then
          ! Uniform nodal grid: x_node is stored GLOBAL (build_mesh_2d_local is a
          ! no-op for uniform meshes), so it must NOT be indexed by the LOCAL i/j
          ! at np>1.  Compute the node position from the global node index instead
          ! (node g is at x_left+(g-1)*dx, g = ix_first_global + i - 1), mirroring
          ! the FVM uniform branch below.  At np=1 (ix_first_global=1) this is
          ! bit-identical to x_node(i,j), so the serial path is unchanged.
          x = cfg % x_left + real(state % decomp_2d % ix_first_global + i - 2, wp) * state % dx
          y = cfg % y_left + real(state % decomp_2d % iy_first_global + j - 2, wp) * state % dy
        else if (state % mesh % uniform) then
          x = cfg % x_left + (real(state % decomp_2d % ix_first_global + i - 1, wp) - 0.5_wp) * state % dx
          y = cfg % y_left + (real(state % decomp_2d % iy_first_global + j - 1, wp) - 0.5_wp) * state % dy
        else
          x = state % mesh % xc(i, j)
          y = state % mesh % yc(i, j)
        end if
        xb = x - xc
        yb = y - yc
        r2 = xb * xb + yb * yb
        du = -fac * yb * exp(0.5_wp * (1.0_wp - r2))
        dv = fac * xb * exp(0.5_wp * (1.0_wp - r2))
        dtemp = -gm1 * beta * beta / (8.0_wp * cfg % gam * pi * pi) * exp(1.0_wp - r2)
        rho = (1.0_wp + dtemp)**(1.0_wp / gm1)
        u = u_inf + du
        v = v_inf + dv
        p = rho**cfg % gam
        state % ub(1, i, j) = rho
        state % ub(2, i, j) = rho * u
        state % ub(3, i, j) = rho * v
        state % ub(4, i, j) = p / gm1 + 0.5_wp * rho * (u * u + v * v)
      end do
    end do
  end subroutine ic_isentropic_vortex

  !> Sedov-Taylor blast: ambient rho=1,p=1e-5; total energy E=1 deposited as high
  !! pressure in cells within r0 = 3.5*dx of the domain centre (cylindrical).
  !!
  !! Energy normalisation uses the DISCRETE cell count (not pi*r0^2) so that
  !! exactly E=1 is deposited on any grid, avoiding the O(15%) shortfall that
  !! arises from rounding when the disk boundary crosses cell centres.
  subroutine ic_sedov(state, cfg)
    type(solver_state_2d_t), intent(inout) :: state
    type(config_2d_t), intent(in) :: cfg
    integer :: i, j, ig, jg
    real(wp) :: xg, yg, xc, yc, r, r0, gm1, p_amb, p_hot, e_tot
    integer :: n_hot_global
    real(wp), parameter :: e_blast = 1.0_wp
    gm1 = cfg % gam - 1.0_wp
    p_amb = 1.0e-5_wp
    xc = 0.5_wp * (cfg % x_left + cfg % x_right)
    yc = 0.5_wp * (cfg % y_left + cfg % y_right)
    ! Use max(dx,dy) so r0 spans at least 3.5 cells in the coarser direction on
    ! anisotropic grids, guaranteeing at least one hot cell regardless of aspect ratio.
    r0 = 3.5_wp * max(state % dx, state % dy)
    ! Count discrete hot cells globally so energy normalisation is exact.
    n_hot_global = 0
    do jg = 1, state % decomp_2d % ny_global
      yg = cfg % y_left + (real(jg, wp) - 0.5_wp) * state % dy
      do ig = 1, state % decomp_2d % nx_global
        xg = cfg % x_left + (real(ig, wp) - 0.5_wp) * state % dx
        if (sqrt((xg - xc)**2 + (yg - yc)**2) < r0) n_hot_global = n_hot_global + 1
      end do
    end do
    if (n_hot_global == 0) &
      call parallel_fatal('ic_sedov: no hot cells (r0 too small for grid aspect ratio)')
    ! deposit E=1 uniformly over the discrete hot cells: p * A_cell * N_hot = gm1 * E
    p_hot = gm1 * e_blast / (real(n_hot_global, wp) * state % dx * state % dy)
    do j = 1, state % ny_local
      jg = state % decomp_2d % iy_first_global + j - 1
      yg = cfg % y_left + (real(jg, wp) - 0.5_wp) * state % dy
      do i = 1, state % nx_local
        ig = state % decomp_2d % ix_first_global + i - 1
        xg = cfg % x_left + (real(ig, wp) - 0.5_wp) * state % dx
        r = sqrt((xg - xc)**2 + (yg - yc)**2)
        e_tot = merge(p_hot, p_amb, r < r0) / gm1   ! u=v=0 → E = p/(g-1)
        state % ub(1, i, j) = 1.0_wp
        state % ub(2, i, j) = 0.0_wp
        state % ub(3, i, j) = 0.0_wp
        state % ub(4, i, j) = e_tot
      end do
    end do
  end subroutine ic_sedov

  !> Toro 2D explosion: circle R=0.4 at the domain centre; inside rho=1,p=1;
  !! outside rho=0.125,p=0.1; u=v=0.
  subroutine ic_toro_explosion(state, cfg)
    type(solver_state_2d_t), intent(inout) :: state
    type(config_2d_t), intent(in) :: cfg
    integer :: i, j, ig, jg
    real(wp) :: xg, yg, xc, yc, r, gm1, rho, p
    gm1 = cfg % gam - 1.0_wp
    xc = 0.5_wp * (cfg % x_left + cfg % x_right)
    yc = 0.5_wp * (cfg % y_left + cfg % y_right)
    do j = 1, state % ny_local
      jg = state % decomp_2d % iy_first_global + j - 1
      yg = cfg % y_left + (real(jg, wp) - 0.5_wp) * state % dy
      do i = 1, state % nx_local
        ig = state % decomp_2d % ix_first_global + i - 1
        xg = cfg % x_left + (real(ig, wp) - 0.5_wp) * state % dx
        r = sqrt((xg - xc)**2 + (yg - yc)**2)
        if (r < 0.4_wp) then
          rho = 1.0_wp; p = 1.0_wp
        else
          rho = 0.125_wp; p = 0.1_wp
        end if
        state % ub(1, i, j) = rho
        state % ub(2, i, j) = 0.0_wp
        state % ub(3, i, j) = 0.0_wp
        state % ub(4, i, j) = p / gm1
      end do
    end do
  end subroutine ic_toro_explosion

  !> Lax-Liu Config 3 four-quadrant Riemann problem (split at the domain centre).
  subroutine ic_config3(state, cfg)
    type(solver_state_2d_t), intent(inout) :: state
    type(config_2d_t), intent(in) :: cfg
    integer :: i, j, ig, jg
    real(wp) :: xg, yg, xc, yc, gm1, rho, u, v, p
    gm1 = cfg % gam - 1.0_wp
    xc = 0.5_wp * (cfg % x_left + cfg % x_right)
    yc = 0.5_wp * (cfg % y_left + cfg % y_right)
    do j = 1, state % ny_local
      jg = state % decomp_2d % iy_first_global + j - 1
      yg = cfg % y_left + (real(jg, wp) - 0.5_wp) * state % dy
      do i = 1, state % nx_local
        ig = state % decomp_2d % ix_first_global + i - 1
        xg = cfg % x_left + (real(ig, wp) - 0.5_wp) * state % dx
        if (xg >= xc .and. yg >= yc) then          ! upper-right
          rho = 1.5_wp; u = 0.0_wp; v = 0.0_wp; p = 1.5_wp
        else if (xg < xc .and. yg >= yc) then       ! upper-left
          rho = 0.5323_wp; u = 1.206_wp; v = 0.0_wp; p = 0.3_wp
        else if (xg < xc .and. yg < yc) then        ! lower-left
          rho = 0.138_wp; u = 1.206_wp; v = 1.206_wp; p = 0.029_wp
        else                                        ! lower-right
          rho = 0.5323_wp; u = 0.0_wp; v = 1.206_wp; p = 0.3_wp
        end if
        state % ub(1, i, j) = rho
        state % ub(2, i, j) = rho * u
        state % ub(3, i, j) = rho * v
        state % ub(4, i, j) = p / gm1 + 0.5_wp * rho * (u * u + v * v)
      end do
    end do
  end subroutine ic_config3

  !> Uniform free-stream from the config reference state (rho_inf,u_inf,v_inf,p_inf).
  !! Used to initialise supersonic inflow cases (e.g. the compression ramp).
  subroutine ic_freestream(state, cfg)
    type(solver_state_2d_t), intent(inout) :: state
    type(config_2d_t), intent(in) :: cfg
    integer :: i, j
    real(wp) :: e_tot
    e_tot = cfg % p_inf / (cfg % gam - 1.0_wp) &
            + 0.5_wp * cfg % rho_inf * (cfg % u_inf**2 + cfg % v_inf**2)
    do j = 1, state % ny_local
      do i = 1, state % nx_local
        state % ub(1, i, j) = cfg % rho_inf
        state % ub(2, i, j) = cfg % rho_inf * cfg % u_inf
        state % ub(3, i, j) = cfg % rho_inf * cfg % v_inf
        state % ub(4, i, j) = e_tot
      end do
    end do
  end subroutine ic_freestream

end module initial_conditions_2d