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,h2o_surfdensity_avg,h2o_soildensity_avg, &
                          q_h2o_ts,q_co2_ts,d_h2oice,d_co2ice)
!-----------------------------------------------------------------------
! NAME
!     load_xios_data
!
! DESCRIPTION
!     Reads yearly and daily XIOS data, computes frost and ice tendencies.
!
! AUTHORS & DATE
!     JB Clement, 2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry,   only: ngrid, 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 tendencies, only: compute_tend
use frost,      only: compute_frost4PCM
use stoppage,   only: stop_clean
use display,    only: print_msg
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, h2o_surfdensity_avg, d_h2oice, d_co2ice, ps_ts, q_h2o_ts, q_co2_ts
real(dp), dimension(:,:,:),   intent(out) :: tsoil_avg, h2o_soildensity_avg
real(dp), dimension(:,:,:,:), intent(out) :: tsoil_ts

! 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(:),                 allocatable    :: num ! To read slope variables
real(dp), dimension(ngrid,nslope,2)          :: min_h2operice, min_co2perice, min_h2ofrost, min_co2frost
real(dp), dimension(ngrid,nsoil,nslope,nday) :: h2o_soildensity_ts
! CODE
! ----
! Initialization
min_h2operice(:,:,:) = 0._dp
min_co2perice(:,:,:) = 0._dp
min_h2ofrost(:,:,:) = 0._dp
min_co2frost(:,:,:) = 0._dp
if (nslope == 1) then ! No slope
    allocate(character(0) :: num)
else ! Using slopes
    allocate(character(8) :: num)
end if

!~~~~~~~~~~~~~~~~~~~~~~~~ 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//'"')
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 = '  '
        write(num,'(i2.2)') islope
        num = '_slope'//num
    end if
    call get_var_nc('co2ice'//num,var_read_2d)          ; call lonlat2vect(nlon,nlat,ngrid,var_read_2d,min_co2frost(:,islope,1))
    call get_var_nc('h2o_ice_s'//num,var_read_2d)       ; call lonlat2vect(nlon,nlat,ngrid,var_read_2d,min_h2ofrost(:,islope,1))
    call get_var_nc('watercap'//num,var_read_2d)        ; call lonlat2vect(nlon,nlat,ngrid,var_read_2d,min_h2operice(:,islope,1))
    call get_var_nc('perennial_co2ice'//num,var_read_2d); call lonlat2vect(nlon,nlat,ngrid,var_read_2d,min_co2perice(:,islope,1))
    call get_var_nc('tsurf'//num,var_read_2d)           ; call lonlat2vect(nlon,nlat,ngrid,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//'"')
call open_nc(xios_yr_name2,'read')

! Allocate and read the variables
call get_var_nc('ps',var_read_2d); call lonlat2vect(nlon,nlat,ngrid,var_read_2d,ps_avg)
do islope = 1,nslope
    if (nslope /= 1) then
        num='  '
        write(num,'(i2.2)') islope
        num = '_slope'//num
    end if
    call get_var_nc('tsurf'//num,var_read_2d)           ; call lonlat2vect(nlon,nlat,ngrid,var_read_2d,tsurf_avg(:,islope))
    call get_var_nc('co2ice'//num,var_read_2d)          ; call lonlat2vect(nlon,nlat,ngrid,var_read_2d,min_co2frost(:,islope,2))
    call get_var_nc('h2o_ice_s'//num,var_read_2d)       ; call lonlat2vect(nlon,nlat,ngrid,var_read_2d,min_h2ofrost(:,islope,2))
    call get_var_nc('watercap'//num,var_read_2d)        ; call lonlat2vect(nlon,nlat,ngrid,var_read_2d,min_h2operice(:,islope,2))
    call get_var_nc('perennial_co2ice'//num,var_read_2d); call lonlat2vect(nlon,nlat,ngrid,var_read_2d,min_co2perice(:,islope,2))
    if (do_soil) then
        call get_var_nc('soiltemp'//num,var_read_3d)
        do isoil = 1,nsoil_PCM
            call lonlat2vect(nlon,nlat,ngrid,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
        call get_var_nc('waterdensity_surface'//num,var_read_2d); call lonlat2vect(nlon,nlat,ngrid,var_read_2d,h2o_surfdensity_avg(:,islope))
    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//'"')
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(nlon,nlat,ngrid,var_read_3d(:,:,iday),ps_ts(:,iday))
end do
call get_var_nc('h2o_layer1',var_read_3d)
do iday = 1,nday
    call lonlat2vect(nlon,nlat,ngrid,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(nlon,nlat,ngrid,var_read_3d(:,:,iday),q_co2_ts(:,iday))
end do
if (do_soil) then
    do islope = 1,nslope
        if (nslope /= 1) then
            num='  '
            write(num,'(i2.2)') islope
            num = '_slope'//num
        end if
        call get_var_nc('soiltemp'//num,var_read_4d)
        do iday = 1,nday
            do isoil = 1,nsoil_PCM
                call lonlat2vect(nlon,nlat,ngrid,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
        call get_var_nc('waterdensity_soil'//num,var_read_4d)
        do iday = 1,nday
            do isoil = 1,nsoil_PCM
                call lonlat2vect(nlon,nlat,ngrid,var_read_4d(:,:,isoil,iday),h2o_soildensity_ts(:,isoil,islope,iday))
            end do
            do isoil = nsoil_PCM + 1,nsoil
                h2o_soildensity_ts(:,isoil,islope,iday) = h2o_soildensity_ts(:,nsoil_PCM,islope,iday) ! Explicit initialization because dimension with size nsoil > nsoil_PCM
            end do
        end do
    end do
    h2o_soildensity_avg(:,:,:) = sum(h2o_soildensity_ts,4)/nday
end if
deallocate(var_read_3d,var_read_4d,num)

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

!~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
! Compute frost from yearly minima
call compute_frost4PCM(min_h2ofrost(:,:,2),min_co2frost(:,:,2))

!~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
! Compute ice tendencies from yearly minima
call print_msg('> Computing surface ice tendencies')
call compute_tend(min_h2operice + min_h2ofrost,d_h2oice)
call print_msg('H2O ice tendencies [kg/m2/yr] (min|max): '//real2str(minval(d_h2oice))//' | '//real2str(maxval(d_h2oice)))
call compute_tend(min_co2perice + min_co2frost,d_co2ice)
call print_msg('CO2 ice tendencies [kg/m2/yr] (min|max): '//real2str(minval(d_co2ice))//' | '//real2str(maxval(d_co2ice)))

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

END MODULE xios_data