!> Output management for VTK visualization and diagnostics (XML VTU format).
!!
!! This module handles the generation of simulation results in modern XML-based 
!! VTK format (.vtu) and CSV-based global diagnostics. It manages the 
!! creation of the output directory, writing mesh summaries, and 
!! generating PVD collection files for time-series visualization in ParaView.
module mod_output
   use mod_kinds, only : rk, zero, path_len, fatal_error
   use mod_input, only : case_params_t
   use mod_mesh_types, only : mesh_t
   use mod_mpi_flow, only : flow_mpi_t, flow_gather_owned_scalar_root, &
                            flow_gather_owned_matrix_root
   use mod_fields, only : flow_fields_t
   use mod_flow_projection, only : solver_stats_t
   use mod_species, only : species_fields_t
   use mod_energy, only : energy_fields_t
   use mod_transport_properties, only : transport_properties_t
   implicit none

   private

   public :: prepare_output, write_diagnostics_header, write_diagnostics_row
   public :: write_vtu_unstructured, write_mesh_summary, write_pvd_collection

contains

   !> Creates the output directory specified in the case parameters.
   subroutine prepare_output(params, flow)
      type(case_params_t), intent(in) :: params
      type(flow_mpi_t), intent(in) :: flow

      integer :: exitstat
      character(len=path_len + 16) :: command

      if (flow%rank /= 0) return

      command = 'mkdir -p '//trim(params%output_dir)
      call execute_command_line(trim(command), exitstat=exitstat)

      if (exitstat /= 0) then
         call fatal_error('output', 'failed to create output directory')
      end if
   end subroutine prepare_output


   !> Writes the CSV header for global simulation diagnostics.
   subroutine write_diagnostics_header(params, flow)
      type(case_params_t), intent(in) :: params
      type(flow_mpi_t), intent(in) :: flow

      integer :: unit_id
      character(len=path_len + 32) :: filename

      if (flow%rank /= 0 .or. .not. params%write_diagnostics) return

      filename = trim(params%output_dir)//'/diagnostics.csv'
      open(newunit=unit_id, file=trim(filename), status='replace', action='write')

      write(unit_id,'(a)') 'step,time,dt,max_divergence,rms_divergence,net_boundary_flux,'// &
                           'kinetic_energy,pressure_iterations,pressure_iterations_total,'// &
                           'pressure_iterations_max,pressure_iterations_avg,pressure_solve_count,'// &
                           'pressure_residual,cfl,wall_time,max_velocity,total_mass,min_species_y'

      close(unit_id)
   end subroutine write_diagnostics_header


   !> Appends a new row of diagnostic data to the CSV file.
   subroutine write_diagnostics_row(params, flow, step, time, stats)
      type(case_params_t), intent(in) :: params
      type(flow_mpi_t), intent(in) :: flow
      integer, intent(in) :: step
      real(rk), intent(in) :: time
      type(solver_stats_t), intent(in) :: stats

      integer :: unit_id
      character(len=path_len + 32) :: filename

      if (flow%rank /= 0 .or. .not. params%write_diagnostics) return

      filename = trim(params%output_dir)//'/diagnostics.csv'
      open(newunit=unit_id, file=trim(filename), status='old', position='append', action='write')

      write(unit_id,'(i0,a,es16.8,a,es16.8,a,es16.8,a,es16.8,a,es16.8,a,es16.8,a,'// &
                    'i0,a,i0,a,i0,a,es16.8,a,i0,a,es16.8,a,es16.8,a,es16.8,a,'// &
                    'es16.8,a,es16.8,a,es16.8)') &
            step, ',', time, ',', params%dt, ',', stats%max_divergence, ',', &
            stats%rms_divergence, ',', stats%net_boundary_flux, ',', &
            stats%kinetic_energy, ',', stats%pressure_iterations, ',', &
            stats%pressure_iterations_total, ',', stats%pressure_iterations_max, ',', &
            stats%pressure_iterations_avg, ',', stats%pressure_solve_count, ',', &
            stats%pressure_residual, ',', stats%cfl, ',', stats%wall_time, ',', &
            stats%max_velocity, ',', stats%total_mass, ',', stats%min_species_y

      close(unit_id)
   end subroutine write_diagnostics_row


   !> Writes a human-readable summary of the mesh connectivity and patches.
   subroutine write_mesh_summary(params, flow, mesh)
      type(case_params_t), intent(in) :: params
      type(flow_mpi_t), intent(in) :: flow
      type(mesh_t), intent(in) :: mesh

      integer :: unit_id, p
      character(len=path_len + 32) :: filename

      if (flow%rank /= 0) return

      filename = trim(params%output_dir)//'/mesh_summary.txt'
      open(newunit=unit_id, file=trim(filename), status='replace', action='write')

      write(unit_id,'(a,i0)') 'points: ', mesh%npoints
      write(unit_id,'(a,i0)') 'cells: ', mesh%ncells
      write(unit_id,'(a,i0)') 'faces: ', mesh%nfaces
      write(unit_id,'(a,i0)') 'patches: ', mesh%npatches

      do p = 1, mesh%npatches
         write(unit_id,'(i0,1x,a,1x,i0)') mesh%patches(p)%id, trim(mesh%patches(p)%name), mesh%patches(p)%nfaces
      end do

      close(unit_id)
   end subroutine write_mesh_summary


   !> Writes the full flow field to an XML Unstructured Grid file (.vtu).
   !!
   !! Fields included:
   !! - Velocity (Cell Vector)
   !! - Pressure (Cell Scalar)
   !! - Divergence (Cell Scalar)
   !! - Species Mass Fractions (Cell Scalars, if enabled)
   subroutine write_vtu_unstructured(params, flow, mesh, fields, species, energy, transport, step)
      type(case_params_t), intent(in) :: params
      type(flow_mpi_t), intent(inout) :: flow
      type(mesh_t), intent(in) :: mesh
      type(flow_fields_t), intent(in) :: fields
      type(species_fields_t), intent(in) :: species
      type(energy_fields_t), intent(in) :: energy
      type(transport_properties_t), intent(in) :: transport
      integer, intent(in) :: step

      integer :: unit_id
      integer :: p, c, n, m, k, n_owned
      character(len=path_len + 64) :: filename, local_filename
      integer, allocatable :: owned_indices(:)

      if (.not. params%write_vtu) return

      ! Count owned cells and collect their indices
      n_owned = 0
      do c = 1, mesh%ncells
         if (flow%owned(c)) n_owned = n_owned + 1
      end do

      allocate(owned_indices(n_owned))
      n_owned = 0
      do c = 1, mesh%ncells
         if (flow%owned(c)) then
            n_owned = n_owned + 1
            owned_indices(n_owned) = c
         end if
      end do

      ! Every rank writes its own piece file
      write(local_filename,'("flow_",i6.6,"_P",i4.4,".vtu")') step, flow%rank
      filename = trim(params%output_dir)//'/'//trim(local_filename)
      open(newunit=unit_id, file=trim(filename), status='replace', action='write')

      write(unit_id,'(a)') '<?xml version="1.0"?>'
      write(unit_id,'(a)') '<VTKFile type="UnstructuredGrid" version="0.1" byte_order="LittleEndian">'
      write(unit_id,'(a)') '  <UnstructuredGrid>'
      write(unit_id,'(a,i0,a,i0,a)') '    <Piece NumberOfPoints="', mesh%npoints, &
                                      '" NumberOfCells="', n_owned, '">'

      ! ------------------------------------------------------------
      ! Correct XML VTK Piece order:
      !
      ! PointData
      ! CellData
      ! Points
      ! Cells
      ! ------------------------------------------------------------

      write(unit_id,'(a)') '      <PointData>'
      write(unit_id,'(a)') '      </PointData>'

      write(unit_id,'(a)') '      <CellData Scalars="pressure" Vectors="velocity">'

      write(unit_id,'(a)') '        <DataArray type="Float64" Name="velocity" NumberOfComponents="3" format="ascii">'
      do n = 1, n_owned
         c = owned_indices(n)
         write(unit_id,'(3(es24.16,1x))') fields%u(1,c), fields%u(2,c), fields%u(3,c)
      end do
      write(unit_id,'(a)') '        </DataArray>'

      write(unit_id,'(a,a,a)') '        <DataArray type="Float64" Name="pressure" format="ascii">'
      do n = 1, n_owned
         c = owned_indices(n)
         write(unit_id,'(es24.16)') fields%p(c)
      end do
      write(unit_id,'(a)') '        </DataArray>'


      write(unit_id,'(a)') '        <DataArray type="Float64" Name="thermo_pressure" format="ascii">'
      do n = 1, n_owned
         write(unit_id,'(es24.16)') params%background_press
      end do
      write(unit_id,'(a)') '        </DataArray>'

      write(unit_id,'(a,a,a)') '        <DataArray type="Float64" Name="divergence" format="ascii">'
      do n = 1, n_owned
         c = owned_indices(n)
         write(unit_id,'(es24.16)') fields%div(c)
      end do
      write(unit_id,'(a)') '        </DataArray>'

      if (params%enable_energy) then
         if (.not. allocated(energy%T) .or. .not. allocated(energy%h) .or. &
             .not. allocated(energy%qrad)) then
            call fatal_error('output', 'energy output requested but energy arrays are not allocated')
         end if

         write(unit_id,'(a)') '        <DataArray type="Float64" Name="temperature" format="ascii">'
         do n = 1, n_owned
            c = owned_indices(n)
            write(unit_id,'(es24.16)') energy%T(c)
         end do
         write(unit_id,'(a)') '        </DataArray>'

         write(unit_id,'(a)') '        <DataArray type="Float64" Name="enthalpy" format="ascii">'
         do n = 1, n_owned
            c = owned_indices(n)
            write(unit_id,'(es24.16)') energy%h(c)
         end do
         write(unit_id,'(a)') '        </DataArray>'

         write(unit_id,'(a)') '        <DataArray type="Float64" Name="qrad" format="ascii">'
         do n = 1, n_owned
            c = owned_indices(n)
            write(unit_id,'(es24.16)') energy%qrad(c)
         end do
         write(unit_id,'(a)') '        </DataArray>'

         if (allocated(energy%cp)) then
            write(unit_id,'(a)') '        <DataArray type="Float64" Name="cp" format="ascii">'
            do n = 1, n_owned
               c = owned_indices(n)
               write(unit_id,'(es24.16)') energy%cp(c)
            end do
            write(unit_id,'(a)') '        </DataArray>'
         end if

         if (allocated(energy%lambda)) then
            write(unit_id,'(a)') '        <DataArray type="Float64" Name="thermal_conductivity" format="ascii">'
            do n = 1, n_owned
               c = owned_indices(n)
               write(unit_id,'(es24.16)') energy%lambda(c)
            end do
            write(unit_id,'(a)') '        </DataArray>'
         end if


         if (allocated(energy%lambda) .and. allocated(energy%cp)) then
            write(unit_id,'(a)') '        <DataArray type="Float64" Name="thermal_diffusivity" format="ascii">'
            do n = 1, n_owned
               c = owned_indices(n)
               write(unit_id,'(es24.16)') energy%lambda(c) / max(params%rho * energy%cp(c), tiny(1.0_rk))
            end do
            write(unit_id,'(a)') '        </DataArray>'
         end if

         if (allocated(energy%rho_thermo)) then
            write(unit_id,'(a)') '        <DataArray type="Float64" Name="rho_thermo" format="ascii">'
            do n = 1, n_owned
               c = owned_indices(n)
               write(unit_id,'(es24.16)') energy%rho_thermo(c)
            end do
            write(unit_id,'(a)') '        </DataArray>'
         end if

      end if

      if (params%enable_species .and. params%nspecies > 0) then

         do k = 1, params%nspecies
            write(unit_id,'(a,a,a)') '        <DataArray type="Float64" Name="Y_', &
               trim(params%species_name(k)), '" format="ascii">'
            do n = 1, n_owned
               c = owned_indices(n)
               write(unit_id,'(es24.16)') species%Y(k,c)
            end do
            write(unit_id,'(a)') '        </DataArray>'
         end do

         write(unit_id,'(a)') '        <DataArray type="Float64" Name="sum_Y" format="ascii">'
         do n = 1, n_owned
            c = owned_indices(n)
            write(unit_id,'(es24.16)') sum(species%Y(:,c))
         end do
         write(unit_id,'(a)') '        </DataArray>'

         if (allocated(transport%diffusivity)) then
            do k = 1, params%nspecies
               write(unit_id,'(a,a,a)') '        <DataArray type="Float64" Name="D_', &
                  trim(params%species_name(k)), '" format="ascii">'
               do n = 1, n_owned
                  c = owned_indices(n)
                  write(unit_id,'(es24.16)') transport%diffusivity(k,c)
               end do
               write(unit_id,'(a)') '        </DataArray>'
            end do
         end if

      end if

      ! Helpful debug scalar so you can color by cell_id in ParaView.
      write(unit_id,'(a)') '        <DataArray type="Int32" Name="cell_id" format="ascii">'
      do n = 1, n_owned
         c = owned_indices(n)
         write(unit_id,'(i0)') c
      end do
      write(unit_id,'(a)') '        </DataArray>'

      write(unit_id,'(a)') '      </CellData>'

      write(unit_id,'(a)') '      <Points>'
      write(unit_id,'(a)') '        <DataArray type="Float64" NumberOfComponents="3" format="ascii">'
      do p = 1, mesh%npoints
         write(unit_id,'(3(es24.16,1x))') mesh%points(1,p), mesh%points(2,p), mesh%points(3,p)
      end do
      write(unit_id,'(a)') '        </DataArray>'
      write(unit_id,'(a)') '      </Points>'

      write(unit_id,'(a)') '      <Cells>'

      write(unit_id,'(a)') '        <DataArray type="Int32" Name="connectivity" format="ascii">'
      do n = 1, n_owned
         c = owned_indices(n)
         write(unit_id,'(8(i0,1x))') (mesh%cells(c)%nodes(m) - 1, m = 1, 8)
      end do
      write(unit_id,'(a)') '        </DataArray>'

      write(unit_id,'(a)') '        <DataArray type="Int32" Name="offsets" format="ascii">'
      do n = 1, n_owned
         write(unit_id,'(i0)') 8 * n
      end do
      write(unit_id,'(a)') '        </DataArray>'

      write(unit_id,'(a)') '        <DataArray type="UInt8" Name="types" format="ascii">'
      do n = 1, n_owned
         write(unit_id,'(i0)') 12
      end do
      write(unit_id,'(a)') '        </DataArray>'

      write(unit_id,'(a)') '      </Cells>'

      write(unit_id,'(a)') '    </Piece>'
      write(unit_id,'(a)') '  </UnstructuredGrid>'
      write(unit_id,'(a)') '</VTKFile>'

      close(unit_id)

      deallocate(owned_indices)

      ! Rank 0 writes the master PVTU file
      if (flow%rank == 0) then
         call write_pvtu_master(params, flow, step)
      end if
   end subroutine write_vtu_unstructured


   !> Writes the master Parallel VTK file (.pvtu) that links the rank pieces.
   subroutine write_pvtu_master(params, flow, step)
      type(case_params_t), intent(in) :: params
      type(flow_mpi_t), intent(in) :: flow
      integer, intent(in) :: step

      integer :: unit_id, r, k
      character(len=path_len + 64) :: filename
      character(len=64) :: piece_name

      write(filename,'(a,"/flow_",i6.6,".pvtu")') trim(params%output_dir), step
      open(newunit=unit_id, file=trim(filename), status='replace', action='write')

      write(unit_id,'(a)') '<?xml version="1.0"?>'
      write(unit_id,'(a)') '<VTKFile type="PUnstructuredGrid" version="0.1" byte_order="LittleEndian">'
      write(unit_id,'(a)') '  <PUnstructuredGrid GhostLevel="0">'

      write(unit_id,'(a)') '    <PPointData>'
      write(unit_id,'(a)') '    </PPointData>'

      write(unit_id,'(a)') '    <PCellData Scalars="pressure" Vectors="velocity">'
      write(unit_id,'(a)') '      <PDataArray type="Float64" Name="velocity" NumberOfComponents="3" format="ascii"/>'
      write(unit_id,'(a)') '      <PDataArray type="Float64" Name="pressure" format="ascii"/>'
      write(unit_id,'(a)') '      <PDataArray type="Float64" Name="thermo_pressure" format="ascii"/>'
      write(unit_id,'(a)') '      <PDataArray type="Float64" Name="divergence" format="ascii"/>'
      if (params%enable_energy) then
         write(unit_id,'(a)') '      <PDataArray type="Float64" Name="temperature" format="ascii"/>'
         write(unit_id,'(a)') '      <PDataArray type="Float64" Name="enthalpy" format="ascii"/>'
         write(unit_id,'(a)') '      <PDataArray type="Float64" Name="qrad" format="ascii"/>'
         write(unit_id,'(a)') '      <PDataArray type="Float64" Name="cp" format="ascii"/>'
         write(unit_id,'(a)') '      <PDataArray type="Float64" Name="thermal_conductivity" format="ascii"/>'
         write(unit_id,'(a)') '      <PDataArray type="Float64" Name="thermal_diffusivity" format="ascii"/>'
         write(unit_id,'(a)') '      <PDataArray type="Float64" Name="rho_thermo" format="ascii"/>'
      end if
      if (params%enable_species .and. params%nspecies > 0) then
         do k = 1, params%nspecies
            write(unit_id,'(a,a,a)') '      <PDataArray type="Float64" Name="Y_', &
               trim(params%species_name(k)), '" format="ascii"/>'
         end do
         write(unit_id,'(a)') '      <PDataArray type="Float64" Name="sum_Y" format="ascii"/>'
         do k = 1, params%nspecies
            write(unit_id,'(a,a,a)') '      <PDataArray type="Float64" Name="D_', &
               trim(params%species_name(k)), '" format="ascii"/>'
         end do
      end if
      write(unit_id,'(a)') '      <PDataArray type="Int32" Name="cell_id" format="ascii"/>'
      write(unit_id,'(a)') '    </PCellData>'

      write(unit_id,'(a)') '    <PPoints>'
      write(unit_id,'(a)') '      <PDataArray type="Float64" NumberOfComponents="3" format="ascii"/>'
      write(unit_id,'(a)') '    </PPoints>'

      do r = 0, flow%nprocs - 1
         write(piece_name,'("flow_",i6.6,"_P",i4.4,".vtu")') step, r
         write(unit_id,'(a,a,a)') '    <Piece Source="', trim(piece_name), '"/>'
      end do

      write(unit_id,'(a)') '  </PUnstructuredGrid>'
      write(unit_id,'(a)') '</VTKFile>'

      close(unit_id)
   end subroutine write_pvtu_master


   !> Writes a PVD collection file to allow ParaView to load time-series data.
   subroutine write_pvd_collection(params, flow, nsteps, output_interval, dt)
      type(case_params_t), intent(in) :: params
      type(flow_mpi_t), intent(in) :: flow
      integer, intent(in) :: nsteps
      integer, intent(in) :: output_interval
      real(rk), intent(in) :: dt

      integer :: unit_id
      integer :: step
      character(len=path_len + 32) :: filename
      character(len=64) :: vtu_name

      if (flow%rank /= 0 .or. .not. params%write_vtu) return

      filename = trim(params%output_dir)//'/flow.pvd'
      open(newunit=unit_id, file=trim(filename), status='replace', action='write')

      write(unit_id,'(a)') '<?xml version="1.0"?>'
      write(unit_id,'(a)') '<VTKFile type="Collection" version="0.1" byte_order="LittleEndian">'
      write(unit_id,'(a)') '  <Collection>'

      call write_dataset_line(unit_id, 0, zero)

      do step = output_interval, nsteps, output_interval
         call write_dataset_line(unit_id, step, real(step, rk) * dt)
      end do

      if (mod(nsteps, output_interval) /= 0) then
         call write_dataset_line(unit_id, nsteps, real(nsteps, rk) * dt)
      end if

      write(unit_id,'(a)') '  </Collection>'
      write(unit_id,'(a)') '</VTKFile>'

      close(unit_id)

   contains

      !> Appends a Dataset entry to the PVD collection file.
      !!
      !! @param unit_id The file unit for the .pvd file.
      !! @param step_id The current simulation time step index.
      !! @param time_value The simulation time at this step.
      subroutine write_dataset_line(unit_id, step_id, time_value)
         integer, intent(in) :: unit_id
         integer, intent(in) :: step_id
         real(rk), intent(in) :: time_value

         character(len=32) :: time_text

         write(vtu_name,'("flow_",i6.6,".pvtu")') step_id
         write(time_text,'(es16.8)') time_value
         time_text = adjustl(time_text)

         write(unit_id,'(a,a,a,a,a)') '    <DataSet timestep="', trim(time_text), &
            '" group="" part="0" file="', trim(vtu_name), '"/>'
      end subroutine write_dataset_line

   end subroutine write_pvd_collection


   !> Internal helper to write a scalar field to a VTU file.
   subroutine write_vtu_cell_scalar(unit_id, name, field)
      integer, intent(in) :: unit_id
      character(len=*), intent(in) :: name
      real(rk), intent(in) :: field(:)

      integer :: i

      write(unit_id,'(a,a,a)') '        <DataArray type="Float64" Name="', trim(name), '" format="ascii">'

      do i = 1, size(field)
         write(unit_id,'(es24.16)') field(i)
      end do

      write(unit_id,'(a)') '        </DataArray>'
   end subroutine write_vtu_cell_scalar


   !> Performs sanity checks on hex connectivity before writing output.
   subroutine validate_hex_connectivity(mesh)
      type(mesh_t), intent(in) :: mesh

      integer :: c, n, node_id

      if (mesh%npoints <= 0) then
         call fatal_error('output', 'mesh has no points')
      end if

      if (mesh%ncells <= 0) then
         call fatal_error('output', 'mesh has no cells')
      end if

      do c = 1, mesh%ncells
         do n = 1, 8
            node_id = mesh%cells(c)%nodes(n)

            if (node_id < 1 .or. node_id > mesh%npoints) then
               call fatal_error('output', 'cell connectivity has node id outside 1..npoints')
            end if
         end do
      end do
   end subroutine validate_hex_connectivity

end module mod_output
