MODULE output
!-----------------------------------------------------------------------
! NAME
!     output
!
! DESCRIPTION
!     Tools to write PEM diagnostic outputs.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! PARAMETERS
! ----------
character(12),                             parameter :: diagdef_name = 'diagevol.def'
integer(di),                               parameter :: dim_ngrid = -1        ! Dimension ID to identify 'ngrid' for "diagevol.nc" logic
logical(k4),                               protected :: is_diagdef = .false.  ! Flag to know if "diagevol.def" exists
logical(k4),                               protected :: is_diagevol = .false. ! Flag to know if "diagevol.nc" exists
character(256), dimension(:), allocatable, protected :: outputs_name          ! Names of output variables defined in "diagevol.def"
integer(di),                               protected :: n_outputs             ! Number of output variables defined in "diagevol.def"
integer(di),                               protected :: output_rate           ! Output rate

! VARIABLES
! ---------
integer(di) :: writing_idt = -1                                                  ! Current writing number of timestep
integer(di) :: itime                                                             ! Current time index to record variables
integer(di) :: dim_nlon, dim_nlat, dim_nalt, dim_interlayer, dim_nsoil, dim_time ! Dimensions ID to write in "diagevol.nc"

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

!=======================================================================
SUBROUTINE ini_output()
!-----------------------------------------------------------------------
! NAME
!     ini_output
!
! DESCRIPTION
!     Initialize the parameters of module 'output'.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!     At first call.
!-----------------------------------------------------------------------

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

! CODE
! ----
n_outputs = 0
if (.not. allocated(outputs_name)) allocate(outputs_name(0))
call read_diagdef()

END SUBROUTINE ini_output
!=======================================================================

!=======================================================================
SUBROUTINE end_output()
!-----------------------------------------------------------------------
! NAME
!     end_output
!
! DESCRIPTION
!     Deallocate output arrays.
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!
!-----------------------------------------------------------------------

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

! CODE
! ----
if (allocated(outputs_name)) deallocate(outputs_name)

END SUBROUTINE end_output
!=======================================================================

!=======================================================================
SUBROUTINE set_output_config(output_rate_in)
!-----------------------------------------------------------------------
! NAME
!     set_output_config
!
! DESCRIPTION
!     Setter for 'output' configuration parameters.
!
! AUTHORS & DATE
!     JB Clement, 02/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use utility, only: int2str
use display, only: print_msg

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

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

! CODE
! ----
output_rate = output_rate_in
call print_msg('output_rate = '//int2str(output_rate))

END SUBROUTINE set_output_config
!=======================================================================

!=======================================================================
SUBROUTINE append_name(var_name)
!-----------------------------------------------------------------------
! NAME
!     append_name
!
! DESCRIPTION
!     Append the variable name read in "diagevol.def" to 'outputs_name'.
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!
!-----------------------------------------------------------------------

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

! ARGUMENTS
! ---------
character(*), intent(in) :: var_name

! LOCAL VARIABLES
! ---------------
character(256), dimension(:), allocatable :: tmp

! CODE
! ----
allocate(tmp(n_outputs + 1))
if (n_outputs > 0) tmp(1:n_outputs) = outputs_name
tmp(n_outputs + 1) = var_name

call move_alloc(tmp,outputs_name)
n_outputs = n_outputs + 1

END SUBROUTINE append_name
!=======================================================================

!=======================================================================
SUBROUTINE read_diagdef()
!-----------------------------------------------------------------------
! NAME
!     read_diagdef
!
! DESCRIPTION
!     Read the variables names defined in "diagevol.def" to ouput.
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!     > If no "diagevol.def", then no output.
!     > If "diagevol.def" is empty, then all the possible coded outputs.
!     > If "diagevol.def" holds a list of variable names, then these are
!     outputted.
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use stoppage, only: stop_clean
use utility,  only: is_id_1st_char, is_id_char, int2str
use display, only: print_msg

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

! LOCAL VARIABLES
! ---------------
integer(di)    :: i, iline, ierr, funit
character(256) :: line

! CODE
! ----
inquire(file = diagdef_name,exist = is_diagdef)
if (is_diagdef) then
    call print_msg('> Reading "'//diagdef_name//'"')
    open(newunit = funit,file = diagdef_name,status = 'old',form = 'formatted',action = 'read',iostat = ierr)
    if (ierr /= 0) call stop_clean(__FILE__,__LINE__,'error opening file "'//diagdef_name//'"!',ierr)
    iline = 1
    do
        ! Read the line
        read(funit,'(a)',iostat = ierr) line
        if (ierr /= 0) exit ! Empty file or end of file
        iline = iline + 1

        ! Examine the line
        line = trim(line)
        if (len(line) == 0) cycle ! Ignore empty lines
        if (.not. is_id_1st_char(line(1:1))) call print_msg('Invalid output name at line '//int2str(iline)//' in '//diagdef_name//'"!')
        do i = 2,len(line)
            if (.not. is_id_char(line(i:i))) call print_msg('Invalid output name at line '//int2str(iline)//' in '//diagdef_name//'"!')
        end do

        ! Fill in the table 'outputs_name'
        call append_name(line)
    end do
    close(funit)
    if (n_outputs == 0) call print_msg('There is a "'//diagdef_name//'" which is is empty. So all the possible coded variables are outputted.')
    ! Create the output file
    call create_diagevol()
else
    call print_msg('There is no "'//diagdef_name//'". So no variable is outputted.')
end if

END SUBROUTINE read_diagdef
!=======================================================================

!=======================================================================
FUNCTION is_output_selected(var_name) RESULT(output_var)
!-----------------------------------------------------------------------
! NAME
!     is_output_selected
!
! DESCRIPTION
!     Read the variables names to output defined in "diagevol.def".
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!     > If no "diagevol.def", then no output.
!     > If "diagevol.def" is empty, then all the possible coded outputs.
!     > If "diagevol.def" holds a list of variable names, then these are
!     outputted.
!-----------------------------------------------------------------------

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

! ARGUMENTS
! ---------
character(*), intent(in) :: var_name
logical(k4)              :: output_var

! LOCAL VARIABLES
! ---------------
integer(di) :: i

! CODE
! ----
output_var = .false.

! If diagdef exists but is empty, output everything
if (n_outputs == 0) then
    output_var = .true.
    return
end if

do i = 1,n_outputs
    if (trim(outputs_name(i)) == trim(var_name)) then
        output_var = .true.
        return
    end if
end do

END FUNCTION is_output_selected
!=======================================================================

!=======================================================================
SUBROUTINE create_diagevol()
!-----------------------------------------------------------------------
! NAME
!     create_diagevol
!
! DESCRIPTION
!     Create a NetCDF file "diagevol.nc" containing the diagnostic
!     variables.
!
! 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 io_netcdf,  only: check_nc, open_nc, close_nc, put_var_nc, diagevol_name, ncid_diagevol
use geometry,   only: dim_init, ngrid, nlon, nlat, nlayer, nsoil, cell_area, longitudes, latitudes, vect2dyngrd
use stoppage,   only: stop_clean
use soil,       only: mlayer
use atmosphere, only: ap, bp
use display,    only: print_msg

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

! LOCAL VARIABLES
! ---------------
integer(di)                           :: varid ! Variable ID
integer(di)                           :: nlon_loc
real(dp), dimension(:,:), allocatable :: var_dyn

! CODE
! ----
! Check if dimensions are well initialized
if (.not. dim_init) call stop_clean(__FILE__,__LINE__,'dimensions are not initilized in the PEM!',1)
if (ngrid == 1) then ! 1D
    nlon_loc = 1
else ! 3D
    nlon_loc = nlon + 1
end if

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

! Define dimensions
call check_nc(nf90_def_dim(ncid_diagevol,'Time',NF90_UNLIMITED,dim_time),'defining dimension Time')
call check_nc(nf90_def_dim(ncid_diagevol,'longitude',nlon_loc,dim_nlon),'defining dimension longitude')
call check_nc(nf90_def_dim(ncid_diagevol,'latitude',nlat,dim_nlat),'defining dimension latitude')
call check_nc(nf90_def_dim(ncid_diagevol,'altitude',nlayer,dim_nalt),'defining dimension altitude')
call check_nc(nf90_def_dim(ncid_diagevol,'interlayer',nlayer + 1,dim_interlayer),'defining dimension interlayer')
call check_nc(nf90_def_dim(ncid_diagevol,'subsurface_layers',nsoil,dim_nsoil),'defining dimension subsurface_layers')

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

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

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

call check_nc(nf90_def_var(ncid_diagevol,'ap',nf90_double,(/dim_interlayer/),varid),'defining variable ap')
call check_nc(nf90_put_att(ncid_diagevol,varid,'title','Hybrid pressure at interlayers'),'putting title attribute for ap')

call check_nc(nf90_def_var(ncid_diagevol,'bp',nf90_double,(/dim_interlayer/),varid),'defining variable bp')
call check_nc(nf90_put_att(ncid_diagevol,varid,'title','Hybrid sigma at interlayers'),'putting title attribute for bp')

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

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

! Global attributes
call check_nc(nf90_put_att(ncid_diagevol,nf90_global,'title','PEM diagnostic file'),'putting global attribute')

! Leave define mode and putting variables defining dimensions
call check_nc(nf90_enddef(ncid_diagevol),'leaving define mode')
allocate(var_dyn(nlon_loc,nlat))
call vect2dyngrd(nlon + 1,nlat,ngrid,longitudes,var_dyn)
call put_var_nc('longitude',var_dyn)
call vect2dyngrd(nlon + 1,nlat,ngrid,latitudes,var_dyn)
call put_var_nc('latitude',var_dyn)
call put_var_nc('ap',ap)
call put_var_nc('bp',bp)
call put_var_nc('soildepth',mlayer)
call vect2dyngrd(nlon + 1,nlat,ngrid,cell_area,var_dyn)
call put_var_nc('cell_area',var_dyn)
deallocate(var_dyn)
call close_nc(diagevol_name)

! File creation done
is_diagevol = .true.

END SUBROUTINE create_diagevol
!=======================================================================

!=======================================================================
SUBROUTINE write_diagevol(var_name,title,units,var,dimids_in)
!-----------------------------------------------------------------------
! NAME
!     write_diagevol
!
! DESCRIPTION
!     Write selected diagnostic variables to NetCDF file "diagevol.nc".
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!     > Output rate controlled by 'output_rate' from "run.def" (PEM).
!     > Variable selection controlled by the file "diagevol.def".
!         - If no "diagevol.def", then no output.
!         - If "diagevol.def" is empty, then all the possible coded
!           outputs are done.
!         - If "diagevol.def" holds a list of variable names, then these
!           are outputted.
!     > If 'var' is an array, then the argument 'dimids_in' is mandatory.
!       'dimids_in' contains the dimension IDs of the variable in order.
!       If the variable holds a dimension 'ngrid', then it must appear
!       in the first position.
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use io_netcdf, only: open_nc, close_nc, def_var_nc, put_var_nc, diagevol_name
use evolution, only: idt, n_yr_run
use stoppage,  only: stop_clean
use geometry,  only: vect2dyngrd, ngrid, nlon, nlat

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

! ARGUMENTS
! ---------
character(*),               intent(in)           :: var_name, title, units
real(dp),    dimension(..), intent(in)           :: var ! Assumed-rank array
integer(di), dimension(:),  intent(in), optional :: dimids_in

! LOCAL VARIABLES
! ---------------
integer(di), dimension(:),       allocatable :: dimids
integer(di)                                  :: idim_ngrid, nlon_loc, ndim, i, j
real(dp),    dimension(:,:),     allocatable :: var_dyn
real(dp),    dimension(:,:,:),   allocatable :: var1d_dyn
real(dp),    dimension(:,:,:,:), allocatable :: var2d_dyn

! CODE
! ----
! If no diagdef file, no output at all
if (.not. is_diagdef) return

! If variable not selected, return
! Empty diagdef file means all the variables are ouputted
if (.not. is_output_selected(var_name)) return

! If output timing not met, return
if (abs(mod(idt,output_rate)) > minieps) return

! If it is the first writing of the current timestep
if (writing_idt /= idt) then
    writing_idt = idt
    call open_nc(diagevol_name,'write',itime) ! Get the current time index
    call put_var_nc('Time',n_yr_run,itime)
end if

! Write the variable
call open_nc(diagevol_name,'write') ! Open diagevol_name only it is not done yet
select rank (var) ! Assumed-rank array removes compile-time rank information which is required by the following subroutines using generic interface overload
    rank(0) ! Scalar
        call def_var_nc(var_name,title,units,(/dim_time/)) ! Define the variable only it is not done yet
        call put_var_nc(var_name,var,itime)
    rank default ! Arrays
        if (.not. present(dimids_in)) call stop_clean(__FILE__,__LINE__,'The argument ''dimids_in'' must be present to output arrays!',1)
        ndim = size(dimids_in)
        if (any(dimids_in == dim_ngrid)) then
            idim_ngrid = findloc(dimids_in,dim_ngrid,dim = 1)
            if (idim_ngrid /= 1) call stop_clean(__FILE__,__LINE__,'The dimension ''ngrid'' must be the first one to output the array!',1)
            allocate(dimids(ndim + 2))
            dimids(1) = dim_nlon
            dimids(2) = dim_nlat
            dimids(3:ndim + 1) = dimids_in(2:ndim)
            dimids(ndim + 2) = dim_time
            if (ngrid == 1) then ! 1D
                nlon_loc = 1
            else ! 3D
                nlon_loc = nlon + 1
            end if
            call def_var_nc(var_name,title,units,dimids) ! Define the variable only it is not done yet
            select rank (var) ! Assumed-rank array removes compile-time rank information which is required by the following subroutines using generic interface overload
                rank(1)
                    allocate(var_dyn(nlon_loc,nlat))
                    call vect2dyngrd(nlon + 1,nlat,ngrid,var,var_dyn)
                    call put_var_nc(var_name,var_dyn,itime)
                    deallocate(var_dyn)
                rank(2)
                    allocate(var1d_dyn(nlon_loc,nlat,size(var,2)))
                    do i = 1,size(var,2)
                        call vect2dyngrd(nlon + 1,nlat,ngrid,var(:,i),var1d_dyn(:,:,i))
                    end do
                    call put_var_nc(var_name,var1d_dyn,itime)
                    deallocate(var1d_dyn)
                rank(3)
                    allocate(var2d_dyn(nlon_loc,nlat,size(var,2),size(var,3)))
                    do i = 1,size(var,2)
                        do j = 1,size(var,3)
                            call vect2dyngrd(nlon + 1,nlat,ngrid,var(:,i,j),var2d_dyn(:,:,i,j))
                        end do
                    end do
                    call put_var_nc(var_name,var2d_dyn,itime)
                    deallocate(var2d_dyn)
                rank default
                    call stop_clean(__FILE__,__LINE__,'unsupported rank for the input array!',1)
            end select
        else ! No 'ngrid' dimension
            allocate(dimids(ndim + 1))
            dimids(1:ndim) = dimids_in(:)
            dimids(ndim + 1) = dim_time
            call def_var_nc(var_name,title,units,dimids) ! Define the variable only it is not done yet
            select rank (var) ! Assumed-rank array removes compile-time rank information which is required by the following subroutines using generic interface overload
                rank(1)
                    call put_var_nc(var_name,var_dyn,itime)
                rank(2)
                    call put_var_nc(var_name,var1d_dyn,itime)
                rank(3)
                    call put_var_nc(var_name,var2d_dyn,itime)
                rank(4)
                    call put_var_nc(var_name,var2d_dyn,itime)
                rank default
                    call stop_clean(__FILE__,__LINE__,'unsupported rank for the input array!',1)
            end select
        end if
        deallocate(dimids)
end select
call close_nc(diagevol_name)

END SUBROUTINE write_diagevol
!=================================================================

END MODULE output