!> Grey, absorbing-emitting, non-scattering P-1 radiation solver model for incident radiation diffusion.
!!
!! [Physical Assumptions]
!! - Grey (single wavenumber/band) absorbing-emitting medium.
!! - Non-scattering; no scattering source is included and Gamma = 1 / (3*kappa).
!!
!! [X-1 Contract]
!! - This module computes and populates only the unreduced partial absorption-like term
!!   kappa*G in source%qrad.
!! - The radiation driver is responsible for performing MPI reductions and then locally
!!   subtracting the emission term 4*kappa*sigma*T^4 on all ranks.
!!
!! [Parallel Spectral Contract]
!! - In a parallel spectral run, only the single rank assigned to wavenumber/band 1
!!   computes the P-1 equation. Other ranks exit immediately returning zero source%qrad,
!!   which is subsequently synchronized globally by the driver using MPI_Allreduce(MPI_SUM).
!!
!! This module solves the discretized finite-volume P-1 equation:
!!   -div(Gamma grad(G)) + kappa*G = 4*kappa*sigma*T^4
!! where Gamma = 1/(3*kappa), and G is the incident radiation.
!! Marshak boundary conditions are applied at walls and open boundaries.
module mod_radiation_p1
   use mod_precision, only : rk, zero, tiny_safe, fatal_error
   use mod_mesh, only : mesh_t
   use mod_input, only : case_params_t
   use mod_bc, only : bc_set_t, patch_type_for_face, face_effective_neighbor, &
                      boundary_temperature, bc_symmetry
   use mod_radiation_types, only : radiation_context_t, radiation_state_t, radiation_source_t
   use mod_linear_solver, only : face_normal_distance
   implicit none

   private
   public :: radiation_compute_p1

   !> Stefan-Boltzmann constant [W/(m^2*K^4)]
   real(rk), parameter :: sigma = 5.670374419e-8_rk

contains

   !> Main entry point for the P-1 radiation model.
   subroutine radiation_compute_p1(mesh, bc, context, state, source, params)
      type(mesh_t), intent(in) :: mesh
      type(bc_set_t), intent(in) :: bc
      type(radiation_context_t), intent(in) :: context
      type(radiation_state_t), intent(in) :: state
      type(radiation_source_t), intent(inout) :: source
      type(case_params_t), intent(in) :: params

      real(rk), allocatable :: diag(:)
      real(rk), allocatable :: coeff(:,:)
      integer, allocatable :: nb(:,:)
      real(rk), allocatable :: rhs(:)
      real(rk), allocatable :: G(:)

      integer :: c, lf, f, neighbor, max_faces, max_iter
      real(rk) :: absorption, gamma, emissivity
      real(rk) :: dist, coeff_val, denom, T_w, rhs_bc
      logical :: is_dir
      logical, save :: p1_bc_warned = .false.
      logical, save :: p1_neg_g_warned = .false.

      ! Initialize the output partial source term.
      source%qrad = zero

      ! In the parallel spectral framework, only the rank assigned to wavenumber 1
      ! computes this single grey band.
      if (context%wn_first > 1 .or. context%wn_last < 1) return

      max_faces = size(mesh%cell_faces, 1)

      allocate(diag(mesh%ncells))
      allocate(coeff(max_faces, mesh%ncells))
      allocate(nb(max_faces, mesh%ncells))
      allocate(rhs(mesh%ncells))
      allocate(G(mesh%ncells))

      diag = zero
      coeff = zero
      nb = 0
      rhs = zero

      absorption = max(params%radiation_absorption_coeff, 1.0e-6_rk)
      emissivity = max(min(params%radiation_emissivity, 1.0_rk), 1.0e-6_rk)
      gamma = 1.0_rk / (3.0_rk * absorption)

      ! Assemble the P-1 finite-volume system.
      do c = 1, mesh%ncells
         diag(c) = diag(c) + absorption * mesh%cells(c)%volume
         rhs(c) = rhs(c) + 4.0_rk * absorption * sigma * &
                            (state%temperature(c)**4) * mesh%cells(c)%volume

         do lf = 1, mesh%ncell_faces(c)
            f = mesh%cell_faces(lf, c)
            neighbor = face_effective_neighbor(mesh, bc, f, c)

            if (neighbor > 0) then
               ! Internal or periodic face.
               dist = face_normal_distance(mesh, bc, f, c, neighbor)
               if (dist <= tiny_safe) then
                  call fatal_error('p1', 'Non-positive face-normal distance in internal P-1 face assembly')
               end if

               coeff_val = gamma * mesh%faces(f)%area / dist
               nb(lf, c) = neighbor
               coeff(lf, c) = coeff_val
               diag(c) = diag(c) + coeff_val
            else
               ! Boundary face.
               if (patch_type_for_face(mesh, bc, f) == bc_symmetry) cycle

               ! Non-periodic wall or open boundary: apply Marshak BC.
               ! The denominator represents boundary resistance:
               !   denom = Delta_x/Gamma + 2*(2 - epsilon)/epsilon
               dist = face_normal_distance(mesh, bc, f, c, 0)
               if (dist <= tiny_safe) then
                  call fatal_error('p1', 'Non-positive face-normal distance in boundary P-1 face assembly')
               end if

               denom = 2.0_rk * (2.0_rk - emissivity) / emissivity + dist / gamma
               if (denom <= tiny_safe) then
                  call fatal_error('p1', 'Non-positive Marshak boundary denominator')
               end if

               coeff_val = mesh%faces(f)%area / denom
               diag(c) = diag(c) + coeff_val

               call boundary_temperature(mesh, bc, f, state%temperature(c), T_w, is_dir)
               if (.not. is_dir .and. .not. p1_bc_warned) then
                  p1_bc_warned = .true.
                  write(*,'(A)') ' [WARNING] mod_radiation_p1: Radiative boundary faces have no prescribed temperature; ' // &
                                    'falling back to cell temperature.'
               end if

               rhs_bc = coeff_val * 4.0_rk * sigma * (T_w**4)
               rhs(c) = rhs(c) + rhs_bc
            end if
         end do
      end do

      ! Initial guess: local blackbody incident radiation.
      do c = 1, mesh%ncells
         G(c) = 4.0_rk * sigma * (state%temperature(c)**4)
      end do

      max_iter = 1000
      call solve_pcg(mesh, diag, coeff, nb, rhs, G, max_iter, 1.0e-8_rk)

      ! Incident radiation should be nonnegative. Warn once if the linear solve produces
      ! an unphysical value, but do not clip silently.
      if (minval(G) < -1.0e-10_rk .and. .not. p1_neg_g_warned) then
         p1_neg_g_warned = .true.
         write(*,'(A,ES12.5)') ' [WARNING] mod_radiation_p1: Negative incident radiation detected: ', minval(G)
      end if

      ! Return only the partial absorption-like contribution. The driver subtracts
      ! 4*kappa*sigma*T^4 after global synchronization.
      do c = 1, mesh%ncells
         source%qrad(c) = absorption * G(c)
      end do

      deallocate(diag, coeff, nb, rhs, G)
   end subroutine radiation_compute_p1

   !> Jacobi-preconditioned Conjugate Gradient solver for the symmetric P-1 system.
   subroutine solve_pcg(mesh, diag, coeff, nb, rhs, G, max_iter, tol)
      type(mesh_t), intent(in) :: mesh
      real(rk), intent(in) :: diag(:)
      real(rk), intent(in) :: coeff(:,:)
      integer, intent(in) :: nb(:,:)
      real(rk), intent(in) :: rhs(:)
      real(rk), intent(inout) :: G(:)
      integer, intent(in) :: max_iter
      real(rk), intent(in) :: tol

      real(rk), allocatable :: r(:), z(:), p(:), ap(:), inv_diag(:)
      real(rk) :: rz, rz_new, alpha, beta, pap, r_norm, r0_norm
      integer :: iter, c
      logical :: converged

      allocate(r(mesh%ncells))
      allocate(z(mesh%ncells))
      allocate(p(mesh%ncells))
      allocate(ap(mesh%ncells))
      allocate(inv_diag(mesh%ncells))

      do c = 1, mesh%ncells
         if (diag(c) <= tiny_safe) then
            call fatal_error('p1', 'Non-positive diagonal encountered in P-1 PCG solver')
         end if
         inv_diag(c) = 1.0_rk / diag(c)
      end do

      ! r = rhs - A*G
      call apply_matrix(mesh, diag, coeff, nb, G, ap)
      r = rhs - ap

      r0_norm = sqrt(dot_product(r, r))
      r_norm = r0_norm
      converged = .false.

      if (r0_norm < tol) then
         converged = .true.
         deallocate(r, z, p, ap, inv_diag)
         return
      end if

      z = inv_diag * r
      p = z
      rz = dot_product(r, z)

      do iter = 1, max_iter
         call apply_matrix(mesh, diag, coeff, nb, p, ap)

         pap = dot_product(p, ap)
         if (abs(pap) < 1.0e-30_rk) then
            write(*,'(A)') ' [WARNING] mod_radiation_p1: PCG breakdown (pAp = 0).'
            exit
         end if

         alpha = rz / pap
         G = G + alpha * p
         r = r - alpha * ap

         r_norm = sqrt(dot_product(r, r))
         if (r_norm < tol * r0_norm .or. r_norm < 1.0e-12_rk) then
            converged = .true.
            exit
         end if

         z = inv_diag * r
         rz_new = dot_product(r, z)
         if (abs(rz) < 1.0e-30_rk) then
            write(*,'(A)') ' [WARNING] mod_radiation_p1: PCG breakdown (rMz = 0).'
            exit
         end if

         beta = rz_new / rz
         p = z + beta * p
         rz = rz_new
      end do

      if (.not. converged) then
         write(*,'(A,I0,A,ES12.5)') ' [WARNING] mod_radiation_p1: PCG failed to converge in ', &
                                      max_iter, ' iterations. Final residual norm: ', r_norm
      end if

      deallocate(r, z, p, ap, inv_diag)
   end subroutine solve_pcg

   !> Sparse matrix-vector product for the assembled P-1 system.
   subroutine apply_matrix(mesh, diag, coeff, nb, x, ax)
      type(mesh_t), intent(in) :: mesh
      real(rk), intent(in) :: diag(:)
      real(rk), intent(in) :: coeff(:,:)
      integer, intent(in) :: nb(:,:)
      real(rk), intent(in) :: x(:)
      real(rk), intent(out) :: ax(:)

      integer :: c, lf, neighbor

      do c = 1, mesh%ncells
         ax(c) = diag(c) * x(c)
         do lf = 1, mesh%ncell_faces(c)
            neighbor = nb(lf, c)
            if (neighbor > 0) then
               ax(c) = ax(c) - coeff(lf, c) * x(neighbor)
            end if
         end do
      end do
   end subroutine apply_matrix

end module mod_radiation_p1
