MODULE climate_rec
!-----------------------------------------------------------------------
! NAME
!     climate_rec
!
! DESCRIPTION
!     Write the restart files to save the climate state.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use numerics, only: dp, di, k4

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

! VARIABLES
! ---------
logical(k4), protected, private :: is_restartevo = .false. ! Flag to know if "restartevo.nc" exists

contains
!+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

!=======================================================================
SUBROUTINE write_restart(ps4PCM,pa4PCM,preff4PCM,q4PCM,teta4PCM,air_mass4PCM)
!-----------------------------------------------------------------------
! NAME
!    write_restart
!
! DESCRIPTION
!    Write the file "restart.nc".
!
! AUTHORS & DATE
!    JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use io_netcdf, only: open_nc, close_nc, put_var_nc, start_name
use geometry,  only: vect2dyngrd, ngrid, nlon, nlat, nlayer
use tracers,   only: qnames, nq
use stoppage,  only: stop_clean
use display,   only: print_msg, LVL_NFO

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

! ARGUMENTS
! ---------
real(dp),                   intent(in) :: pa4PCM, preff4PCM
real(dp), dimension(:),     intent(in) :: ps4PCM
real(dp), dimension(:,:),   intent(in) :: teta4PCM, air_mass4PCM
real(dp), dimension(:,:,:), intent(in) :: q4PCM

! LOCAL VARIABLES
! ---------------
integer(di)                               :: cstat, l, i
real(dp), dimension(nlon + 1,nlat)        :: ps_dyn
real(dp), dimension(nlon + 1,nlat,nlayer) :: var_dyn

! CODE
! ----
! In case of 1D
! ~~~~~~~~~~~~~
if (ngrid == 1) then
    call write_restart1D(ps4PCM,pa4PCM,preff4PCM,teta4PCM,q4PCM)
    return
end if

! In case of 3D
! ~~~~~~~~~~~~~
! Copy "start.nc" into "restart.nc"
call print_msg('> Writing "re'//start_name//'"',LVL_NFO)
call execute_command_line('\cp -f '//start_name//' re'//start_name,cmdstat = cstat)
if (cstat > 0) then
    call stop_clean(__FILE__,__LINE__,'command execution failed!',1)
else if (cstat < 0) then
    call stop_clean(__FILE__,__LINE__,'command execution not supported!',1)
end if

! Rewrite the variables modified by the PEM
call open_nc('re'//start_name,'write')

! Surface pressure
call vect2dyngrd(ps4PCM,ps_dyn)
call put_var_nc('ps',ps_dyn,1)

! Potential temperature
do l = 1,nlayer
    call vect2dyngrd(teta4PCM(:,l),var_dyn(:,:,l))
end do
call put_var_nc('teta',var_dyn,1)

! Air mass
do l = 1,nlayer
    call vect2dyngrd(air_mass4PCM(:,l),var_dyn(:,:,l),extensive = .true.)
end do
call put_var_nc('masse',var_dyn,1)

! Tracers
do i = 1,nq
    do l = 1,nlayer
        call vect2dyngrd(q4PCM(:,l,i),var_dyn(:,:,l))
    end do
    call put_var_nc(qnames(i),var_dyn,1)
end do

! Close
call close_nc('re'//start_name)

END SUBROUTINE write_restart
!=======================================================================

!=======================================================================
SUBROUTINE write_restart1D(ps4PCM,pa4PCM,preff4PCM,teta4PCM,q4PCM)
!-----------------------------------------------------------------------
! NAME
!    write_restart1D
!
! DESCRIPTION
!    Write the file "restart1D.txt".
!
! AUTHORS & DATE
!    JB Clement, 01/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use io_netcdf,  only: start1D_name
use geometry,   only: nlayer
use atmosphere, only: u_PCM, v_PCM
use tracers,    only: nq, qnames
use stoppage,   only: stop_clean
use display,    only: print_msg, LVL_NFO

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

! ARGUMENTS
! ---------
real(dp),                   intent(in) :: pa4PCM, preff4PCM
real(dp), dimension(:),     intent(in) :: ps4PCM
real(dp), dimension(:,:),   intent(in) :: teta4PCM
real(dp), dimension(:,:,:), intent(in) :: q4PCM

! LOCAL VARIABLES
! ---------------
integer(di) :: funit, ierr, i, l

! CODE
! ----
! Write "restart1D.txt"
call print_msg('> Writing "re'//start1D_name//'"',LVL_NFO)
open(newunit = funit,file = 're'//start1D_name,status = "replace",action = "write",iostat = ierr)
if (ierr /= 0) call stop_clean(__FILE__,__LINE__,'error opening file "re'//start1D_name//'"!',ierr)
write(funit,*) 'ps', ps4PCM(1), pa4PCM, preff4PCM
do i = 1,nq
    write(funit,*) qnames(i), (q4PCM(1,l,i), l = 1,nlayer)
end do
write(funit,*) 'u', (u_PCM(1,l), l = 1,nlayer)
write(funit,*) 'v', (v_PCM(1,l), l = 1,nlayer)
write(funit,*) 'teta', (teta4PCM(1,l), l = 1,nlayer)
close(funit)

END SUBROUTINE write_restart1D
!=======================================================================

!=======================================================================
SUBROUTINE write_restartfi(is_h2o_perice,co2_ice4PCM,tsurf4PCM,tsoil4PCM,inertiesoil4PCM,albedo4PCM,emissivity4PCM,flux_geo4PCM,h2oice_depth4PCM)
!-----------------------------------------------------------------------
! NAME
!    write_restartfi
!
! DESCRIPTION
!    Write the file "restartfi.nc".
!
! AUTHORS & DATE
!    JB Clement, 01/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use io_netcdf, only: open_nc, close_nc, put_var_nc, get_dim_nc, get_var_nc, startfi_name
use geometry,  only: ngrid, nslope
use orbit,     only: obliquity, aphelion, perihelion, date_peri
use frost,     only: h2o_frost4PCM, co2_frost4PCM
use stoppage,  only: stop_clean
use display,   only: print_msg, LVL_NFO

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

! ARGUMENTS
! ---------
real(dp),    dimension(:,:),   intent(in) :: co2_ice4PCM, albedo4PCM, emissivity4PCM, tsurf4PCM, flux_geo4PCM, h2oice_depth4PCM
real(dp),    dimension(:,:,:), intent(in) :: tsoil4PCM, inertiesoil4PCM
logical(k4), dimension(:),     intent(in) :: is_h2o_perice

! LOCAL VARIABLES
! ---------------
integer(di)                         :: cstat
integer(di)                         :: nindex ! Size of dimension 'index'
real(dp), dimension(:), allocatable :: controle
real(dp), dimension(ngrid,nslope)   :: watercap

! CODE
! ----
! Copy "startfi.nc" into "restartfi.nc"
call print_msg('> Writing "re'//startfi_name//'"',LVL_NFO)
call execute_command_line('\cp -f '//startfi_name//' re'//startfi_name,cmdstat = cstat)
if (cstat > 0) then
    call stop_clean(__FILE__,__LINE__,'command execution failed!',1)
else if (cstat < 0) then
    call stop_clean(__FILE__,__LINE__,'command execution not supported!',1)
end if

! Load the variable 'controle' to modify it with new values
call open_nc('re'//startfi_name,'read')
call get_dim_nc('index',nindex)
allocate(controle(nindex))
call get_var_nc('controle',controle)
call close_nc('re'//startfi_name)

! Rewrite the variables modified by the PEM
call open_nc('re'//startfi_name,'write')

! Orbital parameters (controle)
controle(18) = obliquity  ! Obliquity
controle(15) = perihelion ! Perihelion
controle(16) = aphelion   ! Aphelion
controle(17) = date_peri  ! Date of perihelion
call put_var_nc('controle',controle)
deallocate(controle)

! In the Mars PCM, only the variation of perennial H2O ice is monitored, not the absolute quantity h2o_ice4PCM
watercap(:,:) = 0._dp

! Other variables that have been modified
call put_var_nc('watercaptag',merge(1._dp,0._dp,is_h2o_perice))
call put_var_nc('watercap',watercap,1)
call put_var_nc('h2o_ice',h2o_frost4PCM,1)
call put_var_nc('co2',co2_frost4PCM,1)
call put_var_nc('perennial_co2ice',co2_ice4PCM,1)
call put_var_nc('tsurf',tsurf4PCM,1)
call put_var_nc('tsoil',tsoil4PCM,1)
call put_var_nc('inertiesoil',inertiesoil4PCM,1)
call put_var_nc('albedo',albedo4PCM,1)
call put_var_nc('emis',emissivity4PCM,1)
call put_var_nc('flux_geo',flux_geo4PCM,1)
call put_var_nc('h2oice_depth',h2oice_depth4PCM)

! Close
call close_nc('re'//startfi_name)

END SUBROUTINE write_restartfi
!=======================================================================

!=======================================================================
SUBROUTINE create_startevo(nb_str_max)
!-----------------------------------------------------------------------
! NAME
!     create_startevo
!
! DESCRIPTION
!     Create a NetCDF file to restart the PEM.
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use io_netcdf,        only: create_nc, close_nc, def_dim_nc, def_var_nc, put_var_nc, startevo_name
use geometry,         only: dim_init, ngrid, nsoil, nslope, longitudes, latitudes, cell_area
use stoppage,         only: stop_clean
use soil,             only: do_soil, mlayer
use sorption,         only: do_sorption
use layered_deposits, only: do_layering
use display,          only: print_msg, LVL_NFO
use utility,          only: int2str

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

! ARGUMENTS
! ---------
integer(di), intent(in) :: nb_str_max

! LOCAL VARIABLES
! ---------------
integer(di) :: dim_ngrid, dim_nsoil, dim_nslope, dim_time, dim_nb_str

! CODE
! ----
! Check if dimensions are well initialized
if (.not. dim_init) call stop_clean(__FILE__,__LINE__,'dimensions are not initilized!',1)

! Create file
call create_nc('re'//startevo_name,'Starting file for the PEM')

! Define dimensions
call def_dim_nc('Time',-1,dim_time)
call def_dim_nc('physical_points',ngrid,dim_ngrid)
call def_dim_nc('subsurface_layers',nsoil,dim_nsoil)
call def_dim_nc('nslope',nslope,dim_nslope)
if (do_layering) then
    call print_msg('nb_str_max = '//int2str(nb_str_max),LVL_NFO)
    call def_dim_nc('nb_str_max',nb_str_max,dim_nb_str)
end if

! Define the variables
call def_var_nc('Time','Year of simulation','Planetary year',(/dim_time/))
call def_var_nc('longitude','Longitudes of the grid','Radian East-West',(/dim_ngrid/))
call def_var_nc('latitude','Latitudes of the grid','Radian North-South',(/dim_ngrid/))
call def_var_nc('cell_area','Cell area','m2',(/dim_ngrid/))
call def_var_nc('h2o_ice','H2O ice','kg/m2',(/dim_ngrid,dim_nslope,dim_time/))
call def_var_nc('co2_ice','CO2 ice','kg/m2',(/dim_ngrid,dim_nslope,dim_time/))
if (do_soil) then
    call def_var_nc('soildepth','Depths of soil layers','m',(/dim_nsoil/))
    call def_var_nc('tsoil','Soil temperature','K',(/dim_ngrid,dim_nsoil,dim_nslope,dim_time/))
    call def_var_nc('TI','Soil thermal inertia','J/m2/K/s^0.5',(/dim_ngrid,dim_nsoil,dim_nslope,dim_time/))
    call def_var_nc('inertiedat','Reference soil thermal inertia','J/m2/K/s^0.5',(/dim_ngrid,dim_nsoil,dim_time/))
    call def_var_nc('icetable_depth','Depth of ice table','m',(/dim_ngrid,dim_nslope,dim_time/))
    call def_var_nc('icetable_thickness','Thickness of ice table','m',(/dim_ngrid,dim_nslope,dim_time/))
    call def_var_nc('ice_porefilling','Subsurface ice pore filling','kg/m2',(/dim_ngrid,dim_nsoil,dim_nslope,dim_time/))
    if (do_sorption) then
        call def_var_nc('h2o_ads_reg','H2O adsorbded in regolith','kg/m2',(/dim_ngrid,dim_nsoil,dim_nslope,dim_time/))
        call def_var_nc('co2_ads_reg','CO2 adsorbded in regolith','kg/m2',(/dim_ngrid,dim_nsoil,dim_nslope,dim_time/))
    end if
end if

if (do_layering) then
    call def_var_nc('stratif_top_elevation','Layering top elevation','m',(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/))
    call def_var_nc('stratif_h_h2oice','Layering H2O ice height','m',(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/))
    call def_var_nc('stratif_h_co2ice','Layering CO2 ice height','m',(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/))
    call def_var_nc('stratif_h_dust','Layering dust height','m',(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/))
    call def_var_nc('stratif_h_pore','Layering pore height','m',(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/))
    call def_var_nc('stratif_poreice_coef1','Layering pore-ice coefficient #1','-',(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/))
    call def_var_nc('stratif_poreice_coef2','Layering pore-ice coefficient #2','-',(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/))
    call def_var_nc('stratif_poreice_coef3','Layering pore-ice coefficient #3','-',(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/))
    call def_var_nc('stratif_poreice_coef4','Layering pore-ice coefficient #4','-',(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/))
    call def_var_nc('stratif_date','Layering stratum formation date','Planetary year',(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/))
end if

! Fill the variables related to dimensions
call put_var_nc('longitude',longitudes)
call put_var_nc('latitude',latitudes)
call put_var_nc('cell_area',cell_area)
call put_var_nc('soildepth',mlayer)

! File creation done
call close_nc('re'//startevo_name)
is_restartevo = .true.

END SUBROUTINE create_startevo
!=======================================================================

!=======================================================================
SUBROUTINE write_restartevo(h2o_ice,co2_ice,tsoil,TI,icetable_depth,icetable_thickness,ice_porefilling,h2o_ads_reg,co2_ads_reg,layerings_map)
!-----------------------------------------------------------------------
! NAME
!    write_restartevo
!
! DESCRIPTION
!    Write the file "restartevo.nc".
!
! AUTHORS & DATE
!    JB Clement, 01/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use io_netcdf,        only: open_nc, close_nc, put_var_nc, startevo_name
use geometry,         only: ngrid, nslope
use evolution,        only: pem_ini_date, n_yr_sim
use layered_deposits, only: layering, do_layering, map2array, get_nb_str_max, ncoef_poreice
use soil,             only: do_soil, inertiedat
use sorption,         only: do_sorption
use ice_table,        only: icetable_equilibrium, icetable_dynamic
use stoppage,         only: stop_clean
use display,          only: print_msg, LVL_NFO

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

! ARGUMENTS
! ---------
real(dp),       dimension(:,:),   intent(in) :: h2o_ice, co2_ice, icetable_depth, icetable_thickness
real(dp),       dimension(:,:,:), intent(in) :: tsoil, TI, ice_porefilling, h2o_ads_reg, co2_ads_reg
type(layering), dimension(:,:),   intent(in) :: layerings_map

! LOCAL VARIABLES
! ---------------
integer(di)                               :: itime ! Current NetCDF time index
integer(di)                               :: nb_str_max
real(dp), dimension(:,:,:,:), allocatable :: layerings_array

! CODE
! ----
! Create the file
nb_str_max = get_nb_str_max(layerings_map)
call create_startevo(nb_str_max)

call print_msg('> Writing "re'//startevo_name//'"',LVL_NFO)
if (.not. is_restartevo) call stop_clean(__FILE__,__LINE__,'The file"'//startevo_name//'" has not been created',1)

! Writing time counter
call open_nc('re'//startevo_name,'write',itime)
call put_var_nc('Time',pem_ini_date + n_yr_sim,itime)

! Writing other variables
call put_var_nc('h2o_ice',h2o_ice,itime)
call put_var_nc('co2_ice',co2_ice,itime)

if (do_soil) then
    call put_var_nc('tsoil',tsoil,itime)
    call put_var_nc('TI',TI,itime)
    call put_var_nc('inertiedat',inertiedat,itime)
    call put_var_nc('icetable_depth',icetable_depth,itime)
    if (icetable_equilibrium) then
        call put_var_nc('icetable_thickness',icetable_thickness,itime)
    else if (icetable_dynamic) then
        call put_var_nc('ice_porefilling',ice_porefilling,itime)
    end if
    if (do_sorption) then
        call put_var_nc('h2o_ads_reg',h2o_ads_reg,itime)
        call put_var_nc('co2_ads_reg',co2_ads_reg,itime)
    end if
end if

if (do_layering) then
    allocate(layerings_array(ngrid,nslope,nb_str_max,6 + ncoef_poreice))
    call map2array(layerings_map,layerings_array)
    call put_var_nc('stratif_top_elevation',layerings_array(:,:,:,1),itime)
    call put_var_nc('stratif_h_co2ice',layerings_array(:,:,:,2),itime)
    call put_var_nc('stratif_h_h2oice',layerings_array(:,:,:,3),itime)
    call put_var_nc('stratif_h_dust',layerings_array(:,:,:,4),itime)
    call put_var_nc('stratif_h_pore',layerings_array(:,:,:,5),itime)
    call put_var_nc('stratif_poreice_coef1',layerings_array(:,:,:,6),itime)
    call put_var_nc('stratif_poreice_coef2',layerings_array(:,:,:,7),itime)
    call put_var_nc('stratif_poreice_coef3',layerings_array(:,:,:,8),itime)
    call put_var_nc('stratif_poreice_coef4',layerings_array(:,:,:,9),itime)
    call put_var_nc('stratif_date',layerings_array(:,:,:,10),itime)
    deallocate(layerings_array)
end if

! Close
call close_nc('re'//startevo_name)

END SUBROUTINE write_restartevo
!=======================================================================

END MODULE climate_rec