!> Finite-Volume transport solver for chemical species mass fractions.
!!
!! This module implements non-reacting finite-volume transport for \(N\)
!! chemical species mass fractions \(Y_k\). The solver supports:
!! 1. **Upwind Advection**: A 1st-order stable scheme for robust transport 
!!    of sharp scalar gradients.
!! 2. **Corrected Diffusion**: Diffusive fluxes are explicitly corrected 
!!    to ensure the net mass flux sums to zero \(\sum \mathbf{j}_k = 0\).
!! 3. **Mass Conservation**: Enforces \(\sum Y_k = 1\) and boundedness \([0, 1]\)
!!    after every timestep.  Chemical reactions and source terms are currently
!!    disabled in the supported path.
!! 4. **MPI Synchronization**: Efficiently gathers owned-cell updates into 
!!    the globally replicated mesh field.
module mod_species
   use mod_precision, only : rk, zero, one, tiny_safe, fatal_error, name_len, lowercase
   use mod_mesh, only : mesh_t
   use mod_input, only : case_params_t
   use mod_bc, only : bc_set_t, bc_periodic, patch_type_for_face, face_effective_neighbor, boundary_species
   use mod_flow_fields, only : flow_fields_t
   use mod_mpi_flow, only : flow_mpi_t, flow_exchange_cell_matrix
   use mod_flow_projection, only : face_normal_distance
   use mod_cantera, only : transport_properties_t
   use mod_reconstruction, only : compute_scalar_gradients, scalar_face_value
   use mod_species_composition, only : parse_named_composition, composition_string_is_set
   implicit none

   private

   public :: species_fields_t
   public :: initialize_species, finalize_species, advance_species_transport

   !> Container for multi-species mass fraction fields.
   type :: species_fields_t
      integer :: nspecies = 0                       !< Total number of transport species \(N_s\).
      real(rk), allocatable :: Y(:,:)               !< Current mass fractions \(Y_k\) (nspecies, ncells).
      real(rk), allocatable :: Y_old(:,:)           !< Mass fractions from previous step \(n\).
      real(rk), allocatable :: rhs_old(:,:)         !< Previous transport RHS for optional AB2 time integration.
      logical :: rhs_old_valid = .false.            !< True once rhs_old contains a previous transport RHS.
      character(len=name_len), allocatable :: names(:) !< Array of species names (e.g., "H2", "O2", "N2").
   end type species_fields_t

contains

   !> Populates species fields with initial mass fractions and handles naming.
   !!
   !! Performs name-based matching between the transport registry and 
   !! namelist initial conditions. Normalizes the initial mixture to 
   !! ensure the physical constraint \(\sum Y_k = 1\) is met at \(t=0\).
   !!
   !! @param mesh The computational mesh.
   !! @param params Input configuration.
   !! @param species The fields to initialize.
   subroutine initialize_species(mesh, params, species)
      type(mesh_t), intent(in) :: mesh
      type(case_params_t), intent(in) :: params
      type(species_fields_t), intent(inout) :: species

      integer :: c, k, j
      real(rk) :: sum_Y
      real(rk) :: init_mixture(params%nspecies)
      character(len=name_len) :: target_name

      call finalize_species(species)

      species%nspecies = params%nspecies
      if (species%nspecies <= 0) return

      allocate(species%Y(species%nspecies, mesh%ncells))
      allocate(species%Y_old(species%nspecies, mesh%ncells))
      allocate(species%rhs_old(species%nspecies, mesh%ncells))
      allocate(species%names(species%nspecies))

      species%names = params%species_name(1:species%nspecies)

      ! Prefer named initial composition strings because Cantera may replace
      ! the input species ordering with the full mechanism species ordering.
      init_mixture = zero
      if (composition_string_is_set(params%initial_composition)) then
         call parse_named_composition(params%initial_composition, species%names, species%nspecies, &
                                      init_mixture, 'species initial_composition', &
                                      normalize=.true., require_sum_positive=.true.)
      else
         ! Legacy path: match namelist initial_Y entries by the species_name list
         ! supplied in &species_input.  This remains safer than assuming numeric
         ! indices are unchanged after Cantera species discovery.
         do j = 1, params%namelist_nspecies
            target_name = trim(lowercase(params%namelist_species_name(j)))
            if (len_trim(target_name) == 0) cycle
            do k = 1, species%nspecies
               if (trim(lowercase(species%names(k))) == target_name) then
                  init_mixture(k) = params%initial_Y(j)
                  exit
               end if
            end do
         end do
      end if

      ! Normalize the initial mixture vector
      sum_Y = sum(init_mixture)
      if (sum_Y > zero) then
         init_mixture = init_mixture / sum_Y
      else
         ! Fallback: If no IC specified, set the mixture to 100% of the first species.
         if (species%nspecies > 0) init_mixture(1) = one
      end if

      do c = 1, mesh%ncells
         species%Y(:, c) = init_mixture
      end do

      species%Y_old = species%Y
      species%rhs_old = zero
      species%rhs_old_valid = .false.
   end subroutine initialize_species


   !> Safely deallocates species fields and names.
   subroutine finalize_species(species)
      type(species_fields_t), intent(inout) :: species

      if (allocated(species%Y)) deallocate(species%Y)
      if (allocated(species%Y_old)) deallocate(species%Y_old)
      if (allocated(species%rhs_old)) deallocate(species%rhs_old)
      if (allocated(species%names)) deallocate(species%names)
      species%nspecies = 0
      species%rhs_old_valid = .false.
   end subroutine finalize_species


   !> Performs one explicit Euler step for non-reacting species transport.
   !!
   !! This routine evaluates:
   !! - **Constant-density advection**: `fields%face_flux` [m^3/s] advects
   !!   \(Y_k\) directly.
   !! - **Variable-density advection**: `fields%mass_flux` [kg/s] advances
   !!   conservative \(\rho Y_k\), then divides by the active density.
   !! - **Diffusive fluxes**: constant-density mode uses \(D_k \nabla Y_k\);
   !!   variable-density mode uses \(\rho_f D_k \nabla Y_k\).
   !! - **Correction velocity**: subtracts \(Y_k \sum J_{diff,k}\) at each face
   !!   so the corrected diffusive species flux has zero net mass flux.
   !! - **Bounding & Normalization**: Clamps \(Y_k \in [0, 1]\) and enforces \(\sum Y_k = 1\) locally.
   !! Reactions are not applied here; `enable_reactions` is rejected by the
   !! supported variable-density path.
   !!
   !! @param mesh The computational mesh.
   !! @param flow MPI decomposition data for synchronization.
   !! @param bc Boundary condition settings.
   !! @param params Simulation parameters (dt, etc).
   !! @param fields Flow field (velocity/face fluxes).
   !! @param species Mass fraction fields to update.
   !! @param transport Physical properties (diffusivities).
   subroutine advance_species_transport(mesh, flow, bc, params, fields, species, transport)
      type(mesh_t), intent(in) :: mesh
      type(flow_mpi_t), intent(inout) :: flow
      type(bc_set_t), intent(in) :: bc
      type(case_params_t), intent(in) :: params
      type(flow_fields_t), intent(in) :: fields
      type(species_fields_t), intent(inout) :: species
      type(transport_properties_t), intent(in) :: transport

      real(rk), allocatable :: dY(:), diff_flux(:), adv_flux(:), Y_face_lin(:), rhs_current(:)
      real(rk), allocatable :: gradY(:,:,:)
      real(rk) :: flux, face_area, dist, sum_diff_flux
      real(rk) :: rho_face, rho_cell, rho_old_cell
      real(rk) :: Y_cell, Y_other, Y_face
      real(rk) :: diff_face, rhs_update
      real(rk) :: c_curr, c_old
      integer :: c, f, fid, neighbor
      integer :: k
      logical :: is_dirichlet, use_variable_density_species, use_ab2

      if (species%nspecies <= 0) return

      use_ab2 = trim(lowercase(params%species_time_scheme)) == 'ab2' .or. &
                trim(lowercase(params%species_time_scheme)) == 'adams_bashforth2' .or. &
                trim(lowercase(params%species_time_scheme)) == 'adams-bashforth2'
      c_curr = 1.5_rk
      c_old = -0.5_rk
      if (use_ab2 .and. species%rhs_old_valid .and. params%dt_old > tiny_safe) then
         c_curr = 1.0_rk + 0.5_rk * (params%dt / params%dt_old)
         c_old = -0.5_rk * (params%dt / params%dt_old)
      end if
      use_variable_density_species = params%enable_variable_density
      if (use_variable_density_species) then
         if (.not. allocated(fields%mass_flux)) then
            call fatal_error('species', 'variable-density species transport requires fields%mass_flux')
         end if
         if (.not. allocated(transport%rho)) then
            call fatal_error('species', 'variable-density species transport requires transport%rho')
         end if
         if (.not. allocated(transport%rho_old)) then
            call fatal_error('species', 'variable-density species transport requires transport%rho_old')
         end if
      end if

      allocate(dY(species%nspecies))
      allocate(diff_flux(species%nspecies))
      allocate(adv_flux(species%nspecies))
      allocate(Y_face_lin(species%nspecies))
      allocate(rhs_current(species%nspecies))
      allocate(gradY(3, mesh%ncells, species%nspecies))

      species%Y_old = species%Y
      do k = 1, species%nspecies
         call compute_scalar_gradients(mesh, species%Y_old(k, :), gradY(:, :, k))
      end do

      ! Iterate through MPI-owned cells
      do c = flow%first_cell, flow%last_cell
         dY = zero

         do f = 1, mesh%ncell_faces(c)
            fid = mesh%cell_faces(f,c)
            if (mesh%faces(fid)%owner == c) then
               if (use_variable_density_species) then
                  flux = fields%mass_flux(fid)
               else
                  flux = fields%face_flux(fid)
               end if
            else
               if (use_variable_density_species) then
                  flux = -fields%mass_flux(fid)
               else
                  flux = -fields%face_flux(fid)
               end if
            end if
            neighbor = face_effective_neighbor(mesh, bc, fid, c)

            face_area = mesh%faces(fid)%area
            dist = face_normal_distance(mesh, bc, fid, c, neighbor)

            sum_diff_flux = zero
            do k = 1, species%nspecies
               Y_cell = species%Y_old(k, c)
               
               if (neighbor == 0) then
                  call boundary_species(mesh, bc, fid, k, Y_cell, Y_other, is_dirichlet)
               else
                  Y_other = species%Y_old(k, neighbor)
                  is_dirichlet = .true.
               end if

               ! 1. Advective flux.  The selected scalar scheme can be
               ! first-order upwind, central, bounded central, or compact
               ! limited-linear/MUSCL reconstruction.  flux is oriented
               ! outward from the current cell.
               Y_face = scalar_face_value(mesh, species%Y_old(k, :), gradY(:, :, k), &
                                          params%species_convection_scheme, params%scalar_limiter, &
                                          c, neighbor, fid, flux, Y_other, is_dirichlet)
               adv_flux(k) = -flux * Y_face

               ! 2. Diffusive flux using central difference
               diff_flux(k) = zero
               if (neighbor /= 0 .or. is_dirichlet) then
                  if (neighbor == 0) then
                     diff_face = transport%diffusivity(k, c)
                  else
                     diff_face = 0.5_rk * (transport%diffusivity(k, c) + transport%diffusivity(k, neighbor))
                  end if
                  if (use_variable_density_species) then
                     if (neighbor == 0) then
                        rho_face = transport%rho(c)
                     else
                        rho_face = 0.5_rk * (transport%rho(c) + transport%rho(neighbor))
                     end if
                     diff_flux(k) = rho_face * diff_face * (Y_other - Y_cell) / dist * face_area
                  else
                     diff_flux(k) = diff_face * (Y_other - Y_cell) / dist * face_area
                  end if
               end if
               
               sum_diff_flux = sum_diff_flux + diff_flux(k)
               Y_face_lin(k) = 0.5_rk * (Y_cell + Y_other)
            end do

            ! 3. Apply Correction Velocity to ensure mass conservation: sum(j_k) = 0
            do k = 1, species%nspecies
               dY(k) = dY(k) + adv_flux(k) + (diff_flux(k) - Y_face_lin(k) * sum_diff_flux)
            end do
         end do

         ! Explicit timestep update. Constant-density mode updates Y directly.
         ! Variable-density mode updates rho*Y using mass fluxes, then divides
         ! by the active cell density for this split substep.
         rho_cell = one
         rho_old_cell = one
         if (use_variable_density_species) then
            rho_cell = max(transport%rho(c), tiny_safe)
            rho_old_cell = max(transport%rho_old(c), tiny_safe)
         end if

         rhs_current = dY / mesh%cells(c)%volume

         do k = 1, species%nspecies
            if (use_ab2 .and. species%rhs_old_valid) then
               rhs_update = c_curr * rhs_current(k) + c_old * species%rhs_old(k, c)
            else
               rhs_update = rhs_current(k)
            end if

            if (use_variable_density_species) then
               species%Y(k,c) = (rho_old_cell * species%Y_old(k,c) + params%dt * rhs_update) / rho_cell
            else
               species%Y(k,c) = species%Y_old(k,c) + params%dt * rhs_update
            end if
            
         end do

         species%rhs_old(:, c) = rhs_current

         ! Keep the local composition on the bounded simplex.  This still is a
         ! fallback repair for explicit transport overshoots, but it is less
         ! destructive than scaling every species after clipping: the sum defect
         ! is preferentially put into/removes from the dominant bath component,
         ! preserving trace-species ratios whenever possible.
         call repair_cell_mass_fractions(species%Y(:, c), species%nspecies)
      end do

      deallocate(dY)
      deallocate(diff_flux)
      deallocate(adv_flux)
      deallocate(Y_face_lin)
      deallocate(rhs_current)
      deallocate(gradY)

      species%rhs_old_valid = .true.

      ! Synchronize updated species ghosts for the next transport/property step.
      call flow_exchange_cell_matrix(flow, species%Y)

   end subroutine advance_species_transport


   !> Keeps the local species mass fractions on the bounded physical simplex.
   !!
   !! Implements a simplex projection to correct transport defects and roundoff
   !! errors without unphysically scaling trace species. Defects are preferentially
   !! corrected in the dominant species.
   !!
   !! @param Y_cell Current mass fractions to be repaired.
   !! @param nspecies Total number of species in Y_cell.
   subroutine repair_cell_mass_fractions(Y_cell, nspecies)
      real(rk), intent(inout) :: Y_cell(:) !< Current mass fractions \(Y_k\)
      integer, intent(in) :: nspecies !< Number of species \(N_s\)

      integer :: k !< Species index loop counter
      integer :: k_best !< Index of the species with the highest mass fraction
      integer :: iter !< Loop iteration counter for convergence safety
      real(rk) :: sum_Y !< Sum of mass fractions
      real(rk) :: defect !< The discrepancy from unity constraint: 1 - sum(Y)
      real(rk) :: amount !< The incremental correction applied to the dominant species
      real(rk) :: best_value !< The maximum species mass fraction found
      real(rk), parameter :: simplex_tol = 1.0e-12_rk !< Numerical tolerance for constraint convergence


      if (nspecies <= 0) return

      do k = 1, nspecies
         Y_cell(k) = max(zero, min(one, Y_cell(k)))
      end do

      sum_Y = sum(Y_cell(1:nspecies))
      if (sum_Y <= tiny_safe) then
         Y_cell(1:nspecies) = zero
         Y_cell(1) = one
         return
      end if

      defect = one - sum_Y
      iter = 0
      do while (abs(defect) > simplex_tol .and. iter < 2*nspecies)
         iter = iter + 1
         k_best = 0
         best_value = -one
         
         ! Find the dominant species in this cell (highest mass fraction)
         ! to absorb both surplus and deficit, preserving trace species.
         do k = 1, nspecies
            if (Y_cell(k) > best_value) then
               best_value = Y_cell(k)
               k_best = k
            end if
         end do
         
         if (k_best <= 0) exit
         
         if (defect > zero) then
            amount = min(defect, one - Y_cell(k_best))
            Y_cell(k_best) = Y_cell(k_best) + amount
         else
            amount = min(-defect, Y_cell(k_best))
            Y_cell(k_best) = Y_cell(k_best) - amount
         end if
         
         defect = one - sum(Y_cell(1:nspecies))
      end do

      if (abs(defect) > 10.0_rk*simplex_tol) then
         sum_Y = sum(Y_cell(1:nspecies))
         if (sum_Y > tiny_safe) then
            Y_cell(1:nspecies) = Y_cell(1:nspecies) / sum_Y
         else
            Y_cell(1:nspecies) = zero
            Y_cell(1) = one
         end if
      end if
   end subroutine repair_cell_mass_fractions


end module mod_species

