MODULE xios_data
!-----------------------------------------------------------------------
! NAME
!     xios_data
!
! DESCRIPTION
!     Read XIOS output data and process it for PEM initialization.
!
! AUTHORS & DATE
!     JB Clement, 2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

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

!=======================================================================
SUBROUTINE load_xios_data(ps_avg,ps_ts,tsurf_avg,tsurf_avg_yr1,tsoil_avg,tsoil_ts,q_h2o_ts,q_co2_ts,min_yrPCM_h2ofrost,min_yrPCM_co2frost)
!-----------------------------------------------------------------------
! NAME
!     load_xios_data
!
! DESCRIPTION
!     Reads yearly and daily XIOS data, computes frost and ice tendencies.
!
! AUTHORS & DATE
!     JB Clement, 2025 & 04/2026
!     C. Metz, 04/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry,  only: nslope, nsoil, nsoil_PCM, nday, lonlat2vect, nlon, nlat
use io_netcdf, only: xios_day_name2, xios_yr_name1, xios_yr_name2, open_nc, close_nc, get_var_nc, get_dim_nc
use soil,      only: do_soil
use frost,     only: compute_frost4PCM
use stoppage,  only: stop_clean
use display,   only: print_msg, LVL_NFO
use utility,   only: real2str

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

! ARGUMENTS
! ---------
real(dp), dimension(:),       intent(out) :: ps_avg
real(dp), dimension(:,:),     intent(out) :: tsurf_avg, tsurf_avg_yr1, ps_ts, q_h2o_ts, q_co2_ts
real(dp), dimension(:,:,:),   intent(out) :: tsoil_avg
real(dp), dimension(:,:,:,:), intent(out) :: tsoil_ts
real(dp), dimension(:,:,:),   intent(out) ::  min_yrPCM_h2ofrost, min_yrPCM_co2frost

! LOCAL VARIABLES
! ---------------
logical(k4)                               :: here
integer(di)                               :: islope, isoil, iday
real(dp), dimension(:,:),     allocatable :: var_read_2d
real(dp), dimension(:,:,:),   allocatable :: var_read_3d
real(dp), dimension(:,:,:,:), allocatable :: var_read_4d
character(8)                              :: num ! Slope suffix to read variables

! CODE
! ----
! Initialization
min_yrPCM_h2ofrost(:,:,:) = 0._dp
min_yrPCM_co2frost(:,:,:) = 0._dp

!~~~~~~~~~~~~~~~~~~~~~~~~ Year 1 - Yearly data ~~~~~~~~~~~~~~~~~~~~~~~~~
inquire(file = xios_yr_name1,exist = here)
if (.not. here) call stop_clean(__FILE__,__LINE__,'cannot find required file "'//xios_yr_name1//'"!',1)

! Open the NetCDF file of XIOS outputs
call print_msg('> Reading "'//xios_yr_name1//'"',LVL_NFO)
call open_nc(xios_yr_name1,'read')

! Allocate and read the variables
allocate(var_read_2d(nlon,nlat),var_read_3d(nlon,nlat,nsoil_PCM))
do islope = 1,nslope
    if (nslope == 1) then
        num = ''
    else
        write(num,'("_slope",i2.2)') islope
    end if
    call get_var_nc('co2_ice_surf'//trim(num),var_read_2d); call lonlat2vect(var_read_2d,min_yrPCM_co2frost(:,islope,1))
    call get_var_nc('h2o_ice_surf'//trim(num),var_read_2d); call lonlat2vect(var_read_2d,min_yrPCM_h2ofrost(:,islope,1))
    call get_var_nc('tsurf'//trim(num),var_read_2d)       ; call lonlat2vect(var_read_2d,tsurf_avg_yr1(:,islope))
end do

! Close the NetCDF file of XIOS outputs
call close_nc(xios_yr_name1)

!~~~~~~~~~~~~~~~~~~~~~~~~ Year 2 - Yearly data ~~~~~~~~~~~~~~~~~~~~~~~~~
inquire(file = xios_yr_name2,exist = here)
if (.not. here) call stop_clean(__FILE__,__LINE__,'cannot find required file "'//xios_yr_name2//'"!',1)
! Open the NetCDF file of XIOS outputs
call print_msg('> Reading "'//xios_yr_name2//'"',LVL_NFO)
call open_nc(xios_yr_name2,'read')

! Allocate and read the variables
call get_var_nc('ps',var_read_2d); call lonlat2vect(var_read_2d,ps_avg)
do islope = 1,nslope
    if (nslope == 1) then
        num = ''
    else
        write(num,'("_slope",i2.2)') islope
    end if
    call get_var_nc('tsurf'//trim(num),var_read_2d)       ; call lonlat2vect(var_read_2d,tsurf_avg(:,islope))
    call get_var_nc('co2_ice_surf'//trim(num),var_read_2d); call lonlat2vect(var_read_2d,min_yrPCM_co2frost(:,islope,2))
    call get_var_nc('h2o_ice_surf'//trim(num),var_read_2d); call lonlat2vect(var_read_2d,min_yrPCM_h2ofrost(:,islope,2))
    if (do_soil) then
        call get_var_nc('tsoil'//trim(num),var_read_3d)
        do isoil = 1,nsoil_PCM
            call lonlat2vect(var_read_3d(:,:,isoil),tsoil_avg(:,isoil,islope))
        end do
        do isoil = nsoil_PCM + 1,nsoil
            tsoil_avg(:,isoil,islope) = tsoil_avg(:,nsoil_PCM,islope) ! Explicit initialization because dimension with size nsoil > nsoil_PCM
        end do
    end if
end do

! Close the NetCDF file of XIOS outputs
call close_nc(xios_yr_name2)

!~~~~~~~~~~~~~~~~~~~~~~~~~ Year 2 - Daily data ~~~~~~~~~~~~~~~~~~~~~~~~~
inquire(file = xios_day_name2,exist = here)
if (.not. here) call stop_clean(__FILE__,__LINE__,'cannot find required file "'//xios_day_name2//'"!',1)
! Open the NetCDF file of XIOS outputs
call print_msg('> Reading "'//xios_day_name2//'"',LVL_NFO)
call open_nc(xios_day_name2,'read')

! Allocate and read the variables
deallocate(var_read_2d,var_read_3d)
allocate(var_read_3d(nlon,nlat,nday),var_read_4d(nlon,nlat,nsoil_PCM,nday))
call get_var_nc('ps',var_read_3d)
do iday = 1,nday
    call lonlat2vect(var_read_3d(:,:,iday),ps_ts(:,iday))
end do
call get_var_nc('h2o_layer1',var_read_3d)
do iday = 1,nday
    call lonlat2vect(var_read_3d(:,:,iday),q_h2o_ts(:,iday))
end do
call get_var_nc('co2_layer1',var_read_3d)
do iday = 1,nday
    call lonlat2vect(var_read_3d(:,:,iday),q_co2_ts(:,iday))
end do
if (do_soil) then
    do islope = 1,nslope
        if (nslope == 1) then
            num = ''
        else
            write(num,'("_slope",i2.2)') islope
        end if
        call get_var_nc('tsoil'//trim(num),var_read_4d)
        do iday = 1,nday
            do isoil = 1,nsoil_PCM
                call lonlat2vect(var_read_4d(:,:,isoil,iday),tsoil_ts(:,isoil,islope,iday))
            end do
            do isoil = nsoil_PCM + 1,nsoil
                tsoil_ts(:,isoil,islope,iday) = tsoil_ts(:,nsoil_PCM,islope,iday) ! Explicit initialization because dimension with size nsoil > nsoil_PCM
            end do
        end do
    end do
end if
deallocate(var_read_3d,var_read_4d)

! Close the NetCDF file of XIOS outputs
call close_nc(xios_day_name2)

!~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
! Compute frost from yearly minima
where (min_yrPCM_h2ofrost(:,:,:) < 0._dp) min_yrPCM_h2ofrost(:,:,:) = 0._dp ! Enforcing positivity
where (min_yrPCM_co2frost(:,:,:) < 0._dp) min_yrPCM_co2frost(:,:,:) = 0._dp ! Enforcing positivity
call compute_frost4PCM(min_yrPCM_h2ofrost(:,:,2),min_yrPCM_co2frost(:,:,2))

END SUBROUTINE load_xios_data
!=======================================================================

END MODULE xios_data