// 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 "../image_analysis/scale_merge/AndersonMixing.h"

namespace {
    // A noiseless sweep of 24 frames whose reflections each span three consecutive frames, scaled by the
    // alternating least squares of the per-frame scaling loop: the reference intensities from the
    // scales, the scales from the references, the gauge pinned at a geometric mean of 1. A frame is
    // tied only to its neighbours, so a smooth error of the scales over the sweep decays by a few
    // percent a round - the slow linear convergence of the real loop.
    struct SlowScalingLoop {
        static constexpr int n_frames = 24, span = 3;
        std::vector<int> first;   // reflection h is on frames first[h] .. first[h] + span - 1
        std::vector<double> I_true, G_true;

        SlowScalingLoop() {
            for (int h = 0; h + span <= n_frames; ++h)
                for (int rep = 0; rep < 3; ++rep) {
                    first.push_back(h);
                    I_true.push_back(1.0 + 0.5 * std::sin(1.7 * h + rep));
                }
            for (int f = 0; f < n_frames; ++f)
                G_true.push_back(1.0 + 0.4 * std::sin(6.283 * f / n_frames) + 0.2 * std::cos(2.1 * f));
        }

        std::vector<double> Round(const std::vector<double> &g) const {
            std::vector<double> num(n_frames, 0.0), den(n_frames, 0.0), out(n_frames);
            for (size_t h = 0; h < first.size(); ++h) {
                double mean = 0.0;
                for (int k = 0; k < span; ++k) mean += G_true[first[h] + k] * I_true[h] / g[first[h] + k];
                mean /= span;
                for (int k = 0; k < span; ++k) {
                    num[first[h] + k] += G_true[first[h] + k] * I_true[h] * mean;
                    den[first[h] + k] += mean * mean;
                }
            }
            double s = 0.0;
            for (int f = 0; f < n_frames; ++f) { out[f] = num[f] / den[f]; s += std::log(out[f]); }
            for (int f = 0; f < n_frames; ++f) out[f] /= std::exp(s / n_frames);
            return out;
        }

        // Rounds to a step below `tolerance`, and the largest |log G - log G_true| there (gauge-pinned).
        std::pair<int, double> Run(int depth, double tolerance) const {
            std::vector<double> g(n_frames, 1.0), log_g(n_frames), log_g_new(n_frames), w(n_frames, 1.0);
            AndersonMixing anderson(depth);
            int rounds = 0;
            for (double step = INFINITY; step >= tolerance && rounds < 100000; ++rounds) {
                std::vector<double> g_new = Round(g);
                double ss = 0.0;
                for (int f = 0; f < n_frames; ++f) {
                    log_g[f] = std::log(g[f]);
                    log_g_new[f] = std::log(g_new[f]);
                    ss += (log_g_new[f] - log_g[f]) * (log_g_new[f] - log_g[f]);
                }
                step = std::sqrt(ss / n_frames);
                if (depth > 0) {
                    const std::vector<double> next = anderson.Next(log_g, log_g_new, w);
                    double s = 0.0;
                    for (double v : next) s += v;
                    for (int f = 0; f < n_frames; ++f) g_new[f] = std::exp(next[f] - s / n_frames);
                }
                g = g_new;
            }
            double s = 0.0, err = 0.0;
            for (double v : G_true) s += std::log(v);
            for (int f = 0; f < n_frames; ++f)
                err = std::max(err, std::fabs(std::log(g[f]) - (std::log(G_true[f]) - s / n_frames)));
            return {rounds, err};
        }
    };
}

TEST_CASE("AndersonMixing_SlowScalingLoop", "[RotationScale]") {
    const SlowScalingLoop loop;
    const auto [plain_rounds, plain_err] = loop.Run(0, 1e-6);
    const auto [aa_rounds, aa_err] = loop.Run(5, 1e-6);
    // The plain loop creeps: hundreds of rounds, and stopped on its step it still sits a hundred
    // times the tolerance from the answer.
    CHECK(plain_rounds > 300);
    CHECK(plain_err > 5e-5);
    // Accelerated: the same fixed point, in a tenth of the rounds, and closer to it at the same stop.
    CHECK(aa_rounds * 10 < plain_rounds);
    CHECK(aa_err < 0.5 * plain_err);
}

TEST_CASE("AndersonMixing_ZeroWeightTakesThePlainStep", "[RotationScale]") {
    // x -> 0.9 x on two components, the second not part of the fit.
    AndersonMixing anderson(5);
    std::vector<double> x{1.0, 1.0};
    const std::vector<double> w{1.0, 0.0};
    for (int it = 0; it < 4; ++it) {
        const std::vector<double> fx{0.9 * x[0], 0.9 * x[1]};
        const std::vector<double> next = anderson.Next(x, fx, w);
        CHECK(next[1] == fx[1]);
        x = next;
    }
    // A linear map on the fitted component is solved by the first mixed step.
    CHECK(std::fabs(x[0]) < 1e-12);
    CHECK(anderson.Restarts() == 0);
}

TEST_CASE("AndersonMixing_RestartsWhenTheResidualGrows", "[RotationScale]") {
    AndersonMixing anderson(5);
    const std::vector<double> w{1.0};
    CHECK(anderson.Next({1.0}, {0.5}, w)[0] == 0.5);    // no history: the plain step
    anderson.Next({0.5}, {0.25}, w);                      // residual fell: history kept
    CHECK(anderson.Restarts() == 0);
    CHECK(anderson.Next({0.0}, {2.0}, w)[0] == 2.0);    // residual grew: history dropped, plain step
    CHECK(anderson.Restarts() == 1);
}