!------------------------------------------------------------------------------------------------
! Copyright (c) 2021 - 2025, Oleg Shatrov
! All rights reserved.
! This file is part of dtFFT library.

! dtFFT is free software: you can redistribute it and/or modify
! it under the terms of the GNU General Public License as published by
! the Free Software Foundation, either version 3 of the License, or
! (at your option) any later version.

! dtFFT is distributed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU General Public License for more details.

! You should have received a copy of the GNU General Public License
! along with this program.  If not, see <https://www.gnu.org/licenses/>.
!------------------------------------------------------------------------------------------------
#include "dtfft_config.h"
module dtfft_backend_cufftmp_m
!! cuFFTMp GPU Backend [[backend_cufftmp]]
use iso_fortran_env
use iso_c_binding
use dtfft_abstract_backend, only: abstract_backend, backend_helper
use dtfft_interface_nvshmem
use dtfft_interface_cuda_runtime
use dtfft_interface_cufft
use dtfft_errors
use dtfft_parameters
use dtfft_pencil, only: pencil
use dtfft_utils
#include "_dtfft_mpi.h"
#include "_dtfft_cuda.h"
#include "_dtfft_private.h"
implicit none
private
public :: backend_cufftmp

    type :: Box3D
    !! cuFFTMp Box
        integer(c_long_long) :: lower(3)   !! Lower box boundaries
        integer(c_long_long) :: upper(3)   !! Upper box boundaries
        integer(c_long_long) :: strides(3) !! Strides in memory
    end type Box3D

    type, extends(abstract_backend) :: backend_cufftmp
    !! cuFFTMp GPU Backend
    private
        type(cufftReshapeHandle) :: plan
    contains
        procedure :: create_private => create
        procedure :: execute_private => execute
        procedure :: destroy_private => destroy
    end type backend_cufftmp

contains

    subroutine create(self, helper, base_storage)
    !! Creates cuFFTMp GPU Backend
        class(backend_cufftmp), intent(inout)   :: self               !! cuFFTMp GPU Backend
        type(backend_helper),   intent(in)      :: helper             !! Backend helper
        integer(int64),         intent(in)      :: base_storage       !! Number of bytes to store single element
        type(Box3D) :: inbox, outbox  !! Reshape boxes
        type(pencil), pointer :: in, out
        type(c_ptr) :: c_comm
        integer(int64) :: aux_size
        logical :: is_transpose

        is_transpose = .false.
        if ( is_valid_transpose_type(helper%transpose_type) ) then
            is_transpose = .true.
            select case (helper%transpose_type%val)
            case (DTFFT_TRANSPOSE_X_TO_Y%val)
                in => helper%pencils(1)
                out => helper%pencils(2)
            case (DTFFT_TRANSPOSE_Y_TO_X%val)
                in => helper%pencils(2)
                out => helper%pencils(1)
            case (DTFFT_TRANSPOSE_Y_TO_Z%val)
                in => helper%pencils(2)
                out => helper%pencils(3)
            case (DTFFT_TRANSPOSE_Z_TO_Y%val)
                in => helper%pencils(3)
                out => helper%pencils(2)
            case (DTFFT_TRANSPOSE_X_TO_Z%val)
                in => helper%pencils(1)
                out => helper%pencils(3)
            case (DTFFT_TRANSPOSE_Z_TO_X%val)
                in => helper%pencils(3)
                out => helper%pencils(1)
            case default
                INTERNAL_ERROR("backend_cufftmp: unknown `transpose_type`")
            end select
        else
            select case (helper%reshape_type%val)
            case ( DTFFT_RESHAPE_X_BRICKS_TO_PENCILS%val)
                in => helper%pencils(1)
                out => helper%pencils(3)
            case ( DTFFT_RESHAPE_X_PENCILS_TO_BRICKS%val )
                in => helper%pencils(3)
                out => helper%pencils(1)
            case ( DTFFT_RESHAPE_Z_BRICKS_TO_PENCILS%val )
                in => helper%pencils(2)
                out => helper%pencils(4)
            case ( DTFFT_RESHAPE_Z_PENCILS_TO_BRICKS%val )
                in => helper%pencils(4)
                out => helper%pencils(2)
            case default
                INTERNAL_ERROR("backend_cufftmp: unknown `transpose_type`")
            endselect
        endif

        if (in%rank == 3) then
            if ( is_transpose ) then
                if (any(helper%transpose_type == [DTFFT_TRANSPOSE_X_TO_Z, DTFFT_TRANSPOSE_Y_TO_X, DTFFT_TRANSPOSE_Z_TO_Y])) then
                    inbox%lower = [in%starts(2), in%starts(1), in%starts(3)]
                    inbox%upper = [in%starts(2) + in%counts(2), in%starts(1) + in%counts(1), in%starts(3) + in%counts(3)]
                    inbox%strides = [in%counts(1) * in%counts(3), in%counts(3), 1]
                else
                    inbox%lower = [in%starts(1), in%starts(3), in%starts(2)]
                    inbox%upper = [in%starts(1) + in%counts(1), in%starts(3) + in%counts(3), in%starts(2) + in%counts(2)]
                    inbox%strides = [in%counts(2) * in%counts(3), in%counts(2), 1]
                end if
            else
                inbox%lower = [in%starts(3), in%starts(2), in%starts(1)]
                inbox%upper = [in%starts(3) + in%counts(3), in%starts(2) + in%counts(2), in%starts(1) + in%counts(1)]
                inbox%strides = [in%counts(1) * in%counts(2), in%counts(1), 1]
            endif

            outbox%lower = [out%starts(3), out%starts(2), out%starts(1)]
            outbox%upper = [out%starts(3) + out%counts(3), out%starts(2) + out%counts(2), out%starts(1) + out%counts(1)]
            outbox%strides = [out%counts(1) * out%counts(2), out%counts(1), 1]
        else
            if ( is_transpose ) then
                inbox%lower = [0, in%starts(1), in%starts(2)]
                inbox%upper = [1, in%starts(1) + in%counts(1), in%starts(2) + in%counts(2)]
                inbox%strides = [in%counts(1) * in%counts(2), in%counts(2), 1]
            else
                inbox%lower = [0, in%starts(2), in%starts(1)]
                inbox%upper = [1, in%starts(2) + in%counts(2), in%starts(1) + in%counts(1)]
                inbox%strides = [in%counts(1) * in%counts(2), in%counts(1), 1]
            endif

            outbox%lower = [0, out%starts(2), out%starts(1)]
            outbox%upper = [1, out%starts(2) + out%counts(2), out%starts(1) + out%counts(1)]
            outbox%strides = [out%counts(1) * out%counts(2), out%counts(1), 1]
        end if

        CUFFT_CALL( cufftMpCreateReshape(self%plan) )
        c_comm = Comm_f2c(GET_MPI_VALUE(helper%comms(1)))
        CUFFT_CALL( cufftMpAttachReshapeComm(self%plan, CUFFT_COMM_MPI, c_comm) )
        CUFFT_CALL( cufftMpMakeReshape(self%plan, base_storage, 3, inbox%lower, inbox%upper, outbox%lower, outbox%upper, inbox%strides, outbox%strides) )
        CUFFT_CALL( cufftMpGetReshapeSize(self%plan, aux_size) )
        self%aux_bytes = max(aux_size, self%aux_bytes)
    end subroutine create

    subroutine execute(self, in, out, stream, aux, error_code)
    !! Executes cuFFTMp GPU Backend
        class(backend_cufftmp),             intent(inout)   :: self       !! cuFFTMp GPU Backend
        real(real32),   target, contiguous, intent(inout)   :: in(:)      !! Send pointer
        real(real32),   target, contiguous, intent(inout)   :: out(:)     !! Recv pointer
        type(dtfft_stream_t),               intent(in)      :: stream     !! Main execution CUDA stream
        real(real32),   target, contiguous, intent(inout)   :: aux(:)     !! Aux pointer
        integer(int32),                     intent(out)     :: error_code !! Error code
        integer(int32) :: ierr

        ! call nvshmemx_sync_all_on_stream(stream)
        CUDA_CALL( cudaStreamSynchronize(stream) )
        call MPI_Barrier(self%comm, ierr)
        CUFFT_CALL( cufftMpExecReshapeAsync(self%plan, c_loc(out), c_loc(in), c_loc(aux), stream) )
        error_code = DTFFT_SUCCESS
    end subroutine execute

    subroutine destroy(self)
    !! Destroys cuFFTMp GPU Backend
        class(backend_cufftmp), intent(inout) :: self        !! cuFFTMp GPU Backend

        CUFFT_CALL( cufftMpDestroyReshape(self%plan) )
    end subroutine destroy
end module dtfft_backend_cufftmp_m
