! -----------------------------------------------------------------------------
! Copyright © 2025 Alexandre Gauvain
! Contributor(s): Alexandre Gauvain
! Email:       alexandre.gauvain@lmd.ipsl.fr or alexandre.gauvain.ag@gmail.com
! Affiliation: LMD/IPSL, CNRS, Sorbonne Université, Paris, France
! Created:     2023-09-03
! Modified:    2025-06-10
!
! This software is a computer program whose purpose is to model hydrological
! processes at planetary scale, including water flow, storage, climate
! interactions, and groundwater.
!
! This software is governed by the CeCILL license under
! French law and abiding by the rules of distribution of free software.
! You can use, modify and/or redistribute the software under the terms of
! the CeCILL license as circulated by CEA, CNRS and INRIA
! at the following URL "http://www.cecill.info".
!
! As a counterpart to the access to the source code and rights to copy,
! modify and redistribute granted by the license, users are provided only
! with a limited warranty and the software's author, the holder of the
! economic rights, and the successive licensors have only limited liability.
!
! In this respect, the user's attention is drawn to the risks associated
! with loading, using, modifying and/or developing or reproducing the
! software by the user in light of its specific status of free software,
! that may mean that it is complicated to manipulate, and that also therefore
! means that it is reserved for developers and experienced professionals
! having in-depth computer knowledge. Users are therefore encouraged to load
! and test the software's suitability as regards their requirements in
! conditions enabling the security of their systems and/or data to be ensured
! and, more generally, to use and operate it in the same conditions as
! regards security.
!
! The fact that you are presently reading this means that you have had
! knowledge of the CeCILL license and that you accept its terms.
!
! -----------------------------------------------------------------------------

module hydro_database
!  use Arguments
  use display, only: print_msg, LVL_NFO
  use netcdf
  use display, only: print_msg, LVL_NFO
  use, intrinsic :: iso_fortran_env, only : real64, real128
  implicit none
!===============================================================================
! Hydrological Database (INPUT)
!===============================================================================
  integer, parameter :: dp = real64, qp = real128
  ! Longitudes
  integer :: lonlen
  real, allocatable :: longitude(:)
  ! Latitudes
  integer :: latlen
  real, allocatable :: latitude(:)
  ! Depressions
  integer:: deplen, IDdep, size_leaf, ndep, nleaf
  ! Function between water volume and water surface elevation
  double precision, allocatable :: lake_vol(:,:), lake_area(:,:), lake_elev(:,:)
  integer :: lake_dislen
  double precision, allocatable :: lake_perc(:)
  ! Graph table
  integer, allocatable :: child_list(:), lchild_list(:), rchild_list(:), parent_list(:)
  integer, allocatable :: pit_cell_list(:), out_cell_list(:), sibling_list(:), geolink_list(:)
  integer, allocatable :: imin_list(:), imax_list(:), jmin_list(:), jmax_list(:)
  double precision, allocatable :: Vmax_list(:), Vmin_list(:), lake_lat_list(:), lake_long_list(:)
  double precision, allocatable :: Zmin_list(:), Zmax_list(:), Amax_list(:), watershed_area_list(:)
  ! Maps informations
  integer,allocatable :: label(:,:)
  real(real128) :: planet_area
  real(real128), allocatable :: cell_area(:)
  character(256) :: infile

contains
  subroutine load_hydro_database(infile_in)
    implicit none
    ! NetCDF stuff
    character(256), intent(in) :: infile_in
    integer :: i
    integer :: status ! NetCDF return code
    integer :: ncid ! NetCDF file ID
    integer :: dimid ! NetCDF dimension ID
    integer :: varid ! NetCDF variable ID

    ! 1. Load input datasets
    ! call PrintHeader("Hydrological Database")
    infile = infile_in

    ! 1.1. Open input datafile
    status=nf90_open(infile,NF90_NOWRITE,ncid)
    call print_msg('> Loading hydrological database',LVL_NFO)
    ! 1.2 Identify/load dimensions
    ! Longitude
    status=nf90_inq_dimid(ncid,"lon",dimid)
    status=nf90_inquire_dimension(ncid,dimid,len=lonlen)
    write(*,*) "lonlen                 =",lonlen
    ! Latitude
    status=nf90_inq_dimid(ncid,"lat",dimid)
    status=nf90_inquire_dimension(ncid,dimid,len=latlen)
    write(*,*) "latlen                 =",latlen
    ! Depressions
    status=nf90_inq_dimid(ncid,"depression",dimid) ! Find depression dimension
    status=nf90_inquire_dimension(ncid,dimid,len=deplen) ! Find depression length
    write(*,*) "deplen                 =",deplen
    ! lake_dis
    status=nf90_inq_dimid(ncid,"lake",dimid) ! Find vol2cell dimension
    status=nf90_inquire_dimension(ncid,dimid,len=lake_dislen) ! Find vol2cell length
    write(*,*) "lake_dislen            =",lake_dislen
    ! 1.3 Allocate dimensions and read variables
    ! 1.3.1 Latitudes and Longitudes
    ! Longitude
    allocate(longitude(lonlen),stat=status) ! Allocate longitude()
    status=nf90_inq_varid(ncid,"longitude",varid) ! Find longitude varid
    status=nf90_get_var(ncid,varid,longitude)! read longitude
    write(*,*) 'lonlen                 =',lonlen
    write(*,*) 'longitude(1)           =',longitude(1)
    write(*,*) 'longitude(lonlen)      =',longitude(lonlen)
    ! Latitude
    allocate(latitude(latlen),stat=status) ! Allocate latitude()
    status=nf90_inq_varid(ncid,"latitude",varid) ! Find latitude varid
    status=nf90_get_var(ncid,varid,latitude) ! Read latitude
    write(*,*) 'latlen                 =',latlen
    write(*,*) 'latitude(1)            =',latitude(1)
    write(*,*) 'latitude(2)            =',latitude(2)
    write(*,*) 'latitude(latlen)       =',latitude(latlen)
    ! 1.3.2 Graph table
    ! Child
    allocate(child_list(deplen),stat=status) ! Allocate child()
    status=nf90_inq_varid(ncid,"child",varid) ! Find child varid
    status=nf90_get_var(ncid,varid,child_list) ! Load and read child
    write(*,*) 'child_list(1)          =',child_list(1)
    write(*,*) 'child_list(deplen)     =',child_list(deplen)
    ! Parent
    allocate(parent_list(deplen),stat=status) ! Allocate parent()
    status=nf90_inq_varid(ncid,"parent",varid) ! Find parent varid
    status=nf90_get_var(ncid,varid,parent_list) ! Load and read parent
    write(*,*) 'parent_list(1)         =',parent_list(1)
    write(*,*) 'parent_list(deplen)    =',parent_list(deplen)
    ! lchild
    allocate(lchild_list(deplen),stat=status) ! Allocate lchild()
    status=nf90_inq_varid(ncid,"lchild",varid) ! Find lchild varid
    status=nf90_get_var(ncid,varid,lchild_list) ! Load and read lchild
    write(*,*) 'lchild_list(1)         =',lchild_list(1)
    write(*,*) 'lchild_list(deplen)    =',lchild_list(deplen)
    ! rchild
    allocate(rchild_list(deplen),stat=status) ! Allocate rchild()
    status=nf90_inq_varid(ncid,"rchild",varid) ! Find rchild varid
    status=nf90_get_var(ncid,varid,rchild_list) ! Load and read rchild
    write(*,*) 'rchild_list(1)         =',rchild_list(1)
    write(*,*) 'rchild_list(deplen)    =',rchild_list(deplen)
    ! Vmax
    allocate(Vmax_list(deplen),stat=status) ! Allocate Vmax()
    status=nf90_inq_varid(ncid,"Vmax",varid) ! Find Vmax varid
    status=nf90_get_var(ncid,varid,Vmax_list) ! Load and read Vmax
    write(*,*) 'Vmax_list(1)           =',Vmax_list(1)
    write(*,*) 'Vmax_list(deplen)      =',Vmax_list(deplen)
    ! Vmin
    allocate(Vmin_list(deplen),stat=status) ! Allocate Vmax()
    status=nf90_inq_varid(ncid,"Vmin",varid) ! Find Vmax varid
    status=nf90_get_var(ncid,varid,Vmin_list) ! Load and read Vmax
    write(*,*) 'Vmin_list(1)           =',Vmin_list(1)
    write(*,*) 'Vmin_list(deplen)      =',Vmin_list(deplen)
    ! pit_cell
    allocate(pit_cell_list(deplen),stat=status) ! Allocate pit_cell()
    status=nf90_inq_varid(ncid,"pit_cell",varid) ! Find pit_cell varid
    status=nf90_get_var(ncid,varid,pit_cell_list) ! Load and read pit_cell
    write(*,*) 'pit_cell_list(1)       =',pit_cell_list(1)
    write(*,*) 'pit_cell_list(deplen)  =',pit_cell_list(deplen)
    ! out_cell
    allocate(out_cell_list(deplen),stat=status) ! Allocate out_cell()
    status=nf90_inq_varid(ncid,"out_cell",varid) ! Find out_cell varid
    status=nf90_get_var(ncid,varid,out_cell_list) ! Load and pread out_cell
    write(*,*) 'out_cell_list(1)       =',out_cell_list(1)
    write(*,*) 'out_cell_list(deplen)  =',out_cell_list(deplen)
    ! lake_long
    allocate(lake_long_list(deplen),stat=status) ! Allocate lake_long()
    status=nf90_inq_varid(ncid,"lake_lon",varid) ! Find lake_long varid
    status=nf90_get_var(ncid,varid,lake_long_list) ! Load and read lake_long
    write(*,*) 'lake_long_list(1)      =',lake_long_list(1)
    write(*,*) 'lake_long_list(deplen) =',lake_long_list(deplen)
    ! lake_lat
    allocate(lake_lat_list(deplen),stat=status) ! Allocate lake_lat()
    status=nf90_inq_varid(ncid,"lake_lat",varid) ! Find lake_lat varid
    status=nf90_get_var(ncid,varid,lake_lat_list) ! Load and read lake_lat
    write(*,*) 'lake_lat_list(1)       =',lake_lat_list(1)
    write(*,*) 'lake_lat_list(deplen)  =',lake_lat_list(deplen)
    ! Zmin
    allocate(Zmin_list(deplen),stat=status) ! Allocate Zmin()
    status=nf90_inq_varid(ncid,"Zmin",varid) ! Find Zmin varid
    status=nf90_get_var(ncid,varid,Zmin_list) ! Load and read Zmin
    write(*,*) 'Zmin_list(1)           =',Zmin_list(1)
    write(*,*) 'Zmin_list(deplen)      =',Zmin_list(deplen)
    ! Zmax
    allocate(Zmax_list(deplen),stat=status) ! Allocate Zmax()
    status=nf90_inq_varid(ncid,"Zmax",varid) ! Find Zmax varid
    status=nf90_get_var(ncid,varid,Zmax_list) ! Load and read Zmax
    write(*,*) 'Zmax_list(1)           =',Zmax_list(1)
    write(*,*) 'Zmax_list(deplen)      =',Zmax_list(deplen)
    ! sibling
    allocate(sibling_list(deplen),stat=status) ! Allocate sibling()
    status=nf90_inq_varid(ncid,"sibling",varid) ! Find sibling varid
    status=nf90_get_var(ncid,varid,sibling_list) ! Load and read sibling
    write(*,*) 'sibling_list(1)        =',sibling_list(1)
    write(*,*) 'sibling_list(deplen)   =',sibling_list(deplen)
    ! geolink
    allocate(geolink_list(deplen),stat=status) ! Allocate geolink()
    status=nf90_inq_varid(ncid,"geolink",varid) ! Find geolink varid
    status=nf90_get_var(ncid,varid,geolink_list) ! Load and read geolink
    write(*,*) 'geolink_list(1)        =',geolink_list(1)
    write(*,*) 'geolink_list(deplen)   =',geolink_list(deplen)
    ! watershed_area
    allocate(watershed_area_list(deplen),stat=status) ! Allocate watershed_area()
    status=nf90_inq_varid(ncid,"watershed_area",varid) ! Find watershed_area varid
    status=nf90_get_var(ncid,varid,watershed_area_list) ! Load and read watershed_area
    write(*,*) 'watershed_area_list(1) =',watershed_area_list(1)
    ! 1.3.3 Maps informations
    ! ID Label of depression
    allocate(label(lonlen,latlen),stat=status) ! Allocate label()
    status=nf90_inq_varid(ncid,"label",varid) ! Find label varid
    status=nf90_get_var(ncid,varid,label) ! Load and read label
    write(*,*) 'shape(label)           =',shape(label)

   ! 1.3.4 Index to zoom on leaf into maps informations
    ! Imin
    allocate(imin_list(deplen),stat=status) ! Allocate imin
    status=nf90_inq_varid(ncid,"imin",varid) ! Find imin varid
    status=nf90_get_var(ncid,varid,imin_list) ! Load and read imin
    write(*,*) 'imin_list(1)           =',imin_list(1)
    write(*,*) 'imin_list(deplen)      =',imin_list(deplen)
    ! Imax
    allocate(imax_list(deplen),stat=status) ! Allocate imax
    status=nf90_inq_varid(ncid,"imax",varid) ! Find imax varid
    status=nf90_get_var(ncid,varid,imax_list) ! Load and read imax
    write(*,*) 'imax_list(1)           =',imax_list(1)
    write(*,*) 'imax_list(deplen)      =',imax_list(deplen)
    ! Jmin
    allocate(jmin_list(deplen),stat=status) ! Allocate jmin
    status=nf90_inq_varid(ncid,"jmin",varid) ! Find jmin varid
    status=nf90_get_var(ncid,varid,jmin_list) ! Load and read jmin
    write(*,*) 'jmin_list(1)           =',jmin_list(1)
    write(*,*) 'jmin_list(deplen)      =',jmin_list(deplen)
    ! Jmax
    allocate(jmax_list(deplen),stat=status) ! Allocate jmax
    status=nf90_inq_varid(ncid,"jmax",varid) ! Find jmax varid
    status=nf90_get_var(ncid,varid,jmax_list) ! Load and read jmax
    write(*,*) 'jmax_list(1)           =',jmax_list(1)
    write(*,*) 'jmax_list(deplen)      =',jmax_list(deplen)
    ! 1.3.5 Read and load function between water volume and water surface elevation
    ! Lake volume discretisation
    allocate(lake_perc(lake_dislen),stat=status) ! Allocate lake_perc()
    status=nf90_inq_varid(ncid,"lake_perc",varid) ! Find lake_perc varid
    status=nf90_get_var(ncid,varid,lake_perc) ! Load and read lake_perc
    write(*,*) 'shape(lake_perc)       =',shape(lake_perc)
    ! Lake elevation
    allocate(lake_elev(lake_dislen, deplen),stat=status) ! Allocate lake_elev()
    status=nf90_inq_varid(ncid,"lake_elev",varid) ! Find lake_elev varid
    status=nf90_get_var(ncid,varid,lake_elev) ! Load and read lake_elev
    write(*,*) 'shape(lake_elev)       =',shape(lake_elev)
    ! Lake volume
    allocate(lake_vol(lake_dislen, deplen),stat=status) ! Allocate lake_vol()
    status=nf90_inq_varid(ncid,"lake_vol",varid) ! Find lake_vol varid
    status=nf90_get_var(ncid,varid,lake_vol) ! Load and read lake_vol
    write(*,*) 'shape(lake_vol)        =',shape(lake_vol)
    ! lake surface
    allocate(lake_area(lake_dislen, deplen),stat=status) ! Allocate lake_area()
    status=nf90_inq_varid(ncid,"lake_area",varid) ! Find lake_area varid
    status=nf90_get_var(ncid,varid,lake_area) ! Load and read lake_area
    write(*,*) 'shape(lake_area)       =',shape(lake_area)

    ! Nleaf
    nleaf = maxval(label)
    write(*,*) 'nleaf                  =',nleaf
    ! Size_leaf
    size_leaf = sum(merge(1,0,child_list<=nleaf))
    write(*,*) 'size_leaf              =',size_leaf
    ! Ndep
    ndep = size(child_list)
    write(*,*) 'ndep                   =',ndep
    ! IDdep
    IDdep = maxval(child_list)+1

    allocate(Amax_list(deplen),stat=status)
    Amax_list = lake_area(lake_dislen, :)

    ! 2. Close NetCDF file
    status = nf90_close(ncid)

    !planet area
    planet_area = 0.0_real128
    do i = 1, ndep
      if (child_list(i) <= nleaf) then
        planet_area = planet_area + watershed_area_list(i)
      end if
    end do
    write(*,*) 'planet_area            =',planet_area

    ! Cell area
    allocate(cell_area(latlen),stat=status)
    call get_cells_area(latitude, longitude, cell_area)
    !write(*,*) 'cell_area              =', sum(cell_area*lonlen)
  end subroutine load_hydro_database

  subroutine get_cells_area(latitude, longitude, area)
    implicit none

    ! Inputs
    real, intent(in)  :: latitude(:)
    real, intent(in)  :: longitude(:)

    ! Output
    real(real128), intent(inout) :: area(:)

    ! Local variables
    real(real128) :: delta_lat, delta_lon
    real(real128) :: R
    real(real128) :: pi

    ! Grid spacing (degrees)
    delta_lat = abs(latitude(1)  - latitude(2))
    delta_lon = abs(longitude(1) - longitude(2))

    ! Radius of Mars (km)
    R  = 3389.5
    pi = acos(-1.0)
    ! Cell area computation (m²)
    area = (pi * R * 1000.0_real128 / 180.0_real128 )**2 * &
                         delta_lat * delta_lon * &
                         cos(latitude * pi / 180.0_real128 )
end subroutine get_cells_area

  subroutine load_ID_map(imin, imax, jmin, jmax, dep_label, dep_elev, nx, ny)
    implicit none
    ! NetCDF stuff
    integer :: status ! NetCDF return code
    integer :: ncid   ! NetCDF file ID
    integer :: varid  ! NetCDF variable ID
    integer, intent(out) :: nx, ny ! map dimensions
    integer, dimension(2) :: start, count
    integer, intent(in) :: imin, imax, jmin, jmax
    integer, allocatable :: dep_label(:,:)
    integer, allocatable, intent(out) :: dep_elev(:,:)

    nx = imax - imin + 1
    ny = jmax - jmin + 1

    start = (/imin, jmin/)
    count = (/nx, ny/)

    ! 1. Load input datasets
    ! 1.1. Open input datafile
    status=nf90_open(infile,NF90_NOWRITE,ncid)
    ! 1.2 Read elevation and label variable
    allocate(dep_label(nx,ny),stat=status) ! Allocate label()
    status=nf90_inq_varid(ncid,"label",varid) ! Find label varid
    status=nf90_get_var(ncid,varid,dep_label,start=start,count=count) ! Load and read label
    allocate(dep_elev(nx,ny),stat=status) ! Allocate elevation()
    status=nf90_inq_varid(ncid,"elevation",varid) ! Find elevation varid
    status=nf90_get_var(ncid,varid,dep_elev,start=start,count=count) ! Load and read elevation
    ! 2. Close NetCDF file
    status = nf90_close(ncid)

    !label(:, :) = label(:, nx:1:-1)

  end subroutine load_ID_map

  subroutine clear_database_variables()
    implicit none
    if (allocated(lchild_list)) deallocate(lchild_list)
    if (allocated(rchild_list)) deallocate(rchild_list)
    if (allocated(parent_list)) deallocate(parent_list)
    if (allocated(sibling_list)) deallocate(sibling_list)
    if (allocated(geolink_list)) deallocate(geolink_list)
    if (allocated(imin_list)) deallocate(imin_list)
    if (allocated(imax_list)) deallocate(imax_list)
    if (allocated(jmin_list)) deallocate(jmin_list)
    if (allocated(jmax_list)) deallocate(jmax_list)
    !if (allocated(Vmax_list)) deallocate(Vmax_list)
    if (allocated(Vmin_list)) deallocate(Vmin_list)
    if (allocated(lake_lat_list)) deallocate(lake_lat_list)
    if (allocated(lake_long_list)) deallocate(lake_long_list)
    if (allocated(Zmin_list)) deallocate(Zmin_list)
    if (allocated(Zmax_list)) deallocate(Zmax_list)
    if (allocated(Amax_list)) deallocate(Amax_list)
    if (allocated(watershed_area_list)) deallocate(watershed_area_list)
    if (allocated(out_cell_list)) deallocate(out_cell_list)
    if (allocated(pit_cell_list)) deallocate(pit_cell_list)
    if (allocated(lake_perc)) deallocate(lake_perc)
    if (allocated(lake_area)) deallocate(lake_area)
    if (allocated(lake_elev)) deallocate(lake_elev)
    if (allocated(lake_vol)) deallocate(lake_vol)
  end subroutine clear_database_variables

  function load_diaghydro(init_file, deplen, Z_list, V_list, A_list, Vout_list, Vin_list, active_list) result(status)
    character(len=*), intent(in) :: init_file
    integer, intent(in) :: deplen
    real(real128), intent(inout) :: Z_list(deplen), V_list(deplen), A_list(deplen), Vout_list(deplen), Vin_list(deplen)
    integer, intent(inout) :: active_list(deplen)
    double precision :: Z(deplen), V(deplen), A(deplen), Vout(deplen), Vin(deplen)

    ! NetCDF stuff
    integer :: status ! NetCDF return code
    integer :: ncid   ! NetCDF file ID
    integer :: varid  ! NetCDF variable ID

    ! 1. Load input datasets
    write(*,*) "###################################"
    write(*,*) "###### Load Diaghydro File ########"
    write(*,*) "###################################"
    ! 1.1. Open input datafile
    status=nf90_open(init_file,NF90_NOWRITE,ncid)
    write(*,*) "Reading input file: ",trim(init_file)

    write(*,*) "deplen=",deplen

    status=nf90_inq_varid(ncid,"water_volume",varid)
    status=nf90_get_var(ncid,varid,V)
    V_list = V

    status=nf90_inq_varid(ncid,"water_elevation",varid)
    status=nf90_get_var(ncid,varid,Z)
    Z_list = Z

    status=nf90_inq_varid(ncid,"water_area",varid)
    status=nf90_get_var(ncid,varid,A)
    A_list = A

    status=nf90_inq_varid(ncid,"water_active",varid)
    status=nf90_get_var(ncid,varid,active_list)
    write(*,*) "Active depressions loaded", sum(active_list)

    status=nf90_inq_varid(ncid,"water_outflow",varid)
    status=nf90_get_var(ncid,varid,Vout)
    Vout_list = Vout

    status=nf90_inq_varid(ncid,"water_inflow",varid)
    status=nf90_get_var(ncid,varid,Vin)
    Vin_list = Vin
  end function load_diaghydro

  function load_diaghydro_iteration(init_file, deplen, Z_list, V_list, A_list, Vout_list, Vin_list, active_list) result(status)
    !use :: TreeNodeModule
    character(len=*), intent(in) :: init_file
    integer, intent(in) :: deplen
    real(real128), intent(inout) :: Z_list(deplen), V_list(deplen), A_list(deplen), Vout_list(deplen), Vin_list(deplen)
    integer, intent(out) :: active_list(deplen)
    integer :: status

    ! NetCDF stuff
    integer :: ncid, varid_count, varid_start, varid_ids, varid_idx
    integer :: varid_V, varid_Z, varid_A, varid_Vout, varid_Vin
    integer ::  last_idx, start_index, count, i
    integer :: dimid_time, len_time
    integer, allocatable :: depression_ids(:), depression_idx(:)
    double precision, allocatable :: V_file(:), Z_file(:), A_file(:)
    double precision, allocatable :: Vout_file(:), Vin_file(:)
    integer, allocatable :: start_index_arr(:), count_arr(:)
    integer :: child
    !type(TreeNode), intent(inout) :: tree(IDdep)


    ! 1. Load input datasets
    write(*,*) "###################################"
    write(*,*) "###### Load Diaghydro File ########"
    write(*,*) "###################################"

    ! 1.1. Open input datafile
    status=nf90_open(init_file,NF90_NOWRITE,ncid)
    write(*,*) "Reading input file: ",trim(init_file)

    write(*,*) "deplen=",deplen

    active_list = 0

    ! Find the number of time steps (iterations)
    status = nf90_inq_dimid(ncid, "time", dimid_time)
    status = nf90_inquire_dimension(ncid, dimid_time, len=len_time)
    last_idx = len_time
    write(*,*) "last_idx=", last_idx

    ! Get variable IDs for start_index and count
    allocate(start_index_arr(len_time), count_arr(len_time))
    status = nf90_inq_varid(ncid, "start_index", varid_start)
    status = nf90_inq_varid(ncid, "count", varid_count)
    status = nf90_get_var(ncid, varid_start, start_index_arr)
    status = nf90_get_var(ncid, varid_count, count_arr)
    start_index = start_index_arr(last_idx)
    count = count_arr(last_idx)
    deallocate(start_index_arr, count_arr)

    write(*,*) "start_index=", start_index

    ! Get variable IDs for the data

    allocate(V_file(count), Z_file(count), A_file(count), Vout_file(count), Vin_file(count), depression_ids(count), depression_idx(count))

    status = nf90_inq_varid(ncid, "depression_ids", varid_ids)
    status = nf90_get_var(ncid, varid_ids, depression_ids, start=[start_index], count=[count])
    write(*,*) "depression_ids loaded"

    status = nf90_inq_varid(ncid, "depression_idx", varid_idx)
    status = nf90_get_var(ncid, varid_idx, depression_idx, start=[start_index], count=[count])
    write(*,*) "depression_idx loaded"

    status = nf90_inq_varid(ncid, "water_volume", varid_V)
    status = nf90_get_var(ncid, varid_V, V_file, start=[start_index], count=[count])
    write(*,*) "water_volume loaded"

    status = nf90_inq_varid(ncid, "water_elevation", varid_Z)
    status = nf90_get_var(ncid, varid_Z, Z_file, start=[start_index], count=[count])
    write(*,*) "water_elevation loaded"

    status = nf90_inq_varid(ncid, "water_area", varid_A)
    status = nf90_get_var(ncid, varid_A, A_file, start=[start_index], count=[count])
    write(*,*) "water_area loaded"

    status = nf90_inq_varid(ncid, "water_outflow", varid_Vout)
    status = nf90_get_var(ncid, varid_Vout, Vout_file, start=[start_index], count=[count])
    write(*,*) "water_outflow loaded"

    status = nf90_inq_varid(ncid, "water_inflow", varid_Vin)
    status = nf90_get_var(ncid, varid_Vin, Vin_file, start=[start_index], count=[count])
    write(*,*) "water_inflow loaded"

    do i = 1, count
      child = depression_idx(i)

      V_list(child) = V_file(i)
      Z_list(child) = Z_file(i)
      A_list(child) = A_file(i)
      Vout_list(child) = Vout_file(i)
      Vin_list(child) = Vin_file(i)
      active_list(child) = 1
      write(*,*) i/count *100, "%"
    end do
  end function load_diaghydro_iteration


  function load_diaghydro_active(init_file, deplen, Z_list, V_list, A_list, Vout_list, Vin_list, active_list) result(status)
    !use :: TreeNodeModule
    character(len=*), intent(in) :: init_file
    integer, intent(in) :: deplen
    real(real128), intent(inout) :: Z_list(deplen), V_list(deplen), A_list(deplen), Vout_list(deplen), Vin_list(deplen)
    integer, intent(out) :: active_list(deplen)
    integer :: status

    ! NetCDF stuff
    integer :: ncid, varid_ids, varid_idx
    integer :: varid_V, varid_Z, varid_A, varid_Vout, varid_Vin
    integer :: dimid_time, i
    integer, allocatable :: depression_ids(:), depression_idx(:)
    double precision, allocatable :: V_file(:), Z_file(:), A_file(:)
    double precision, allocatable :: Vout_file(:), Vin_file(:)
    integer :: child, dep
    !type(TreeNode), intent(inout) :: tree(IDdep)

    ! 1.1. Open input datafile
    status=nf90_open(init_file,NF90_NOWRITE,ncid)
    write(*,*) "Initialization from file: ",trim(init_file)
    write(*,*) "deplen=",deplen

    active_list = 0

    ! Find the number of time steps (iterations)
    status = nf90_inq_dimid(ncid, "depression", dimid_time)
    status = nf90_inquire_dimension(ncid, dimid_time, len=dep)
    write(*,*) "depression=", dep

    ! Get variable IDs for the data
    allocate(V_file(dep), Z_file(dep), A_file(dep), Vout_file(dep), Vin_file(dep), depression_ids(dep), depression_idx(dep))

    status = nf90_inq_varid(ncid, "depression_ids", varid_ids)
    status = nf90_get_var(ncid, varid_ids, depression_ids)

    status = nf90_inq_varid(ncid, "depression_idx", varid_idx)
    status = nf90_get_var(ncid, varid_idx, depression_idx)

    status = nf90_inq_varid(ncid, "water_volume", varid_V)
    status = nf90_get_var(ncid, varid_V, V_file)

    status = nf90_inq_varid(ncid, "water_elevation", varid_Z)
    status = nf90_get_var(ncid, varid_Z, Z_file)

    status = nf90_inq_varid(ncid, "water_area", varid_A)
    status = nf90_get_var(ncid, varid_A, A_file)

    status = nf90_inq_varid(ncid, "water_outflow", varid_Vout)
    status = nf90_get_var(ncid, varid_Vout, Vout_file)

    status = nf90_inq_varid(ncid, "water_inflow", varid_Vin)
    status = nf90_get_var(ncid, varid_Vin, Vin_file)

    do i = 1, dep
      child = depression_idx(i)
      V_list(child) = V_file(i)
      Z_list(child) = Z_file(i)
      A_list(child) = A_file(i)
      Vout_list(child) = Vout_file(i)
      Vin_list(child) = Vin_file(i)
      active_list(child) = 1
    end do
  end function load_diaghydro_active

end module hydro_database