// 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) {
    // Two transforms at a time, each on a workspace of its own. A workspace is padded to twice the
    // detector each way, so on a large detector the four workspaces and three spectra that the seven
    // transforms used to hold at once were gigabytes. In order: the forward transforms of the image
    // and the mask; the two inverse ones that read the image's spectrum, which is then handed back;
    // the squared image's forward transform beside the mask's own inverse; and the last inverse.
    // 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.
    const int64_t ny = BeamCenterFFTPadSize(2 * h), nx = BeamCenterFFTPadSize(2 * w);
    Rfft2 ws0(ny, nx), ws1(ny, nx);

    std::vector<std::complex<float>> A, M;
    {
        auto fA = std::async(std::launch::async, [&] { return ws0.Forward(a, h, w); });
        M = ws1.Forward(m, h, w);
        A = fA.get();
    }

    BeamCenterConvSurfaces2D out;
    {
        auto fC = std::async(std::launch::async, [&] { return ws0.InverseProduct(A, A, 2 * h, 2 * w); });
        out.S = ws1.InverseProduct(A, M, 2 * h, 2 * w);
        out.C = fC.get();
    }
    std::vector<std::complex<float>>().swap(A);

    std::vector<float> a2(a.size());
    for (size_t i = 0; i < a2.size(); i++)
        a2[i] = a[i] * a[i];
    {
        auto fD = std::async(std::launch::async, [&] { return ws1.InverseProduct(M, M, 2 * h, 2 * w); });
        const auto A2 = ws0.Forward(a2, h, w);
        out.Q = ws0.InverseProduct(A2, M, 2 * h, 2 * w);
        out.D = fD.get();
    }
    return out;
}