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

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

#include "../image_analysis/rotation_indexer/RotationIndexer.h"
#include "../image_analysis/bragg_prediction/BraggPrediction.h"

TEST_CASE("RotationIndexer") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
         .BeamX_pxl(1000)
         .BeamY_pxl(1000)
         .PoniRot1_rad(0.01)
         .PoniRot2_rad(0.02)
         .DetectorDistance_mm(200)
    .ImagesPerTrigger(50);

    IndexingSettings settings;
    settings.Algorithm(IndexingAlgorithmEnum::Auto);   // rotation: FFT with a GPU, FFTW without

    settings.RotationIndexing(true).RotationIndexingAngularStride_deg(1.0).RotationIndexingMinAngularRange_deg(30.0);
    exp_i.ImportIndexingSettings(settings);

    // Base lattice (non-pathological)
    CrystalLattice latt_base(40, 50, 80, 90, 90, 90);
    latt_base = latt_base.Multiply(RotMatrix(2.0, Coord(sqrt(3)/3,sqrt(3)/3,sqrt(3)/3)));

    // Rotation axis: around X with 1 deg per image
    GoniometerAxis axis("omega", 0.0f, 1.0f, Coord(1,0,0), std::nullopt);
    exp_i.Goniometer(axis);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.3,
        .ewald_dist_cutoff = 0.002
    };

    IndexerThreadPool indexer_thread_pool(exp_i.GetIndexingSettings());
    RotationIndexer indexer(exp_i, indexer_thread_pool);

    BraggPrediction prediction;

    int cnt = 0;

    // Predict reflections for images at 0-30 deg.
    for (int img = 0; img < 50; ++img) {
        std::vector<SpotToSave> spots;
        // For a rotated image, per-image lattice is obtained as Multiply(rot.transpose())
        const float angle_deg = axis.GetAngle_deg(img) + axis.GetWedge_deg() / 2.0f;
        const RotMatrix rot = axis.GetTransformationAngle(angle_deg);
        const CrystalLattice latt_img = latt_base.Multiply(rot.transpose());

        const auto n = prediction.Calc(exp_i, latt_img, prediction_settings);
        for (int i = 0; i < n; ++i) {
            const auto& r = prediction.GetReflections().at(i);
            SpotToSave s{};
            s.x = r.predicted_x;
            s.y = r.predicted_y;
            s.image = img;       // provide image index for rotation-aware refinement
            s.intensity = 1.0f;  // minimal positive value
            s.phi = angle_deg;
            s.ice_ring = false;
            s.indexed = true;
            spots.push_back(s);
        }

        indexer.ProcessImage(img, spots);
        if (img == 30)
            indexer.RunIndexing();

        auto result = indexer.GetLattice();
        if (result.has_value())
            cnt++;
    }

    CHECK(cnt == 20);
    // An indexer that ran records no error; only one that threw does. A caller reporting "no lattice"
    // to the user tells the two apart on this.
    CHECK_FALSE(indexer.GetIndexerError().has_value());

    auto ret = indexer.GetLattice();
    REQUIRE(ret.has_value());
    auto uc = ret->lattice.GetUnitCell();
    auto uc_ref = latt_base.GetUnitCell();
    REQUIRE(std::fabs(uc.a - uc_ref.a) < 0.1);
    REQUIRE(std::fabs(uc.b - uc_ref.b) < 0.1);
    REQUIRE(std::fabs(uc.c - uc_ref.c) < 0.1);
    REQUIRE(std::fabs(uc.alpha - uc_ref.alpha) < 0.1);
    REQUIRE(std::fabs(uc.beta - uc_ref.beta) < 0.1);
    REQUIRE(std::fabs(uc.gamma - uc_ref.gamma) < 0.1);
    CHECK(ret->search_result.centering == 'P');
    CHECK(ret->search_result.system == gemmi::CrystalSystem::Orthorhombic);
}

// RefineConstrained is what a caller reaches for when it has ADOPTED a symmetry the indexing never
// refined under - the intensities confirm a two-fold the spot positions never offered - and so holds
// a cell whose metric is still the free fit's. Give it such a cell: the lattice this crystal indexes
// on, sheared so that alpha is a degree off, which is what that situation looks like. The constraint
// has to take it back to a cell the group can describe, and the spots have to be happier for it.
TEST_CASE("RotationIndexer::RefineConstrained puts a free metric back on its class") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
         .BeamX_pxl(1000)
         .BeamY_pxl(1000)
         .DetectorDistance_mm(200)
         .ImagesPerTrigger(50);

    IndexingSettings settings;
    settings.Algorithm(IndexingAlgorithmEnum::Auto);   // rotation: FFT with a GPU, FFTW without
    settings.RotationIndexing(true).RotationIndexingAngularStride_deg(1.0).RotationIndexingMinAngularRange_deg(30.0);
    exp_i.ImportIndexingSettings(settings);

    const CrystalLattice latt_base =
        CrystalLattice(40, 50, 80, 90, 105, 90).Multiply(RotMatrix(2.0, Coord(sqrt(3)/3, sqrt(3)/3, sqrt(3)/3)));

    GoniometerAxis axis("omega", 0.0f, 1.0f, Coord(1, 0, 0), std::nullopt);
    exp_i.Goniometer(axis);

    BraggPredictionSettings prediction_settings{ .high_res_A = 1.3, .ewald_dist_cutoff = 0.002 };
    IndexerThreadPool indexer_thread_pool(exp_i.GetIndexingSettings());
    RotationIndexer indexer(exp_i, indexer_thread_pool);
    BraggPrediction prediction;

    for (int img = 0; img < 50; ++img) {
        std::vector<SpotToSave> spots;
        const float angle_deg = axis.GetAngle_deg(img) + axis.GetWedge_deg() / 2.0f;
        const CrystalLattice latt_img = latt_base.Multiply(axis.GetTransformationAngle(angle_deg).transpose());
        const auto n = prediction.Calc(exp_i, latt_img, prediction_settings);
        for (int i = 0; i < n; ++i) {
            const auto &r = prediction.GetReflections().at(i);
            SpotToSave s{};
            s.x = r.predicted_x;
            s.y = r.predicted_y;
            s.image = img;
            s.intensity = 1.0f;
            s.phi = angle_deg;
            s.ice_ring = false;
            s.indexed = true;
            spots.push_back(s);
        }
        indexer.ProcessImage(img, spots);
        if (img == 30)
            indexer.RunIndexing();
    }

    REQUIRE(indexer.GetLattice().has_value());

    // Shear c along b: the orientation and two of the axes are untouched, and alpha - which the class
    // fixes at 90 - moves by about a degree. A free refinement that has walked into a class leaves
    // exactly this, a cell the group cannot describe standing in the group's own setting.
    const CrystalLattice sheared(latt_base.Vec0(), latt_base.Vec1(),
                                 latt_base.Vec2() + latt_base.Vec1() * 0.02f);
    CHECK(std::fabs(sheared.GetUnitCell().alpha - 90.0) > 0.5);

    const auto refit = indexer.RefineConstrained(sheared, gemmi::CrystalSystem::Monoclinic);
    REQUIRE(refit.has_value());
    const auto uc = refit->lattice.GetUnitCell();
    CHECK(uc.alpha == Catch::Approx(90.0).margin(1e-3));
    CHECK(uc.gamma == Catch::Approx(90.0).margin(1e-3));
    // ...and it is the cell the crystal has, not merely a cell obeying the constraint.
    CHECK(uc.a == Catch::Approx(40.0).margin(0.2));
    CHECK(uc.b == Catch::Approx(50.0).margin(0.2));
    CHECK(uc.c == Catch::Approx(80.0).margin(0.2));
    CHECK(uc.beta == Catch::Approx(105.0).margin(0.2));
    // The spots decide whether a caller keeps it, so the fractions have to be the real comparison:
    // the sheared cell indexes worse than the one the constraint brings back.
    CHECK(refit->indexed_fraction > refit->indexed_fraction_before);
    CHECK(refit->indexed_fraction > 0.5f);
}

// Index a synthetic sweep recorded on a detector tilted by true_tilt_deg beyond the tilt the indexer
// is handed, with reflections to res_A.
static std::optional<RotationIndexerResult> IndexOnTiltedDetector(double true_tilt_deg, float res_A) {
    constexpr double header_rot2_rad = 0.02;
    DiffractionExperiment exp_header;
    exp_header.IncidentEnergy_keV(WVL_1A_IN_KEV)
              .BeamX_pxl(1000)
              .BeamY_pxl(1000)
              .PoniRot1_rad(0.01)
              .PoniRot2_rad(header_rot2_rad)
              .DetectorDistance_mm(200)
              .ImagesPerTrigger(50);

    IndexingSettings settings;
    settings.Algorithm(IndexingAlgorithmEnum::Auto);   // rotation: FFT with a GPU, FFTW without
    settings.RotationIndexing(true).RotationIndexingAngularStride_deg(1.0).RotationIndexingMinAngularRange_deg(30.0);
    exp_header.ImportIndexingSettings(settings);

    GoniometerAxis axis("omega", 0.0f, 1.0f, Coord(1, 0, 0), std::nullopt);
    exp_header.Goniometer(axis);

    // The detector the spots were actually recorded on.
    DiffractionExperiment exp_true = exp_header;
    exp_true.PoniRot2_rad(header_rot2_rad + true_tilt_deg * PI / 180.0);

    const CrystalLattice latt_base =
        CrystalLattice(40, 50, 80, 90, 90, 90).Multiply(RotMatrix(2.0, Coord(sqrt(3)/3, sqrt(3)/3, sqrt(3)/3)));

    BraggPredictionSettings prediction_settings{ .high_res_A = res_A, .ewald_dist_cutoff = 0.002 };
    IndexerThreadPool indexer_thread_pool(exp_header.GetIndexingSettings());
    RotationIndexer indexer(exp_header, indexer_thread_pool);
    BraggPrediction prediction;

    for (int img = 0; img < 50; ++img) {
        std::vector<SpotToSave> spots;
        const float angle_deg = axis.GetAngle_deg(img) + axis.GetWedge_deg() / 2.0f;
        const CrystalLattice latt_img = latt_base.Multiply(axis.GetTransformationAngle(angle_deg).transpose());
        const auto n = prediction.Calc(exp_true, latt_img, prediction_settings);
        for (int i = 0; i < n; ++i) {
            const auto &r = prediction.GetReflections().at(i);
            SpotToSave s{};
            s.x = r.predicted_x;
            s.y = r.predicted_y;
            s.image = img;
            s.intensity = 1.0f;
            s.phi = angle_deg;
            s.ice_ring = false;
            s.indexed = true;
            spots.push_back(s);
        }
        indexer.ProcessImage(img, spots);
        if (img == 30)
            indexer.RunIndexing();
    }
    // Round-trip through ForceResult - how a canonical pass takes over the result of the scheme
    // indexer that found the lattice, and what the report then reads - so what comes back is what a
    // run sees, the tilt walk included.
    const auto found = indexer.GetLattice();
    if (!found)
        return {};
    RotationIndexer forced(exp_header, indexer_thread_pool);
    forced.ForceResult(*found);
    return forced.GetLattice();
}

// The detector tilt is refined freely, and a fit that walks it further from where it started than a
// mounted detector can be off square by (ROT_TILT_PRIOR_DEG) is refused and made again with the tilt
// held. The prior must not touch a tilt a mounting can have: on a detector tilted half a degree
// beyond the tilt the indexer is handed, with reflections to 2.5 A, the fit finds it, keeps it and
// reports no refusal.
TEST_CASE("RotationIndexer keeps a tilt a mounting can have") {
    const auto ret = IndexOnTiltedDetector(0.5, 2.5f);
    REQUIRE(ret.has_value());
    CHECK_FALSE(ret->tilt_walk.has_value());
    CHECK((ret->geom.GetPoniRot2_rad() - 0.02) * 180.0 / PI == Catch::Approx(0.5).margin(0.05));
    CHECK(ret->geom.GetPoniRot1_rad() == Catch::Approx(0.01).margin(1e-3));
    const auto uc = ret->lattice.GetUnitCell();
    CHECK(uc.a == Catch::Approx(40.0).margin(0.3));
    CHECK(uc.b == Catch::Approx(50.0).margin(0.3));
    CHECK(uc.c == Catch::Approx(80.0).margin(0.5));
}

// The other side of the prior, and the unidentifiability that makes it necessary: the same detector
// tilted 1.5 deg beyond the handed tilt, but with reflections only to 6 A, where the keystone the
// fit could read the tilt off is a fraction of a pixel. The free fit does not find 1.5 deg - it runs
// away to about 4 deg (measured 3.9; at 8 A it reaches 37), because at that 2theta reach the tilt is
// a whole-pattern shift the beam centre imitates and nothing pins its size. That walk is past the
// prior, so the lattice is refined again with the tilt held and the two are judged on the spots they
// index; the result records the walk, both counts, the verdict, and a geometry that matches it. The
// lattice is not checked: a 1.5 deg detector error on 6 A data already puts the FFT on a different
// cell before any fit.
TEST_CASE("RotationIndexer judges a tilt no mounting can have on the spots", "[portable]") {
    const auto ret = IndexOnTiltedDetector(1.5, 6.0f);
    REQUIRE(ret.has_value());
    REQUIRE(ret->tilt_walk.has_value());
    const auto &tw = *ret->tilt_walk;
    CHECK(std::hypot(tw.tilt_rad[0] - 0.01, tw.tilt_rad[1] - 0.02) * 180.0 / PI > 1.0);
    CHECK(tw.spots_walked > 0.0f);
    CHECK(tw.spots_held > 0.0f);
    CHECK(tw.refused == !(tw.spots_walked > tw.spots_held + std::sqrt(tw.spots_held)));
    if (tw.refused) {
        CHECK(ret->geom.GetPoniRot1_rad() == Catch::Approx(0.01).margin(1e-7));
        CHECK(ret->geom.GetPoniRot2_rad() == Catch::Approx(0.02).margin(1e-7));
    } else {
        CHECK(ret->geom.GetPoniRot1_rad() == Catch::Approx(tw.tilt_rad[0]).margin(1e-7));
        CHECK(ret->geom.GetPoniRot2_rad() == Catch::Approx(tw.tilt_rad[1]).margin(1e-7));
    }
}