MODULE clim_state_rec
!-----------------------------------------------------------------------
! NAME
!     clim_state_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
use io_netcdf, only: open_nc, close_nc, check_nc, put_var_nc, get_dim_nc, get_var_nc, start_name, start1D_name, startfi_name, startpem_name

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

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

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

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

! DEPENDENCIES
! ------------
use geometry, only: vect2dyngrd, ngrid, nlon, nlat, nlayer
use tracers,  only: qnames, nq
use stoppage, only: stop_clean
use display,  only: print_msg

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

! ARGUMENTS
! ---------
real(dp),                   intent(in) :: pa4PCM, preff4PCM
real(dp), dimension(:),     intent(in) :: ps4PCM
real(dp), dimension(:,:),   intent(in) :: tsurf4PCM, 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,tsurf4PCM,q4PCM)
    return
end if

! In case of 3D
! ~~~~~~~~~~~~~
! Copy "start.nc" into "restart.nc"
call print_msg('> Writing "re'//start_name//'"')
call execute_command_line('cp '//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(nlon + 1,nlat,ngrid,ps4PCM,ps_dyn)
call put_var_nc('ps',ps_dyn)

! Potential temperature
call vect2dyngrd(nlon + 1,nlat,ngrid,teta4PCM,var_dyn)
call put_var_nc('teta',var_dyn)

! Air mass
call vect2dyngrd(nlon + 1,nlat,ngrid,air_mass4PCM,var_dyn)
call put_var_nc('masse',var_dyn)

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

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

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

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

! DEPENDENCIES
! ------------
use geometry,   only: nlayer, nslope
use atmosphere, only: teta_PCM, u_PCM, v_PCM
use tracers,    only: nq, qnames
use stoppage,   only: stop_clean
use display,    only: print_msg

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

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

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

! CODE
! ----
! Write "restart1D.txt"
call print_msg('> Writing "re'//start1D_name//'"')
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', (tsurf4PCM(1,j), j = 1,nslope), (teta_PCM(1,l), l = 1,nlayer)
close(funit)

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

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

! DEPENDENCIES
! ------------
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

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

! ARGUMENTS
! ---------
real(dp),    dimension(:,:),   intent(in) :: h2o_ice4PCM, co2_ice4PCM, albedo4PCM, emissivity4PCM, tsurf4PCM, flux_geo4PCM
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

! CODE
! ----
! Copy "startfi.nc" into "restartfi.nc"
call print_msg('> Writing "re'//startfi_name//'"')
call execute_command_line('cp '//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)

! Variables that have been modified
call put_var_nc('watercaptag',merge(1.,0.,is_h2o_perice))
call put_var_nc('watercap',h2o_ice4PCM,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)

! h2oice_depth
! lag_co2_ice
! zdqsdif_ssi_tot
! d_coef

call close_nc('re'//startfi_name)

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

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

! DEPENDENCIES
! ------------
use netcdf,           only: nf90_double, nf90_clobber, nf90_unlimited, nf90_global, &
                            nf90_create, nf90_enddef, nf90_def_dim, nf90_def_var, nf90_put_att
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
use utility,          only: int2str

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

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

! LOCAL VARIABLES
! ---------------
integer(di) :: ncid  ! File ID
integer(di) :: varid ! Variable ID
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 in the PEM!',1)

! Create file
call print_msg('> Creating "re'//startpem_name//'"')
call check_nc(nf90_create('re'//startpem_name,nf90_clobber,ncid),'creating re'//startpem_name) ! Enter define mode

! Define dimensions
call check_nc(nf90_def_dim(ncid,'Time',nf90_unlimited,dim_time),'defining dimension Time')
call check_nc(nf90_def_dim(ncid,'physical_points',ngrid,dim_ngrid),'defining dimension physical_points')
call check_nc(nf90_def_dim(ncid,'subsurface_layers',nsoil,dim_nsoil),'defining dimension subsurface_layers')
call check_nc(nf90_def_dim(ncid,'nslope',nslope,dim_nslope),'defining dimension nslope')
if (do_layering) then
    call print_msg('nb_str_max = '//int2str(nb_str_max))
    call check_nc(nf90_def_dim(ncid,'nb_str_max',nb_str_max,dim_nb_str),'defining dimension nb_str_max')
end if

! Define variables
call check_nc(nf90_def_var(ncid,'Time',nf90_double,(/dim_time/),varid),'defining variable Time')
call check_nc(nf90_put_att(ncid,varid,'title','Year of simulation'),'putting title attribute for Time')
call check_nc(nf90_put_att(ncid,varid,'units','Planetary year'),'putting units attribute for Time')

call check_nc(nf90_def_var(ncid,'longitude',nf90_double,(/dim_ngrid/),varid),'defining variable longitude')
call check_nc(nf90_put_att(ncid,varid,'title','Longitudes of the grid'),'putting title attribute for longitude')

call check_nc(nf90_def_var(ncid,'latitude',nf90_double,(/dim_ngrid/),varid),'defining variable latitude')
call check_nc(nf90_put_att(ncid,varid,'title','Latitudes of the grid'),'putting title attribute for latitude')

call check_nc(nf90_def_var(ncid,'cell_area',nf90_double,(/dim_ngrid/),varid),'defining variable cell_area')
call check_nc(nf90_put_att(ncid,varid,'title','Cell area'),'putting title attribute for cell_area')

call check_nc(nf90_def_var(ncid,'h2o_ice',nf90_double,(/dim_ngrid,dim_nslope,dim_time/),varid),'defining variable h2o_ice')
call check_nc(nf90_put_att(ncid,varid,'title','H2O ice'),'putting title attribute for h2o_ice')

call check_nc(nf90_def_var(ncid,'co2_ice',nf90_double,(/dim_ngrid,dim_nslope,dim_time/),varid),'defining variable co2_ice')
call check_nc(nf90_put_att(ncid,varid,'title','CO2 ice'),'putting title attribute for co2_ice')

if (do_soil) then
    call check_nc(nf90_def_var(ncid,'soildepth',nf90_double,(/dim_nsoil/),varid),'defining variable soildepth')
    call check_nc(nf90_put_att(ncid,varid,'title','Depths of soil layers'),'putting title attribute for soildepth')

    call check_nc(nf90_def_var(ncid,'tsoil',nf90_double,(/dim_ngrid,dim_nsoil,dim_nslope,dim_time/),varid),'defining variable tsoil')
    call check_nc(nf90_put_att(ncid,varid,'title','Soil temperature'),'putting title attribute for tsoil')

    call check_nc(nf90_def_var(ncid,'TI',nf90_double,(/dim_ngrid,dim_nsoil,dim_nslope,dim_time/),varid),'defining variable inertiedat')
    call check_nc(nf90_put_att(ncid,varid,'title','Thermal inertie of PEM'),'putting title attribute for TI')

    call check_nc(nf90_def_var(ncid,'inertiedat',nf90_double,(/dim_ngrid,dim_nsoil,dim_time/),varid),'defining variable longitude')
    call check_nc(nf90_put_att(ncid,varid,'title','Soil thermal inertia'),'putting title attribute for inertiedat')

    call check_nc(nf90_def_var(ncid,'icetable_depth',nf90_double,(/dim_ngrid,dim_nslope,dim_time/),varid),'defining variable icetable_depth')
    call check_nc(nf90_put_att(ncid,varid,'title','Depth of ice table'),'putting title attribute for icetable_depth')

    call check_nc(nf90_def_var(ncid,'icetable_thickness',nf90_double,(/dim_ngrid,dim_nslope,dim_time/),varid),'defining variable icetable_thickness')
    call check_nc(nf90_put_att(ncid,varid,'title','Thickness of ice table'),'putting title attribute for icetable_thickness')

    call check_nc(nf90_def_var(ncid,'ice_porefilling',nf90_double,(/dim_ngrid,dim_nsoil,dim_nslope,dim_time/),varid),'defining variable ice_porefilling')
    call check_nc(nf90_put_att(ncid,varid,'title','Subsurface ice pore filling'),'putting title attribute for ice_porefilling')

    if (do_sorption) then
        call check_nc(nf90_def_var(ncid,'h2o_ads_reg',nf90_double,(/dim_ngrid,dim_nsoil,dim_nslope,dim_time/),varid),'defining variable h2o_ads_reg')
        call check_nc(nf90_put_att(ncid,varid,'title','Mass of H2O adsorbded in the regolith'),'putting title attribute for h2o_ads_reg')

        call check_nc(nf90_def_var(ncid,'co2_ads_reg',nf90_double,(/dim_ngrid,dim_nsoil,dim_nslope,dim_time/),varid),'defining variable co2_ads_reg')
        call check_nc(nf90_put_att(ncid,varid,'title','Mass of CO2 adsorbded in the regolith'),'putting title attribute for co2_ads_reg')
    end if
end if

if (do_layering) then
    call check_nc(nf90_def_var(ncid,'stratif_top_elevation',nf90_double,(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/),varid),'defining variable stratif_top_elevation')
    call check_nc(nf90_put_att(ncid,varid,'title','Layering top elevation'),'putting title attribute for stratif_top_elevation')

    call check_nc(nf90_def_var(ncid,'stratif_h_h2oice',nf90_double,(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/),varid),'defining variable stratif_h_h2oice')
    call check_nc(nf90_put_att(ncid,varid,'title','Layering H2O ice height'),'putting title attribute for stratif_h_h2oice')

    call check_nc(nf90_def_var(ncid,'stratif_h_co2ice',nf90_double,(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/),varid),'defining variable longitude')
    call check_nc(nf90_put_att(ncid,varid,'title','Layering CO2 ice height'),'putting title attribute for stratif_h_co2ice')

    call check_nc(nf90_def_var(ncid,'stratif_h_dust',nf90_double,(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/),varid),'defining variable stratif_h_co2ice')
    call check_nc(nf90_put_att(ncid,varid,'title','Layering dust height'),'putting title attribute for stratif_h_dust')

    call check_nc(nf90_def_var(ncid,'stratif_h_pore',nf90_double,(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/),varid),'defining variable stratif_h_pore')
    call check_nc(nf90_put_att(ncid,varid,'title','Layering pore height'),'putting title attribute for stratif_h_pore')

    call check_nc(nf90_def_var(ncid,'stratif_poreice_volfrac',nf90_double,(/dim_ngrid,dim_nb_str,dim_nslope,dim_time/),varid),'defining variable stratif_poreice_volfrac')
    call check_nc(nf90_put_att(ncid,varid,'title','Layering ice pore volume fraction'),'putting title attribute for stratif_poreice_volfrac')
end if
! Global attributes
call check_nc(nf90_put_att(ncid,nf90_global,'title','PEM start file'),'putting global attribute')

! Leave define mode and putting variables defining dimensions
call check_nc(nf90_enddef(ncid),'leaving define mode')
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)
call close_nc('re'//startpem_name)

! File creation done
is_restartpem = .true.

END SUBROUTINE create_nc_pem
!=======================================================================

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

! DEPENDENCIES
! ------------
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
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

! 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 ! Time index to record variables
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_nc_pem(nb_str_max)

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

! Writing time counter
call open_nc('re'//startpem_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))
    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_volfrac',layerings_array(:,:,:,6),itime)
    deallocate(layerings_array)
end if

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

END SUBROUTINE write_restartpem
!=======================================================================

END MODULE clim_state_rec