// 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 <cmath>
#include <random>

#include "../common/JFJochMath.h"
#include "../image_analysis/geom_refinement/GeometryRefiner.h"
#include "../image_analysis/bragg_prediction/BraggPrediction.h"

// A detector of four modules (the quadrants around the beam), each displaced by a known rigid move: the
// bundle adjustment with module placement must find the moves, and the fit without it must be worse.
TEST_CASE("RefineGlobalGeometry_ModulePlacement") {
    DiffractionExperiment exp;
    exp.IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(1000).BeamY_pxl(1000).DetectorDistance_mm(150);
    const DiffractionGeometry geom = exp.GetDiffractionGeometry();

    // Moves that hold none of what the gauge takes away: zero in sum, in turn and outwards.
    std::vector<ModuleShift> truth(4);
    const double centre[4][2] = {{500, 500}, {1500, 500}, {500, 1500}, {1500, 1500}};
    const double move[4][3] = {{2.0, -2.0, 0.2}, {2.0, 2.0, -0.2}, {-2.0, -2.0, 0.1}, {-2.0, 2.0, -0.1}};
    for (int m = 0; m < 4; m++) {
        truth[m].x0 = centre[m][0];
        truth[m].y0 = centre[m][1];
        truth[m].dx = move[m][0];
        truth[m].dy = move[m][1];
        truth[m].rot_rad = move[m][2] * PI / 180.0;
    }
    auto module_of = [](double x, double y) { return (x < 1000 ? 0 : 1) + (y < 1000 ? 0 : 2); };

    const CrystalLattice latt(50, 60, 70, 90, 90, 90);
    BraggPredictionSettings settings{.high_res_A = 2.0, .ewald_dist_cutoff = 0.001};
    std::mt19937 rng(7);
    std::uniform_real_distribution<float> u(-1.0f, 1.0f);

    std::vector<GeomRefineFrame> frames;
    for (int f = 0; f < 60; f++) {
        const Coord axis = Coord(u(rng), u(rng), u(rng)).Normalize();
        const CrystalLattice rotated = latt.Multiply(RotMatrix(static_cast<float>(PI) * u(rng), axis));
        BraggPrediction prediction;
        const int n = prediction.Calc(exp, rotated, settings);
        GeomRefineFrame frame;
        frame.lattice = rotated;
        frame.ordinal = f;
        for (int i = 0; i < n; i++) {
            const auto &r = prediction.GetReflections().at(i);
            const int m = module_of(r.predicted_x, r.predicted_y);
            // where the image shows a spot whose true place is the prediction: the inverse of the move
            const auto &t = truth[m];
            const double c = std::cos(t.rot_rad), s = std::sin(t.rot_rad);
            const double vx = r.predicted_x - t.x0 - t.dx, vy = r.predicted_y - t.y0 - t.dy;
            GeomRefineSpot spot;
            spot.x = static_cast<float>(t.x0 + c * vx + s * vy);
            spot.y = static_cast<float>(t.y0 - s * vx + c * vy);
            spot.h = r.h;
            spot.k = r.k;
            spot.l = r.l;
            spot.module = static_cast<int16_t>(m);
            frame.spots.push_back(spot);
        }
        if (frame.spots.size() >= 20)
            frames.push_back(frame);
    }
    REQUIRE(frames.size() > 30);

    GeometryRefinerSettings st;
    st.crystal_system = gemmi::CrystalSystem::Orthorhombic;
    st.num_threads = 2;
    for (const auto &t: truth) {
        ModuleShift start;
        start.x0 = t.x0;
        start.y0 = t.y0;
        st.modules.push_back(start);
    }

    const auto held = RefineGlobalGeometry(geom, latt.GetUnitCell(), frames, st);
    REQUIRE(held.ok);

    st.refine_modules = true;
    const auto placed = RefineGlobalGeometry(geom, latt.GetUnitCell(), frames, st);
    REQUIRE(placed.ok);
    CHECK(placed.median_residual_px < 0.1);
    CHECK(held.median_residual_px > 5 * placed.median_residual_px);

    for (int m = 0; m < 4; m++) {
        CHECK(placed.modules[m].dx == Catch::Approx(truth[m].dx).margin(0.1));
        CHECK(placed.modules[m].dy == Catch::Approx(truth[m].dy).margin(0.1));
        CHECK(placed.modules[m].rot_rad * 180.0 / PI == Catch::Approx(move[m][2]).margin(0.01));
    }
    CHECK(placed.beam_x_px == Catch::Approx(1000).margin(0.1));
    CHECK(placed.beam_y_px == Catch::Approx(1000).margin(0.1));
    CHECK(placed.distance_mm == Catch::Approx(150).margin(0.05));
}