// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// 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-FileCopyrightText: The Eigen Authors
// SPDX-License-Identifier: MPL-2.0

#ifndef EIGEN_FFT_IMKLFFT_IMPL_H
#define EIGEN_FFT_IMKLFFT_IMPL_H

#include <mkl_dfti.h>

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

#include <complex>
#include <memory>

namespace Eigen {
namespace internal {
namespace imklfft {

#define RUN_OR_ASSERT(EXPR, ERROR_MSG)                    \
  {                                                       \
    MKL_LONG status = (EXPR);                             \
    eigen_assert(status == DFTI_NO_ERROR && (ERROR_MSG)); \
  };

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

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

/*
 * Parameters:
 * precision: enum, Precision of the transform: DFTI_SINGLE or DFTI_DOUBLE.
 * forward_domain: enum, Forward domain of the transform: DFTI_COMPLEX or
 * DFTI_REAL. dimension: MKL_LONG Dimension of the transform. sizes: MKL_LONG if
 * dimension = 1. Length of the transform for a one-dimensional transform. sizes:
 * Array of type MKL_LONG otherwise. Lengths of each dimension for a
 * multi-dimensional transform.
 */
inline void configure_descriptor(std::shared_ptr<DFTI_DESCRIPTOR>& handl, enum DFTI_CONFIG_VALUE precision,
                                 enum DFTI_CONFIG_VALUE forward_domain, MKL_LONG dimension, MKL_LONG* sizes) {
  eigen_assert(dimension == 1 || dimension == 2 && "Transformation dimension must be less than 3.");

  DFTI_DESCRIPTOR_HANDLE res = nullptr;
  if (dimension == 1) {
    RUN_OR_ASSERT(DftiCreateDescriptor(&res, precision, forward_domain, dimension, *sizes),
                  "DftiCreateDescriptor failed.")
    handl.reset(res, [](DFTI_DESCRIPTOR_HANDLE handle) { DftiFreeDescriptor(&handle); });
    if (forward_domain == DFTI_REAL) {
      // Set CCE storage
      RUN_OR_ASSERT(DftiSetValue(handl.get(), DFTI_CONJUGATE_EVEN_STORAGE, DFTI_COMPLEX_COMPLEX),
                    "DftiSetValue failed.")
    }
  } else {
    RUN_OR_ASSERT(DftiCreateDescriptor(&res, precision, DFTI_COMPLEX, dimension, sizes), "DftiCreateDescriptor failed.")
    handl.reset(res, [](DFTI_DESCRIPTOR_HANDLE handle) { DftiFreeDescriptor(&handle); });
  }

  RUN_OR_ASSERT(DftiSetValue(handl.get(), DFTI_PLACEMENT, DFTI_NOT_INPLACE), "DftiSetValue failed.")
  RUN_OR_ASSERT(DftiCommitDescriptor(handl.get()), "DftiCommitDescriptor failed.")
}

template <typename T>
struct plan {};

template <>
struct plan<float> {
  typedef float scalar_type;
  typedef MKL_Complex8 complex_type;

  std::shared_ptr<DFTI_DESCRIPTOR> m_plan;

  plan() = default;

  enum DFTI_CONFIG_VALUE precision = DFTI_SINGLE;

  inline void forward(complex_type* dst, complex_type* src, MKL_LONG nfft) {
    if (m_plan == 0) {
      configure_descriptor(m_plan, precision, DFTI_COMPLEX, 1, &nfft);
    }
    RUN_OR_ASSERT(DftiComputeForward(m_plan.get(), src, dst), "DftiComputeForward failed.")
  }

  inline void inverse(complex_type* dst, complex_type* src, MKL_LONG nfft) {
    if (m_plan == 0) {
      configure_descriptor(m_plan, precision, DFTI_COMPLEX, 1, &nfft);
    }
    RUN_OR_ASSERT(DftiComputeBackward(m_plan.get(), src, dst), "DftiComputeBackward failed.")
  }

  inline void forward(complex_type* dst, scalar_type* src, MKL_LONG nfft) {
    if (m_plan == 0) {
      configure_descriptor(m_plan, precision, DFTI_REAL, 1, &nfft);
    }
    RUN_OR_ASSERT(DftiComputeForward(m_plan.get(), src, dst), "DftiComputeForward failed.")
  }

  inline void inverse(scalar_type* dst, complex_type* src, MKL_LONG nfft) {
    if (m_plan == 0) {
      configure_descriptor(m_plan, precision, DFTI_REAL, 1, &nfft);
    }
    RUN_OR_ASSERT(DftiComputeBackward(m_plan.get(), src, dst), "DftiComputeBackward failed.")
  }

  inline void forward2(complex_type* dst, complex_type* src, int n0, int n1) {
    if (m_plan == 0) {
      MKL_LONG sizes[2] = {n0, n1};
      configure_descriptor(m_plan, precision, DFTI_COMPLEX, 2, sizes);
    }
    RUN_OR_ASSERT(DftiComputeForward(m_plan.get(), src, dst), "DftiComputeForward failed.")
  }

  inline void inverse2(complex_type* dst, complex_type* src, int n0, int n1) {
    if (m_plan == 0) {
      MKL_LONG sizes[2] = {n0, n1};
      configure_descriptor(m_plan, precision, DFTI_COMPLEX, 2, sizes);
    }
    RUN_OR_ASSERT(DftiComputeBackward(m_plan.get(), src, dst), "DftiComputeBackward failed.")
  }
};

template <>
struct plan<double> {
  typedef double scalar_type;
  typedef MKL_Complex16 complex_type;

  std::shared_ptr<DFTI_DESCRIPTOR> m_plan;

  plan() = default;

  enum DFTI_CONFIG_VALUE precision = DFTI_DOUBLE;

  inline void forward(complex_type* dst, complex_type* src, MKL_LONG nfft) {
    if (m_plan == 0) {
      configure_descriptor(m_plan, precision, DFTI_COMPLEX, 1, &nfft);
    }
    RUN_OR_ASSERT(DftiComputeForward(m_plan.get(), src, dst), "DftiComputeForward failed.")
  }

  inline void inverse(complex_type* dst, complex_type* src, MKL_LONG nfft) {
    if (m_plan == 0) {
      configure_descriptor(m_plan, precision, DFTI_COMPLEX, 1, &nfft);
    }
    RUN_OR_ASSERT(DftiComputeBackward(m_plan.get(), src, dst), "DftiComputeBackward failed.")
  }

  inline void forward(complex_type* dst, scalar_type* src, MKL_LONG nfft) {
    if (m_plan == 0) {
      configure_descriptor(m_plan, precision, DFTI_REAL, 1, &nfft);
    }
    RUN_OR_ASSERT(DftiComputeForward(m_plan.get(), src, dst), "DftiComputeForward failed.")
  }

  inline void inverse(scalar_type* dst, complex_type* src, MKL_LONG nfft) {
    if (m_plan == 0) {
      configure_descriptor(m_plan, precision, DFTI_REAL, 1, &nfft);
    }
    RUN_OR_ASSERT(DftiComputeBackward(m_plan.get(), src, dst), "DftiComputeBackward failed.")
  }

  inline void forward2(complex_type* dst, complex_type* src, int n0, int n1) {
    if (m_plan == 0) {
      MKL_LONG sizes[2] = {n0, n1};
      configure_descriptor(m_plan, precision, DFTI_COMPLEX, 2, sizes);
    }
    RUN_OR_ASSERT(DftiComputeForward(m_plan.get(), src, dst), "DftiComputeForward failed.")
  }

  inline void inverse2(complex_type* dst, complex_type* src, int n0, int n1) {
    if (m_plan == 0) {
      MKL_LONG sizes[2] = {n0, n1};
      configure_descriptor(m_plan, precision, DFTI_COMPLEX, 2, sizes);
    }
    RUN_OR_ASSERT(DftiComputeBackward(m_plan.get(), src, dst), "DftiComputeBackward failed.")
  }
};

template <typename Scalar_>
struct imklfft_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, /*real_io=*/false, dst, src).forward(complex_cast(dst), complex_cast(src), nfft);
  }

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

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

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

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

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

 private:
  std::map<int64_t, plan<Scalar>> m_plans;

  // Pack (real_io, inplace, aligned) into 3 contiguous low bits of the cache
  // key.  real_io distinguishes plans built with DFTI_REAL from those built
  // with DFTI_COMPLEX so reusing the same FFT object across real-input and
  // complex-input transforms doesn't return a cached plan of the wrong domain.
  // MKL DFTI descriptors are bidirectional, so forward and inverse share
  // slots intentionally.
  static int64_t plan_flags(bool real_io, void* dst, const void* src) {
    int inplace = (dst == src) ? 1 : 0;
    int aligned = ((reinterpret_cast<size_t>(src) & 15) | (reinterpret_cast<size_t>(dst) & 15)) == 0 ? 1 : 0;
    return (int(real_io) << 2) | (inplace << 1) | aligned;
  }

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

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

#undef RUN_OR_ASSERT

}  // namespace imklfft
}  // namespace internal
}  // namespace Eigen

#endif  // EIGEN_FFT_IMKLFFT_IMPL_H