// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifdef PADDLE_WITH_XPU_FFT
#include "paddle/phi/kernels/complex_kernel.h"
#include "fft/cuComplex.h"
#include "paddle/phi/backends/xpu/enforce_xpu.h"
#include "paddle/phi/common/type_traits.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/expand_kernel.h"
#include "paddle/phi/kernels/funcs/common_infer_shape_functions.h"
#include "paddle/phi/kernels/xpu/xpu_api_wrapper.h"

namespace xfft_internal::xpu {
// just for declaration here, the real implementation is in libcufft.so
template <typename T, typename TComplex>
int combine_as_complex(
    const XPUStream stream, int N, const T* real, const T* imag, TComplex* out);
template <>
int combine_as_complex(const XPUStream stream,
                       int N,
                       const float* real,
                       const float* imag,
                       float2* out);
template <>
int combine_as_complex(const XPUStream stream,
                       int N,
                       const double* real,
                       const double* imag,
                       double2* out);

template <typename TComplex, typename T>
int complex_spilt(
    const XPUStream stream, int N, const TComplex* in, T* real, T* imag);
template <>
int complex_spilt(
    const XPUStream stream, int N, const float2* in, float* real, float* imag);
template <>
int complex_spilt(const XPUStream stream,
                  int N,
                  const double2* in,
                  double* real,
                  double* imag);

template <typename T>  // T supports float2, double2
int Conj(const XPUStream stream, int N, const T* input, T* output);
}  // namespace xfft_internal::xpu

namespace phi {
template <typename T, typename Context>
void ConjKernel(const Context& dev_ctx,
                const DenseTensor& x,
                DenseTensor* out) {
  if (out->numel() == 0) {
    dev_ctx.template Alloc<T>(out);
    return;
  }
  dev_ctx.template Alloc<T>(out);
  if (std::is_same_v<T, phi::complex64>) {
    int r = xfft_internal::xpu::Conj(
        dev_ctx.x_context()->xpu_stream,
        x.numel(),
        reinterpret_cast<const cuFloatComplex*>(x.data<T>()),
        reinterpret_cast<cuFloatComplex*>(out->data<T>()));
    PADDLE_ENFORCE_XPU_SUCCESS(r);
  } else if (std::is_same_v<T, phi::complex128>) {
    int r = xfft_internal::xpu::Conj(
        dev_ctx.x_context()->xpu_stream,
        x.numel(),
        reinterpret_cast<const cuDoubleComplex*>(x.data<T>()),
        reinterpret_cast<cuDoubleComplex*>(out->data<T>()));
    PADDLE_ENFORCE_XPU_SUCCESS(r);
  } else {
    using XPUType = typename XPUCopyTypeTrait<T>::Type;
    const auto* input_data = x.data<T>();
    int r = xpu::copy<XPUType>(dev_ctx.x_context(),
                               reinterpret_cast<const XPUType*>(input_data),
                               reinterpret_cast<XPUType*>(out->data<T>()),
                               x.numel());
    PADDLE_ENFORCE_XDNN_SUCCESS(r, "copy");
  }
}

template <typename T, typename Context>
void RealKernel(const Context& dev_ctx,
                const DenseTensor& x,
                DenseTensor* out) {
  using XPUComplexType =
      typename XPUComplexTypeTrait<phi::dtype::Real<T>>::Type;
  if (out->numel() == 0) {
    dev_ctx.template Alloc<phi::dtype::Real<T>>(out);
    return;
  }
  dev_ctx.template Alloc<phi::dtype::Real<T>>(out);
  // The allocation of imag here is redundant and could be optimized.
  DenseTensor imag;
  imag.Resize(x.dims());
  dev_ctx.template Alloc<phi::dtype::Real<T>>(&imag);
  int r = xfft_internal::xpu::complex_spilt(
      dev_ctx.x_context()->xpu_stream,
      out->numel(),
      reinterpret_cast<const XPUComplexType*>(x.data<T>()),
      out->data<phi::dtype::Real<T>>(),
      imag.data<phi::dtype::Real<T>>());
  PADDLE_ENFORCE_XPU_SUCCESS(r);
}

template <typename T, typename Context>
void ImagKernel(const Context& dev_ctx,
                const DenseTensor& x,
                DenseTensor* out) {
  using XPUComplexType =
      typename XPUComplexTypeTrait<phi::dtype::Real<T>>::Type;
  if (out->numel() == 0) {
    dev_ctx.template Alloc<phi::dtype::Real<T>>(out);
    return;
  }
  dev_ctx.template Alloc<phi::dtype::Real<T>>(out);
  // The allocation of ‘real’ here is redundant and could be optimized.
  DenseTensor real;
  real.Resize(x.dims());
  dev_ctx.template Alloc<phi::dtype::Real<T>>(&real);
  int r = xfft_internal::xpu::complex_spilt(
      dev_ctx.x_context()->xpu_stream,
      out->numel(),
      reinterpret_cast<const XPUComplexType*>(x.data<T>()),
      real.data<phi::dtype::Real<T>>(),
      out->data<phi::dtype::Real<T>>());
  PADDLE_ENFORCE_XPU_SUCCESS(r);
}

template <typename T, typename Context>
void ComplexKernel(const Context& dev_ctx,
                   const DenseTensor& x,
                   const DenseTensor& y,
                   DenseTensor* out) {
  using C = phi::dtype::complex<T>;
  using XPUComplexType = typename XPUComplexTypeTrait<T>::Type;
  if (out->numel() == 0) {
    dev_ctx.template Alloc<C>(out);
    return;
  }
  auto x_dims = x.dims();
  auto y_dims = y.dims();
  auto out_dims = funcs::BroadcastTwoDims(x_dims, y_dims);
  std::vector<int64_t> out_dims_vec = vectorize(out_dims);

  DenseTensor broadcasted_x, broadcasted_y;
  const T* x_data = nullptr;
  const T* y_data = nullptr;

  if (x_dims == out_dims) {
    x_data = x.data<T>();
  } else {
    broadcasted_x.Resize(out_dims);
    dev_ctx.template Alloc<T>(&broadcasted_x);
    ExpandKernel<T, Context>(
        dev_ctx, x, phi::IntArray(out_dims_vec), &broadcasted_x);
    x_data = broadcasted_x.data<T>();
  }

  if (y_dims == out_dims) {
    y_data = y.data<T>();
  } else {
    broadcasted_y.Resize(out_dims);
    dev_ctx.template Alloc<T>(&broadcasted_y);
    ExpandKernel<T, Context>(
        dev_ctx, y, phi::IntArray(out_dims_vec), &broadcasted_y);
    y_data = broadcasted_y.data<T>();
  }

  dev_ctx.template Alloc<C>(out);
  int r = xfft_internal::xpu::combine_as_complex(
      dev_ctx.x_context()->xpu_stream,
      out->numel(),
      x_data,
      y_data,
      reinterpret_cast<XPUComplexType*>(out->data<C>()));
  PADDLE_ENFORCE_XPU_SUCCESS(r);
}
}  // namespace phi

PD_REGISTER_KERNEL(conj,
                   XPU,
                   ALL_LAYOUT,
                   phi::ConjKernel,
                   bool,
                   int,
                   int64_t,
                   float,
                   double,
                   phi::float16,
                   phi::bfloat16,
                   phi::complex64,
                   phi::complex128) {}

PD_REGISTER_KERNEL(
    real, XPU, ALL_LAYOUT, phi::RealKernel, phi::complex64, phi::complex128) {
  kernel->OutputAt(0).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
}

PD_REGISTER_KERNEL(
    imag, XPU, ALL_LAYOUT, phi::ImagKernel, phi::complex64, phi::complex128) {
  kernel->OutputAt(0).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
}

PD_REGISTER_KERNEL(
    complex, XPU, ALL_LAYOUT, phi::ComplexKernel, float, double) {
  kernel->OutputAt(0).SetDataType(phi::dtype::ToComplex(kernel_key.dtype()));
}
#endif  // PADDLE_WITH_XPU_FFT