!= 配列の分割／拡張
!
!= Division and expansion of arrays
!
! Authors::   Hiroki KASHIMURA, Yoshiyuki O. TAKAHASHI, Shin-ichi Takehiro
! Version::
! Tag Name::  $Name:  $
! Copyright:: Copyright (C) GFD Dennou Club, 2013-2026. All rights reserved.
! License::   See COPYRIGHT[link:../../../COPYRIGHT]
!
module sltt_extarr

  != 配列の分割／拡張
  !
  != Division and expansion of arrays (for ISPACK3/spml2)
  !
  !== 配列書法
  !
  !   * 通常変数 : xyz_Var(0:imax-1,1:jmax,1:kmax)
  !
  !   * 拡張変数 : xyz_ExtVar(0:imax_global-1,jexmin:jexmax,1:kc)
  !
  !
  !---------------------------------------------------------------------
  ! 種別型パラメタ
  ! Kind type parameter
  !
  use dc_types, only: DP, &      ! 倍精度実数型. Double precision.
    &                 STRING     ! 文字列.       Strings.

  use mpi_wrapper, only : myrank

  ! 格子点設定
  ! Grid points settings
  !
#ifdef LIB_MPI
  use gridset, only:       &
    &                imax, & ! 経度格子点数.
                             ! Number of grid points in longitude
    &                jmax, & ! 緯度格子点数.
                             ! Number of grid points in latitude
    &                kmax, & ! 鉛直層数.
                             ! Number of vertical level
    &                lmax, & ! スペクトルデータの配列サイズ
                             ! Size of array for spectral data
    &          imax_global   ! 経度格子点数 (全球).
                             ! Number of grid points in longitude on whole globe
  use ua_mpi_module_base, only : nproc_h, iproc_h, mpi_comm_h, &
       & xvb_pva, kc
  use mpi
#else
  use gridset, only:       &
    &                imax, & ! 経度格子点数.
                             ! Number of grid points in longitude
    &                jmax, & ! 緯度格子点数.
                             ! Number of grid points in latitude
    &                kmax, & ! 鉛直層数.
                             ! Number of vertical level
    &           kc => kmax,& ! 並列鉛直層数.
                             ! Number of vertical level for parallel compuation
    &                lmax, & ! スペクトルデータの配列サイズ
                             ! Size of array for spectral data
    &          imax_global   ! 経度格子点数 (全球).
                             ! Number of grid points in longitude on whole globe
#endif
  
  ! 組成に関わる配列の設定
  ! Settings of array for atmospheric composition
  !
  use composition, only : ncmax

  use sltt_const, only : PIx2, PIH, jexmin, jexmax, dtjw

  implicit none
  
  private

  ! 非公開変数
  ! Private variables
  !
  logical, save :: sltt_extarr_inited = .false.
                              ! 初期設定フラグ.
                              ! Initialization flag

  public :: SLTTExtArrInit
  public :: SLTTExtArrf
  public :: SLTTExtArr


  character(*), parameter:: module_name = 'sltt_extarr'
                              ! モジュールの名称.
                              ! Module name
  character(*), parameter:: version = &
    & '$Name:  $' // &
    & '$Id: sltt_extarr.f90,v 1.4 2026/05/21 22:00:00 takepiro Exp $'
                              ! モジュールのバージョン
                              ! Module version

contains

  !---------------------------------------------------------------------
  ! モジュール初期化
  ! 
  subroutine SLTTExtArrInit(   &
    & x_Lon, y_Lat,            & ! (in )
    & x_ExtLon, y_ExtLat       & ! (out)
    & )

    ! メッセージ出力
    ! Message output
    !
    use dc_message, only: MessageNotify

    !
    ! MPI
    !
    use mpi

    use constants0, only : PI
!    use axesset   , only : y_Lat
    use sltt_const, only : PIx2, jexmin, jexmax

    use ua_mpi_module_base, only : nproc_h, iproc_h, mpi_comm_h, &
         nproc_v, iproc_v, mpi_comm_v, ic

    real(DP), intent(in ) :: x_Lon   (   0:imax-1  )
    real(DP), intent(in ) :: y_Lat   (   1:jmax    )
    real(DP), intent(out) :: x_ExtLon(0:imax_global-1)
    real(DP), intent(out) :: y_ExtLat(jexmin:jexmax)

    !
    ! 作業変数
    !
    integer, allocatable :: isp(:), iep(:), icp(:)
    
    real(DP), allocatable :: a_sendbuf(:)
    real(DP), allocatable :: a_recvbufm(:)
    real(DP), allocatable :: a_recvbufp(:)
    integer :: ireqm, ireqp
    integer, allocatable  :: istat(:)
    
    integer :: ip, ipm, ipp, ipv
    integer :: ierr, ierrm, ierrp, ierrw
    integer :: itagm, itagp

    integer :: ireq_sm, ireq_rm
    logical :: sent_m, recvd_m
    integer :: ireq_sp, ireq_rp
    logical :: sent_p, recvd_p
    
    integer :: i, j

    if ( sltt_extarr_inited ) return

    allocate(a_sendbuf(dtjw))
    allocate(a_recvbufm(dtjw))
    allocate(a_recvbufp(dtjw))
    allocate(istat(MPI_STATUS_SIZE))

!!$#if defined(AXISYMMETRY) || defined(AXISYMMETRY_SJPACK)
!!$    do i = iexmin, 0-1
!!$      x_ExtLon(i) = x_Lon(0) - ( PIx2 * dble(-i) - x_Lon(0) )
!!$    end do
!!$    do i = 0, imax-1
!!$      x_ExtLon(i) = x_Lon(i)
!!$    end do
!!$    x_ExtLon(imax-1+1) = PIx2
!!$    do i = imax-1+1+1, iexmax
!!$      x_ExtLon(i) = PIx2 + ( PIx2 * dble(i-(imax-1+1)) - x_Lon(0) )
!!$    end do
#if defined(LIB_MPI)
    allocate(isp(0:nproc_v-1))
    allocate(iep(0:nproc_v-1))
    allocate(icp(0:nproc_v-1))

    call MPI_ALLGATHER(ic,  1, MPI_INTEGER, &
                       icp, 1, MPI_INTEGER, &
                       MPI_COMM_V, ierr)
    isp(0) = 0 ; iep(0) = icp(0)-1
    do ipv = 1, nproc_v-1
       isp(ipv) = isp(ipv-1) + icp(ipv-1)
       iep(ipv) = isp(ipv) + icp(ipv) -1
    end do
    
    call MPI_ALLGATHERV(x_Lon,    imax, MPI_REAL8,      &
                       x_ExtLon, icp,  isp, MPI_REAL8, &
                       MPI_COMM_V, ierr)
    deallocate(isp)
    deallocate(iep)
    deallocate(icp)
#else
    x_ExtLon = x_Lon
#endif

    do j=1,jmax
       y_ExtLat(j) = y_Lat(j)
    end do

#if defined(LIB_MPI)
    !------------------------------
    ! 南側
    sent_m = .false.
    recvd_m = .false.

    ! --- 通信の発行ループ ---
    do ip = 1, nproc_h-1
       ipm = modulo(ip-1+nproc_h, nproc_h)
       itagm = ipm*10 + ip

       ! 送信側：データをバッファに詰めて非同期送信を開始
       if ( ipm == iproc_h ) then
          a_sendbuf = y_Lat(jmax-dtjw+1:jmax)
          call MPI_ISEND(a_sendbuf, dtjw, MPI_REAL8, ip, itagm, mpi_comm_h, ireq_sm, ierrm)
          sent_m = .true.
       endif

       ! 受信側：非同期受信を開始（ここでは待たない）
       if ( ip == iproc_h ) then
          call MPI_IRECV(a_recvbufm, dtjw, MPI_REAL8, ipm, itagm, mpi_comm_h, ireq_rm, ierrm)
          recvd_m = .true.
       endif
    end do

    ! --- 完了待機とデータの格納 ---

    ! 送信完了待機（バッファの再利用を安全にするため）
    if (sent_m) then
       call MPI_WAIT(ireq_sm, istat, ierrw)
    endif

    ! 受信完了待機と格納
    if (recvd_m) then
       call MPI_WAIT(ireq_rm, istat, ierrw)
       y_ExtLat(-dtjw+1:0) = a_recvbufm
    endif
#endif

#if defined(LIB_MPI)    
    !------------------------------
    ! 北側
    ! リクエスト変数を送信(s)と受信(r)で個別に定義
    sent_p = .false.
    recvd_p = .false.

    ! --- 1. すべての通信（ISEND / IRECV）を一斉に発行 ---
    do ip = 0, nproc_h-2
       ipp = modulo(ip+1, nproc_h)
       itagp = ipp*10 + ip + 100

       ! 送信側：バッファから非同期送信を開始
       if ( ipp == iproc_h ) then
          a_sendbuf = y_Lat(1:dtjw) 
          call MPI_ISEND(a_sendbuf, dtjw, MPI_REAL8, ip, itagp, mpi_comm_h, ireq_sp, ierrp)
          sent_p = .true.
       endif

       ! 受信側：非同期受信を開始
       if ( ip == iproc_h ) then
          call MPI_IRECV(a_recvbufp, dtjw, MPI_REAL8, ipp, itagp, mpi_comm_h, ireq_rp, ierrp)
          recvd_p = .true.
       endif
    end do

    ! --- 2. 完了待機とデータの格納（ループの外で実施） ---

    ! 送信完了を確認
    if (sent_p) then
       call MPI_WAIT(ireq_sp, istat, ierrw)
    endif

    ! 受信完了を確認し、データを格納
    if (recvd_p) then
       call MPI_WAIT(ireq_rp, istat, ierrw)
       y_ExtLat(jmax+1:jmax+dtjw) = a_recvbufp
    endif
#endif

    !------------------------------
    ! 南極
    if ( iproc_h == 0 ) then
       y_ExtLat(-dtjw+1:0) &
            = -PIH - ( y_Lat(dtjw:1:-1) - ( -PIH ) )
    endif
       
    !------------------------------
    ! 北極
    if ( iproc_h == nproc_h-1 .OR. nproc_h==0 ) then
       y_ExtLat(jmax+1:jmax+dtjw) &
            = PIH + ( PIH - y_Lat(jmax:jmax-dtjw+1:-1) )
    endif

#if defined(LIB_MPI)       
    deallocate(a_sendbuf)
    deallocate(a_recvbufm)
    deallocate(a_recvbufp)
    deallocate(istat)
#endif
    
  end subroutine SLTTExtArrInit

  !------------------------------------------------------------------------
  ! 配列拡張 (xyzf)
  !
  subroutine SLTTExtArrf(       &
    & xyzf_Data, PM,            & ! (in)
    & xyzf_ExtData              & ! (out)
    & )

    real(DP), intent(in ) :: xyzf_Data    (0:imax-1, 1:jmax, 1:kmax, 1:ncmax)
    real(DP), intent(in ) :: PM   ! 配列拡張する際、極ごえ後に符号が変わる場合は -1.0を与える。そうでない場合は1.0を与える。
    real(DP), intent(out) :: xyzf_ExtData (0:imax_global-1, jexmin:jexmax, 1:kc, 1:ncmax)

    integer :: n
    integer :: ierr

    do n=1,ncmax
       call SLTTExtArr(                          &
            & xyzf_Data(:,:,:,n), PM,            & ! (in)
            & xyzf_ExtData(:,:,:,n)              & ! (out)
            & )
    end do

  end subroutine SLTTExtArrf
  
  !------------------------------------------------------------------------
  ! 配列拡張 (xyz)
  !
  subroutine SLTTExtArr(       &
    & xyz_Data, PM,            & ! (in)
    & xyz_ExtData              & ! (out)
    & )
    
    real(DP), intent(in ) :: xyz_Data    (0:imax-1, 1:jmax, 1:kmax)
    real(DP), intent(in ) :: PM   ! 配列拡張する際、極ごえ後に符号が変わる場合は -1.0を与える。そうでない場合は1.0を与える。
    real(DP), intent(out) :: xyz_ExtData (0:imax_global-1, jexmin:jexmax, 1:kc)
   
    !
    ! 作業変数
    !
#ifdef LIB_MPI
    real(DP) :: xaz_sendbuf(imax_global,dtjw,kc)
    real(DP) :: xaz_recvbufm(imax_global,dtjw,kc)
    real(DP) :: xaz_recvbufp(imax_global,dtjw,kc)
    integer  :: istat(MPI_STATUS_SIZE)
    integer  :: ireqm, ireqp
    
    integer :: ip, ipm, ipp
    integer :: ierrm, ierrp, ierrw
    integer :: itagm, itagp
    integer :: icount
#endif
    
    integer :: i, ind, j

    integer :: ireq_s, ireq_r
    logical :: sent, recvd
    
    !------------------------------
    ! 経度方向拡張, 鉛直分割
#ifdef LIB_MPI
    xyz_ExtData(:,1:jmax,:) = xvb_pva(xyz_Data)
#else
    xyz_ExtData(:,1:jmax,:) = xyz_Data
#endif

    !------------------------------
    ! 南側
#ifdef LIB_MPI
    sent = .false.
    recvd = .false.

    do ip = 1, nproc_h-1
       ipm = modulo(ip-1+nproc_h, nproc_h)
       itagm = ipm*10 + ip
       icount = imax_global * dtjw * kc

       ! --- 送信処理 ---
       if ( ipm == iproc_h ) then
          xaz_sendbuf = xyz_ExtData(:, jmax-dtjw+1:jmax, :)
          call MPI_ISEND(xaz_sendbuf, icount, MPI_REAL8, ip, itagm, mpi_comm_h, ireq_s, ierrm)
          sent = .true.
       endif

       ! --- 受信処理 ---
       if ( ip == iproc_h ) then
          ! WAITはここで行わず、まず通信の開始だけを行う
          call MPI_IRECV(xaz_recvbufm, icount, MPI_REAL8, ipm, itagm, mpi_comm_h, ireq_r, ierrm)
          recvd = .true.
       endif
    end do

    ! --- 通信完了の待機（ループの外で実施） ---
    if (sent) then
       call MPI_WAIT(ireq_s, istat, ierrw)
    endif

    if (recvd) then
       call MPI_WAIT(ireq_r, istat, ierrw)
       ! 受信完了を確認してからデータを格納
       xyz_ExtData(:, -dtjw+1:0, :) = xaz_recvbufm
    endif
#endif
    
    !------------------------------
    ! 北側
#ifdef LIB_MPI
    sent = .false.
    recvd = .false.
    do ip = 0, nproc_h-2
       ipp = modulo(ip+1, nproc_h)
       itagp = ipp*10 + ip + 100
       icount = imax_global * dtjw * kc

       ! --- 送信処理 ---
       if ( ipp == iproc_h ) then
          xaz_sendbuf = xyz_ExtData(:, 1:dtjw, :) 
          call MPI_ISEND(xaz_sendbuf, icount, MPI_REAL8, ip, itagp, mpi_comm_h, ireq_s, ierrp)
          sent = .true.
       endif

       ! --- 受信処理 ---
       if ( ip == iproc_h ) then
          call MPI_IRECV(xaz_recvbufp, icount, MPI_REAL8, ipp, itagp, mpi_comm_h, ireq_r, ierrp)
          recvd = .true.
       endif
    end do

    ! --- 通信完了の待機 ---
    if (sent) then
       call MPI_WAIT(ireq_s, istat, ierrw)
    endif

    if (recvd) then
       call MPI_WAIT(ireq_r, istat, ierrw)
       ! 受信完了後にデータをコピー
       xyz_ExtData(:, jmax+1:jmax+dtjw, :) = xaz_recvbufp
    endif
#endif
    
    !------------------------------
    ! 南極
#ifdef LIB_MPI
    if ( iproc_h == 0 ) then
       do j=1,dtjw
          do i = 0, imax_global-1
             ind = modulo(i + imax_global/2, imax_global)
             xyz_ExtData(i,1-j,:) = PM * xyz_ExtData(ind,j,:)
          end do
       end do
    endif
#else
    do j=1,dtjw
       do i = 0, imax_global-1
          ind = modulo(i + imax_global/2, imax_global)
          xyz_ExtData(i,1-j,:) = PM * xyz_ExtData(ind,j,:)
       end do
    end do
#endif
    
    !------------------------------
    ! 北極
#ifdef LIB_MPI
    if ( iproc_h == nproc_h-1 .or. nproc_h == 0 ) then
       do j=1,dtjw
          do i = 0, imax_global-1
             ind = modulo(i + imax_global/2, imax_global)
             xyz_ExtData(i,jmax+j,:) = PM * xyz_ExtData(ind,jmax+1-j,:)
          end do
       end do
    endif
#else
    do j=1,dtjw
       do i = 0, imax_global-1
          ind = modulo(i + imax_global/2, imax_global)
          xyz_ExtData(i,jmax+j,:) = PM * xyz_ExtData(ind,jmax+1-j,:)
       end do
    end do
#endif
       
  end subroutine SLTTExtArr
  
end module sltt_extarr