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

#include "../image_analysis/scale_merge/AnisotropyAnalysis.h"
#include "gemmi/symmetry.hpp"
#include "gemmi/unitcell.hpp"

namespace {
    gemmi::UnitCell Cell(double a, double b, double c, double al, double be, double ga) {
        gemmi::UnitCell out;
        out.set(a, b, c, al, be, ga);
        return out;
    }

    // A synthetic merge: Wilson-distributed intensities with an isotropic fall-off and a known
    // deviatoric anisotropy on top, over every hkl inside the resolution limit. The "structure
    // factor" is a hash of the index, so the set is reproduced bit for bit.
    std::vector<MergedReflection> SyntheticMerge(const gemmi::UnitCell &cell, const gemmi::SpaceGroup *sg,
                                                 double d_min, double b_iso,
                                                 const gemmi::SMat33<double> &b_dev) {
        const gemmi::GroupOps ops = sg->operations();
        const gemmi::ReciprocalAsu asu(sg);
        std::vector<MergedReflection> out;
        const int hmax = static_cast<int>(cell.a / d_min) + 1;
        const int kmax = static_cast<int>(cell.b / d_min) + 1;
        const int lmax = static_cast<int>(cell.c / d_min) + 1;
        for (int h = -hmax; h <= hmax; ++h)
            for (int k = -kmax; k <= kmax; ++k)
                for (int l = -lmax; l <= lmax; ++l) {
                    const gemmi::Op::Miller hkl{{h, k, l}};
                    if ((h == 0 && k == 0 && l == 0) || !asu.is_in(hkl) || ops.is_systematically_absent(hkl))
                        continue;
                    const double d2 = cell.calculate_1_d2(hkl);
                    if (d2 <= 0 || 1.0 / std::sqrt(d2) < d_min)
                        continue;
                    // Wilson draw from a hash of the index: deterministic, and spanning a realistic range.
                    const uint32_t seed = static_cast<uint32_t>(h * 73856093 ^ k * 19349663 ^ l * 83492791);
                    const double u = ((seed * 2654435761u) >> 8) / static_cast<double>(1 << 24);
                    const double wilson = -std::log(std::max(1e-6, u));
                    // The tensor is applied in the Cartesian frame, s = F^T h.
                    const gemmi::Vec3 s = cell.frac.mat.left_multiply(gemmi::Vec3(h, k, l));
                    const double aniso = -0.5 * b_dev.r_u_r(s);
                    MergedReflection r;
                    r.h = h; r.k = k; r.l = l;
                    r.d = static_cast<float>(1.0 / std::sqrt(d2));
                    r.I = static_cast<float>(1000.0 * wilson * std::exp(-0.5 * b_iso * d2 + aniso));
                    r.sigma = static_cast<float>(0.02 * std::fabs(r.I) + 1.0);
                    out.push_back(r);
                }
        return out;
    }
}

// The number of free deviatoric anisotropy parameters is fixed by the Laue class alone. This is the
// self-test of the constraint basis: 5 / 3 / 2 / 1 / 1 / 0 for triclinic / monoclinic / orthorhombic /
// tetragonal / trigonal-hexagonal / cubic, and nothing else is possible.
TEST_CASE("Anisotropy free-parameter count", "[anisotropy]") {
    struct Case {
        const char *space_group;
        gemmi::UnitCell cell;
        int free_directions;
    };
    const std::vector<Case> cases{
        {"P 1",       Cell(51, 62, 73, 84.0, 95.0, 103.0), 5},
        {"P 1 21 1",  Cell(51, 62, 73, 90.0, 95.0, 90.0),  3},
        {"C 1 2 1",   Cell(91, 62, 73, 90.0, 105.0, 90.0), 3},
        {"P 21 21 21", Cell(51, 62, 73, 90.0, 90.0, 90.0), 2},
        {"I 2 2 2",   Cell(51, 62, 73, 90.0, 90.0, 90.0),  2},
        {"P 43 21 2", Cell(79, 79, 38, 90.0, 90.0, 90.0),  1},   // lysozyme, the field's test specimen
        {"P 31 2 1",  Cell(62, 62, 91, 90.0, 90.0, 120.0), 1},
        {"R 3 :H",    Cell(78, 78, 33, 90.0, 90.0, 120.0), 1},
        {"P 63",      Cell(62, 62, 91, 90.0, 90.0, 120.0), 1},
        {"I 2 3",     Cell(78, 78, 78, 90.0, 90.0, 90.0),  0},
        {"F 4 3 2",   Cell(78, 78, 78, 90.0, 90.0, 90.0),  0},
    };
    for (const auto &c : cases) {
        const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name(c.space_group);
        REQUIRE(sg != nullptr);
        std::vector<MergedReflection> merged(1);
        merged[0].h = 1; merged[0].k = 0; merged[0].l = 0;
        merged[0].I = 1.0f; merged[0].sigma = 1.0f; merged[0].d = 10.0f;
        const auto result = AnalyzeAnisotropy(merged, {}, c.cell, sg);
        INFO(c.space_group);
        CHECK(result.n_free_parameters == c.free_directions);
    }
}

// A refined cell need not obey its space group's metric constraints exactly, and rugnux writes the
// unconstrained refined cell. An angle a hundredth of a degree off 90 must not create an extra
// anisotropy direction.
TEST_CASE("Anisotropy free-parameter count with an off-metric refined cell", "[anisotropy]") {
    std::vector<MergedReflection> merged(1);
    merged[0].h = 1; merged[0].k = 0; merged[0].l = 0;
    merged[0].I = 1.0f; merged[0].sigma = 1.0f; merged[0].d = 10.0f;
    CHECK(AnalyzeAnisotropy(merged, {}, Cell(51, 62, 73, 90.02, 105.0, 89.97),
                            gemmi::find_spacegroup_by_name("C 1 2 1")).n_free_parameters == 3);
    CHECK(AnalyzeAnisotropy(merged, {}, Cell(51.0, 62.0, 73.0, 89.98, 90.03, 90.01),
                            gemmi::find_spacegroup_by_name("I 2 2 2")).n_free_parameters == 2);
    CHECK(AnalyzeAnisotropy(merged, {}, Cell(78.01, 77.99, 78.02, 90.01, 89.99, 90.0),
                            gemmi::find_spacegroup_by_name("I 2 3")).n_free_parameters == 0);
}

// A cubic crystal has no free deviatoric parameter, so its deltaB is exactly zero by symmetry - not
// small, not measured, zero - and the verdict is a statement about symmetry rather than about data.
TEST_CASE("Anisotropy is exactly zero in a cubic Laue class", "[anisotropy]") {
    const gemmi::UnitCell cell = Cell(78, 78, 78, 90, 90, 90);
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("I 2 3");
    const auto merged = SyntheticMerge(cell, sg, 2.5, 20.0, {8.0, -4.0, -4.0, 0.0, 0.0, 0.0});
    REQUIRE(merged.size() > 1000);
    const auto result = AnalyzeAnisotropy(merged, {}, cell, sg);
    CHECK(result.n_free_parameters == 0);
    CHECK(result.delta_b == 0.0);
    CHECK(result.verdict == AnisotropyVerdict::NotDetected);
}

// The tensor itself: put a known deviatoric B into a tetragonal merge and read it back. The
// tetragonal Laue class leaves one free direction, along c*, and its magnitude is what deltaB means.
TEST_CASE("Anisotropy tensor is recovered from a synthetic merge", "[anisotropy]") {
    const gemmi::UnitCell cell = Cell(79, 79, 38, 90, 90, 90);
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 43 21 2");
    // Uniaxial about c, deltaB = B_zz - B_xx = 15 A^2.
    const gemmi::SMat33<double> b_dev{-5.0, -5.0, 10.0, 0.0, 0.0, 0.0};
    const auto merged = SyntheticMerge(cell, sg, 2.0, 20.0, b_dev);
    REQUIRE(merged.size() > 2000);
    const auto result = AnalyzeAnisotropy(merged, {}, cell, sg);
    REQUIRE(result.n_free_parameters == 1);
    REQUIRE(result.n_cells > 0);
    CHECK(result.delta_b == Catch::Approx(15.0).margin(1.5));
    // c* is the weak direction here, so the largest principal value points along z.
    CHECK(std::fabs(result.eigenvector[0][2]) == Catch::Approx(1.0).margin(0.05));
    // A pure Debye-Waller fall-off is a straight line through the origin in s^2.
    CHECK(result.shape == AnisotropyShape::Linear);
    // With no unmerged observations the systematic-error scale cannot be measured, and the verdict
    // says so rather than falling back on a counting-statistics error bar.
    CHECK(result.verdict == AnisotropyVerdict::CannotDetermine);
    CHECK_FALSE(result.refusal.empty());
}

// An isotropic merge must not produce a tensor, whatever the Laue class allows.
TEST_CASE("Anisotropy of an isotropic synthetic merge is small", "[anisotropy]") {
    const gemmi::UnitCell cell = Cell(79, 79, 38, 90, 90, 90);
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 43 21 2");
    const auto merged = SyntheticMerge(cell, sg, 2.0, 20.0, {0.0, 0.0, 0.0, 0.0, 0.0, 0.0});
    REQUIRE(merged.size() > 2000);
    const auto result = AnalyzeAnisotropy(merged, {}, cell, sg);
    REQUIRE(result.n_cells > 0);
    CHECK(std::fabs(result.delta_b) < 1.0);
}

namespace {
    // One partial of reflection (h,k,l) on image `frame`, with the fields ScaledObservations reads.
    Reflection Partial(int h, int k, int l, float frame, float I, float partiality) {
        Reflection r{};
        r.h = h; r.k = k; r.l = l;
        r.image_number = frame;
        r.d = 3.0f;
        r.I = I;
        r.sigma = 10.0f;
        r.rlp = 1.0f;
        r.partiality = partiality;
        return r;
    }
}

// ScaledObservations assembles rotation partials into fulls, and - because the floor downstream only
// ever reads a sample of the unique reflections - it may hand back a sample of them rather than all.
// This pins the assembly on a run far below the sampling cap, where the sample is everything: the
// partials of one reflection over consecutive frames become one full on the merge's own scale, a gap
// wider than the combine's splits the run in two, an event that caught too little of its rocking
// curve is dropped, and an image with no fitted scale contributes nothing.
TEST_CASE("Rotation partials are assembled into scaled fulls", "[anisotropy]") {
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 43 21 2");
    std::vector<IntegrationOutcome> outcomes(8);
    for (auto &o : outcomes)
        o.image_scale_g = 2.0f;
    // (1,2,3) is measured over frames 0-2, then again over frames 6-7 after a four-frame gap.
    outcomes[0].reflections.push_back(Partial(1, 2, 3, 0.0f, 100.0f, 0.25f));
    outcomes[1].reflections.push_back(Partial(1, 2, 3, 1.0f, 300.0f, 0.50f));
    outcomes[2].reflections.push_back(Partial(1, 2, 3, 2.0f, 100.0f, 0.25f));
    outcomes[6].reflections.push_back(Partial(1, 2, 3, 6.0f, 200.0f, 0.40f));
    outcomes[7].reflections.push_back(Partial(1, 2, 3, 7.0f, 300.0f, 0.60f));
    // A second reflection whose two partials sum to less than min_partiality: not a measurement.
    outcomes[3].reflections.push_back(Partial(4, 5, 6, 3.0f, 50.0f, 0.10f));
    outcomes[4].reflections.push_back(Partial(4, 5, 6, 4.0f, 50.0f, 0.15f));
    // A third on an image with no fitted per-image scale.
    outcomes[5].image_scale_g.reset();
    outcomes[5].reflections.push_back(Partial(7, 8, 9, 5.0f, 500.0f, 1.00f));

    const auto obs = ScaledObservations(outcomes, /*rotation=*/true, sg);
    REQUIRE(obs.size() == 2);
    for (const auto &o : obs) {
        CHECK(o.h == 1);
        CHECK(o.k == 2);
        CHECK(o.l == 3);
    }
    // Both events are complete, so their partialities sum to 1 and the divisor leaves them alone;
    // what is left is the summed intensity over the per-image scale.
    CHECK(obs[0].I == Catch::Approx(250.0));    // (100 + 300 + 100) / 2
    CHECK(obs[1].I == Catch::Approx(250.0));    // (200 + 300) / 2
    // A still is a whole measurement of its reflection, so there is nothing to assemble and nothing
    // to sample: each partial stands or falls on its own partiality, and only two clear the floor.
    const auto stills = ScaledObservations(outcomes, /*rotation=*/false, sg);
    CHECK(stills.size() == 2);
}