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

#include <catch2/catch_all.hpp>

#include <algorithm>
#include <cmath>
#include <cstring>
#include <map>
#include <random>
#include <vector>

#include "../common/CorrelationCoefficient.h"
#include "../common/CUDAWrapper.h"
#include "../common/Logger.h"
#include "../image_analysis/scale_merge/RotationScaleMerge.h"

namespace {

// A synthetic crystal measured as pre-combined fulls on several sweeps: a cubic 20 A cell, Wilson
// intensities at B = 15 A^2 to 1.6 A, each sweep at its own scale and relative B, Poisson-like noise.
struct SweepSpec {
    int n_frames;
    double scale;
    double relative_b;                // at the sweep's first frame
    double relative_b_end = NAN;      // at its last; NaN: the same B throughout
};

struct Synthetic {
    std::vector<IntegrationOutcome> frames;
    std::map<std::array<int, 3>, double> I_true;   // by the Friedel-merged index
};

constexpr double CELL = 20.0;

std::array<int, 3> FriedelKey(int h, int k, int l) {
    const std::array<int, 3> a{h, k, l}, b{-h, -k, -l};
    return std::max(a, b);
}

Synthetic MakeSweeps(const std::vector<SweepSpec> &sweeps, uint32_t seed) {
    std::mt19937 rng(seed);
    std::vector<std::array<int, 3>> hkl;
    for (int h = -12; h <= 12; ++h)
        for (int k = -12; k <= 12; ++k)
            for (int l = -12; l <= 12; ++l) {
                const double d = CELL / std::sqrt(static_cast<double>(h * h + k * k + l * l));
                if ((h || k || l) && d >= 1.6) hkl.push_back({h, k, l});
            }
    Synthetic out;
    std::exponential_distribution<double> wilson(1.0);
    for (const auto &m: hkl) {
        const auto key = FriedelKey(m[0], m[1], m[2]);
        if (!out.I_true.contains(key)) {
            const double s2 = (m[0] * m[0] + m[1] * m[1] + m[2] * m[2]) / (4.0 * CELL * CELL);
            out.I_true[key] = 2000.0 * std::exp(-2.0 * 15.0 * s2) * (0.05 + wilson(rng));
        }
    }
    std::uniform_int_distribution<size_t> pick(0, hkl.size() - 1);
    std::normal_distribution<double> noise(0.0, 1.0);
    int frame = 0;
    for (const auto &sweep: sweeps)
        for (int f = 0; f < sweep.n_frames; ++f, ++frame) {
            IntegrationOutcome o;
            o.latt = CrystalLattice(UnitCell{static_cast<float>(CELL), static_cast<float>(CELL),
                                             static_cast<float>(CELL), 90, 90, 90});
            for (int j = 0; j < 400; ++j) {
                const auto &m = hkl[pick(rng)];
                const double s2 = (m[0] * m[0] + m[1] * m[1] + m[2] * m[2]) / (4.0 * CELL * CELL);
                const double b = std::isfinite(sweep.relative_b_end)
                                 ? sweep.relative_b + (sweep.relative_b_end - sweep.relative_b) * f / (sweep.n_frames - 1)
                                 : sweep.relative_b;
                const double expected = out.I_true.at(FriedelKey(m[0], m[1], m[2])) * sweep.scale
                                        * std::exp(-2.0 * b * s2);
                const double sigma = std::sqrt(expected + 20.0);
                Reflection r{};
                r.h = m[0]; r.k = m[1]; r.l = m[2];
                r.image_number = static_cast<float>(frame);
                r.d = static_cast<float>(1.0 / std::sqrt(4.0 * s2));
                r.I = static_cast<float>(expected + sigma * noise(rng));
                r.sigma = static_cast<float>(sigma);
                // The combine rebuilds a full's variance as var_bkg + I: what is not the intensity's own
                r.var_bkg = static_cast<float>(sigma * sigma - std::max(0.0, static_cast<double>(r.I)));
                r.bkg = 20.0f;
                r.prescaling_corr = 1.0f;
                r.partiality = 1.0f;
                r.zeta = 1.0f;
                r.image_scale_corr = 1.0f;
                r.predicted_x = 500.0f;
                r.predicted_y = 500.0f;
                o.reflections.push_back(r);
            }
            out.frames.push_back(std::move(o));
        }
    return out;
}

DiffractionExperiment SyntheticExperiment(bool correction_surfaces) {
    DiffractionExperiment x;
    x.IncidentEnergy_keV(12.4f);
    x.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(1, 0, 0), {}));
    x.SetUnitCell(UnitCell{static_cast<float>(CELL), static_cast<float>(CELL), static_cast<float>(CELL), 90, 90, 90});
    ScalingSettings s;
    s.ScaleFulls(true).CorrectionSurfaces(correction_surfaces).ResolutionCutoff(ResolutionCutoffMethod::Off);
    x.ImportScalingSettings(s);
    return x;
}

std::vector<SweepSegment> Segments(const std::vector<SweepSpec> &sweeps) {
    std::vector<SweepSegment> out;
    int32_t first = 0;
    for (const auto &s: sweeps) {
        out.push_back({first, s.n_frames});
        first += s.n_frames;
    }
    return out;
}

RotationScaleMerge::Result Merge(Synthetic &data, const DiffractionExperiment &x,
                                 const std::vector<SweepSegment> &segments, bool export_fulls = false,
                                 bool *on_gpu = nullptr) {
    Logger logger("RotationScaleMergeSegmentsTest");
    const UnitCell cell{static_cast<float>(CELL), static_cast<float>(CELL), static_cast<float>(CELL), 90, 90, 90};
    RotationScaleMerge engine(x, data.frames, cell, 20, 4, logger);
    if (!segments.empty())
        engine.SetSweepSegments(segments);
    engine.SetExportScaledFulls(export_fulls);
    engine.Ingest();
    if (on_gpu)
        *on_gpu = engine.OnGPU();
    return engine.Run(false, true, false);
}

// The median over a frame range of the scaled fulls' ratio to the true intensity (the strong ones).
double MedianRatio(const RotationScaleMerge::Result &r, const Synthetic &data, int f0, int f1) {
    std::vector<double> ratio;
    for (const auto &o: r.scaled_fulls)
        if (o.image_number >= static_cast<float>(f0) && o.image_number < static_cast<float>(f1)) {
            const double t = data.I_true.at(FriedelKey(o.h, o.k, o.l));
            if (t > 500.0)
                ratio.push_back(o.I / t);
        }
    REQUIRE(ratio.size() > 20);
    std::nth_element(ratio.begin(), ratio.begin() + ratio.size() / 2, ratio.end());
    return ratio[ratio.size() / 2];
}

}  // namespace

TEST_CASE("RotationScaleMergeSegments_PerSweepScaleRecovered", "[RotationScale]") {
    const std::vector<SweepSpec> sweeps{{60, 1.0, 0.0}, {60, 0.5, 0.0}};
    auto data = MakeSweeps(sweeps, 11);
    const auto r = Merge(data, SyntheticExperiment(false), Segments(sweeps), true);
    REQUIRE(r.segments.size() == 2);
    CHECK(r.segments[1].scale / r.segments[0].scale == Catch::Approx(0.5).epsilon(0.03));
    CHECK(r.statistics.overall.cc_half > 0.99);

    // No coupling across the boundary: the frames on either side of it are put on the common scale as
    // well as the middle of their sweeps are.
    const double middle = MedianRatio(r, data, 20, 40);
    CHECK(MedianRatio(r, data, 57, 60) / middle == Catch::Approx(1.0).epsilon(0.03));
    CHECK(MedianRatio(r, data, 60, 63) / middle == Catch::Approx(1.0).epsilon(0.03));
    CHECK(MedianRatio(r, data, 80, 100) / middle == Catch::Approx(1.0).epsilon(0.03));
}

TEST_CASE("RotationScaleMergeSegments_DoseOrderFromTheRelativeB", "[RotationScale]") {
    // Given first, the sweep taken second: its B is 8 A^2 higher. The order comes from the data.
    const std::vector<SweepSpec> sweeps{{60, 1.0, 8.0}, {60, 1.0, 0.0}, {60, 1.0, 4.0}};
    auto data = MakeSweeps(sweeps, 12);
    const auto r = Merge(data, SyntheticExperiment(true), Segments(sweeps));
    REQUIRE(r.segments.size() == 3);
    CHECK(r.segments[1].dose_rank == 0);
    CHECK(r.segments[2].dose_rank == 1);
    CHECK(r.segments[0].dose_rank == 2);
    CHECK(r.segments[1].relative_b == Catch::Approx(0.0).margin(1e-9));
    CHECK(r.segments[0].relative_b == Catch::Approx(8.0).margin(1.0));
    CHECK(r.segments[2].relative_b == Catch::Approx(4.0).margin(1.0));
    // Flat within each sweep: one B per sweep predicts the held-out frames better than a dose-linear ramp
    CHECK(r.segment_decay_model == "per-sweep B");
}

TEST_CASE("RotationScaleMergeSegments_DoseLinearDecayAcrossSweeps", "[RotationScale]") {
    // One crystal damaged steadily: the sweep given first was taken second, its B going on from where the
    // other's ended. Along the dose order the B rises in a straight line, which the ramp describes.
    const std::vector<SweepSpec> sweeps{{60, 1.0, 6.0, 12.0}, {60, 1.0, 0.0, 6.0}};
    auto data = MakeSweeps(sweeps, 15);
    const auto r = Merge(data, SyntheticExperiment(true), Segments(sweeps));
    REQUIRE(r.segments.size() == 2);
    CHECK(r.segments[1].dose_rank == 0);
    CHECK(r.segments[0].dose_rank == 1);
    CHECK(r.segments[0].relative_b == Catch::Approx(6.0).margin(1.5));
    CHECK(r.segment_decay_model == "dose-linear");
}

TEST_CASE("RotationScaleMergeSegments_SegmentsMustCoverTheFrames", "[RotationScale]") {
    const std::vector<SweepSpec> sweeps{{10, 1.0, 0.0}, {10, 1.0, 0.0}};
    auto data = MakeSweeps(sweeps, 13);
    const auto x = SyntheticExperiment(false);
    CHECK_THROWS_AS(Merge(data, x, {{0, 10}, {11, 9}}), JFJochException);
    CHECK_THROWS_AS(Merge(data, x, {{0, 10}}), JFJochException);
}

TEST_CASE("RotationScaleMergeSegments_NoSegmentsIsTheEngineAsBefore", "[RotationScale]") {
    // The same synthetic sweep merged without segments gives, to the bit, what the engine gave before
    // segments existed: FNV-1a over the merged I and sigma, recorded on the engine of rc181 e56b19abb -
    // on the CPU, and on the GPU for the card model it was recorded on (the device path matches the host
    // path to the bit only where every reduction does; this one carries its own, so it is checked against
    // its own record).
    const std::vector<SweepSpec> sweeps{{40, 1.0, 0.0}, {40, 0.7, 3.0}};
    auto data = MakeSweeps(sweeps, 14);
    bool on_gpu = false;
    const auto r = Merge(data, SyntheticExperiment(true), {}, false, &on_gpu);
    uint64_t hash = 1469598103934665603ULL;
    auto mix = [&hash](float v) {
        uint32_t u;
        std::memcpy(&u, &v, sizeof u);
        for (int b = 0; b < 4; ++b) {
            hash ^= (u >> (8 * b)) & 0xffu;
            hash *= 1099511628211ULL;
        }
    };
    for (const auto &m: r.merged) {
        mix(m.I);
        mix(m.sigma);
    }
    CHECK(r.segments.empty());
    CHECK(r.merged.size() == 4106);
    if (!on_gpu)
        CHECK(hash == 3376995616949734726ULL);
    else if (get_gpu_description().find("RTX 5080") != std::string::npos)
        CHECK(hash == 6594155396455850264ULL);
}

// A sweep in which no rocking event is captured whole enough to make a full - as in one exposed 0.2 deg out
// of every 0.6 deg, where each reflection is caught on one frame and only in part. On the GPU the combine
// returned no fulls, the scale-fulls step then fell through to the host loop, which slices a per-frame CSR
// that only the host combine builds, and the process died on SIGSEGV instead of reporting the empty merge.
TEST_CASE("RotationScaleMerge_NoFullsIsAnErrorNotACrash", "[RotationScale]") {
    const std::vector<SweepSpec> sweeps{{60, 1.0, 0.0}};
    auto data = MakeSweeps(sweeps, 17);
    auto x = SyntheticExperiment(false);
    // More than any rocking event can capture: the combine makes no full.
    ScalingSettings s = x.GetScalingSettings();
    s.MinCapturedFraction(1000.0);
    x.ImportScalingSettings(s);
    CHECK_THROWS_AS(Merge(data, x, {}), JFJochException);
}

// The reported CC1/2 is the plain Pearson of the half-set means, every reflection counted alike, as
// XDS, Aimless and DIALS report it; the cc_weight-weighted form is internal (cc_half_internal). A sweep
// with a weak stretch makes the weights differ from 1, so the two are not the same number.
TEST_CASE("RotationScaleMerge_ReportedCCHalfIsThePlainPearson", "[RotationScale]") {
    const std::vector<SweepSpec> sweeps{{60, 1.0, 0.0}, {60, 0.1, 0.0}};
    auto data = MakeSweeps(sweeps, 19);
    const auto r = Merge(data, SyntheticExperiment(false), {});

    CorrelationCoefficient plain, weighted;
    int down_weighted = 0;
    for (const auto &m: r.merged)
        if (std::isfinite(m.I) && std::isfinite(m.sigma) && m.sigma > 0.0f
            && std::isfinite(m.I_half[0]) && std::isfinite(m.I_half[1])) {
            plain.Add(m.I_half[0], m.I_half[1]);
            weighted.Add(m.I_half[0], m.I_half[1], m.cc_weight);
            if (m.cc_weight < 0.9f) ++down_weighted;
        }
    REQUIRE(down_weighted > 100);
    CHECK(r.statistics.overall.cc_half == Catch::Approx(plain.GetCC()).epsilon(1e-9));
    CHECK(r.statistics.overall.cc_half_internal == Catch::Approx(weighted.GetCC()).epsilon(1e-9));
    CHECK(std::fabs(r.statistics.overall.cc_half - r.statistics.overall.cc_half_internal) > 1e-6);
}