// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "BeamCenterFFTCPU.h"

#include <algorithm>
#include <complex>
#include <cstring>
#include <future>
#include <memory>
#include <mutex>
#include <stdexcept>

#include <fftw3.h>

#include "../../common/FFTWPlannerLock.h"

namespace {

// One real-to-complex 2D FFT workspace of fixed padded size, reused for every forward transform.
class Rfft2 {
    int64_t ny, nx, nxc;
    std::vector<float> real;
    fftwf_plan fwd = nullptr;
    fftwf_plan bwd = nullptr;

public:
    std::vector<std::complex<float>> spectrum;

    Rfft2(int64_t ny_, int64_t nx_) : ny(ny_), nx(nx_), nxc(nx_ / 2 + 1) {
        real.resize(static_cast<size_t>(ny * nx));
        spectrum.resize(static_cast<size_t>(ny * nxc));
        std::unique_lock lock(FFTWPlannerMutex());
        fwd = fftwf_plan_dft_r2c_2d(static_cast<int>(ny), static_cast<int>(nx), real.data(),
                                    reinterpret_cast<fftwf_complex *>(spectrum.data()),
                                    FFTW_ESTIMATE);
        bwd = fftwf_plan_dft_c2r_2d(static_cast<int>(ny), static_cast<int>(nx),
                                    reinterpret_cast<fftwf_complex *>(spectrum.data()),
                                    real.data(), FFTW_ESTIMATE);
        if (!fwd || !bwd)
            throw std::runtime_error("BeamCenterFFT: fftwf plan failed");
    }
    Rfft2(const Rfft2 &) = delete;
    Rfft2 &operator=(const Rfft2 &) = delete;
    ~Rfft2() {
        std::unique_lock lock(FFTWPlannerMutex());
        if (fwd)
            fftwf_destroy_plan(fwd);
        if (bwd)
            fftwf_destroy_plan(bwd);
    }

    // Zero-pad `src` (height x width) into the workspace and return its spectrum.
    std::vector<std::complex<float>> Forward(const std::vector<float> &src, int64_t height,
                                             int64_t width) {
        std::fill(real.begin(), real.end(), 0.0f);
        for (int64_t y = 0; y < height; y++)
            std::memcpy(&real[static_cast<size_t>(y * nx)], &src[static_cast<size_t>(y * width)],
                        sizeof(float) * static_cast<size_t>(width));
        fftwf_execute(fwd);
        return spectrum;   // copy: the workspace is reused
    }

    // Inverse-transform the elementwise product u*v, normalised, cropped to out_h x out_w.
    std::vector<float> InverseProduct(const std::vector<std::complex<float>> &u,
                                      const std::vector<std::complex<float>> &v, int64_t out_h,
                                      int64_t out_w) {
        const float norm = 1.0f / (static_cast<float>(ny) * static_cast<float>(nx));
        for (size_t i = 0; i < spectrum.size(); i++)
            spectrum[i] = u[i] * v[i];
        fftwf_execute(bwd);
        std::vector<float> out(static_cast<size_t>(out_h * out_w));
        for (int64_t y = 0; y < out_h; y++)
            for (int64_t x = 0; x < out_w; x++)
                out[static_cast<size_t>(y * out_w + x)] =
                        real[static_cast<size_t>(y * nx + x)] * norm;
        return out;
    }
};

} // namespace

BeamCenterConvSurfaces2D BeamCenterFFTCPU::PointSurfaces(const std::vector<float> &a,
                                                         const std::vector<float> &m, int64_t h,
                                                         int64_t w) {
    // The three forward transforms are independent of each other, and so are the four inverse ones,
    // so each group runs at the same time, one workspace per transform. A workspace is built exactly
    // as a single one was - its own buffers and its own plan for them - so every transform is the one
    // it was, bit for bit; the four inverse ones reuse the three forward workspaces and one more.
    const int64_t ny = BeamCenterFFTPadSize(2 * h), nx = BeamCenterFFTPadSize(2 * w);
    std::vector<std::unique_ptr<Rfft2>> fft;
    for (int i = 0; i < 4; i++)
        fft.push_back(std::make_unique<Rfft2>(ny, nx));

    std::vector<float> a2(a.size());
    for (size_t i = 0; i < a2.size(); i++)
        a2[i] = a[i] * a[i];
    auto fA = std::async(std::launch::async, [&] { return fft[0]->Forward(a, h, w); });
    auto fM = std::async(std::launch::async, [&] { return fft[1]->Forward(m, h, w); });
    auto fA2 = std::async(std::launch::async, [&] { return fft[2]->Forward(a2, h, w); });
    const auto A = fA.get();
    const auto M = fM.get();
    const auto A2 = fA2.get();

    BeamCenterConvSurfaces2D out;
    auto fC = std::async(std::launch::async, [&] { return fft[0]->InverseProduct(A, A, 2 * h, 2 * w); });
    auto fS = std::async(std::launch::async, [&] { return fft[1]->InverseProduct(A, M, 2 * h, 2 * w); });
    auto fQ = std::async(std::launch::async, [&] { return fft[2]->InverseProduct(A2, M, 2 * h, 2 * w); });
    auto fD = std::async(std::launch::async, [&] { return fft[3]->InverseProduct(M, M, 2 * h, 2 * w); });
    out.C = fC.get();
    out.S = fS.get();
    out.Q = fQ.get();
    out.D = fD.get();
    return out;
}