// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2009 Mark Borgerding mark a borgerding net
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0

#ifndef EIGEN_FFT_FFTW_IMPL_H
#define EIGEN_FFT_FFTW_IMPL_H

// IWYU pragma: private
#include "./InternalHeaderCheck.h"

#include <memory>
#include <mutex>

namespace Eigen {

namespace internal {

// FFTW uses non-const arguments,
// so const_cast is needed for all the args it uses.
//
// This should be safe as long as
// 1. we use FFTW_ESTIMATE for all our planning
//       see the FFTW docs section 4.3.2 "Planner Flags"
// 2. fftw_complex is compatible with std::complex
//    This assumes std::complex<T> layout is array of size 2 with real,imag
template <typename T>
inline T *fftw_cast(const T *p) {
  return const_cast<T *>(p);
}

inline fftw_complex *fftw_cast(const std::complex<double> *p) {
  return const_cast<fftw_complex *>(reinterpret_cast<const fftw_complex *>(p));
}

inline fftwf_complex *fftw_cast(const std::complex<float> *p) {
  return const_cast<fftwf_complex *>(reinterpret_cast<const fftwf_complex *>(p));
}

inline fftwl_complex *fftw_cast(const std::complex<long double> *p) {
  return const_cast<fftwl_complex *>(reinterpret_cast<const fftwl_complex *>(p));
}

// The FFTW planner is not thread-safe: fftw_execute and its new-array variants,
// which is what this backend runs transforms through, are the only entry points
// that may be called concurrently (FFTW manual, "Thread safety"), so plan
// creation and destruction are serialized through this mutex. A template
// static gives the header-only definition; std::mutex is
// constexpr-constructible, so the mutex is ready before any thread starts.
// The planner state it stands in for is one per process, so the mutex must be
// too: the explicit default visibility is what keeps the definition from being
// bound locally under -fvisibility=hidden, where each library planning through
// Eigen would get a mutex of its own and serialize nothing between them. The
// module documentation covers what no header can reach, which needs FFTW's own
// fftw_make_planner_thread_safe().
#if EIGEN_HAS_ATTRIBUTE(visibility) && !EIGEN_OS_WIN
#define EIGEN_FFTW_PLANNER_MUTEX_VISIBILITY __attribute__((visibility("default")))
#else
#define EIGEN_FFTW_PLANNER_MUTEX_VISIBILITY
#endif

template <typename Dummy = void>
struct fftw_planner_lock {
  static EIGEN_FFTW_PLANNER_MUTEX_VISIBILITY std::mutex mutex;
};
template <typename Dummy>
EIGEN_FFTW_PLANNER_MUTEX_VISIBILITY std::mutex fftw_planner_lock<Dummy>::mutex;

inline std::mutex &fftw_planner_mutex() { return fftw_planner_lock<>::mutex; }

template <typename PlanFactory>
inline decltype(auto) fftw_make_plan(PlanFactory factory) {
  std::lock_guard<std::mutex> lock(fftw_planner_mutex());
  return factory();
}

template <typename T>
struct fftw_plan {};

template <>
struct fftw_plan<float> {
  typedef float scalar_type;
  typedef fftwf_complex complex_type;
  std::shared_ptr<fftwf_plan_s> m_plan;
  fftw_plan() = default;

  void set_plan(fftwf_plan p) {
    m_plan.reset(p, [](fftwf_plan plan) {
      std::lock_guard<std::mutex> lock(fftw_planner_mutex());
      fftwf_destroy_plan(plan);
    });
  }
  inline void fwd(complex_type *dst, complex_type *src, int nfft) {
    if (!m_plan)
      set_plan(fftw_make_plan(
          [&] { return fftwf_plan_dft_1d(nfft, src, dst, FFTW_FORWARD, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftwf_execute_dft(m_plan.get(), src, dst);
  }
  inline void inv(complex_type *dst, complex_type *src, int nfft) {
    if (!m_plan)
      set_plan(fftw_make_plan(
          [&] { return fftwf_plan_dft_1d(nfft, src, dst, FFTW_BACKWARD, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftwf_execute_dft(m_plan.get(), src, dst);
  }
  inline void fwd(complex_type *dst, scalar_type *src, int nfft) {
    if (!m_plan)
      set_plan(
          fftw_make_plan([&] { return fftwf_plan_dft_r2c_1d(nfft, src, dst, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftwf_execute_dft_r2c(m_plan.get(), src, dst);
  }
  inline void inv(scalar_type *dst, complex_type *src, int nfft) {
    if (!m_plan)
      set_plan(
          fftw_make_plan([&] { return fftwf_plan_dft_c2r_1d(nfft, src, dst, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftwf_execute_dft_c2r(m_plan.get(), src, dst);
  }

  inline void fwd2(complex_type *dst, complex_type *src, int n0, int n1) {
    if (!m_plan)
      set_plan(fftw_make_plan(
          [&] { return fftwf_plan_dft_2d(n0, n1, src, dst, FFTW_FORWARD, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftwf_execute_dft(m_plan.get(), src, dst);
  }
  inline void inv2(complex_type *dst, complex_type *src, int n0, int n1) {
    if (!m_plan)
      set_plan(fftw_make_plan(
          [&] { return fftwf_plan_dft_2d(n0, n1, src, dst, FFTW_BACKWARD, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftwf_execute_dft(m_plan.get(), src, dst);
  }
};
template <>
struct fftw_plan<double> {
  typedef double scalar_type;
  typedef fftw_complex complex_type;
  std::shared_ptr<fftw_plan_s> m_plan;
  fftw_plan() = default;

  void set_plan(::fftw_plan p) {
    m_plan.reset(p, [](::fftw_plan plan) {
      std::lock_guard<std::mutex> lock(fftw_planner_mutex());
      fftw_destroy_plan(plan);
    });
  }
  inline void fwd(complex_type *dst, complex_type *src, int nfft) {
    if (!m_plan)
      set_plan(fftw_make_plan(
          [&] { return fftw_plan_dft_1d(nfft, src, dst, FFTW_FORWARD, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftw_execute_dft(m_plan.get(), src, dst);
  }
  inline void inv(complex_type *dst, complex_type *src, int nfft) {
    if (!m_plan)
      set_plan(fftw_make_plan(
          [&] { return fftw_plan_dft_1d(nfft, src, dst, FFTW_BACKWARD, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftw_execute_dft(m_plan.get(), src, dst);
  }
  inline void fwd(complex_type *dst, scalar_type *src, int nfft) {
    if (!m_plan)
      set_plan(
          fftw_make_plan([&] { return fftw_plan_dft_r2c_1d(nfft, src, dst, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftw_execute_dft_r2c(m_plan.get(), src, dst);
  }
  inline void inv(scalar_type *dst, complex_type *src, int nfft) {
    if (!m_plan)
      set_plan(
          fftw_make_plan([&] { return fftw_plan_dft_c2r_1d(nfft, src, dst, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftw_execute_dft_c2r(m_plan.get(), src, dst);
  }
  inline void fwd2(complex_type *dst, complex_type *src, int n0, int n1) {
    if (!m_plan)
      set_plan(fftw_make_plan(
          [&] { return fftw_plan_dft_2d(n0, n1, src, dst, FFTW_FORWARD, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftw_execute_dft(m_plan.get(), src, dst);
  }
  inline void inv2(complex_type *dst, complex_type *src, int n0, int n1) {
    if (!m_plan)
      set_plan(fftw_make_plan(
          [&] { return fftw_plan_dft_2d(n0, n1, src, dst, FFTW_BACKWARD, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftw_execute_dft(m_plan.get(), src, dst);
  }
};
template <>
struct fftw_plan<long double> {
  typedef long double scalar_type;
  typedef fftwl_complex complex_type;
  std::shared_ptr<fftwl_plan_s> m_plan;
  fftw_plan() = default;

  void set_plan(fftwl_plan p) {
    m_plan.reset(p, [](fftwl_plan plan) {
      std::lock_guard<std::mutex> lock(fftw_planner_mutex());
      fftwl_destroy_plan(plan);
    });
  }
  inline void fwd(complex_type *dst, complex_type *src, int nfft) {
    if (!m_plan)
      set_plan(fftw_make_plan(
          [&] { return fftwl_plan_dft_1d(nfft, src, dst, FFTW_FORWARD, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftwl_execute_dft(m_plan.get(), src, dst);
  }
  inline void inv(complex_type *dst, complex_type *src, int nfft) {
    if (!m_plan)
      set_plan(fftw_make_plan(
          [&] { return fftwl_plan_dft_1d(nfft, src, dst, FFTW_BACKWARD, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftwl_execute_dft(m_plan.get(), src, dst);
  }
  inline void fwd(complex_type *dst, scalar_type *src, int nfft) {
    if (!m_plan)
      set_plan(
          fftw_make_plan([&] { return fftwl_plan_dft_r2c_1d(nfft, src, dst, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftwl_execute_dft_r2c(m_plan.get(), src, dst);
  }
  inline void inv(scalar_type *dst, complex_type *src, int nfft) {
    if (!m_plan)
      set_plan(
          fftw_make_plan([&] { return fftwl_plan_dft_c2r_1d(nfft, src, dst, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftwl_execute_dft_c2r(m_plan.get(), src, dst);
  }
  inline void fwd2(complex_type *dst, complex_type *src, int n0, int n1) {
    if (!m_plan)
      set_plan(fftw_make_plan(
          [&] { return fftwl_plan_dft_2d(n0, n1, src, dst, FFTW_FORWARD, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftwl_execute_dft(m_plan.get(), src, dst);
  }
  inline void inv2(complex_type *dst, complex_type *src, int n0, int n1) {
    if (!m_plan)
      set_plan(fftw_make_plan(
          [&] { return fftwl_plan_dft_2d(n0, n1, src, dst, FFTW_BACKWARD, FFTW_ESTIMATE | FFTW_PRESERVE_INPUT); }));
    fftwl_execute_dft(m_plan.get(), src, dst);
  }
};

template <typename Scalar_>
struct fftw_impl {
  typedef Scalar_ Scalar;
  typedef std::complex<Scalar> Complex;

  inline void clear() { m_plans.clear(); }

  // complex-to-complex forward FFT
  inline void fwd(Complex *dst, const Complex *src, int nfft) {
    get_plan(nfft, false, /*real_io=*/false, dst, src).fwd(fftw_cast(dst), fftw_cast(src), nfft);
  }

  // real-to-complex forward FFT
  inline void fwd(Complex *dst, const Scalar *src, int nfft) {
    get_plan(nfft, false, /*real_io=*/true, dst, src).fwd(fftw_cast(dst), fftw_cast(src), nfft);
  }

  // 2-d complex-to-complex
  inline void fwd2(Complex *dst, const Complex *src, int n0, int n1) {
    get_plan(n0, n1, false, /*real_io=*/false, dst, src).fwd2(fftw_cast(dst), fftw_cast(src), n0, n1);
  }

  // inverse complex-to-complex
  inline void inv(Complex *dst, const Complex *src, int nfft) {
    get_plan(nfft, true, /*real_io=*/false, dst, src).inv(fftw_cast(dst), fftw_cast(src), nfft);
  }

  // half-complex to scalar
  inline void inv(Scalar *dst, const Complex *src, int nfft) {
    get_plan(nfft, true, /*real_io=*/true, dst, src).inv(fftw_cast(dst), fftw_cast(src), nfft);
  }

  // 2-d complex-to-complex
  inline void inv2(Complex *dst, const Complex *src, int n0, int n1) {
    get_plan(n0, n1, true, /*real_io=*/false, dst, src).inv2(fftw_cast(dst), fftw_cast(src), n0, n1);
  }

 protected:
  typedef fftw_plan<Scalar> PlanData;

  typedef Eigen::numext::int64_t int64_t;

  typedef std::map<int64_t, PlanData> PlanMap;

  PlanMap m_plans;

  // Pack (inverse, real_io, inplace, aligned) into 4 contiguous low bits of
  // the cache key.  real_io distinguishes r2c/c2r from c2c so that reusing
  // the same FFT object across real-input and complex-input transforms
  // doesn't return a cached plan of the wrong kind.
  static int64_t plan_flags(bool inverse, bool real_io, void *dst, const void *src) {
    bool inplace = (dst == src);
    bool aligned = ((reinterpret_cast<size_t>(src) & 15) | (reinterpret_cast<size_t>(dst) & 15)) == 0;
    return (inverse << 3) | (real_io << 2) | (inplace << 1) | aligned;
  }

  inline PlanData &get_plan(int nfft, bool inverse, bool real_io, void *dst, const void *src) {
    int64_t key = ((nfft << 4) | plan_flags(inverse, real_io, dst, src)) << 1;
    return m_plans[key];
  }

  inline PlanData &get_plan(int n0, int n1, bool inverse, bool real_io, void *dst, const void *src) {
    int64_t key = (((((int64_t)n0) << 31) | (n1 << 4) | plan_flags(inverse, real_io, dst, src)) << 1) + 1;
    return m_plans[key];
  }
};

}  // end namespace internal

}  // end namespace Eigen

#endif  // EIGEN_FFT_FFTW_IMPL_H