! -----------------------------------------------------------------------------
! 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_diaghydro
  use netcdf
  use hydro_database, only: ndep, deplen, child_list
  use, intrinsic ::iso_fortran_env, only : real64, real128
  implicit none
  !=============================================================================
  ! Diaghydro file (OUTPUT)
  !=============================================================================
  real(real128), allocatable :: Z_list(:), V_list(:), A_list(:), Vout_list(:), Vin_list(:)
  real(real128), allocatable :: Vout_list_prev(:), Vout_relative_error(:), V_list_prev(:), V_relative_error(:), V_list_new(:), V_list_old(:)
  real(real128), allocatable :: A_list_prev(:), A_relative_error(:)
  real(real128), allocatable :: Z_list_prev(:), Z_relative_error(:)
  real(real128), allocatable :: Vevap_list(:), Vinj_list(:)
  real(real128), allocatable :: Qin(:), Qout(:)
  integer, allocatable :: active_list(:)

  double precision, allocatable, public :: Z_sub(:), V_sub(:), A_sub(:), Vout_sub(:), Vin_sub(:), selected_ids_sub(:)
  integer, allocatable, public          :: active_sub(:)

  double precision, allocatable, public :: Z_out(:), V_out(:), A_out(:), Vout_out(:), Vin_out(:), ids_out(:)
  integer, allocatable, public :: active_out(:), idx_out(:)

  contains
  subroutine allocate_outputs(deplen)
    implicit none
    integer, intent(in) :: deplen

    ! Allocate output variables
    allocate(Z_list(deplen), V_list(deplen), A_list(deplen), Vout_list(deplen), Vin_list(deplen))
    allocate(Vout_list_prev(deplen), Vout_relative_error(deplen), V_list_prev(deplen), V_relative_error(deplen))
    allocate(V_list_new(deplen), V_list_old(deplen))
    allocate(A_list_prev(deplen), A_relative_error(deplen))
    allocate(Z_list_prev(deplen), Z_relative_error(deplen))
    allocate(Vevap_list(deplen), Vinj_list(deplen))
    allocate(Qin(deplen), Qout(deplen))
    allocate(active_list(deplen))

    ! Initialize output variables
    Z_list = 0.0
    V_list = 0.0
    A_list = 0.0
    Vout_list = 0.0
    Vin_list = 0.0

  end subroutine allocate_outputs

  function save_output_model(outfile) result(status)
    implicit none
    character(len=*), intent(in) :: outfile

    ! NetCDF stuff
    integer :: status ! NetCDF return code
    integer :: ncid ! NetCDF file ID
    integer :: dimid_dep
    integer :: varid_V, varid_Z, varid_A, varid_active, varid_Vout, varid_Vin

    status = nf90_create(outfile, NF90_CLOBBER, ncid)
    status = nf90_def_dim(ncid, "depression", deplen, dimid_dep)
    status = nf90_def_var(ncid, "water_volume", NF90_FLOAT, dimid_dep, varid_V)
    status = nf90_def_var(ncid, "water_elevation", NF90_FLOAT, dimid_dep, varid_Z)
    status = nf90_def_var(ncid, "water_area", NF90_FLOAT, dimid_dep, varid_A)
    status = nf90_def_var(ncid, "water_active", NF90_INT, dimid_dep, varid_active)
    status = nf90_def_var(ncid, "water_outflow", NF90_DOUBLE, dimid_dep, varid_Vout)
    status = nf90_def_var(ncid, "water_inflow", NF90_DOUBLE, dimid_dep, varid_Vin)
    status = nf90_enddef(ncid)
    status = nf90_put_var(ncid, varid_V, real(V_list,8))
    status = nf90_put_var(ncid, varid_Z, real(Z_list,8))
    status = nf90_put_var(ncid, varid_A, real(A_list,8))
    status = nf90_put_var(ncid, varid_active, real(active_list,8))
    status = nf90_put_var(ncid, varid_Vout, real(Vout_list,8))
    status = nf90_put_var(ncid, varid_Vin, real(Vin_list,8))
    status = nf90_close(ncid)

    write(*,*) "Output saved ! :", outfile
  end function save_output_model

  function save_output_subregion(outfile, Z_list, V_list, A_list, Vout_list, Vin_list, &
    active_list) result(status)
    implicit none

    character(len=*), intent(in)  :: outfile
    double precision, intent(in)  :: Z_list(:), V_list(:), A_list(:), Vout_list(:), Vin_list(:)
    integer, intent(in)           :: active_list(:)
    integer                       :: status
    integer                       :: ncid, dimid_dep
    integer                       :: varid_V, varid_Z, varid_A, varid_active, varid_Vout, varid_Vin
    integer                       :: i
    integer                       :: varid_ids, nsel, counter

    nsel = count(active_list == 1)
    counter = 1

    do i = 1, ndep
      if(active_list(i) == 1) then
        Z_sub(counter) = Z_list(i)
        V_sub(counter) = V_list(i)
        A_sub(counter) = A_list(i)
        Vout_sub(counter) = Vout_list(i)
        Vin_sub(counter) = Vin_list(i)
        active_sub(counter) = active_list(i)
        selected_ids_sub(counter) = child_list(i)
        counter = counter + 1
      end if
    end do

    status = nf90_create(outfile, NF90_CLOBBER, ncid)
    status = nf90_def_dim(ncid, "depression", nsel, dimid_dep)
    status = nf90_def_var(ncid, "water_volume", NF90_FLOAT, dimid_dep, varid_V)
    status = nf90_def_var(ncid, "water_elevation", NF90_FLOAT, dimid_dep, varid_Z)
    status = nf90_def_var(ncid, "water_area", NF90_FLOAT, dimid_dep, varid_A)
    status = nf90_def_var(ncid, "water_active", NF90_INT, dimid_dep, varid_active)
    status = nf90_def_var(ncid, "water_outflow", NF90_DOUBLE, dimid_dep, varid_Vout)
    status = nf90_def_var(ncid, "water_inflow", NF90_DOUBLE, dimid_dep, varid_Vin)
    status = nf90_def_var(ncid, "selected_ids", NF90_INT, dimid_dep, varid_ids)
    status = nf90_enddef(ncid)
    status = nf90_put_var(ncid, varid_V, V_sub(1:nsel))
    status = nf90_put_var(ncid, varid_Z, Z_sub(1:nsel))
    status = nf90_put_var(ncid, varid_A, A_sub(1:nsel))
    status = nf90_put_var(ncid, varid_active, active_sub(1:nsel))
    status = nf90_put_var(ncid, varid_Vout, Vout_sub(1:nsel))
    status = nf90_put_var(ncid, varid_Vin, Vin_sub(1:nsel))
    status = nf90_put_var(ncid, varid_ids, selected_ids_sub(1:nsel))
    status = nf90_close(ncid)

    write(*,*) "Output subregion file is saved ! :", outfile
  end function save_output_subregion

  function save_output_active(outfile) result(status)
    implicit none
    character(len=*), intent(in) :: outfile
    integer :: status
    integer :: ncid, dimid_dep
    integer :: varid_V, varid_Z, varid_A, varid_Vout, varid_Vin
    integer :: varid_ids, varid_idx, nsel, counter, i

    nsel = count(active_list == 1)
    allocate(Z_out(nsel), V_out(nsel), A_out(nsel), Vout_out(nsel), Vin_out(nsel), ids_out(nsel), idx_out(nsel))
    counter  = 1
    do i = 1, ndep
      if (active_list(i) == 1) then
        Z_out(counter) = real(Z_list(i),8)
        V_out(counter) = real(V_list(i),8)
        A_out(counter) = real(A_list(i),8)
        Vout_out(counter) = real(Vout_list(i),8)
        Vin_out(counter) = real(Vin_list(i),8)
        ids_out(counter) = child_list(i)
        idx_out(counter) = i
        counter = counter + 1
      end if
    end do

    status = nf90_create(outfile, NF90_CLOBBER, ncid)
    status = nf90_def_dim(ncid, "depression", nsel, dimid_dep)
    status = nf90_def_var(ncid, "water_volume", NF90_DOUBLE, dimid_dep, varid_V)
    status = nf90_def_var(ncid, "water_elevation", NF90_DOUBLE, dimid_dep, varid_Z)
    status = nf90_def_var(ncid, "water_area", NF90_DOUBLE, dimid_dep, varid_A)
    status = nf90_def_var(ncid, "water_outflow", NF90_DOUBLE, dimid_dep, varid_Vout)
    status = nf90_def_var(ncid, "water_inflow", NF90_DOUBLE, dimid_dep, varid_Vin)
    status = nf90_def_var(ncid, "depression_ids", NF90_INT, dimid_dep, varid_ids)
    status = nf90_def_var(ncid, "depression_idx", NF90_INT, dimid_dep, varid_idx)
    status = nf90_enddef(ncid)
    status = nf90_put_var(ncid, varid_V, V_out)
    status = nf90_put_var(ncid, varid_Z, Z_out)
    status = nf90_put_var(ncid, varid_A, A_out)
    status = nf90_put_var(ncid, varid_Vout, Vout_out)
    status = nf90_put_var(ncid, varid_Vin, Vin_out)
    status = nf90_put_var(ncid, varid_ids, ids_out)
    status = nf90_put_var(ncid, varid_idx, idx_out)
    status = nf90_close(ncid)
    deallocate(Z_out, V_out, A_out, Vout_out, Vin_out, ids_out, idx_out)

    write(*,*) "Output saved : ", trim(outfile)
  end function save_output_active

  function save_output_ragged_active_iteration(outfile, iter_idx, Z_list, &
    V_list, A_list, Vout_list, Vin_list, pos_record, it) result(status)
    implicit none
    character(len=*), intent(in) :: outfile
    integer, intent(in) :: iter_idx
    !integer, intent(in) :: depression_ids(nsel)
    double precision, intent(in) :: Z_list(:), V_list(:), A_list(:), Vout_list(:), Vin_list(:)
    integer, intent(inout) :: pos_record
    integer,intent(in) :: it
    integer, allocatable :: ids_out(:)
    integer :: i, counter, nsel

    ! NetCDF variables
    integer :: ncid, dim_time_id, dim_record_id
    integer :: var_ids, var_idx, var_Z, var_V, var_A, var_Vout, var_Vin
    integer :: var_start_index, var_count, var_iteration
    integer :: status
    logical :: file_exists

    ! Filter on active depressions
    nsel = count(active_list == 1)
    allocate(Z_out(nsel), V_out(nsel), A_out(nsel), Vout_out(nsel), Vin_out(nsel), ids_out(nsel), idx_out(nsel))
    counter  = 1
    do i = 1, ndep
      if (active_list(i) == 1) then
        Z_out(counter) = Z_list(i)
        V_out(counter) = V_list(i)
        A_out(counter) = A_list(i)
        Vout_out(counter) = Vout_list(i)
        Vin_out(counter) = Vin_list(i)
        ids_out(counter) = child_list(i)
        idx_out(counter) = i
        counter = counter + 1
      end if
    end do


    inquire(file=outfile, exist=file_exists)                                                      ! Check if the file exists
    if (.not. file_exists) then                                                                   ! if does not exist, create a new one
      status = nf90_create(outfile, IOR(NF90_NETCDF4, NF90_CLOBBER), ncid)
      status = nf90_def_dim(ncid, "time", NF90_UNLIMITED, dim_time_id)      ! Define unlimited dimension for time
      status = nf90_def_dim(ncid, "record", NF90_UNLIMITED, dim_record_id)                        ! Define unlimited dimension for records
      status = nf90_def_var(ncid, "depression_ids", NF90_INT, [dim_record_id], var_ids)           ! Define variable for depression ID
      status = nf90_def_var(ncid, "depression_idx", NF90_INT, [dim_record_id], var_idx)           ! Define variable for depression index
      status = nf90_def_var(ncid, "water_volume", NF90_DOUBLE, [dim_record_id], var_V)            ! Define variable for water volume
      status = nf90_def_var(ncid, "water_elevation", NF90_DOUBLE, [dim_record_id], var_Z)         ! Define variable for water elevation
      status = nf90_def_var(ncid, "water_area", NF90_DOUBLE, [dim_record_id], var_A)              ! Define variable for water area
      status = nf90_def_var(ncid, "water_outflow", NF90_DOUBLE, [dim_record_id], var_Vout)        ! Define variable for water outflow
      status = nf90_def_var(ncid, "water_inflow", NF90_DOUBLE, [dim_record_id], var_Vin)          ! Define variable for water inflow
      status = nf90_def_var(ncid, "start_index", NF90_INT, [dim_time_id], var_start_index)        ! Define variable for start index
      status = nf90_def_var(ncid, "count", NF90_INT, [dim_time_id], var_count)                    ! Define variable for count of depressions in this iteration
      status = nf90_def_var(ncid, "real_iteration", NF90_INT, [dim_time_id], var_iteration)            ! Define variable for iteration index
      status = nf90_enddef(ncid)
    else
      status = nf90_open(outfile, NF90_WRITE, ncid)                                               ! Open existing NetCDF file
      status = nf90_inq_varid(ncid, "depression_ids", var_ids)
      status = nf90_inq_varid(ncid, "depression_idx", var_idx)
      status = nf90_inq_varid(ncid, "water_volume", var_V)
      status = nf90_inq_varid(ncid, "water_elevation", var_Z)
      status = nf90_inq_varid(ncid, "water_area", var_A)
      status = nf90_inq_varid(ncid, "water_outflow", var_Vout)
      status = nf90_inq_varid(ncid, "water_inflow", var_Vin)
      status = nf90_inq_varid(ncid, "start_index", var_start_index)
      status = nf90_inq_varid(ncid, "count", var_count)
      status = nf90_inq_varid(ncid, "real_iteration", var_iteration)                                 ! Get the variable ID for iteration index
    end if

    ! Write data for this iteration
    status = nf90_put_var(ncid, var_ids, ids_out, start=[pos_record+1], count=[nsel]) ! Write depression IDs
    status = nf90_put_var(ncid, var_idx, idx_out, start=[pos_record+1], count=[nsel]) ! Write depression index
    status = nf90_put_var(ncid, var_V, V_out, start=[pos_record+1], count=[nsel])                      ! Write water volume
    status = nf90_put_var(ncid, var_Z, Z_out, start=[pos_record+1], count=[nsel])                      ! Write water elevation
    status = nf90_put_var(ncid, var_A, A_out, start=[pos_record+1], count=[nsel])                      ! Write water area
    status = nf90_put_var(ncid, var_Vout, Vout_out, start=[pos_record+1], count=[nsel])                ! Write water outflow
    status = nf90_put_var(ncid, var_Vin, Vin_out, start=[pos_record+1], count=[nsel])                  ! Write water inflow
    status = nf90_put_var(ncid, var_start_index, [pos_record+1], start=[iter_idx])                     ! Write start index for this iteration
    status = nf90_put_var(ncid, var_count, [nsel], start=[iter_idx])                                   ! Write count for this iteration
    status = nf90_put_var(ncid, var_iteration, [it], start=[iter_idx])                                 ! Write iteration index
    status = nf90_close(ncid)
    pos_record = pos_record + nsel         ! Update the position for the next record
    deallocate(Z_out, V_out, A_out, Vout_out, Vin_out, ids_out, idx_out)
    write(*,*) "Output ragged active iteration file is saved ! :", outfile
  end function save_output_ragged_active_iteration

  subroutine deallocate_outputs()
    implicit none
    if (allocated(Vevap_list)) deallocate(Vevap_list)
    if (allocated(Vinj_list)) deallocate(Vinj_list)
    if (allocated(Qin)) deallocate(Qin)
    if (allocated(Qout)) deallocate(Qout)
    if (allocated(Z_list)) deallocate(Z_list)
    if (allocated(V_list)) deallocate(V_list)
    if (allocated(A_list)) deallocate(A_list)
    if (allocated(Vout_list)) deallocate(Vout_list)
    if (allocated(Vin_list)) deallocate(Vin_list)
    if (allocated(V_list_prev)) deallocate(V_list_prev)
    if (allocated(V_relative_error)) deallocate(V_relative_error)
    if (allocated(Vout_list_prev)) deallocate(Vout_list_prev)
    if (allocated(Vout_relative_error)) deallocate(Vout_relative_error)
    if (allocated(A_list_prev)) deallocate(A_list_prev)
    if (allocated(A_relative_error)) deallocate(A_relative_error)
    if (allocated(Z_list_prev)) deallocate(Z_list_prev)
    if (allocated(Z_relative_error)) deallocate(Z_relative_error)
    if (allocated(active_list)) deallocate(active_list)
  end subroutine deallocate_outputs
end module hydro_diaghydro