// 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"

namespace {

// Frames of a 50/60/70 A orthorhombic crystal in random orientations on a detector of four modules (the
// quadrants around the beam at 1000, 1000), each displaced by the rigid move `truth` gives it: the spots are
// where the image shows them, each indexed as predicted.
std::vector<GeomRefineFrame> FramesOnMovedModules(const DiffractionExperiment &exp, const CrystalLattice &latt,
                                                  const std::vector<ModuleShift> &truth) {
    auto module_of = [](double x, double y) { return (x < 1000 ? 0 : 1) + (y < 1000 ? 0 : 2); };
    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);
    }
    return frames;
}

}  // namespace

// 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;
    }
    const CrystalLattice latt(50, 60, 70, 90, 90, 90);
    const std::vector<GeomRefineFrame> frames = FramesOnMovedModules(exp, latt, truth);
    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));

    // The distance held (it shares a block with the beam): it stays exactly where it was given.
    GeometryRefinerSettings held_distance = st;
    held_distance.refine_distance = false;
    const auto fixed = RefineGlobalGeometry(geom, latt.GetUnitCell(), frames, held_distance);
    REQUIRE(fixed.ok);
    CHECK(fixed.distance_mm == geom.GetDetectorDistance_mm());
    CHECK(fixed.modules[0].dx == Catch::Approx(truth[0].dx).margin(0.1));

    // Only the orientations moved, on the modules as placed: the spots fit as well as in the placement.
    GeometryRefinerSettings ev = st;
    ev.refine_modules = false;
    ev.refine_beam_distance = false;
    ev.modules = placed.modules;
    const auto scored = RefineGlobalGeometry(geom, latt.GetUnitCell(), frames, ev);
    REQUIRE(scored.ok);
    CHECK(scored.median_residual_px < 0.1);
    CHECK(scored.distance_mm == geom.GetDetectorDistance_mm());
}

// The coarse search ahead of the placement: a module off by several pixels is found near where it is, and the
// search gives the same answer on any number of threads.
TEST_CASE("ScanModuleShifts_FindsShiftedModule") {
    DiffractionExperiment exp;
    exp.IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(1000).BeamY_pxl(1000).DetectorDistance_mm(150);
    const DiffractionGeometry geom = exp.GetDiffractionGeometry();

    std::vector<ModuleShift> truth(4), given(4);
    const double centre[4][2] = {{500, 500}, {1500, 500}, {500, 1500}, {1500, 1500}};
    for (int m = 0; m < 4; m++) {
        truth[m].x0 = given[m].x0 = centre[m][0];
        truth[m].y0 = given[m].y0 = centre[m][1];
    }
    truth[1].dx = 6.0;
    truth[1].dy = -4.0;

    const CrystalLattice latt(50, 60, 70, 90, 90, 90);
    std::vector<GeomRefineFrame> frames = FramesOnMovedModules(exp, latt, truth);
    REQUIRE(frames.size() > 30);
    // The module's spots did not index where it was given: the search sees them as unindexed.
    for (auto &f: frames)
        for (auto &s: f.spots)
            if (s.module == 1) {
                s.h = s.k = s.l = 0;
                f.unindexed.push_back(s);
            }
    for (auto &f: frames)
        std::erase_if(f.spots, [](const GeomRefineSpot &s) { return s.module == 1; });

    const auto one = ScanModuleShifts(geom, frames, given, 25.0, 1);
    const auto many = ScanModuleShifts(geom, frames, given, 25.0, 4);
    REQUIRE(one.size() == 4);
    // Coarse: a spot is on its lattice within 0.15 of an index, here ~4 px of the spacing, so the search lands
    // within that of the true shift and the placement takes it from there.
    CHECK(one[1].dx == Catch::Approx(6.0).margin(4.0));
    CHECK(one[1].dy == Catch::Approx(-4.0).margin(4.0));
    CHECK(one[1].on_lattice_at_best > 2 * one[1].on_lattice_at_zero);
    for (int m = 0; m < 4; m++) {
        CHECK(many[m].dx == one[m].dx);
        CHECK(many[m].dy == one[m].dy);
        CHECK(many[m].on_lattice_at_best == one[m].on_lattice_at_best);
        CHECK(many[m].on_lattice_at_zero == one[m].on_lattice_at_zero);
    }
}

// A tilted detector, one panel: started from the untilted geometry with the beam where the beam really lands
// (the direct beam, which is what a beam-centre measurement gives), the refinement with the tilt must find it
// and put the PONI where the tilt has it - 10 px from the direct beam here, further than the beam's 8 px box
// would reach in PONI terms - and the fit with the tilt held must be worse. With the image mirrored and turned
// as well, the direct beam has to be taken through the image orientation.
TEST_CASE("RefineGlobalGeometry_Tilt") {
    for (const auto &orientation : {DetectorOrientation(), DetectorOrientation(true, 1)}) {
        DiffractionExperiment exp;
        exp.IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(1000).BeamY_pxl(1000).DetectorDistance_mm(150)
           .PoniRot1_rad(0.3 * PI / 180.0).PoniRot2_rad(-0.2 * PI / 180.0);
        exp.Detector().ImageOrientation(orientation);
        const DiffractionGeometry truth = exp.GetDiffractionGeometry();
        const auto [direct_x, direct_y] = truth.GetDirectBeam_pxl();
        REQUIRE(std::hypot(direct_x - 1000, direct_y - 1000) > 10);

        // One module covering the whole image: the frames' spots, each where the tilted detector shows it.
        ModuleShift whole;
        whole.x0 = 1000;
        whole.y0 = 1000;
        const CrystalLattice latt(50, 60, 70, 90, 90, 90);
        std::vector<GeomRefineFrame> frames = FramesOnMovedModules(exp, latt, {whole, whole, whole, whole});
        REQUIRE(frames.size() > 30);
        for (auto &f: frames)
            for (auto &s: f.spots)
                s.module = -1;

        DiffractionGeometry nominal = truth;
        nominal.PoniRot1_rad(0).PoniRot2_rad(0).BeamX_pxl(direct_x).BeamY_pxl(direct_y);
        GeometryRefinerSettings st;
        st.crystal_system = gemmi::CrystalSystem::Orthorhombic;
        st.num_threads = 2;

        const auto held = RefineGlobalGeometry(nominal, latt.GetUnitCell(), frames, st);
        REQUIRE(held.ok);
        CHECK(held.rot1_rad == 0);
        CHECK(held.rot2_rad == 0);

        st.refine_tilt = true;
        const auto tilted = RefineGlobalGeometry(nominal, latt.GetUnitCell(), frames, st);
        REQUIRE(tilted.ok);
        CHECK(tilted.rot1_rad * 180.0 / PI == Catch::Approx(0.3).margin(0.01));
        CHECK(tilted.rot2_rad * 180.0 / PI == Catch::Approx(-0.2).margin(0.01));
        CHECK(tilted.beam_x_px == Catch::Approx(1000).margin(0.2));
        CHECK(tilted.beam_y_px == Catch::Approx(1000).margin(0.2));
        CHECK(tilted.distance_mm == Catch::Approx(150).margin(0.05));
        CHECK(tilted.median_residual_px < 0.1);
        CHECK(tilted.capped_rms_px < held.capped_rms_px);
    }
}