!> Grey, absorbing-emitting, non-scattering Discrete Ordinates Method (DOM) radiation solver model.
!!
!! [Physical Assumptions]
!! - Grey (single wavenumber/band) absorbing-emitting medium.
!! - Non-scattering; no scattering source is included in the transport equations.
!!
!! [Boundary Condition Assumptions & Limitations]
!! - Wall/open-boundary inflow currently assumes black diffuse behavior, epsilon = 1.0:
!!     I_in = (sigma/pi)*T_w^4.
!! - Input emissivity values below 1 are detected and warned about, but diffuse-gray
!!   reflection is not implemented in this local independent-ordinate solve.
!! - Symmetry boundaries are simplified/skipped. A true DOM symmetry condition should
!!   reflect the ordinate direction through the symmetry plane.
!!
!! [X-1 Contract]
!! - This module computes and populates only the unreduced local partial absorption-like
!!   term kappa*G_local in source%qrad.
!! - The radiation driver is responsible for performing the MPI reduction and then locally
!!   subtracting the isotropic emission term 4*kappa*sigma*T^4 on all ranks.
!!
!! This module solves the discretized finite-volume grey DOM transport equations:
!!   s_m . grad(I_m) + kappa*I_m = kappa*I_b
!! where I_b = sigma/pi*T^4.
!! Angular decomposition is used to distribute directions over MPI processes.
module mod_radiation_dom
   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
   implicit none

   private
   public :: radiation_compute_dom

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

contains

   !> Main entry point for the DOM radiation model.
   subroutine radiation_compute_dom(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 :: s_x(:), s_y(:), s_z(:), w(:)
      real(rk), allocatable :: diag(:)
      real(rk), allocatable :: coeff(:,:)
      integer, allocatable :: nb(:,:)
      real(rk), allocatable :: rhs(:)
      real(rk), allocatable :: I_m(:)
      real(rk), allocatable :: G_local(:)

      integer :: n_dirs, m, c, lf, f, neighbor, max_faces, max_iter
      real(rk) :: absorption, emissivity_input, I_b_val
      real(rk) :: u, area, T_w, rhs_bc
      real(rk) :: n_out(3)
      logical :: is_dir
      logical, save :: dom_bc_warned = .false.
      logical, save :: dom_neg_intensity_warned = .false.
      logical, save :: dom_emissivity_warned = .false.

      source%qrad = zero

      if (context%rad_size <= 0) then
         call fatal_error('dom', 'Invalid radiation angular communicator size')
      end if

      absorption = max(params%radiation_absorption_coeff, 1.0e-6_rk)
      emissivity_input = max(min(params%radiation_emissivity, 1.0_rk), 1.0e-6_rk)
      if (emissivity_input < 0.999999_rk .and. .not. dom_emissivity_warned) then
         dom_emissivity_warned = .true.
         write(*,'(A)') ' [WARNING] mod_radiation_dom: DOM currently assumes black diffuse boundaries; ' // &
                           'radiation_emissivity is ignored.'
      end if

      call get_quadrature(params%radiation_dom_quadrature, n_dirs, s_x, s_y, s_z, w)

      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(I_m(mesh%ncells))
      allocate(G_local(mesh%ncells))

      G_local = zero

      ! Solve each locally assigned angular direction.
      do m = 1, n_dirs
         if (mod(m - 1, context%rad_size) /= context%rad_rank) cycle

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

         ! Assemble direction-m upwind finite-volume transport system.
         do c = 1, mesh%ncells
            diag(c) = diag(c) + absorption * mesh%cells(c)%volume
            I_b_val = (sigma / pi) * (state%temperature(c)**4)
            rhs(c) = rhs(c) + absorption * I_b_val * mesh%cells(c)%volume

            do lf = 1, mesh%ncell_faces(c)
               f = mesh%cell_faces(lf, c)
               area = mesh%faces(f)%area

               if (mesh%faces(f)%owner == c) then
                  n_out = mesh%faces(f)%normal
               else
                  n_out = -mesh%faces(f)%normal
               end if

               u = s_x(m)*n_out(1) + s_y(m)*n_out(2) + s_z(m)*n_out(3)
               neighbor = face_effective_neighbor(mesh, bc, f, c)

               if (neighbor > 0) then
                  if (u >= zero) then
                     ! Outflow face: local cell contributes to diagonal.
                     diag(c) = diag(c) + u * area
                  else
                     ! Inflow face: upwind value is neighbor intensity.
                     nb(lf, c) = neighbor
                     coeff(lf, c) = -u * area
                  end if
               else
                  if (patch_type_for_face(mesh, bc, f) == bc_symmetry) cycle

                  if (u >= zero) then
                     diag(c) = diag(c) + u * area
                  else
                     ! Black diffuse wall/open-boundary inflow.
                     call boundary_temperature(mesh, bc, f, state%temperature(c), T_w, is_dir)
                     if (.not. is_dir .and. .not. dom_bc_warned) then
                        dom_bc_warned = .true.
                        write(*,'(A)') ' [WARNING] mod_radiation_dom: Boundary temperature missing; ' // &
                                          'falling back to cell temperature.'
                     end if

                     rhs_bc = -u * area * (sigma / pi) * (T_w**4)
                     rhs(c) = rhs(c) + rhs_bc
                  end if
               end if
            end do
         end do

         do c = 1, mesh%ncells
            I_m(c) = (sigma / pi) * (state%temperature(c)**4)
         end do

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

         if (minval(I_m) < -1.0e-10_rk .and. .not. dom_neg_intensity_warned) then
            dom_neg_intensity_warned = .true.
            write(*,'(A,ES12.5)') ' [WARNING] mod_radiation_dom: Negative DOM ordinate intensity detected: ', minval(I_m)
         end if

         G_local = G_local + w(m) * I_m
      end do

      ! Return only local kappa*G contribution. Driver reduces over angular ranks and
      ! subtracts 4*kappa*sigma*T^4.
      do c = 1, mesh%ncells
         source%qrad(c) = absorption * G_local(c)
      end do

      deallocate(diag, coeff, nb, rhs, I_m, G_local)
      deallocate(s_x, s_y, s_z, w)
   end subroutine radiation_compute_dom

   !> Symmetric level-style quadrature lookup/generator.
   subroutine get_quadrature(quad_type, n_dirs, s_x, s_y, s_z, w)
      character(len=*), intent(in) :: quad_type
      integer, intent(out) :: n_dirs
      real(rk), allocatable, intent(out) :: s_x(:), s_y(:), s_z(:), w(:)

      integer :: L, i, j, k, sx, sy, sz, idx
      real(rk), allocatable :: mu(:)
      real(rk) :: w_ijk

      select case (trim(adjustl(quad_type)))
      case ('s2', 'S2')
         n_dirs = 8
         L = 1
         allocate(mu(L))
         mu(1) = 0.5773502691896258_rk
      case ('s4', 'S4')
         n_dirs = 24
         L = 2
         allocate(mu(L))
         mu(1) = 0.3500212_rk
         mu(2) = 0.8688903_rk
      case ('s6', 'S6')
         n_dirs = 48
         L = 3
         allocate(mu(L))
         mu(1) = 0.2666355_rk
         mu(2) = 0.6815076_rk
         mu(3) = 0.9261808_rk
      case ('s8', 'S8')
         n_dirs = 80
         L = 4
         allocate(mu(L))
         mu(1) = 0.2182179_rk
         mu(2) = 0.5773503_rk
         mu(3) = 0.7867958_rk
         mu(4) = 0.9511897_rk
      case default
         call fatal_error('dom', 'Unsupported DOM quadrature type')
      end select

      allocate(s_x(n_dirs), s_y(n_dirs), s_z(n_dirs), w(n_dirs))
      idx = 0

      ! Generate combinations in the first octant where i+j+k=L+2, then reflect
      ! to all octants.
      do i = 1, L
         do j = 1, L
            do k = 1, L
               if (i + j + k /= L + 2) cycle

               w_ijk = quadrature_octant_weight(trim(adjustl(quad_type)), i, j, k)

               do sx = -1, 1, 2
                  do sy = -1, 1, 2
                     do sz = -1, 1, 2
                        idx = idx + 1
                        if (idx > n_dirs) then
                           call fatal_error('dom', 'DOM quadrature generated too many directions')
                        end if
                        s_x(idx) = real(sx, rk) * mu(i)
                        s_y(idx) = real(sy, rk) * mu(j)
                        s_z(idx) = real(sz, rk) * mu(k)
                        w(idx) = w_ijk * (pi / 2.0_rk)
                     end do
                  end do
               end do
            end do
         end do
      end do

      deallocate(mu)

      if (idx /= n_dirs) then
         call fatal_error('dom', 'DOM quadrature generated an unexpected number of directions')
      end if

      call validate_quadrature(n_dirs, s_x, s_y, s_z, w)
   end subroutine get_quadrature

   !> Returns first-octant normalized weights. Full-space weights are w_ijk*pi/2.
   real(rk) function quadrature_octant_weight(quad_type, i, j, k) result(w_ijk)
      character(len=*), intent(in) :: quad_type
      integer, intent(in) :: i, j, k

      select case (quad_type)
      case ('s2', 'S2')
         w_ijk = 1.0_rk
      case ('s4', 'S4')
         w_ijk = 0.333333333333_rk
      case ('s6', 'S6')
         if ((i == 1 .and. j == 1 .and. k == 3) .or. &
             (i == 1 .and. j == 3 .and. k == 1) .or. &
             (i == 3 .and. j == 1 .and. k == 1)) then
            w_ijk = 0.1761263_rk
         else
            w_ijk = 0.1572071_rk
         end if
      case ('s8', 'S8')
         if ((i == 1 .and. j == 1 .and. k == 4) .or. &
             (i == 1 .and. j == 4 .and. k == 1) .or. &
             (i == 4 .and. j == 1 .and. k == 1)) then
            w_ijk = 0.1209877_rk
         else if (i == 2 .and. j == 2 .and. k == 2) then
            w_ijk = 0.0925926_rk
         else
            w_ijk = 0.0907407_rk
         end if
      case default
         call fatal_error('dom', 'Unsupported DOM quadrature type in weight lookup')
      end select
   end function quadrature_octant_weight

   !> Runtime quadrature consistency checks.
   subroutine validate_quadrature(n_dirs, s_x, s_y, s_z, w)
      integer, intent(in) :: n_dirs
      real(rk), intent(in) :: s_x(:), s_y(:), s_z(:), w(:)

      integer :: m
      real(rk) :: w_sum, length
      real(rk) :: mxx, myy, mzz, mxy, mxz, myz
      real(rk), parameter :: tol_quad = 1.0e-5_rk

      w_sum = zero
      mxx = zero
      myy = zero
      mzz = zero
      mxy = zero
      mxz = zero
      myz = zero

      do m = 1, n_dirs
         length = sqrt(s_x(m)**2 + s_y(m)**2 + s_z(m)**2)
         if (abs(length - 1.0_rk) > tol_quad) then
            call fatal_error('dom', 'DOM quadrature direction vector does not have unit length')
         end if
         if (w(m) <= zero) then
            call fatal_error('dom', 'DOM quadrature has non-positive weights')
         end if

         w_sum = w_sum + w(m)
         mxx = mxx + w(m) * s_x(m)**2
         myy = myy + w(m) * s_y(m)**2
         mzz = mzz + w(m) * s_z(m)**2
         mxy = mxy + w(m) * s_x(m) * s_y(m)
         mxz = mxz + w(m) * s_x(m) * s_z(m)
         myz = myz + w(m) * s_y(m) * s_z(m)
      end do

      if (abs(w_sum - 4.0_rk*pi) > tol_quad) then
         call fatal_error('dom', 'DOM quadrature weights do not sum to 4*pi')
      end if

      if (abs(mxx - 4.0_rk*pi/3.0_rk) > tol_quad .or. &
          abs(myy - 4.0_rk*pi/3.0_rk) > tol_quad .or. &
          abs(mzz - 4.0_rk*pi/3.0_rk) > tol_quad) then
         call fatal_error('dom', 'DOM quadrature second moments are not isotropic')
      end if

      if (abs(mxy) > tol_quad .or. abs(mxz) > tol_quad .or. abs(myz) > tol_quad) then
         call fatal_error('dom', 'DOM quadrature cross moments are not zero')
      end if
   end subroutine validate_quadrature

   !> Jacobi-preconditioned BiCGStab solver for the unsymmetric DOM transport system.
   subroutine solve_pbicgstab(mesh, diag, coeff, nb, rhs, x, 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) :: x(:)
      integer, intent(in) :: max_iter
      real(rk), intent(in) :: tol

      real(rk), allocatable :: r(:), r0_star(:), p(:), v(:), s(:), t(:)
      real(rk), allocatable :: y(:), z(:), ax(:), inv_diag(:)
      real(rk) :: rho_prev, rho_curr, alpha, omega, beta, r0_norm, r_norm, denom, tt
      integer :: iter, c
      logical :: converged

      allocate(r(mesh%ncells))
      allocate(r0_star(mesh%ncells))
      allocate(p(mesh%ncells))
      allocate(v(mesh%ncells))
      allocate(s(mesh%ncells))
      allocate(t(mesh%ncells))
      allocate(y(mesh%ncells))
      allocate(z(mesh%ncells))
      allocate(ax(mesh%ncells))
      allocate(inv_diag(mesh%ncells))

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

      call apply_matrix(mesh, diag, coeff, nb, x, ax)
      r = rhs - ax
      r0_star = r
      r0_norm = sqrt(dot_product(r, r))
      r_norm = r0_norm
      converged = .false.

      if (r0_norm < tol) then
         converged = .true.
         deallocate(r, r0_star, p, v, s, t, y, z, ax, inv_diag)
         return
      end if

      p = zero
      v = zero
      rho_prev = 1.0_rk
      alpha = 1.0_rk
      omega = 1.0_rk

      do iter = 1, max_iter
         rho_curr = dot_product(r0_star, r)
         if (abs(rho_curr) < 1.0e-30_rk) then
            write(*,'(A)') ' [WARNING] mod_radiation_dom: BiCGStab breakdown (rho = 0).'
            exit
         end if

         if (iter == 1) then
            p = r
         else
            if (abs(omega) < 1.0e-30_rk) then
               write(*,'(A)') ' [WARNING] mod_radiation_dom: BiCGStab breakdown (omega = 0).'
               exit
            end if
            beta = (rho_curr / rho_prev) * (alpha / omega)
            p = r + beta * (p - omega * v)
         end if

         ! Left Jacobi preconditioning: y = M^{-1} p.
         y = inv_diag * p
         call apply_matrix(mesh, diag, coeff, nb, y, v)

         denom = dot_product(r0_star, v)
         if (abs(denom) < 1.0e-30_rk) then
            write(*,'(A)') ' [WARNING] mod_radiation_dom: BiCGStab breakdown (denom = 0).'
            exit
         end if

         alpha = rho_curr / denom
         s = r - alpha * v

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

         ! z = M^{-1} s.
         z = inv_diag * s
         call apply_matrix(mesh, diag, coeff, nb, z, t)

         tt = dot_product(t, t)
         if (abs(tt) < 1.0e-30_rk) then
            write(*,'(A)') ' [WARNING] mod_radiation_dom: BiCGStab breakdown (t.t = 0).'
            exit
         end if

         omega = dot_product(t, s) / tt
         x = x + alpha * y + omega * z
         r = s - omega * t

         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

         rho_prev = rho_curr
      end do

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

      deallocate(r, r0_star, p, v, s, t, y, z, ax, inv_diag)
   end subroutine solve_pbicgstab

   !> Sparse matrix-vector product for the assembled DOM transport 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_dom
