PROGRAM pem
!-----------------------------------------------------------------------
! NAME
!     pem
!
! DESCRIPTION
!     Main entry point for the Planetary Evolution Model (PEM).
!
! AUTHORS & DATE
!     R. Vandemeulebrouck, 22/07/2022 with r2778 & r2779
!     L. Lange, 22/07/2022
!     JB Clement, 2023-2025
!     C. Metz, 04/2026
!
! NOTES
!     Ownership criterion for declarations:
!       - Declare in module "planet" variables that are persistent climate
!         state or persistent references used across multiple time steps.
!       - Declare in program "pem.F90" transient workflow/control varaibles,
!         one-shot initialization buffers and per-iteration working buffers.
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
! LMDZ.COMMON modules
use job,             only: timelimit, antetime, timewall
use program_options, only: parse_args
! PEM modules
use allocation,      only: ini_allocation, end_allocation
use atmosphere,      only: ps_PCM, evolve_pressure, CO2cond_ps_PCM
use backup,          only: save_clim_state, backup_rate
use climate_init,    only: read_start, read_startfi, read_startevo
use config,          only: read_rundef, read_display_config
use display,         only: print_ini, print_end, print_msg, LVL_NFO, LVL_WRN
use evolution,       only: n_yr_run, n_yr_sim, ntot_yr_sim, nmax_yr_run, dt, idt, r_plnt2earth_yr, pem_ini_date
use geometry,        only: ngrid, nslope, nsoil_PCM, nsoil, cell_area, total_surface
use maths,           only: pi
use numerics,        only: dp, qp, di, li, eps, eps_qp
use orbit,           only: evo_orbit, read_orbitpm, compute_maxyr_orbit, update_orbit
use output,          only: write_diagevo, dim_ngrid, dim_nsoil
use planet
use physics,         only: g
use slopes,          only: subslope_dist, def_slope_mean
use soil,            only: do_soil, set_soil, TI
use soil_temp,       only: tsoil_PCM, shift_tsoil2surf, evolve_soil_temp
use stopping_crit,   only: stopFlags, stopping_crit_pressure, stopping_crit_h2oice, stopping_crit_co2ice
use surface,         only: emissivity_PCM
use surf_ice,        only: evolve_co2ice, evolve_h2oice, balance_h2o_fluxes
use surf_temp,       only: tsurf_PCM, adapt_tsurf2disappearedice
use tendencies,      only: compute_tendice, evolve_tend_co2ice
use tracers,         only: adapt_tracers2pressure
use utility,         only: real2str
use workflow_status, only: i_pem_run, read_workflow_status, update_workflow_status
use xios_data,       only: load_xios_data

! DECLARATION
! -----------
implicit none

! LOCAL VARIABLES
! ---------------
! Utility-related:
integer(li)  :: cr     ! Number of clock ticks per second (count rate)
integer(li)  :: c1, c2 ! Counts of processor clock
character(8) :: num    ! Slope suffix to ouput variables
integer(di)  :: i, islope
! Ice-related:
real(dp), dimension(:,:,:), allocatable :: minPCM_h2ofrost  ! Minimum of H2O frost over the last PCM year [kg/m2]
real(dp), dimension(:,:,:), allocatable :: minPCM_co2frost  ! Minimum of CO2 frost over the last PCM year [kg/m2]
! Surface-related:
real(dp), dimension(:,:), allocatable :: tsurf_avg_yr1 ! Average surface temperature of the second to last PCM run [K]
real(dp), dimension(:,:), allocatable :: zshift_surf   ! Elevation shift for the surface [m]
! Evolution-related:
real(dp)        :: nmax_yr_runorb ! Maximum number of years for the run due to orbital parameters
type(stopFlags) :: stopcrit       ! Stopping criteria
! Balance-related
real(qp) :: totmass_co2ice, totmass_atmco2         ! Current total CO2 masses (surface ice|atmospheric)
real(qp) :: totmass_co2ice_ini, totmass_atmco2_ini ! Initial total CO2 masses (surface ice|atmospheric)
real(qp) :: totmass_ini                            ! Initial total CO2 mass [kg]

! CODE
! ----
! Pre-processing step
!~~~~~~~~~~~~~~~~~~~~
! Elapsed time with system clock
call system_clock(count_rate = cr)
call system_clock(c1)

! Parse command-line options
call parse_args()

! Read display configuration
call read_display_config()

! Initialization
! ~~~~~~~~~~~~~~
! Header
call print_ini()

! Allocate module arrays
call ini_allocation()

! Read the duration information of the workflow
call read_workflow_status()

! Read the PEM parameters
call read_rundef()

! Read the orbital parameters
call read_orbitpm(n_yr_sim)

! Read the "start.nc"
call read_start()

! Read the "startfi.nc"
call read_startfi()

! Read the PCM data given by XIOS
call allocate_xios_state()
allocate(tsurf_avg_yr1(ngrid,nslope))
allocate(minPCM_h2ofrost(ngrid,nslope,2))
allocate(minPCM_co2frost(ngrid,nslope,2))

call load_xios_data(ps_avg,ps_ts,tsurf_avg,tsurf_avg_yr1,tsoil_avg,tsoil_ts,q_h2o_ts,q_co2_ts,minPCM_h2ofrost,minPCM_co2frost)

! Initiate soil settings and TI
if (do_soil) call set_soil(TI)

! Compute the deviation from the average
call allocate_deviation_state()
ps_dev(:) = ps_PCM(:) - ps_avg(:)
tsurf_dev(:,:) = tsurf_PCM(:,:) - tsurf_avg(:,:)
tsoil_dev(:,:,:) = tsoil_PCM(:,:,:) - tsoil_avg(:,1:nsoil_PCM,:)

! Compute global surface pressure
ps_avg_glob = sum(cell_area*ps_avg)/total_surface

! Compute the accepted maximum number of years due to orbital parameters (if needed)
call compute_maxyr_orbit(n_yr_sim,nmax_yr_runorb)

! Read the "startevo.nc"
call allocate_startevo_state()
call read_startevo(tsurf_avg_yr1,tsurf_avg,ps_avg_glob_ini,ps_ts,q_co2_ts,q_h2o_ts,h2o_ice,co2_ice,tsoil_avg)
deallocate(tsurf_avg_yr1)

! Compute ice tendencies from yearly minima
call allocate_tendencies()
call print_msg('> Computing surface ice tendencies',LVL_NFO)
call compute_tendice(minPCM_h2ofrost,h2o_ice,d_h2oice)
call print_msg('H2O ice tendencies [kg/m2/y] (min|max): '//real2str(minval(d_h2oice))//' | '//real2str(maxval(d_h2oice)),LVL_NFO)
call compute_tendice(minPCM_co2frost,co2_ice,d_co2ice)
call print_msg('CO2 ice tendencies [kg/m2/y] (min|max): '//real2str(minval(d_co2ice))//' | '//real2str(maxval(d_co2ice)),LVL_NFO)
deallocate(minPCM_h2ofrost,minPCM_co2frost)

! Save initial set-up useful for the next computations
call print_msg('> Saving some initial climate state variables',LVL_NFO)
call allocate_initial_snapshots()
ps_avg_glob_ini = ps_avg_glob
d_co2ice_ini(:,:) = d_co2ice(:,:)
q_co2_ts_ini(:,:) = q_co2_ts(:,:)
h2oice_sublim_coverage_ini = 0._dp
co2ice_sublim_coverage_ini = 0._dp
totmass_co2ice_ini = 0._qp
totmass_atmco2_ini = 0._qp
do i = 1,ngrid
    totmass_atmco2_ini = totmass_atmco2_ini + cell_area(i)*ps_avg(i)/g
    do islope = 1,nslope
        totmass_co2ice_ini = totmass_co2ice_ini + co2_ice(i,islope)*cell_area(i)*subslope_dist(i,islope)/cos(pi*def_slope_mean(islope)/180._dp)
        if (co2_ice(i,islope) > 0._dp) then
            is_co2ice_ini(i,islope) = .true.
            if (d_co2ice(i,islope) < 0._dp) co2ice_sublim_coverage_ini = co2ice_sublim_coverage_ini + cell_area(i)*subslope_dist(i,islope)/cos(pi*def_slope_mean(islope)/180._dp)
        end if
        ! cmetz we'll have to either reword this or get rid of this when we go to warmer planets,
        ! since h2o will have other reasons to go down other than just sublimation
        if (h2o_ice(i,islope) > 0._dp) then
            is_h2oice_ini(i,islope) = .true.
            if (d_h2oice(i,islope) < 0._dp) h2oice_sublim_coverage_ini = h2oice_sublim_coverage_ini + cell_area(i)*subslope_dist(i,islope)/cos(pi*def_slope_mean(islope)/180._dp)
        end if
    end do
end do
call print_msg("Initial global average pressure [Pa] = "//real2str(ps_avg_glob_ini),LVL_NFO)
call print_msg("Initial surface area of sublimating CO2 ice [m2] = "//real2str(co2ice_sublim_coverage_ini),LVL_NFO)
call print_msg("Initial surface area of sublimating H2O ice [m2] = "//real2str(h2oice_sublim_coverage_ini),LVL_NFO)

! Main evolution loop
! ~~~~~~~~~~~~~~~~~~~
call print_msg('',LVL_NFO)
call print_msg('********* Evolution *********',LVL_NFO)
call allocate_loop_state()
n_yr_run = 0
idt = 0
do while (n_yr_run < min(nmax_yr_runorb,nmax_yr_run) .and. n_yr_sim < ntot_yr_sim)
    call print_msg('**** Iteration of the PEM run [plnt y]: '//real2str(n_yr_run + dt),LVL_NFO)
    ! Evolve global surface pressure
    call evolve_pressure(d_co2ice,ps_avg_glob_old,ps_avg_glob,ps_avg)

    ! Adapt tracers according to global surface pressure
    call adapt_tracers2pressure(ps_avg_glob_old,ps_avg_glob,ps_ts,q_h2o_ts,q_co2_ts)

    ! Evolve surface ice
    allocate(zshift_surf(ngrid,nslope))
    zshift_surf(:,:) = 0._dp
    call evolve_h2oice(h2o_ice,d_h2oice,zshift_surf,stopcrit)
    if (stopcrit%stop_code() > 0) then
        deallocate(zshift_surf)
        call print_msg(stopcrit%stop_message(),LVL_NFO)
        exit
    end if
    call evolve_co2ice(co2_ice,d_co2ice,zshift_surf)

    ! Adapt surface temperature if surface ice disappeared
    call adapt_tsurf2disappearedice(co2_ice,is_co2ice_ini,is_co2ice_disappeared,tsurf_avg)
    ! cmetz question for JB why is this commented out because h2o ice doesn't lock the tsurf?
    !call adapt_tsurf2disappearedice(h2o_ice,is_h2oice_ini,tsurf_avg)

    if (do_soil) then
        ! Shift soil temperature to surface
        call shift_tsoil2surf(ngrid,nsoil,nslope,zshift_surf,tsurf_avg,tsoil_avg)

        ! Evolve soil temperature
        call evolve_soil_temp(tsoil_avg,tsurf_avg,tsoil_ts,tsoil_ts_old)
    end if ! do_soil
    deallocate(zshift_surf)

    ! Output the results
    call print_msg('> Standard outputs',LVL_NFO)
    call write_diagevo('ps_avg_glob','Global average pressure','Pa',ps_avg_glob)
    do islope = 1,nslope
        if (nslope == 1) then
            num = ''
        else
            write(num,'("_slope",i2.2)') islope
        end if
        call write_diagevo('h2oice'//trim(num),'H2O ice','kg/m2',h2o_ice(:,islope),(/dim_ngrid/))
        call write_diagevo('co2ice'//trim(num),'CO2 ice','kg/m2',co2_ice(:,islope),(/dim_ngrid/))
        call write_diagevo('d_h2oice'//trim(num),'H2O ice tendency','kg/m2/y',d_h2oice(:,islope),(/dim_ngrid/))
        call write_diagevo('d_co2ice'//trim(num),'CO2 ice tendency','kg/m2/y',d_co2ice(:,islope),(/dim_ngrid/))
        call write_diagevo('tsurf'//trim(num),'Surface temperature','K',tsurf_avg(:,islope),(/dim_ngrid/))
        if (do_soil) then
            call write_diagevo('tsoil_avg'//trim(num),'Soil temperature','K',tsoil_avg(:,:,islope),(/dim_ngrid,dim_nsoil/))
        end if
    end do

    ! Checking mass balance for CO2
    if (abs(CO2cond_ps_PCM - 1._dp) < eps) then
        totmass_co2ice = 0._dp
        totmass_atmco2 = 0._dp
        do i = 1,ngrid
            ! cmetz we need to introduce a co2 mixing ratio
            totmass_atmco2 = totmass_atmco2 + cell_area(i)*ps_avg(i)/g
            do islope = 1,nslope
                totmass_co2ice = totmass_co2ice + co2_ice(i,islope)*cell_area(i)*subslope_dist(i,islope)/cos(pi*def_slope_mean(islope)/180.)
            end do
        end do
        totmass_ini = max(totmass_atmco2_ini + totmass_co2ice_ini,eps_qp) ! To avoid division by 0
        call print_msg('> Relative total CO2 mass balance = '//real2str(100._qp*(totmass_atmco2 + totmass_co2ice - totmass_atmco2_ini - totmass_co2ice_ini)/totmass_ini)//' %',LVL_NFO)
        if (abs((totmass_atmco2 + totmass_co2ice - totmass_atmco2_ini - totmass_co2ice_ini)/totmass_ini) > 0.01_qp) then
            call print_msg('Mass balance is not conserved!',LVL_WRN)
            totmass_ini = max(totmass_atmco2_ini,eps_qp) ! To avoid division by 0
            call print_msg('    Atmospheric CO2 mass balance = '//real2str(100._qp*(totmass_atmco2 - totmass_atmco2_ini)/totmass_ini)//' %',LVL_WRN)
            totmass_ini = max(totmass_co2ice_ini,eps_qp) ! To avoid division by 0
            call print_msg('    CO2 ice mass balance         = '//real2str(100._qp*(totmass_co2ice - totmass_co2ice_ini)/totmass_ini)//' %',LVL_WRN)
        end if
    end if

    ! Evolve the tendencies
    call evolve_tend_co2ice(d_co2ice_ini,co2_ice,emissivity_PCM,q_co2_ts_ini,q_co2_ts,ps_ts,ps_avg_glob_ini,ps_avg_glob,d_co2ice)

    ! Increment time
    n_yr_run = n_yr_run + dt
    n_yr_sim = n_yr_sim + dt
    idt = idt + 1

    ! Save periodic backups of restart files
    if (backup_rate > 0) then
        if (mod(idt,backup_rate) == 0) call save_clim_state(h2o_ice,co2_ice,tsurf_avg,tsurf_dev,tsoil_avg,tsoil_dev,ps_avg,ps_dev,ps_avg_glob,ps_avg_glob_ini,idt)
    end if

    ! Check the stopping criteria
    call print_msg("> Checking the stopping criteria",LVL_NFO)
    call stopping_crit_pressure(ps_avg_glob_ini,ps_avg_glob,stopcrit)
    call stopping_crit_h2oice(h2oice_sublim_coverage_ini,h2o_ice,d_h2oice,stopcrit)
    call stopping_crit_co2ice(co2ice_sublim_coverage_ini,co2_ice,d_co2ice,stopcrit)
    call system_clock(c2)
    if (stopcrit%stop_code() == 0 .and. n_yr_run >= nmax_yr_run) stopcrit%nmax_yr_run_reached = .true.
    if (stopcrit%stop_code() == 0 .and. n_yr_run >= nmax_yr_runorb) stopcrit%nmax_yr_runorb_reached = .true.
    if (stopcrit%stop_code() == 0 .and. n_yr_sim >= ntot_yr_sim) stopcrit%nmax_yr_sim_reached = .true.
    if (stopcrit%stop_code() == 0 .and. timewall .and. real(c2 - c1,dp)/real(cr,dp) >= timelimit - antetime) stopcrit%time_limit_reached = .true.
    if (stopcrit%is_any_set()) then
        call print_msg(stopcrit%stop_message(),LVL_NFO)
        exit
    else
        call print_msg('**** This run has achieved '//real2str(n_yr_run)//' Planetary years.',LVL_NFO)
        call print_msg('**** The workflow has achieved '//real2str(n_yr_sim)//' Planetary years.',LVL_NFO)
        call print_msg('**** This run is continuing!',LVL_NFO)
        call print_msg('****',LVL_NFO)
    end if
end do ! End of the evolution loop

! Finalization
! ~~~~~~~~~~~~
call print_msg('',LVL_NFO)
call print_msg('********* Finalization *********',LVL_NFO)
! Update orbital parameters
if (evo_orbit) call update_orbit(n_yr_sim,n_yr_run)

call save_clim_state(h2o_ice,co2_ice,tsurf_avg,tsurf_dev,tsoil_avg,tsoil_dev,ps_avg,ps_dev,ps_avg_glob,ps_avg_glob_ini)

! Update the duration information of the workflow
call update_workflow_status(n_yr_run,stopcrit%stop_code(),n_yr_sim,ntot_yr_sim)

! Clean-up
call end_allocation()

! Footer
call system_clock(c2)
call print_end(i_pem_run,n_yr_run,n_yr_sim,real((c2 - c1),dp)/real(cr,dp),pem_ini_date,r_plnt2earth_yr)

END PROGRAM pem