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

#include "ModelFFT.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <complex>
#include <map>
#include <mutex>

#include <fftw3.h>

#include "gemmi/fail.hpp"

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

namespace {

// The r2c plan (or, inverse, the c2r plan) for an (nu, nv, nw) map stored u fastest, halving w - the
// layout gemmi's own transforms read and write. FFTW's planner is not thread-safe, so plans are made
// under the process-wide planner lock (which also guards this cache); executing one on new arrays
// (fftwf_execute_dft_r2c / _c2r) is, which is how every caller uses it.
fftwf_plan PlanFor(int nu, int nv, int nw, bool inverse) {
    static std::map<std::array<int, 4>, fftwf_plan> plans;
    std::lock_guard lock(FFTWPlannerMutex());
    const std::array<int, 4> key{nu, nv, nw, inverse};
    if (const auto it = plans.find(key); it != plans.end())
        return it->second;
    float *real = fftwf_alloc_real(static_cast<size_t>(nu) * nv * nw);
    fftwf_complex *cplx = fftwf_alloc_complex(static_cast<size_t>(nu) * nv * (nw / 2 + 1));
    // The last dimension is the one r2c halves; the strides put u fastest on both sides.
    fftwf_iodim dims[3] = {{nu, 1, 1}, {nv, nu, nu}, {nw, nu * nv, nu * nv}};
    fftwf_plan plan = nullptr;
    if (real != nullptr && cplx != nullptr)
        plan = inverse ? fftwf_plan_guru_dft_c2r(3, dims, 0, nullptr, cplx, real, FFTW_ESTIMATE)
                       : fftwf_plan_guru_dft_r2c(3, dims, 0, nullptr, real, cplx, FFTW_ESTIMATE);
    fftwf_free(real);
    fftwf_free(cplx);
    if (plan == nullptr)
        gemmi::fail("ModelFFT: no FFTW plan for the grid");
    plans.emplace(key, plan);
    return plan;
}

}  // namespace

gemmi::FPhiGrid<float> MapToFPhi(const gemmi::Grid<float> &map) {
    if (map.axis_order == gemmi::AxisOrder::ZYX)
        gemmi::fail("MapToFPhi(): ZYX order is not supported");
    gemmi::FPhiGrid<float> hkl;
    hkl.unit_cell = map.unit_cell;
    hkl.spacegroup = map.spacegroup;
    hkl.axis_order = map.axis_order;
    hkl.half_l = true;
    hkl.set_size_without_checking(map.nu, map.nv, map.nw / 2 + 1);
    const float norm = static_cast<float>(map.unit_cell.volume / map.point_count());

    const fftwf_plan plan = PlanFor(map.nu, map.nv, map.nw, false);
    // Buffers from fftwf_malloc: a plan may only be executed on arrays aligned as the ones it was made
    // on were, which the vectors' own allocations do not promise.
    float *in = fftwf_alloc_real(map.data.size());
    fftwf_complex *out = fftwf_alloc_complex(hkl.data.size());
    if (in == nullptr || out == nullptr) {
        fftwf_free(in);
        fftwf_free(out);
        gemmi::fail("MapToFPhi(): cannot allocate the FFT buffers");
    }
    std::copy(map.data.begin(), map.data.end(), in);
    fftwf_execute_dft_r2c(plan, in, out);
    // Conjugated, as gemmi does: FFTW's forward transform carries exp(-2 pi i h.x), and a structure
    // factor is the sum over exp(+2 pi i h.x).
    for (size_t i = 0; i < hkl.data.size(); ++i)
        hkl.data[i] = std::complex<float>(out[i][0] * norm, -out[i][1] * norm);
    fftwf_free(in);
    fftwf_free(out);
    return hkl;
}

void MapFromFPhiExecute(int nu, int nv, int nw, std::complex<float> *in, float *out) {
    static_assert(sizeof(std::complex<float>) == sizeof(fftwf_complex));
    fftwf_execute_dft_c2r(PlanFor(nu, nv, nw, true), reinterpret_cast<fftwf_complex *>(in), out);
}

gemmi::FPhiGrid<float> MapToFPhiInParallel(const gemmi::Grid<float> &map, int reach, size_t nthreads) {
    if (map.axis_order == gemmi::AxisOrder::ZYX)
        gemmi::fail("MapToFPhiInParallel(): ZYX order is not supported");
    const int nu = map.nu, nv = map.nv, nw = map.nw, nwh = nw / 2 + 1, nk = 2 * reach + 1;
    const size_t plane = static_cast<size_t>(nu) * nv;
    float *in = fftwf_alloc_real(map.data.size());
    fftwf_complex *half = fftwf_alloc_complex(plane * nwh);
    if (in == nullptr || half == nullptr) {
        fftwf_free(in);
        fftwf_free(half);
        gemmi::fail("MapToFPhiInParallel(): cannot allocate the FFT buffers");
    }
    // Made once here; executing a plan on new arrays is safe from several threads at once. Unaligned, as a
    // column or a plane starts wherever it does.
    fftwf_plan columns = nullptr, planes = nullptr;
    {
        std::lock_guard lock(FFTWPlannerMutex());
        // The nu columns of one v, w along the stride of a plane.
        fftwf_iodim col_dim{nw, static_cast<int>(plane), static_cast<int>(plane)};
        fftwf_iodim col_many{nu, 1, 1};
        columns = fftwf_plan_guru_dft_r2c(1, &col_dim, 1, &col_many, in, half, FFTW_ESTIMATE | FFTW_UNALIGNED);
        fftwf_iodim plane_dims[2] = {{nv, nu, nu}, {nu, 1, 1}};
        planes = fftwf_plan_guru_dft(2, plane_dims, 0, nullptr, half, half, FFTW_FORWARD, FFTW_ESTIMATE | FFTW_UNALIGNED);
    }
    if (columns == nullptr || planes == nullptr)
        gemmi::fail("MapToFPhiInParallel(): no FFTW plan for the grid");
    std::copy(map.data.begin(), map.data.end(), in);
    ParallelFor(nv, nthreads, [&](int v) {
        fftwf_execute_dft_r2c(columns, in + static_cast<size_t>(v) * nu, half + static_cast<size_t>(v) * nu);
    });
    ParallelFor(reach + 1, nthreads, [&](int l) {
        fftwf_execute_dft(planes, half + l * plane, half + l * plane);
    });

    gemmi::FPhiGrid<float> hkl;
    hkl.unit_cell = map.unit_cell;
    hkl.spacegroup = map.spacegroup;
    hkl.axis_order = map.axis_order;
    hkl.half_l = true;
    hkl.set_size_without_checking(nu, nk, reach + 1);
    const float norm = static_cast<float>(map.unit_cell.volume / map.point_count());
    // Conjugated and scaled as MapToFPhi does it, and cropped as CropHalfL crops it.
    ParallelFor(reach + 1, nthreads, [&](int l) {
        for (int k = -reach; k <= reach; k++) {
            const fftwf_complex *row = half + l * plane + static_cast<size_t>(nu) * gemmi::modulo(k, nv);
            std::complex<float> *out = &hkl.data[static_cast<size_t>(nu) * (gemmi::modulo(k, nk) + static_cast<size_t>(nk) * l)];
            for (int u = 0; u < nu; u++)
                out[u] = std::complex<float>(row[u][0] * norm, -row[u][1] * norm);
        }
    });
    {
        std::lock_guard lock(FFTWPlannerMutex());
        fftwf_destroy_plan(columns);
        fftwf_destroy_plan(planes);
    }
    fftwf_free(in);
    fftwf_free(half);
    return hkl;
}

gemmi::Grid<float> MapFromFPhi(const gemmi::FPhiGrid<float> &hkl) {
    if (hkl.axis_order == gemmi::AxisOrder::ZYX || !hkl.half_l)
        gemmi::fail("MapFromFPhi(): only an XYZ half-l grid is supported");
    gemmi::Grid<float> map;
    map.spacegroup = hkl.spacegroup;
    map.unit_cell = hkl.unit_cell;
    map.set_size(hkl.nu, hkl.nv, 2 * (hkl.nw - 1));
    map.axis_order = hkl.axis_order;
    const float norm = static_cast<float>(1.0 / hkl.unit_cell.volume);

    const fftwf_plan plan = PlanFor(map.nu, map.nv, map.nw, true);
    fftwf_complex *in = fftwf_alloc_complex(hkl.data.size());
    float *out = fftwf_alloc_real(map.data.size());
    if (in == nullptr || out == nullptr) {
        fftwf_free(in);
        fftwf_free(out);
        gemmi::fail("MapFromFPhi(): cannot allocate the FFT buffers");
    }
    // Conjugated, as gemmi does: rho(x) = 1/V sum F exp(-2 pi i h.x), and FFTW's backward transform
    // carries exp(+2 pi i h.x). A missing coefficient (NaN) counts as zero, also as gemmi does.
    for (size_t i = 0; i < hkl.data.size(); ++i) {
        const std::complex<float> x = hkl.data[i];
        in[i][0] = std::isnan(x.imag()) ? 0.0f : x.real();
        in[i][1] = std::isnan(x.imag()) ? 0.0f : -x.imag();
    }
    fftwf_execute_dft_c2r(plan, in, out);
    for (size_t i = 0; i < map.data.size(); ++i)
        map.data[i] = out[i] * norm;
    fftwf_free(in);
    fftwf_free(out);
    return map;
}

// The rows keep their length: cut to 2 reach + 1 as well, the rotation function read them 15% slower on the
// GPU (measured on a large cell).
gemmi::FPhiGrid<float> CropHalfL(const gemmi::FPhiGrid<float> &full, int reach) {
    gemmi::FPhiGrid<float> out;
    out.unit_cell = full.unit_cell;
    out.spacegroup = full.spacegroup;
    out.axis_order = full.axis_order;
    out.half_l = true;
    const int nu = full.nu, n = 2 * reach + 1;
    out.set_size_without_checking(nu, n, reach + 1);
    for (int l = 0; l <= reach; l++)
        for (int k = -reach; k <= reach; k++)
            std::copy_n(&full.data[static_cast<size_t>(nu) * (gemmi::modulo(k, full.nv) + static_cast<size_t>(full.nv) * l)],
                        nu, &out.data[static_cast<size_t>(nu) * (gemmi::modulo(k, n) + static_cast<size_t>(n) * l)]);
    return out;
}