// 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 <mutex>
#include <stdexcept>

#include <fftw3.h>

namespace {

// FFTW planning is not thread-safe (execution is); same idiom as FFTIndexerCPU.
std::mutex &FftwPlanMutex() {
    static std::mutex m;
    return m;
}

// 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(FftwPlanMutex());
        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(FftwPlanMutex());
        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;
    }
};

std::vector<float> Transpose(const std::vector<float> &src, int64_t h, int64_t w) {
    std::vector<float> out(src.size());
    for (int64_t y = 0; y < h; y++)
        for (int64_t x = 0; x < w; x++)
            out[static_cast<size_t>(x * h + y)] = src[static_cast<size_t>(y * w + x)];
    return out;
}

} // namespace

BeamCenterConvSurfaces2D BeamCenterFFTCPU::PointSurfaces(const std::vector<float> &a,
                                                         const std::vector<float> &m, int64_t h,
                                                         int64_t w) {
    Rfft2 fft(BeamCenterFFTPadSize(2 * h), BeamCenterFFTPadSize(2 * w));
    const auto A = fft.Forward(a, h, w);
    const auto M = fft.Forward(m, h, w);
    std::vector<float> a2(a.size());
    for (size_t i = 0; i < a2.size(); i++)
        a2[i] = a[i] * a[i];
    const auto A2 = fft.Forward(a2, h, w);

    BeamCenterConvSurfaces2D out;
    out.C = fft.InverseProduct(A, A, 2 * h, 2 * w);
    out.S = fft.InverseProduct(A, M, 2 * h, 2 * w);
    out.Q = fft.InverseProduct(A2, M, 2 * h, 2 * w);
    out.D = fft.InverseProduct(M, M, 2 * h, 2 * w);
    return out;
}

// Each sequence is correlated with its own mirror and the four accumulators are summed over the
// other coordinate in double complex (the same cancellation argument as the 2D surface), so only
// four inverse transforms are needed however many sequences there are.
BeamCenterConvSurfaces1D BeamCenterFFTCPU::LineSurfaces(const std::vector<float> &a_in,
                                                        const std::vector<float> &m_in, int64_t h,
                                                        int64_t w, BeamCenterMirror mirror) {
    // Mirroring the column coordinate is the same computation on the transposed image.
    const std::vector<float> at = mirror == BeamCenterMirror::Rows ? std::vector<float>()
                                                                   : Transpose(a_in, h, w);
    const std::vector<float> mt = mirror == BeamCenterMirror::Rows ? std::vector<float>()
                                                                   : Transpose(m_in, h, w);
    const std::vector<float> &a = mirror == BeamCenterMirror::Rows ? a_in : at;
    const std::vector<float> &m = mirror == BeamCenterMirror::Rows ? m_in : mt;
    const int64_t nseq = mirror == BeamCenterMirror::Rows ? h : w;
    const int64_t nbatch = mirror == BeamCenterMirror::Rows ? w : h;

    const int64_t nfft = BeamCenterFFTPadSize(2 * nseq);
    const int64_t nc = nfft / 2 + 1;

    std::vector<float> in(static_cast<size_t>(nfft));
    std::vector<std::complex<float>> A(static_cast<size_t>(nc)), M(static_cast<size_t>(nc)),
            A2(static_cast<size_t>(nc));
    std::vector<std::complex<double>> accC(static_cast<size_t>(nc)), accS(static_cast<size_t>(nc)),
            accQ(static_cast<size_t>(nc)), accD(static_cast<size_t>(nc));

    fftwf_plan fwd, bwd;
    {
        std::unique_lock lock(FftwPlanMutex());
        fwd = fftwf_plan_dft_r2c_1d(static_cast<int>(nfft), in.data(),
                                    reinterpret_cast<fftwf_complex *>(A.data()), FFTW_ESTIMATE);
        bwd = fftwf_plan_dft_c2r_1d(static_cast<int>(nfft),
                                    reinterpret_cast<fftwf_complex *>(A.data()), in.data(),
                                    FFTW_ESTIMATE);
    }
    if (!fwd || !bwd)
        throw std::runtime_error("BeamCenterFFT: fftwf 1D plan failed");

    auto forward_into = [&](auto value_of, std::vector<std::complex<float>> &dst) {
        for (int64_t y = 0; y < nseq; y++)
            in[static_cast<size_t>(y)] = value_of(y);
        std::fill(in.begin() + nseq, in.end(), 0.0f);
        fftwf_execute_dft_r2c(fwd, in.data(), reinterpret_cast<fftwf_complex *>(A.data()));
        dst = A;
    };

    for (int64_t x = 0; x < nbatch; x++) {
        forward_into([&](int64_t y) { return a[static_cast<size_t>(y * nbatch + x)]; }, A2);
        std::swap(A, A2);   // A = spectrum of the sequence of a
        std::vector<std::complex<float>> Acol = A;
        forward_into([&](int64_t y) { return m[static_cast<size_t>(y * nbatch + x)]; }, M);
        forward_into(
                [&](int64_t y) {
                    const float v = a[static_cast<size_t>(y * nbatch + x)];
                    return v * v;
                },
                A2);
        for (int64_t i = 0; i < nc; i++) {
            const std::complex<double> ca(Acol[i]), cm(M[i]), ca2(A2[i]);
            accC[i] += ca * ca;
            accS[i] += ca * cm;
            accQ[i] += ca2 * cm;
            accD[i] += cm * cm;
        }
    }

    const double norm = 1.0 / static_cast<double>(nfft);
    auto inverse = [&](const std::vector<std::complex<double>> &acc) {
        for (int64_t i = 0; i < nc; i++)
            A[i] = std::complex<float>(acc[i]);
        fftwf_execute_dft_c2r(bwd, reinterpret_cast<fftwf_complex *>(A.data()), in.data());
        std::vector<double> out(static_cast<size_t>(2 * nseq));
        for (int64_t i = 0; i < 2 * nseq; i++)
            out[static_cast<size_t>(i)] = in[static_cast<size_t>(i)] * norm;
        return out;
    };
    BeamCenterConvSurfaces1D out;
    out.C = inverse(accC);
    out.S = inverse(accS);
    out.Q = inverse(accQ);
    out.D = inverse(accD);
    {
        std::unique_lock lock(FftwPlanMutex());
        fftwf_destroy_plan(fwd);
        fftwf_destroy_plan(bwd);
    }
    return out;
}