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, LVL_NFO
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
allocate(longitudes(ngrid))
allocate(latitudes(ngrid))
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',LVL_NFO)
call print_msg('nlon      = '//int2str(nlon),LVL_NFO)
call print_msg('nlat      = '//int2str(nlat),LVL_NFO)
call print_msg('ngrid     = '//int2str(ngrid),LVL_NFO)
call print_msg('nslope    = '//int2str(nslope),LVL_NFO)
call print_msg('nlayer    = '//int2str(nlayer),LVL_NFO)
call print_msg('nsoil_PCM = '//int2str(nsoil_PCM),LVL_NFO)
call print_msg('nsoil     = '//int2str(nsoil),LVL_NFO)
call print_msg('nday      = '//int2str(nday),LVL_NFO)

! 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(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
! ---------
real(dp), dimension(:,:), intent(in)  :: v_ll
real(dp), dimension(:),   intent(out) :: v_vect

! LOCAL VARIABLES
! ---------------
integer(di) :: ilon, ilat, n_lon, n_lat, n_grd

! CODE
! ----
n_grd = size(v_vect)
n_lon = size(v_ll,1)
n_lat = size(v_ll,2)
if (n_grd == 1) then ! 1D case
    if (n_lon /= 1 .or. n_lat /= 1) call stop_clean(__FILE__,__LINE__,'problem of dimensions!',1)
    v_vect(1) = v_ll(1,1)
else ! 3D
    ! Check
    if (n_grd /= n_lon*(n_lat - 2) + 2) call stop_clean(__FILE__,__LINE__,'problem of dimensions!',1)

    ! Treatment of the poles
    v_vect(1) = v_ll(1,1)
    v_vect(n_grd) = v_ll(1,n_lat)

    ! Treatment of regular points
    do ilat = 2,n_lat - 1
        do ilon = 1,n_lon
            v_vect(1 + ilon + (ilat - 2)*n_lon) = v_ll(ilon,ilat)
        end do
    end do
end if

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

!=======================================================================
SUBROUTINE vect2lonlat(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
! ---------
real(dp), dimension(:),   intent(in)  :: v_vect
real(dp), dimension(:,:), intent(out) :: v_ll

! LOCAL VARIABLES
! ---------------
integer(di) :: ilon, ilat, n_lon, n_lat, n_grd

! CODE
! ----
n_grd = size(v_vect)
n_lon = size(v_ll,1)
n_lat = size(v_ll,2)
if (n_grd == 1) then ! 1D case
    if (n_lon /= 1 .or. n_lat /= 1) call stop_clean(__FILE__,__LINE__,'problem of dimensions!',1)
    v_ll(1,1) = v_vect(1)
else ! 3D
    ! Check
    if (n_grd /= n_lon*(n_lat - 2) + 2) call stop_clean(__FILE__,__LINE__,'problem of dimensions!',1)

    ! Treatment of the poles
    v_ll(:,1) = v_vect(1)
    v_ll(:,n_lat) = v_vect(n_grd)

    ! Treatment of regular points
    do ilat = 2,n_lat - 1
        do ilon = 1,n_lon
            v_ll(ilon,ilat) = v_vect(1 + ilon + (ilat - 2)*n_lon)
        end do
    end do
end if

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

!=======================================================================
SUBROUTINE dyngrd2vect(v_dyn,v_vect,extensive)
!-----------------------------------------------------------------------
! NAME
!     dyngrd2vect
!
! DESCRIPTION
!     Convert data from dynamical lon x lat grid (where 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).
!     For extensive variables, the values at the poles must be averaged
!     over the number of longitudes for computation on the dynamical grid.
!-----------------------------------------------------------------------

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

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

! ARGUMENTS
! ---------
real(dp), dimension(:,:), intent(in)           :: v_dyn
logical(k4),              intent(in), optional :: extensive
real(dp), dimension(:),   intent(out)          :: v_vect

! LOCAL VARIABLES
! ---------------
integer(di) :: ilon, ilat, n_lon, n_lat, n_grd
real(dp)    :: factor

! CODE
! ----
n_grd = size(v_vect)
n_lon = size(v_dyn,1)
n_lat = size(v_dyn,2)
if (n_grd == 1) then ! 1D case
    if (n_lon /= 1 .or. n_lat /= 1) call stop_clean(__FILE__,__LINE__,'problem of dimensions!',1)
    v_vect(1) = v_dyn(1,1)
else ! 3D
    ! Check
    if (n_grd /= (n_lon - 1)*(n_lat - 2) + 2) call stop_clean(__FILE__,__LINE__,'problem of dimensions!',1)

    ! Treatment of the poles
    factor = 1._dp
    if (present(extensive)) then
        ! Correction for polar values of extensive variable when mapped onto the physical grid
        if (extensive) factor = real(n_lon - 1,dp)
    end if
    v_vect(1) = v_dyn(1,1)*factor
    v_vect(n_grd) = v_dyn(1,n_lat)*factor

    ! Treatment of regular points
    do ilat = 2,n_lat - 1
        do ilon = 1,n_lon - 1
            v_vect(1 + ilon + (ilat - 2)*(n_lon - 1)) = v_dyn(ilon,ilat)
        end do
    end do
end if

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

!=======================================================================
SUBROUTINE vect2dyngrd(v_vect,v_dyn,extensive)
!-----------------------------------------------------------------------
! NAME
!     vect2dyngrd
!
! DESCRIPTION
!     Convert data from a vector into lon x lat grid (where poles are 
!     duplicated and longitude -180°/+180° is duplicated).
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!     The longitudes -180 and +180 are duplicated (PCM dynamics).
!     For extensive variables, the values at the poles must be averaged
!     over the number of longitudes for computation on the dynamical grid.
!-----------------------------------------------------------------------

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

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

! ARGUMENTS
! ---------
real(dp), dimension(:),   intent(in)           :: v_vect
logical(k4),              intent(in), optional :: extensive
real(dp), dimension(:,:), intent(out)          :: v_dyn

! LOCAL VARIABLES
! ---------------
integer(di) :: ilon, ilat, n_lon, n_lat, n_grd
real(dp)    :: factor

! CODE
! ----
n_grd = size(v_vect)
n_lon = size(v_dyn,1)
n_lat = size(v_dyn,2)
if (n_grd == 1) then ! 1D case
    if (n_lon /= 1 .or. n_lat /= 1) call stop_clean(__FILE__,__LINE__,'problem of dimensions!',1)
    v_dyn(1,1) = v_vect(1)
else ! 3D
    ! Check
    if (n_grd /= (n_lon - 1)*(n_lat - 2) + 2) call stop_clean(__FILE__,__LINE__,'problem of dimensions!',1)

    ! Treatment of the poles
    factor = 1._dp
    if (present(extensive)) then
        ! Correction for polar values of extensive variable when mapped onto the dynamical grid
        if (extensive) factor = 1._dp/real(n_lon - 1,dp)
    end if
    v_dyn(:,1) = v_vect(1)*factor
    v_dyn(:,n_lat) = v_vect(n_grd)*factor

    ! Treatment of regular points
    do ilat = 2,n_lat - 1
        do ilon = 1,n_lon - 1
            v_dyn(ilon,ilat) = v_vect(1 + ilon + (ilat - 2)*(n_lon - 1))
        end do
        v_dyn(n_lon,ilat) =  v_dyn(1,ilat) ! Treatment of the duplicated longitude
    end do
end if

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

END MODULE geometry