// 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 "../image_analysis/geom_refinement/PostRefine.h"
#include "../common/DiffractionExperiment.h"
#include "../common/Logger.h"

namespace {
    PostRefineResult Measured(double residual, double se) {
        PostRefineResult r;
        r.held_out_before = residual;
        r.held_out_before_se = se;
        return r;
    }

    PostRefineResult Fit(UnitCell before, UnitCell after) {
        PostRefineResult r;
        r.cell_refined = true;
        r.cell_before = before;
        r.cell = after;
        return r;
    }

    UnitCell Cell(float a, float b, float c) {
        return UnitCell{.a = a, .b = b, .c = c, .alpha = 90.0f, .beta = 90.0f, .gamma = 90.0f};
    }

    // The integrated partials of a sweep whose stage turned `true_scale` times the stored angles: the
    // crystal is at rotation true_scale * phi when the file says phi. Each reflection crosses the
    // Ewald sphere where its excitation, p_z + lambda |p|^2 / 2 for p rotated by minus the angle
    // about the spindle, is zero, and is recorded on the frames about that stored angle with a
    // Gaussian rocking curve 0.1 deg wide. No spot positions: the scale is fitted on the angles.
    std::vector<IntegrationOutcome> SimulatedSweep(const CrystalLattice &latt, const GoniometerAxis &axis,
                                                   int images, double lambda, double true_scale) {
        constexpr double D_MIN = 3.0, ROCKING_DEG = 0.1;
        const Coord u = axis.GetAxis().Normalize();
        const double first = axis.GetAngle_deg(0.0f), last = axis.GetAngle_deg(static_cast<float>(images));
        const double inc = axis.GetIncrement_deg();
        std::vector<IntegrationOutcome> out(images);
        for (int h = -20; h <= 20; ++h)
            for (int k = -20; k <= 20; ++k)
                for (int l = -25; l <= 25; ++l) {
                    const Coord p = latt.Astar() * h + latt.Bstar() * k + latt.Cstar() * l;
                    const double p2 = p * p;
                    if (p2 == 0.0 || p2 > 1.0 / (D_MIN * D_MIN))
                        continue;
                    const double up = u * p;
                    const double A = p.z - u.z * up;
                    const double B = u.x * p.y - u.y * p.x;
                    const double C = u.z * up + 0.5 * lambda * p2;
                    const double R = std::hypot(A, B);
                    if (std::fabs(C) >= R)
                        continue;
                    for (const double sign : {-1.0, 1.0}) {
                        // A cos(psi) - B sin(psi) + C = 0, i.e. R cos(psi + atan2(B, A)) = -C
                        double psi = (sign * std::acos(-C / R) - std::atan2(B, A)) * 180.0 / PI;
                        const double phi = std::remainder(psi, 360.0) / true_scale;
                        if (phi < first + 1.0 || phi > last - 1.0)
                            continue;
                        for (int i = 0; i < images; ++i) {
                            const double centre = axis.GetAngle_deg(static_cast<float>(i)) + inc / 2.0;
                            const double x = (centre - phi) / ROCKING_DEG;
                            if (std::fabs(x) > 3.0)
                                continue;
                            Reflection r{};
                            r.h = h; r.k = k; r.l = l;
                            r.image_number = static_cast<float>(i);
                            r.I = static_cast<float>(1000.0 * std::exp(-0.5 * x * x));
                            r.sigma = std::sqrt(r.I) + 1.0f;
                            r.observed_x = r.observed_y = NAN;
                            out[i].reflections.push_back(r);
                        }
                    }
                }
        return out;
    }

    double FittedRotationScale(double true_scale) {
        DiffractionExperiment x(DetJF(1));
        x.IncidentEnergy_keV(12.4);
        const GoniometerAxis axis("omega", -90.0f, 0.1f, Coord(1, 0, 0), {});
        const CrystalLattice latt(40.0f, 50.0f, 60.0f, 90.0f, 90.0f, 90.0f);
        auto outcomes = SimulatedSweep(latt, axis, 1800, x.GetWavelength_A(), true_scale);
        Logger logger("PostRefineTest");
        PostRefineSettings settings;
        settings.refine_geometry = true;
        settings.num_threads = 4;
        return PostRefineRotationGeometry(GatherPostRefineObservations(outcomes, 4, true, logger), axis,
                                          x.GetDiffractionGeometry(), latt, settings, logger)
                .rotation_scale;
    }
}

TEST_CASE("PostRefine_HeldOutResidualFell", "[PostRefine]") {
    // Down by more than the standard error of the difference, sqrt(3^2 + 4^2) = 5.
    CHECK(HeldOutResidualFell(Measured(100.0, 3.0), Measured(94.0, 4.0)));
    // Down, but within that noise.
    CHECK_FALSE(HeldOutResidualFell(Measured(100.0, 3.0), Measured(96.0, 4.0)));
    // Up.
    CHECK_FALSE(HeldOutResidualFell(Measured(100.0, 3.0), Measured(110.0, 4.0)));
    // Not measured on either side.
    CHECK_FALSE(HeldOutResidualFell(PostRefineResult{}, Measured(10.0, 1.0)));
    CHECK_FALSE(HeldOutResidualFell(Measured(100.0, 3.0), PostRefineResult{}));
}

TEST_CASE("PostRefine_ReindexPushesCellBack", "[PostRefine]") {
    // The fit shortens b the most; re-indexing at its geometry returns a longer b: pushed back.
    const auto fit = Fit(Cell(96.9f, 107.9f, 112.9f), Cell(96.8f, 106.6f, 112.9f));
    CHECK(ReindexPushesCellBack(fit, Fit(Cell(96.7f, 107.6f, 112.6f), Cell(96.6f, 106.4f, 112.8f))) == -2);
    // A walk that re-indexing follows (and overshoots a little) is not pushed back.
    const auto walk = Fit(Cell(60.48f, 60.48f, 196.7f), Cell(60.05f, 60.05f, 195.2f));
    CHECK(ReindexPushesCellBack(walk, Fit(Cell(60.00f, 60.00f, 195.1f), Cell(59.8f, 59.8f, 194.2f))) == 0);
    // Nothing committed, or no cell measured by the next pass: no evidence either way.
    PostRefineResult refused = fit;
    refused.cell_refined = false;
    CHECK(ReindexPushesCellBack(refused, Fit(Cell(96.7f, 107.6f, 112.6f), Cell(96.6f, 106.4f, 112.8f))) == 0);
    CHECK(ReindexPushesCellBack(fit, PostRefineResult{}) == 0);
}

TEST_CASE("PostRefine_RotationScale", "[PostRefine]") {
    // Where every frame of the sweep is on the lattice the fit reads the stage's rate itself, and a
    // healthy stage reads as one.
    CHECK(FittedRotationScale(0.97) == Catch::Approx(0.97).margin(0.001));
    CHECK(FittedRotationScale(1.0) == Catch::Approx(1.0).margin(0.001));
}