// Copyright (c) 2023 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 <vector>
#include "paddle/phi/common/type_traits.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/kernels/cpu/elementwise.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"

namespace phi {

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

  dev_ctx.template Alloc<C>(out);

  const size_t x_rank = x.dims().size();
  const size_t out_rank = out->dims().size();

  const size_t n_frames = out->dims()[out_rank - 1];
  const size_t seq_length = x.dims()[x_rank - 1];

  std::vector<int64_t> axes = {1};

  // Frame
  DenseTensor frames;
  DDim frames_dims(out->dims());
  frames_dims.at(axes.back()) = n_fft;
  frames.Resize(frames_dims);
  dev_ctx.template Alloc<T>(&frames);
  funcs::FrameFunctor<Context, T>()(dev_ctx,
                                    &x,
                                    &frames,
                                    seq_length,
                                    n_fft,
                                    n_frames,
                                    hop_length,
                                    /*is_grad*/ false);

  // Window
  DenseTensor frames_w;
  frames_w.Resize(frames_dims);
  dev_ctx.template Alloc<T>(&frames_w);
  funcs::ElementwiseCompute<funcs::MultiplyFunctor<T>, T, T>(
      dev_ctx,
      frames,
      window,
      funcs::MultiplyFunctor<T>(),
      &frames_w,
      axes.back());

  // FFTR2C
  funcs::FFTNormMode normalization;
  if (normalized) {
    normalization = funcs::get_norm_from_string("ortho", true);
  } else {
    normalization = funcs::get_norm_from_string("backward", true);
  }
  funcs::FFTR2CFunctor<Context, T, C> fft_r2c_func;

  if (onesided) {
    fft_r2c_func(dev_ctx, frames_w, out, axes, normalization, true);
  } else {
    DDim onesided_dims(out->dims());
    const int64_t onesided_axis_size = out->dims().at(axes.back()) / 2 + 1;
    onesided_dims.at(axes.back()) = onesided_axis_size;
    DenseTensor onesided_out;
    onesided_out.Resize(onesided_dims);
    dev_ctx.template Alloc<T>(&onesided_out);
    fft_r2c_func(dev_ctx, frames_w, &onesided_out, axes, normalization, true);
    funcs::FFTFillConj<Context, C>(dev_ctx, &onesided_out, out, axes);
  }
}
}  // namespace phi