// Copyright (c) 2024 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.

#pragma once

#include "paddle/common/enforce.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/utils/data_type.h"
#include "paddle/phi/kernels/complex_kernel.h"
#include "paddle/phi/kernels/funcs/broadcast_function.h"
#include "paddle/phi/kernels/funcs/elementwise_base.h"
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
#include "paddle/phi/kernels/funcs/fft.h"
#include "paddle/phi/kernels/funcs/fft_fill_conj.h"
#include "paddle/phi/kernels/funcs/frame_functor.h"
#include "paddle/phi/kernels/funcs/padding.h"

namespace phi {
// Multiply
template <typename T>
using MulFunctor = funcs::MultiplyFunctor<T>;

// It is a common implementation to compute binary calculation with the support
// of broadcast, supporting both CPU and GPU.
// - CPU implementation cannot support the case when x needs broadcast, thus
//   this function need to be called with XxxFunctor and XxxInverseFunctor,
//   like AddFunctor and InverseAddFunctor.
// - GPU implementation supports all the broadcast cases, thus there is no need
//   to define and call with XxxInverseFunctor.
// TODO(liuyiqun): optimize the CPU implementation to support all broadcast
// cases and avoid the need of XxxInverseFunctor.
template <typename Functor, typename Context, typename T, typename OutType = T>
void ElementwiseComputeEx(const Context& dev_ctx,
                          const DenseTensor* x,
                          const DenseTensor* y,
                          int axis,
                          Functor func,
                          DenseTensor* z) {
  dev_ctx.template Alloc<OutType>(z);
  funcs::ElementwiseCompute<Functor, T, OutType>(
      dev_ctx, *x, *y, func, z, axis);
}

template <typename T, typename Context>
void StftGradKernel(const Context& dev_ctx,
                    const DenseTensor& x,
                    const DenseTensor& window,
                    const DenseTensor& out_grad,
                    int n_fft,
                    int hop_length,
                    bool normalized,
                    bool onesided,
                    DenseTensor* x_grad) {
  using C = dtype::complex<T>;

  const auto* dy = &out_grad;
  auto* dx = x_grad;
  dev_ctx.template Alloc<T>(x_grad);

  const size_t dy_rank = dy->dims().size();
  const size_t dx_rank = dx->dims().size();

  const size_t n_frames = dy->dims()[dy_rank - 1];
  const size_t seq_length = dx->dims()[dx_rank - 1];

  std::vector<int64_t> axes = {1};
  DenseTensor d_frames_w;
  DDim d_frames_dims(dy->dims());
  d_frames_dims.at(axes.back()) = n_fft;
  d_frames_w.Resize(d_frames_dims);
  dev_ctx.template Alloc<T>(&d_frames_w);

  DenseTensor complex_d_frames_w;
  complex_d_frames_w.Resize(d_frames_dims);
  dev_ctx.template Alloc<C>(&complex_d_frames_w);

  // dy -> d_frames_w
  funcs::FFTNormMode normalization;
  if (normalized) {
    normalization = funcs::get_norm_from_string("ortho", true);
  } else {
    normalization = funcs::get_norm_from_string("backward", true);
  }
  funcs::FFTC2CFunctor<Context, C, C> fft_c2c_func;

  if (!onesided) {
    fft_c2c_func(dev_ctx, *dy, &complex_d_frames_w, axes, normalization, false);
  } else {
    DenseTensor full_dy;
    full_dy.Resize(d_frames_dims);
    dev_ctx.template Alloc<C>(&full_dy);
    const int64_t zero_length_64 =
        full_dy.dims().at(axes.back()) - dy->dims().at(axes.back());
    PADDLE_ENFORCE_LE_INT_MAX(zero_length_64, "stft zero padding length");
    auto zero_length = static_cast<int>(zero_length_64);
    auto rank = dy->dims().size();

    std::vector<int> pads(rank * 2, 0);
    pads[axes.back() * 2 + 1] = zero_length;

    funcs::PaddingFunctor<Context, C>(
        rank, dev_ctx, pads, static_cast<C>(0), *dy, &full_dy);
    fft_c2c_func(
        dev_ctx, full_dy, &complex_d_frames_w, axes, normalization, false);
  }
  RealKernel<C>(dev_ctx, complex_d_frames_w, &d_frames_w);

  // d_frames_w -> d_frames
  DenseTensor d_frames;
  d_frames.Resize(d_frames_dims);
  dev_ctx.template Alloc<T>(&d_frames);
  const DenseTensor d_frames_w_const = d_frames_w;
  ElementwiseComputeEx<MulFunctor<T>, Context, T>(dev_ctx,
                                                  &d_frames_w_const,
                                                  &window,
                                                  axes.back(),
                                                  MulFunctor<T>(),
                                                  &d_frames);

  // d_frames -> dx
  funcs::FrameFunctor<Context, T>()(dev_ctx,
                                    &d_frames,
                                    dx,
                                    seq_length,
                                    n_fft,
                                    n_frames,
                                    hop_length,
                                    /*is_grad*/ true);
}
}  // namespace phi