!> Rank-independent restart file support.
!!
!! Restart files store global-cell-indexed arrays in canonical mesh order.  The
!! file is therefore independent of the MPI rank count used to write it: a later
!! run with a different contiguous cell decomposition can read the same global
!! arrays and continue from the stored step/time.
module mod_restart
   use mpi_f08
   use mod_precision, only : rk, zero, path_len, fatal_error, output_unit
   use mod_mesh, only : mesh_t
   use mod_input, only : case_params_t
   use mod_mpi_flow, only : flow_mpi_t, flow_gather_owned_scalar_root, flow_gather_owned_matrix_root
   use mod_flow_fields, only : flow_fields_t
   use mod_species, only : species_fields_t
   use mod_energy, only : energy_fields_t
   use mod_cantera, only : transport_properties_t
   implicit none

   private

   public :: read_restart_file
   public :: write_restart_file

   integer, parameter :: restart_version = 2
   character(len=32), parameter :: restart_magic = 'LMRH_RESTART_V2'

contains

   !> Write a restart file if enabled in &restart_input.
   subroutine write_restart_file(mesh, flow, params, fields, species, energy, transport, step, time, chemistry_accum_dt)
      type(mesh_t), intent(in) :: mesh
      type(flow_mpi_t), intent(inout) :: flow
      type(case_params_t), intent(in) :: params
      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
      real(rk), intent(in) :: time
      real(rk), intent(in), optional :: chemistry_accum_dt

      character(len=path_len) :: restart_dir, filename
      character(len=16) :: step_text
      integer :: ierr

      if (.not. params%write_restart) return
      if (params%restart_interval <= 0) return
      if (mod(step, params%restart_interval) /= 0 .and. step /= params%nsteps) return

      call restart_directory(params, restart_dir)
      if (flow%rank == 0) then
         call execute_command_line('mkdir -p ' // trim(restart_dir))
      end if
      call MPI_Barrier(flow%comm, ierr)
      if (ierr /= MPI_SUCCESS) call fatal_error('restart', 'MPI_Barrier failed before restart write')

      write(step_text, '(i8.8)') step
      filename = trim(restart_dir) // '/restart_' // trim(step_text) // '.rst'
      if (present(chemistry_accum_dt)) then
         call write_restart_to_path(trim(filename), mesh, flow, params, fields, species, energy, transport, step, time, chemistry_accum_dt)
      else
         call write_restart_to_path(trim(filename), mesh, flow, params, fields, species, energy, transport, step, time, zero)
      end if

      if (flow%rank == 0) then
         if (trim(params%terminal_detail) == 'full') then
            write(output_unit,'(a,a)') 'wrote restart: ', trim(filename)
         end if
         call prune_restart_files(params, trim(restart_dir))
      end if
   end subroutine write_restart_file


   !> Read a restart file if restart_from_file is enabled.  All ranks read the
   !! same global file; because arrays are global, this naturally supports a new
   !! MPI rank count without scatter logic.
   subroutine read_restart_file(mesh, flow, params, fields, species, energy, transport, step, time, chemistry_accum_dt)
      type(mesh_t), intent(in) :: mesh
      type(flow_mpi_t), intent(inout) :: flow
      type(case_params_t), intent(inout) :: params
      type(flow_fields_t), intent(inout) :: fields
      type(species_fields_t), intent(inout) :: species
      type(energy_fields_t), intent(inout) :: energy
      type(transport_properties_t), intent(inout) :: transport
      integer, intent(inout) :: step
      real(rk), intent(inout) :: time
      real(rk), intent(inout), optional :: chemistry_accum_dt

      character(len=32) :: magic
      integer :: unit_id, ios, version
      integer :: file_ncells, file_nfaces, file_nspecies, file_step, writer_nprocs
      integer :: has_species, has_energy, has_variable_density, rhs_old_valid_int
      integer :: species_rhs_valid_int, energy_rhs_valid_int
      real(rk) :: file_time, file_dt, file_chemistry_accum_dt
      real(rk), allocatable :: tmp_matrix(:,:)
      real(rk), allocatable :: energy_diag(:)
      logical :: load_species, load_energy

      if (.not. params%restart_from_file) return
      if (len_trim(params%restart_file) == 0) then
         call fatal_error('restart', 'restart_from_file=.true. requires restart_file')
      end if

      open(newunit=unit_id, file=trim(params%restart_file), access='stream', form='unformatted', &
           status='old', action='read', iostat=ios)
      if (ios /= 0) call fatal_error('restart', 'could not open restart_file: '//trim(params%restart_file))

      read(unit_id, iostat=ios) magic
      if (ios /= 0) call fatal_error('restart', 'failed reading restart magic')
      if (trim(magic) /= trim(restart_magic)) then
         call fatal_error('restart', 'restart file magic mismatch; expected '//trim(restart_magic))
      end if

      read(unit_id) version
      if (version /= restart_version) call fatal_error('restart', 'unsupported restart version')
      read(unit_id) file_ncells, file_nfaces, file_nspecies, file_step, writer_nprocs
      read(unit_id) file_time, file_dt, file_chemistry_accum_dt
      read(unit_id) has_species, has_energy, has_variable_density, rhs_old_valid_int, &
                    species_rhs_valid_int, energy_rhs_valid_int

      if (file_ncells /= mesh%ncells) call fatal_error('restart', 'restart ncells does not match current mesh')
      if (file_nfaces /= mesh%nfaces) call fatal_error('restart', 'restart nfaces does not match current mesh')
      if (params%enable_species .and. file_nspecies > params%nspecies) then
         call fatal_error('restart', 'restart nspecies exceeds current case nspecies; cannot load restart')
      end if
      if (params%enable_species .and. flow%rank == 0 .and. file_nspecies /= params%nspecies) then
         write(output_unit,'(a,i0,a,i0,a)') &
            'restart: WARNING: file nspecies=', file_nspecies, &
            ' < current nspecies=', params%nspecies, &
            '; remaining species initialized to zero'
      end if
      if (has_energy == 1 .and. .not. params%enable_energy) then
         call fatal_error('restart', 'restart contains energy state but current case has enable_energy=.false.')
      end if
      if (has_species == 1 .and. .not. params%enable_species) then
         call fatal_error('restart', 'restart contains species state but current case has enable_species=.false.')
      end if
      if (has_variable_density == 1 .and. .not. params%enable_variable_density) then
         call fatal_error('restart', 'restart was written from variable-density mode but current case is not variable-density')
      end if

      read(unit_id) fields%u
      read(unit_id) fields%u_old
      read(unit_id) fields%p
      read(unit_id) fields%phi
      read(unit_id) fields%divergence_source
      read(unit_id) fields%projection_rho
      read(unit_id) fields%projection_divergence_source
      read(unit_id) fields%rhs_old
      read(unit_id) fields%face_flux
      read(unit_id) fields%mass_flux

      read(unit_id) transport%rho
      read(unit_id) transport%rho_old
      read(unit_id) transport%mu
      read(unit_id) transport%nu
      read(unit_id) transport%lambda

      if (file_nspecies > 0) then
         allocate(tmp_matrix(file_nspecies, mesh%ncells))
         read(unit_id) tmp_matrix
         if (allocated(transport%diffusivity)) then
            transport%diffusivity(1:min(file_nspecies, params%nspecies), :) = &
               tmp_matrix(1:min(file_nspecies, params%nspecies), :)
         end if
         deallocate(tmp_matrix)
      end if

      load_species = has_species == 1 .and. file_nspecies > 0
      if (load_species) then
         allocate(tmp_matrix(file_nspecies, mesh%ncells))
         read(unit_id) tmp_matrix
         species%Y(1:min(file_nspecies, params%nspecies), :) = &
            tmp_matrix(1:min(file_nspecies, params%nspecies), :)
         if (allocated(species%Y_old)) species%Y_old = species%Y
         read(unit_id) tmp_matrix
         if (allocated(species%rhs_old)) then
            species%rhs_old(1:min(file_nspecies, params%nspecies), :) = &
               tmp_matrix(1:min(file_nspecies, params%nspecies), :)
         end if
         species%rhs_old_valid = species_rhs_valid_int /= 0
         deallocate(tmp_matrix)
      else
         species%rhs_old_valid = .false.
      end if

      load_energy = has_energy == 1
      if (load_energy) then
         read(unit_id) energy%T
         read(unit_id) energy%T_old
         read(unit_id) energy%h
         read(unit_id) energy%h_old
         read(unit_id) energy%rhs_old
         energy%rhs_old_valid = energy_rhs_valid_int /= 0
         read(unit_id) energy%qrad
         read(unit_id) energy%cp
         read(unit_id) energy%lambda
         read(unit_id) energy%rho_thermo
         allocate(energy_diag(12))
         read(unit_id) energy_diag
         call restore_energy_diagnostics(energy, energy_diag)
         deallocate(energy_diag)
      end if

      close(unit_id)

      fields%rhs_old_valid = rhs_old_valid_int /= 0
      step = file_step
      time = file_time
      if (params%restart_use_file_dt) params%dt = file_dt
      if (params%restart_nsteps_are_additional) params%nsteps = file_step + params%nsteps
      if (present(chemistry_accum_dt)) chemistry_accum_dt = file_chemistry_accum_dt

      if (flow%rank == 0) then
         write(output_unit,'(a,a,a,i0,a,es13.5,a,es13.5,a,es13.5,a,i0,a,i0,a,l1,a,l1)') &
            'read restart: ', trim(params%restart_file), ' step=', step, &
            ' time=', time, ' file_dt=', file_dt, ' run_dt=', params%dt, &
            ' nsteps=', params%nsteps, ' writer_ranks=', writer_nprocs, &
            ' use_file_dt=', params%restart_use_file_dt, &
            ' nsteps_additional=', params%restart_nsteps_are_additional
      end if
   end subroutine read_restart_file


   subroutine write_restart_to_path(filename, mesh, flow, params, fields, species, energy, transport, step, time, chemistry_accum_dt)
      character(len=*), intent(in) :: filename
      type(mesh_t), intent(in) :: mesh
      type(flow_mpi_t), intent(inout) :: flow
      type(case_params_t), intent(in) :: params
      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
      real(rk), intent(in) :: time
      real(rk), intent(in) :: chemistry_accum_dt

      integer :: unit_id, ios
      integer :: has_species, has_energy, has_variable_density, rhs_old_valid_int
      integer :: species_rhs_valid_int, energy_rhs_valid_int
      real(rk), allocatable :: root_scalar(:)
      real(rk), allocatable :: root_matrix(:,:)
      real(rk), allocatable :: root_u(:,:)
      real(rk), allocatable :: face_global(:)
      real(rk) :: energy_diag(12)

      allocate(root_scalar(mesh%ncells))
      allocate(root_u(3, mesh%ncells))
      allocate(face_global(mesh%nfaces))
      if (params%nspecies > 0) allocate(root_matrix(params%nspecies, mesh%ncells))

      has_species = merge(1, 0, params%enable_species .and. species%nspecies > 0 .and. allocated(species%Y))
      has_energy = merge(1, 0, params%enable_energy .and. energy%initialized)
      has_variable_density = merge(1, 0, params%enable_variable_density)
      rhs_old_valid_int = merge(1, 0, fields%rhs_old_valid)
      species_rhs_valid_int = merge(1, 0, species%rhs_old_valid)
      energy_rhs_valid_int = merge(1, 0, energy%rhs_old_valid)

      if (flow%rank == 0) then
         open(newunit=unit_id, file=trim(filename), access='stream', form='unformatted', &
              status='replace', action='write', iostat=ios)
         if (ios /= 0) call fatal_error('restart', 'could not open restart output: '//trim(filename))

         write(unit_id) restart_magic
         write(unit_id) restart_version
         write(unit_id) mesh%ncells, mesh%nfaces, params%nspecies, step, flow%nprocs
         write(unit_id) time, params%dt, chemistry_accum_dt
         write(unit_id) has_species, has_energy, has_variable_density, rhs_old_valid_int, &
                        species_rhs_valid_int, energy_rhs_valid_int
      end if

      call flow_gather_owned_matrix_root(flow, fields%u, root_u)
      if (flow%rank == 0) write(unit_id) root_u
      call flow_gather_owned_matrix_root(flow, fields%u_old, root_u)
      if (flow%rank == 0) write(unit_id) root_u
      call flow_gather_owned_scalar_root(flow, fields%p, root_scalar)
      if (flow%rank == 0) write(unit_id) root_scalar
      call flow_gather_owned_scalar_root(flow, fields%phi, root_scalar)
      if (flow%rank == 0) write(unit_id) root_scalar
      call flow_gather_owned_scalar_root(flow, fields%divergence_source, root_scalar)
      if (flow%rank == 0) write(unit_id) root_scalar
      call flow_gather_owned_scalar_root(flow, fields%projection_rho, root_scalar)
      if (flow%rank == 0) write(unit_id) root_scalar
      call flow_gather_owned_scalar_root(flow, fields%projection_divergence_source, root_scalar)
      if (flow%rank == 0) write(unit_id) root_scalar
      call flow_gather_owned_matrix_root(flow, fields%rhs_old, root_u)
      if (flow%rank == 0) write(unit_id) root_u

      call gather_owned_face_scalar(flow, mesh%nfaces, fields%face_flux, face_global)
      if (flow%rank == 0) write(unit_id) face_global
      call gather_owned_face_scalar(flow, mesh%nfaces, fields%mass_flux, face_global)
      if (flow%rank == 0) write(unit_id) face_global

      call flow_gather_owned_scalar_root(flow, transport%rho, root_scalar)
      if (flow%rank == 0) write(unit_id) root_scalar
      call flow_gather_owned_scalar_root(flow, transport%rho_old, root_scalar)
      if (flow%rank == 0) write(unit_id) root_scalar
      call flow_gather_owned_scalar_root(flow, transport%mu, root_scalar)
      if (flow%rank == 0) write(unit_id) root_scalar
      call flow_gather_owned_scalar_root(flow, transport%nu, root_scalar)
      if (flow%rank == 0) write(unit_id) root_scalar
      call flow_gather_owned_scalar_root(flow, transport%lambda, root_scalar)
      if (flow%rank == 0) write(unit_id) root_scalar

      if (params%nspecies > 0) then
         if (allocated(transport%diffusivity)) then
            call flow_gather_owned_matrix_root(flow, transport%diffusivity, root_matrix)
         else
            root_matrix = zero
         end if
         if (flow%rank == 0) write(unit_id) root_matrix
      end if

      if (has_species == 1) then
         call flow_gather_owned_matrix_root(flow, species%Y, root_matrix)
         if (flow%rank == 0) write(unit_id) root_matrix
         call flow_gather_owned_matrix_root(flow, species%rhs_old, root_matrix)
         if (flow%rank == 0) write(unit_id) root_matrix
      end if

      if (has_energy == 1) then
         call flow_gather_owned_scalar_root(flow, energy%T, root_scalar)
         if (flow%rank == 0) write(unit_id) root_scalar
         call flow_gather_owned_scalar_root(flow, energy%T_old, root_scalar)
         if (flow%rank == 0) write(unit_id) root_scalar
         call flow_gather_owned_scalar_root(flow, energy%h, root_scalar)
         if (flow%rank == 0) write(unit_id) root_scalar
         call flow_gather_owned_scalar_root(flow, energy%h_old, root_scalar)
         if (flow%rank == 0) write(unit_id) root_scalar
         call flow_gather_owned_scalar_root(flow, energy%rhs_old, root_scalar)
         if (flow%rank == 0) write(unit_id) root_scalar
         call flow_gather_owned_scalar_root(flow, energy%qrad, root_scalar)
         if (flow%rank == 0) write(unit_id) root_scalar
         call flow_gather_owned_scalar_root(flow, energy%cp, root_scalar)
         if (flow%rank == 0) write(unit_id) root_scalar
         call flow_gather_owned_scalar_root(flow, energy%lambda, root_scalar)
         if (flow%rank == 0) write(unit_id) root_scalar
         call flow_gather_owned_scalar_root(flow, energy%rho_thermo, root_scalar)
         if (flow%rank == 0) write(unit_id) root_scalar

         call pack_energy_diagnostics(energy, energy_diag)
         if (flow%rank == 0) write(unit_id) energy_diag
      end if

      if (flow%rank == 0) close(unit_id)

      if (allocated(root_scalar)) deallocate(root_scalar)
      if (allocated(root_u)) deallocate(root_u)
      if (allocated(root_matrix)) deallocate(root_matrix)
      if (allocated(face_global)) deallocate(face_global)
   end subroutine write_restart_to_path


   subroutine gather_owned_face_scalar(flow, nfaces, face_field, global_field)
      type(flow_mpi_t), intent(in) :: flow
      integer, intent(in) :: nfaces
      real(rk), intent(in) :: face_field(:)
      real(rk), intent(out) :: global_field(:)

      real(rk), allocatable :: local_field(:)
      integer :: i, f, ierr

      if (size(face_field) < nfaces .or. size(global_field) < nfaces) then
         call fatal_error('restart', 'face field size mismatch during restart write')
      end if

      allocate(local_field(nfaces))
      local_field = zero
      if (allocated(flow%owned_faces)) then
         do i = 1, size(flow%owned_faces)
            f = flow%owned_faces(i)
            if (f >= 1 .and. f <= nfaces) local_field(f) = face_field(f)
         end do
      end if

      call MPI_Allreduce(local_field, global_field, nfaces, MPI_DOUBLE_PRECISION, MPI_SUM, flow%comm, ierr)
      if (ierr /= MPI_SUCCESS) call fatal_error('restart', 'MPI_Allreduce failed gathering face field')
      deallocate(local_field)
   end subroutine gather_owned_face_scalar


   subroutine restart_directory(params, restart_dir)
      type(case_params_t), intent(in) :: params
      character(len=path_len), intent(out) :: restart_dir
      character(len=path_len) :: requested

      requested = trim(params%restart_output_dir)
      if (len_trim(requested) == 0) requested = 'restart'

      if (requested(1:1) == '/') then
         restart_dir = requested
      else if (index(trim(requested), trim(params%output_dir)//'/') == 1 .or. trim(requested) == trim(params%output_dir)) then
         restart_dir = requested
      else
         restart_dir = trim(params%output_dir) // '/' // trim(requested)
      end if
   end subroutine restart_directory


   subroutine prune_restart_files(params, restart_dir)
      type(case_params_t), intent(in) :: params
      character(len=*), intent(in) :: restart_dir

      integer :: exitstat, first_to_delete
      character(len=32) :: first_to_delete_text
      character(len=3*path_len + 256) :: command

      if (params%restart_keep <= 0) return

      ! Keep the most recently written restart files in this restart directory.
      ! This mirrors the existing POSIX-shell style used for output directory
      ! creation; paths with spaces are not supported elsewhere in the solver.
      first_to_delete = params%restart_keep + 1
      write(first_to_delete_text, '(i0)') first_to_delete
      command = 'if ls ' // trim(restart_dir) // '/restart_*.rst >/dev/null 2>&1; then ' // &
                'ls -1t ' // trim(restart_dir) // '/restart_*.rst | tail -n +' // &
                trim(adjustl(first_to_delete_text)) // ' | xargs -r rm -f; fi'
      call execute_command_line(trim(command), exitstat=exitstat)
      if (exitstat /= 0) call fatal_error('restart', 'failed pruning old restart files in: '//trim(restart_dir))
   end subroutine prune_restart_files


   subroutine pack_energy_diagnostics(energy, diag)
      type(energy_fields_t), intent(in) :: energy
      real(rk), intent(out) :: diag(12)

      diag = zero
      diag(1) = energy%last_conductive_boundary_flux_out
      diag(2) = real(energy%last_conductive_boundary_flux_available, rk)
      diag(3) = energy%cumulative_boundary_rho_h_advective_flux_out
      diag(4) = energy%cumulative_boundary_rho_h_conductive_flux_out
      diag(5) = energy%cumulative_qrad_integral
      diag(6) = energy%cumulative_rho_species_hdiff_integral
      diag(7) = real(energy%cumulative_energy_budget_available, rk)
      diag(8) = energy%last_energy_update_delta_rate_integral
      diag(9) = energy%last_energy_update_rhs_integral
      diag(10) = energy%last_energy_update_balance_defect
      diag(11) = energy%relative_last_energy_update_balance_defect
      diag(12) = energy%last_operator_consistent_rho_h_integral
   end subroutine pack_energy_diagnostics


   subroutine restore_energy_diagnostics(energy, diag)
      type(energy_fields_t), intent(inout) :: energy
      real(rk), intent(in) :: diag(12)

      energy%last_conductive_boundary_flux_out = diag(1)
      energy%last_conductive_boundary_flux_available = nint(diag(2))
      energy%cumulative_boundary_rho_h_advective_flux_out = diag(3)
      energy%cumulative_boundary_rho_h_conductive_flux_out = diag(4)
      energy%cumulative_qrad_integral = diag(5)
      energy%cumulative_rho_species_hdiff_integral = diag(6)
      energy%cumulative_energy_budget_available = nint(diag(7))
      energy%last_energy_update_delta_rate_integral = diag(8)
      energy%last_energy_update_rhs_integral = diag(9)
      energy%last_energy_update_balance_defect = diag(10)
      energy%relative_last_energy_update_balance_defect = diag(11)
      energy%last_operator_consistent_rho_h_integral = diag(12)
   end subroutine restore_energy_diagnostics

end module mod_restart
