main Program

program main
  use, intrinsic :: iso_fortran_env, only: error_unit
  use bem_kinds, only: dp, i32
  use bem_version, only: beach_build_info, beach_version
  use bem_types, only: sim_stats, mesh_type, injection_state
  use bem_mpi, only: mpi_context, mpi_initialize, mpi_shutdown, mpi_is_root, mpi_world_barrier, mpi_world_size
  use bem_performance_profile, only: perf_configure_from_env, perf_set_output_context, perf_region_begin, &
                                     perf_region_end, perf_write_outputs, perf_region_program_total, perf_region_load_or_init, &
                                     perf_region_history_open, perf_region_write_results, perf_region_write_checkpoint
  use bem_simulator, only: run_absorption_insulator
  use bem_restart, only: load_restart_checkpoint, write_rng_state_file, write_macro_residuals_file
  use bem_periodic_checkpoint, only: resolve_latest_checkpoint_dir
  use bem_checkpoint_contract, only: publish_checkpoint_manifest
  use bem_output_writer, only: open_history_writer, open_potential_history_writer, open_top_reference_history_writer, &
                               open_matching_plane_history_writer, print_run_summary, write_result_files, ensure_output_dir
  use bem_app_config, only: app_config, default_app_config, load_app_config, build_mesh_from_config, &
                            seed_particles_from_config
  use bem_mesh, only: prepare_periodic2_collision_mesh
  use bem_charge_ledger, only: charge_ledger_type
  use bem_electrostatic_snapshot, only: electrostatic_diagnostics_type
  use bem_matching_plane_response_provider, only: matching_plane_response_provider_type, &
                                                  matching_plane_provider_ok
  implicit none

  type(mesh_type) :: mesh
  type(app_config) :: app
  type(sim_stats) :: stats
  type(sim_stats) :: initial_stats
  type(injection_state) :: inject_state
  type(charge_ledger_type) :: charge_ledger
  type(electrostatic_diagnostics_type) :: electrostatic_diagnostics
  type(mpi_context) :: mpi
  integer :: history_unit
  integer :: potential_history_unit
  integer :: top_reference_history_unit
  integer :: matching_plane_history_unit
  logical :: history_opened, potential_history_opened, top_reference_history_opened, matching_plane_history_opened, resumed
  real(dp) :: perf_t0, perf_program_t0
  real(dp), allocatable :: mesh_potential_v(:)

  call handle_early_cli()

  call perf_configure_from_env()
  call perf_region_begin(perf_region_program_total, perf_program_t0)
  call mpi_initialize(mpi)
  call perf_region_begin(perf_region_load_or_init, perf_t0)
  call load_or_init_run_state(app, mesh, initial_stats, inject_state, charge_ledger, resumed, mpi)
  call perf_region_end(perf_region_load_or_init, perf_t0)
  call perf_set_output_context(trim(app%output_dir), app%write_output)
  if (mpi_is_root(mpi)) then
    call perf_region_begin(perf_region_history_open, perf_t0)
    call open_history_writer(app, resumed, history_opened, history_unit)
    call open_potential_history_writer(app, resumed, potential_history_opened, potential_history_unit)
    call open_top_reference_history_writer( &
      app, resumed, top_reference_history_opened, top_reference_history_unit &
      )
    call open_matching_plane_history_writer( &
      app, resumed, matching_plane_history_opened, matching_plane_history_unit &
      )
    call perf_region_end(perf_region_history_open, perf_t0)
  else
    history_opened = .false.
    history_unit = -1
    potential_history_opened = .false.
    potential_history_unit = -1
    top_reference_history_opened = .false.
    top_reference_history_unit = -1
    matching_plane_history_opened = .false.
    matching_plane_history_unit = -1
  end if

  if (history_opened) then
    if (app%write_mesh_potential) then
      if (potential_history_opened) then
        call run_absorption_insulator( &
          mesh, app, stats, history_unit=history_unit, history_stride=app%history_stride, initial_stats=initial_stats, &
          inject_state=inject_state, mpi=mpi, mesh_potential_v=mesh_potential_v, &
          potential_history_unit=potential_history_unit, &
          top_reference_history_unit=top_reference_history_unit, charge_ledger=charge_ledger, &
          electrostatic_diagnostics=electrostatic_diagnostics, &
          matching_plane_history_unit=matching_plane_history_unit &
          )
      else
        call run_absorption_insulator( &
          mesh, app, stats, history_unit=history_unit, history_stride=app%history_stride, initial_stats=initial_stats, &
          inject_state=inject_state, mpi=mpi, mesh_potential_v=mesh_potential_v, charge_ledger=charge_ledger, &
          electrostatic_diagnostics=electrostatic_diagnostics, &
          matching_plane_history_unit=matching_plane_history_unit &
          )
      end if
    else
      if (potential_history_opened) then
        call run_absorption_insulator( &
          mesh, app, stats, history_unit=history_unit, history_stride=app%history_stride, initial_stats=initial_stats, &
          inject_state=inject_state, mpi=mpi, &
          potential_history_unit=potential_history_unit, &
          top_reference_history_unit=top_reference_history_unit, charge_ledger=charge_ledger, &
          electrostatic_diagnostics=electrostatic_diagnostics, &
          matching_plane_history_unit=matching_plane_history_unit &
          )
      else
        call run_absorption_insulator( &
          mesh, app, stats, history_unit=history_unit, history_stride=app%history_stride, initial_stats=initial_stats, &
          inject_state=inject_state, mpi=mpi, charge_ledger=charge_ledger, &
          electrostatic_diagnostics=electrostatic_diagnostics, &
          matching_plane_history_unit=matching_plane_history_unit &
          )
      end if
    end if
    close (history_unit)
  else
    if (app%write_mesh_potential) then
      if (potential_history_opened) then
        call run_absorption_insulator( &
          mesh, app, stats, initial_stats=initial_stats, inject_state=inject_state, mpi=mpi, &
          mesh_potential_v=mesh_potential_v, &
          potential_history_unit=potential_history_unit, &
          top_reference_history_unit=top_reference_history_unit, charge_ledger=charge_ledger, &
          electrostatic_diagnostics=electrostatic_diagnostics, &
          matching_plane_history_unit=matching_plane_history_unit &
          )
      else
        call run_absorption_insulator( &
          mesh, app, stats, initial_stats=initial_stats, inject_state=inject_state, mpi=mpi, &
          mesh_potential_v=mesh_potential_v, charge_ledger=charge_ledger, &
          electrostatic_diagnostics=electrostatic_diagnostics, &
          matching_plane_history_unit=matching_plane_history_unit &
          )
      end if
    else
      if (potential_history_opened) then
        call run_absorption_insulator( &
          mesh, app, stats, initial_stats=initial_stats, inject_state=inject_state, mpi=mpi, &
          potential_history_unit=potential_history_unit, &
          top_reference_history_unit=top_reference_history_unit, charge_ledger=charge_ledger, &
          electrostatic_diagnostics=electrostatic_diagnostics, &
          matching_plane_history_unit=matching_plane_history_unit &
          )
      else
        call run_absorption_insulator( &
          mesh, app, stats, initial_stats=initial_stats, inject_state=inject_state, mpi=mpi, &
          charge_ledger=charge_ledger, electrostatic_diagnostics=electrostatic_diagnostics, &
          matching_plane_history_unit=matching_plane_history_unit &
          )
      end if
    end if
  end if
  if (potential_history_opened) close (potential_history_unit)
  if (top_reference_history_opened) close (top_reference_history_unit)
  if (matching_plane_history_opened) close (matching_plane_history_unit)

  if (mpi_is_root(mpi)) call print_run_summary(mesh, stats)

  if (app%write_output) then
    call ensure_output_dir(app%output_dir)
    if (mpi_is_root(mpi)) then
      call perf_region_begin(perf_region_write_results, perf_t0)
      if (allocated(mesh_potential_v)) then
        call write_result_files( &
          trim(app%output_dir), mesh, stats, app, mpi_world_size=mpi_world_size(mpi), mesh_potential_v=mesh_potential_v, &
          charge_ledger=charge_ledger, electrostatic_diagnostics=electrostatic_diagnostics &
          )
      else
        call write_result_files( &
          trim(app%output_dir), mesh, stats, app, mpi_world_size=mpi_world_size(mpi), charge_ledger=charge_ledger, &
          electrostatic_diagnostics=electrostatic_diagnostics &
          )
      end if
      call perf_region_end(perf_region_write_results, perf_t0)
    end if
    ! rootが旧完了manifestを無効化してから、全rankでrestart stateの公開へ進む。
    call mpi_world_barrier(mpi)
    call perf_region_begin(perf_region_write_checkpoint, perf_t0)
    call write_rng_state_file(trim(app%output_dir), mpi=mpi)
    call write_macro_residuals_file(trim(app%output_dir), inject_state, mpi=mpi)
    call mpi_world_barrier(mpi)
    if (mpi_is_root(mpi)) then
      call publish_checkpoint_manifest( &
        trim(app%output_dir), stats%batches, mpi_world_size(mpi), allocated(inject_state%macro_residual), .true. &
        )
    end if
    call mpi_world_barrier(mpi)
    call perf_region_end(perf_region_write_checkpoint, perf_t0)
    if (mpi_is_root(mpi)) print '(a,a)', 'results written to ', trim(app%output_dir)
  end if
  call perf_region_end(perf_region_program_total, perf_program_t0)
  call perf_write_outputs(mpi)
  call mpi_shutdown(mpi)

contains

  !> MPI 初期化や計算状態の構築前に完結する CLI option を処理する。
  subroutine handle_early_cli()
    character(len=256) :: arg
    character(len=:), allocatable :: config_path
    integer :: path_length

    if (command_argument_count() < 1) return
    call get_command_argument(1, arg)
    select case (trim(arg))
    case ('--version', '-V')
      print '(a)', trim(beach_version)
      stop
    case ('--build-info')
      print '(a)', beach_build_info
      stop
    case ('--check-config')
      if (command_argument_count() /= 2) error stop 'usage: beach --check-config beach.toml'
      call get_command_argument(2, length=path_length)
      allocate (character(len=path_length) :: config_path)
      call get_command_argument(2, config_path)
      call default_app_config(app)
      call load_app_config(config_path, app)
      print '(a,a)', 'config=', config_path
      print '(a)', 'checks=toml,semantic'
      print '(a)', 'status=ok'
      stop
    case ('--help', '-h')
      print '(a)', 'usage: beach [beach.toml]'
      print '(a)', '       beach --check-config beach.toml'
      print '(a)', '       beach --version'
      print '(a)', '       beach --build-info'
      stop
    end select
  end subroutine handle_early_cli

  !> 設定読込・メッシュ構築・再開判定・乱数初期化をまとめて行う。
  !! @param[out] app 読み込み・既定値適用後のアプリ設定。
  !! @param[out] mesh 構築した三角形メッシュ。
  !! @param[out] initial_stats 再開時に引き継ぐ初期統計（新規実行時はゼロ）。
  !! @param[out] inject_state 種別ごとの注入残差状態。
  !! @param[out] resumed チェックポイントから再開した場合に `.true.`。
  subroutine load_or_init_run_state(app, mesh, initial_stats, inject_state, charge_ledger, resumed, mpi)
    type(app_config), intent(out) :: app
    type(mesh_type), intent(out) :: mesh
    type(sim_stats), intent(out) :: initial_stats
    type(injection_state), intent(out) :: inject_state
    type(charge_ledger_type), intent(out) :: charge_ledger
    logical, intent(out) :: resumed
    type(mpi_context), intent(in) :: mpi
    character(len=256) :: cfg_path
    character(len=256) :: restart_dir
    character(len=512) :: matching_response_message
    integer(i32) :: matching_response_status
    logical :: has_config
    type(matching_plane_response_provider_type) :: matching_response_provider

    call default_app_config(app)
    call resolve_config_path(cfg_path, has_config)
    if (has_config) then
      call load_app_config(trim(cfg_path), app)
    end if

    call matching_response_provider%initialize( &
      app, mpi, matching_response_status, matching_response_message &
      )
    if (matching_response_status /= matching_plane_provider_ok) then
      error stop 'matching-plane response preflight failed: '//trim(matching_response_message)
    end if

    call build_mesh_from_config(app, mesh)
    call prepare_periodic2_collision_mesh(mesh, app%sim)
    call initialize_injection_state(inject_state, app%n_particle_species)
    initial_stats = sim_stats()
    resumed = .false.
    if (app%resume_output) then
      restart_dir = app%output_dir
      if (len_trim(app%output_restart_from) > 0) restart_dir = app%output_restart_from
      call resolve_latest_checkpoint_dir(trim(restart_dir), restart_dir)
      call load_restart_checkpoint( &
        trim(restart_dir), mesh, initial_stats, resumed, inject_state, &
        mpi=mpi, require_checkpoint=.true., app=app, charge_ledger=charge_ledger &
        )
    end if

    if (resumed) then
      if (mpi_is_root(mpi)) then
        print '(a,a)', 'resuming_from_dir=', trim(restart_dir)
        print '(a,i0)', 'resuming_from_batches=', initial_stats%batches
        print '(a,i0)', 'resuming_from_processed_particles=', initial_stats%processed_particles
      end if
    else
      call seed_particles_from_config(app, mpi=mpi)
    end if
  end subroutine load_or_init_run_state

  !> 実行時設定ファイルの読み込みパスを決定する。
  !! @param[out] path 読み込む設定ファイルパス（未検出時は空文字）。
  !! @param[out] found 設定ファイルが解決できた場合に `.true.`。
  subroutine resolve_config_path(path, found)
    character(len=*), intent(out) :: path
    logical, intent(out) :: found
    logical :: has_primary
    character(len=*), parameter :: primary_config = 'beach.toml'

    path = ''
    found = .false.
    if (command_argument_count() >= 1) then
      call get_command_argument(1, path)
      if (len_trim(path) == 0) error stop 'Config path argument is empty.'
      found = .true.
      return
    end if

    inquire (file=primary_config, exist=has_primary)
    if (has_primary) then
      path = primary_config
      found = .true.
    end if
  end subroutine resolve_config_path

  !> 種数に合わせて注入状態をゼロ初期化する。
  !! @param[out] state 初期化する注入状態。
  !! @param[in] n_species 粒子種数（source残差と6面boundary流入残差の種方向配列長）。
  subroutine initialize_injection_state(state, n_species)
    type(injection_state), intent(out) :: state
    integer, intent(in) :: n_species

    allocate (state%macro_residual(n_species))
    allocate (state%boundary_macro_residual(6, n_species))
    state%macro_residual = 0.0d0
    state%boundary_macro_residual = 0.0d0
  end subroutine initialize_injection_state

end program main