MODULE geometry
!-----------------------------------------------------------------------
! NAME
!     geometry
!
! DESCRIPTION
!     Grid metrics/parameters and basic tools to convert between
!     different grids.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! PARAMETERS
! ----------
integer(di),                         protected :: nlon               ! Number of longitudes
integer(di),                         protected :: nlat               ! Number of latitudes
integer(di),                         protected :: ngrid              ! Number of grid points
integer(di),                         protected :: nlayer             ! Number of atmospheric layers
integer(di),                         protected :: nsoil_PCM          ! Number of soil layers in the PCM
integer(di),                         parameter :: nsoil = 68_di      ! Number of soil layers in the PEM
integer(di),                         protected :: nslope             ! Number of slopes
integer(di),                         protected :: nday               ! Number of sols in a year
real(dp), dimension(:), allocatable, protected :: longitudes         ! Longitudes
real(dp), dimension(:), allocatable, protected :: latitudes          ! Latitudes
real(dp), dimension(:), allocatable, protected :: cell_area          ! Cell area
real(dp),                            protected :: total_surface      ! Total surface of the grid/planet
logical(k4),                             protected :: dim_init = .false. ! Flag to true if dimensions are well initialized

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

!=======================================================================
SUBROUTINE ini_geometry()
!-----------------------------------------------------------------------
! NAME
!     ini_geometry
!
! DESCRIPTION
!     Initialize the parameters of module 'geometry'.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use stoppage,   only: stop_clean
use io_netcdf,  only: startfi_name, xios_day_name1, open_nc, close_nc, get_dim_nc, get_var_nc
use display,    only: print_msg
use utility,    only: int2str

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

! LOCAL VARIABLES
! ---------------
logical(k4) :: here, found

! CODE
! ----
! Reading from "startfi.nc" file
inquire(file = startfi_name,exist = here)
if (.not. here) call stop_clean(__FILE__,__LINE__,'cannot find required file "'//startfi_name//'"!',1)
call open_nc(startfi_name,'read')
call get_dim_nc('physical_points',ngrid)
call get_dim_nc('subsurface_layers',nsoil_PCM)
call get_dim_nc('nlayer',nlayer)
call get_dim_nc('nslope',nslope,found)
if (.not. found) nslope = 1
if (.not. allocated(longitudes)) allocate(longitudes(ngrid))
if (.not. allocated(latitudes)) allocate(latitudes(ngrid))
if (.not. allocated(cell_area)) allocate(cell_area(ngrid))
call get_var_nc('longitude',longitudes)
call get_var_nc('latitude',latitudes)
call get_var_nc('area',cell_area)
call close_nc(startfi_name)

! Reading from XIOS daily output file
inquire(file = xios_day_name1,exist = here)
if (.not. here) call stop_clean(__FILE__,__LINE__,'cannot find required file "'//xios_day_name1//'"!',1)
call open_nc(xios_day_name1,'read')
call get_dim_nc('lon',nlon)
call get_dim_nc('lat',nlat)
call get_dim_nc('time_counter',nday)
call close_nc(xios_day_name1)

! Total surface
total_surface = sum(cell_area)

! State that dimentions are well initialized
dim_init = .true.
call print_msg('> Reading dimensions')
call print_msg('nlon      = '//int2str(nlon))
call print_msg('nlat      = '//int2str(nlat))
call print_msg('ngrid     = '//int2str(ngrid))
call print_msg('nslope    = '//int2str(nslope))
call print_msg('nlayer    = '//int2str(nlayer))
call print_msg('nsoil_PCM = '//int2str(nsoil_PCM))
call print_msg('nsoil     = '//int2str(nsoil))
call print_msg('nday      = '//int2str(nday))

! Check consistency of grids
if (ngrid /= 2 + nlon*(nlat - 2)) call stop_clean(__FILE__,__LINE__,'geometry inconsistency between nlon/nlat and ngrid!',1)
if (nsoil_PCM > nsoil) call stop_clean(__FILE__,__LINE__,'geometry inconsistency between nsoil and nsoil_PCM! nsoil must be greater than nsoil_PCM.',1)

END SUBROUTINE ini_geometry
!=======================================================================

!=======================================================================
SUBROUTINE end_geometry()
!-----------------------------------------------------------------------
! NAME
!     end_geometry
!
! DESCRIPTION
!     Deallocate geometry arrays.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

! CODE
! ----
if (allocated(longitudes)) deallocate(longitudes)
if (allocated(latitudes)) deallocate(latitudes)
if (allocated(cell_area)) deallocate(cell_area)

END SUBROUTINE end_geometry
!=======================================================================

!=======================================================================
SUBROUTINE lonlat2vect(nlon,nlat,ngrid,v_ll,v_vect)
!-----------------------------------------------------------------------
! NAME
!     lonlat2vect
!
! DESCRIPTION
!     Convert data from lon x lat grid (where values at the poles are
!     duplicated) into a vector.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!     The longitudes -180 and +180 are not duplicated like in the PCM
!     dynamics.
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use stoppage, only: stop_clean

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

! ARGUMENTS
!----------
integer(di),                    intent(in)  :: nlon, nlat, ngrid
real(dp), dimension(nlon,nlat), intent(in)  :: v_ll
real(dp), dimension(ngrid),     intent(out) :: v_vect

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

! CODE
!-----
if (ngrid == 1) then ! 1D case
    v_vect(1) = v_ll(1,1)
    return
else ! 3D
    ! Check
    if (ngrid /= nlon*(nlat - 2) + 2) call stop_clean(__FILE__,__LINE__,'problem of dimensions!',1)

    ! Initialization
    v_vect = 0.

    ! Treatment of the poles
    v_vect(1) = v_ll(1,1)
    v_vect(ngrid) = v_ll(1,nlat)

    ! Treatment of regular points
    do j = 2,nlat - 1
        do i = 1,nlon
            v_vect(1 + i + (j - 2)*nlon) = v_ll(i,j)
        end do
    end do
end if

END SUBROUTINE lonlat2vect
!=======================================================================

!=======================================================================
SUBROUTINE vect2lonlat(nlon,nlat,ngrid,v_vect,v_ll)
!-----------------------------------------------------------------------
! NAME
!     vect2lonlat
!
! DESCRIPTION
!     Convert data from a vector into lon x lat grid (where values
!     at the poles are duplicated).
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!     The longitudes -180 and +180 are not duplicated like in the PCM
!     dynamics.
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use stoppage, only: stop_clean

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

! ARGUMENTS
!----------
integer(di),                    intent(in)  :: nlon, nlat, ngrid
real(dp), dimension(ngrid),     intent(in)  :: v_vect
real(dp), dimension(nlon,nlat), intent(out) :: v_ll

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

! CODE
!-----
if (ngrid == 1) then ! 1D case
    v_ll(1,1) = v_vect(1)
    return
else ! 3D
    ! Check
    if (ngrid /= nlon*(nlat - 2) + 2) call stop_clean(__FILE__,__LINE__,'problem of dimensions!',1)

    ! Initialization
    v_ll = 0.

    ! Treatment of the poles
    v_ll(:,1) = v_vect(1)
    v_ll(:,nlat) = v_vect(ngrid)

    ! Treatment of regular points
    do j = 2,nlat - 1
        do i = 1,nlon
            v_ll(i,j) = v_vect(1 + i + (j - 2)*nlon)
        end do
    end do
end if

END SUBROUTINE vect2lonlat
!=======================================================================

!=======================================================================
SUBROUTINE dyngrd2vect(ni,nj,ngrid,v_dyn,v_vect)
!-----------------------------------------------------------------------
! NAME
!     dyngrd2vect
!
! DESCRIPTION
!     Convert data from dynamical lon x lat grid (where values at the
!     poles are duplicated and longitude -180°/+180° is duplicated) into
!     a vector.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!     The longitudes -180 and +180 are duplicated (PCM dynamics).
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use stoppage, only: stop_clean

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

! ARGUMENTS
!----------
integer(di),                intent(in)  :: ni, nj, ngrid
real(dp), dimension(ni,nj), intent(in)  :: v_dyn
real(dp), dimension(ngrid), intent(out) :: v_vect

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

! CODE
!-----
if (ngrid == 1) then ! 1D case
    v_vect(1) = v_dyn(1,1)
    return
else ! 3D
    ! Check
    if (ngrid /= (ni - 1)*(nj - 2) + 2) call stop_clean(__FILE__,__LINE__,'problem of dimensions!',1)

    ! Initialization
    v_vect = 0.

    ! Treatment of the poles
    v_vect(1) = v_dyn(1,1)
    v_vect(ngrid) = v_dyn(1,nj)

    ! Treatment of regular points
    do j = 2,nj - 1
        do i = 1,ni - 1
            v_vect(1 + i + (j - 2)*(ni - 1)) = v_dyn(i,j)
        end do
    end do
end if

END SUBROUTINE dyngrd2vect
!=======================================================================

!=======================================================================
SUBROUTINE vect2dyngrd(ni,nj,ngrid,v_vect,v_dyn)
!-----------------------------------------------------------------------
! NAME
!     vect2dyngrd
!
! DESCRIPTION
!     Convert data from a vector into lon x lat grid (where values at the
!     poles are duplicated and longitude -180°/+180° is duplicated).
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!     The longitudes -180 and +180 are not duplicated like in the PCM
!     dynamics.
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use stoppage, only: stop_clean

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

! ARGUMENTS
!----------
integer(di),                intent(in)  :: ni, nj, ngrid
real(dp), dimension(ngrid), intent(in)  :: v_vect
real(dp), dimension(ni,nj), intent(out) :: v_dyn

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

! CODE
!-----
if (ngrid == 1) then ! 1D case
    v_dyn(1,1) = v_vect(1)
    return
else ! 3D
    ! Check
    if (ngrid /= (ni - 1)*(nj - 2) + 2) call stop_clean(__FILE__,__LINE__,'problem of dimensions!',1)

    ! Initialization
    v_dyn = 0.

    ! Treatment of the poles
    v_dyn(:,1) = v_vect(1)
    v_dyn(:,nj) = v_vect(ngrid)

    ! Treatment of regular points
    do j = 2,nj - 1
        do i = 1,ni - 1
            v_dyn(i,j) = v_vect(1 + i + (j - 2)*(ni - 1))
        end do
        v_dyn(ni,j) =  v_dyn(1,j)
    end do
end if

END SUBROUTINE vect2dyngrd
!=======================================================================

END MODULE geometry