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

#include <cmath>
#include <catch2/catch_all.hpp>

#include "../rugnux/Rugnux.h"
#include "../image_analysis/scale_merge/RotationScaleMerge.h"

namespace {
    // A synthetic sweep whose stage turned `true_scale` times the stored angles. Scored at a scale k,
    // the validation spots stay on the lattice as long as the angles track the rotation, and the share
    // that does falls off with the relative rate error; a wrong spindle angle keeps half a percent.
    struct SyntheticSweep {
        double true_scale;
        int64_t spots = 35000;
        int index_calls = 0;
        int refit_calls = 0;

        ValidationSpotEvidence IndexAt(float k) {
            ++index_calls;
            const double error = std::fabs(true_scale / k - 1.0);
            const double on = 0.9 * std::max(0.0, 1.0 - 30.0 * error);
            return ValidationSpotEvidence{spots, std::llround(on * spots), std::llround(0.005 * spots)};
        }

        // The post-refinement at k, relative to k. It reads only 70 % of the error that is left: it
        // sees only the frames the angles at k still track.
        std::optional<double> RefitAt(float k) {
            ++refit_calls;
            return 1.0 + 0.7 * (true_scale / k - 1.0);
        }

        RotationScaleWalk Walk(double first_fit) {
            return WalkRotationScale(first_fit, [this](float k) { return IndexAt(k); },
                                     [this](float k) { return RefitAt(k); }, 8);
        }
    };
}

TEST_CASE("ValidationEvidencePrefers", "[RotationScale]") {
    const ValidationSpotEvidence base{10000, 3000, 50};
    // 1 % more of the spots beyond chance is under the noise of two 30 % shares over 10000 spots
    // (sqrt(2 * 0.3 * 0.7 / 10000) = 0.65 %, times 3.29); 5 % is well over it.
    CHECK_FALSE(ValidationEvidencePrefers(base, {10000, 3100, 50}));
    CHECK(ValidationEvidencePrefers(base, {10000, 3500, 50}));
    // A candidate is judged against its own null: more spots on the lattice bought by a null that
    // rose just as much is no gain.
    CHECK_FALSE(ValidationEvidencePrefers(base, {10000, 3500, 550}));
    // Never against itself, and nothing that scored nothing wins.
    CHECK_FALSE(ValidationEvidencePrefers(base, base));
    CHECK_FALSE(ValidationEvidencePrefers(base, {}));
    CHECK(ValidationEvidencePrefers({}, base));
}

TEST_CASE("WalkRotationScale_ReachesTheFixedPoint", "[RotationScale]") {
    // A stage 3 % slow. The first fit reads 70 % of that; each refit at the adopted scale reads 70 %
    // of what is left, and the walk goes on as long as the validation frames prefer the new scale.
    SyntheticSweep sweep{0.97};
    const auto walk = sweep.Walk(sweep.RefitAt(1.0f).value());
    CHECK(walk.scale == Catch::Approx(0.97).margin(0.001));
    CHECK(walk.scale != 1.0f);
    CHECK(walk.evidence.on_lattice > sweep.IndexAt(1.0f).on_lattice);
    CHECK(sweep.refit_calls > 2);
    CHECK_FALSE(walk.trail.empty());
}

TEST_CASE("WalkRotationScale_StoredAnglesStand", "[RotationScale]") {
    SECTION("A healthy stage: a fit off by noise scores no better than the stored angles") {
        SyntheticSweep sweep{1.0};
        const auto walk = sweep.Walk(1.0002);
        CHECK(walk.scale == 1.0f);
        CHECK(sweep.index_calls == 2);   // the stored angles and the fit, nothing more
        CHECK(sweep.refit_calls == 0);
    }
    SECTION("A real but small error the spots cannot resolve beyond their noise") {
        SyntheticSweep sweep{0.999};
        sweep.spots = 400;
        const auto walk = sweep.Walk(0.9993);
        CHECK(walk.scale == 1.0f);
    }
    SECTION("A fit that tracks something other than the rotation scores worse, and is refused") {
        SyntheticSweep sweep{1.0};
        const auto walk = sweep.Walk(0.98);
        CHECK(walk.scale == 1.0f);
        CHECK(sweep.refit_calls == 0);
    }
    SECTION("A fit of exactly one asks for no probe at all") {
        SyntheticSweep sweep{1.0};
        const auto walk = sweep.Walk(1.0);
        CHECK(walk.scale == 1.0f);
        CHECK(walk.trail.empty());
        CHECK(sweep.index_calls == 0);
    }
}

TEST_CASE("SmoothLogScale_FollowsInformationBridgesGaps", "[RotationScale]") {
    const int n = 60;
    // A ramp is no curvature, so any amount of smoothing keeps it exactly.
    std::vector<double> ramp(n), J(n, 1.0);
    for (int f = 0; f < n; ++f) ramp[f] = -0.1 * f;
    auto x = RotationScaleMerge::SmoothLogScale(ramp, J, 1e6);
    for (int f = 0; f < n; ++f) CHECK(x[f] == Catch::Approx(ramp[f]).margin(1e-6));
    // A stretch with no information is bridged by the straight line through its neighbours.
    std::vector<double> Jgap(J);
    for (int f = 20; f < 40; ++f) Jgap[f] = 0.0;
    x = RotationScaleMerge::SmoothLogScale(ramp, Jgap, 1.0);
    CHECK(x[30] == Catch::Approx(-3.0).margin(1e-6));
    // A frame with far more information than its neighbours keeps its own value.
    std::vector<double> y(n, 0.0), Jone(n, 1.0);
    y[30] = 1.0; Jone[30] = 1e6;
    x = RotationScaleMerge::SmoothLogScale(y, Jone, 10.0);
    CHECK(x[30] == Catch::Approx(1.0).margin(1e-3));
    // Fewer than two frames with information: nothing to smooth against.
    std::vector<double> Jsingle(n, 0.0);
    Jsingle[5] = 1.0;
    CHECK(RotationScaleMerge::SmoothLogScale(y, Jsingle, 1.0) == y);
}

TEST_CASE("ChooseLogScaleSmoothing_SmoothsNoiseFollowsSignal", "[RotationScale]") {
    const int n = 400;
    std::vector<double> J(n, 1.0), noisy(n), step(n);
    // Deterministic noise about a flat scale: the chosen curve is close to flat.
    for (int f = 0; f < n; ++f) noisy[f] = 0.3 * std::sin(12.9898 * f) * std::cos(78.233 * f);
    const double l_noise = RotationScaleMerge::ChooseLogScaleSmoothing(noisy, J, 1);
    const auto flat = RotationScaleMerge::SmoothLogScale(noisy, J, l_noise);
    double rms = 0.0;
    for (double v : flat) rms += v * v;
    CHECK(std::sqrt(rms / n) < 0.05);
    // A precise slow wave is followed.
    for (int f = 0; f < n; ++f) step[f] = 2.0 * std::sin(f / 30.0);
    const double l_wave = RotationScaleMerge::ChooseLogScaleSmoothing(step, J, 1);
    const auto wave = RotationScaleMerge::SmoothLogScale(step, J, l_wave);
    CHECK(wave[47] == Catch::Approx(step[47]).margin(0.01));
    CHECK(l_wave < l_noise);
}