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

#include <gemmi/ccp4.hpp>
#include <gemmi/symmetry.hpp>

#include "../receiver/RugnuxLiveStatus.h"
#include "../receiver/RugnuxStructure.h"

namespace {
    constexpr UnitCell TEST_CELL{30.0f, 32.0f, 34.0f, 90.0f, 90.0f, 90.0f};
    const gemmi::Fractional ATOM_FRACT(0.3, 0.4, 0.55);

    // One carbon atom at ATOM_FRACT, in P 1.
    std::shared_ptr<gemmi::Structure> OneAtomModel() {
        auto st = std::make_shared<gemmi::Structure>();
        st->cell = TEST_CELL;
        st->spacegroup_hm = "P 1";
        gemmi::Model &m = st->models.emplace_back(1);
        gemmi::Chain &ch = m.chains.emplace_back("A");
        gemmi::Residue &r = ch.residues.emplace_back(gemmi::ResidueId{gemmi::SeqId(1, ' '), {}, "ALA"});
        r.het_flag = 'A';
        gemmi::Atom &a = r.atoms.emplace_back();
        a.name = "CA";
        a.element = gemmi::Element("C");
        a.pos = st->cell.orthogonalize(ATOM_FRACT);
        a.b_iso = 20.0f;
        st->setup_cell_images();
        return st;
    }

    // The map coefficients of a point at ATOM_FRACT to 2 A, as ModelValidation lays them out: F = 1 and
    // phi = 360 h.x, so that rho(x) = 1/V sum F exp(-2 pi i h.x) peaks there. The difference map is the same
    // point, negative.
    MtzTable PointCoefficients(bool anomalous = false) {
        MtzTable mtz;
        mtz.unit_cell = TEST_CELL;
        mtz.space_group_name = "P 1";
        for (const char *label: {"H", "K", "L", "FWT", "PHWT", "DELFWT", "PHDELWT", "ANOM", "PHANOM"})
            mtz.columns.push_back({label, 'R', 0, {}});
        const gemmi::UnitCell cell = TEST_CELL;
        const gemmi::ReciprocalAsu asu(gemmi::find_spacegroup_by_name("P 1"));
        for (int h = -15; h <= 15; h++)
            for (int k = -16; k <= 16; k++)
                for (int l = -17; l <= 17; l++) {
                    const gemmi::Miller hkl{h, k, l};
                    if ((h == 0 && k == 0 && l == 0) || !asu.is_in(hkl) || cell.calculate_d(hkl) < 2.0)
                        continue;
                    double phi = 360.0 * (h * ATOM_FRACT.x + k * ATOM_FRACT.y + l * ATOM_FRACT.z);
                    phi -= 360.0 * std::floor(phi / 360.0);
                    // The anomalous map the same point, missing on a fifth of the reflections (no Bijvoet pair).
                    const float nan = std::numeric_limits<float>::quiet_NaN();
                    const bool pair = (h + k + l) % 5 != 0;
                    const float row[] = {static_cast<float>(h), static_cast<float>(k), static_cast<float>(l),
                                         1.0f, static_cast<float>(phi),
                                         1.0f, static_cast<float>(std::fmod(phi + 180.0, 360.0)),
                                         pair ? 1.0f : nan, pair ? static_cast<float>(phi) : nan};
                    for (size_t c = 0; c < mtz.columns.size(); c++)
                        mtz.columns[c].values.push_back(row[c]);
                }
        if (!anomalous)
            mtz.columns.resize(7);
        return mtz;
    }

    // The fractional position of the highest point of a CCP4 map read from `bytes`; lowest where `lowest`.
    gemmi::Fractional Extremum(const std::string &bytes, bool lowest, size_t &points, size_t &cell_points) {
        gemmi::Ccp4<float> map;
        map.read_ccp4_from_memory(bytes.data(), bytes.size(), "test");
        points = map.grid.data.size();
        cell_points = static_cast<size_t>(map.header_i32(8)) * map.header_i32(9) * map.header_i32(10);
        const auto start = map.header_3i32(5);
        const auto sampling = map.header_3i32(8);
        size_t best = 0;
        for (size_t i = 0; i < map.grid.data.size(); i++)
            if (lowest ? map.grid.data[i] < map.grid.data[best] : map.grid.data[i] > map.grid.data[best])
                best = i;
        const size_t u = best % map.grid.nu, v = (best / map.grid.nu) % map.grid.nv,
                w = best / (static_cast<size_t>(map.grid.nu) * map.grid.nv);
        return {static_cast<double>(start[0] + static_cast<int>(u)) / sampling[0],
                static_cast<double>(start[1] + static_cast<int>(v)) / sampling[1],
                static_cast<double>(start[2] + static_cast<int>(w)) / sampling[2]};
    }
}

TEST_CASE("RugnuxStructure_MapPeaksAtTheAtom", "[RugnuxLive]") {
    const RugnuxStructure structure(OneAtomModel(), TEST_CELL, *gemmi::find_spacegroup_by_name("P 1"),
                                    PointCoefficients(true));
    REQUIRE(structure.HasMaps());
    REQUIRE(structure.HasAnomalousMap());

    for (const auto type: {RugnuxMapType::TwoFoFc, RugnuxMapType::FoFc, RugnuxMapType::Anomalous}) {
        const std::string bytes = structure.Map(type);
        REQUIRE(bytes.size() > 1024);
        CHECK(bytes.substr(208, 4) == "MAP ");
        size_t points = 0, cell_points = 0;
        const gemmi::Fractional peak = Extremum(bytes, type == RugnuxMapType::FoFc, points, cell_points);
        // Within one grid step (d_min/3 = 0.67 A, i.e. at most 0.022 of these cells).
        CHECK(peak.x == Catch::Approx(ATOM_FRACT.x).margin(0.025));
        CHECK(peak.y == Catch::Approx(ATOM_FRACT.y).margin(0.025));
        CHECK(peak.z == Catch::Approx(ATOM_FRACT.z).margin(0.025));
        // One unit cell (centred on the model, which is smaller than the cell).
        CHECK(points == cell_points);
        // Cached: the same bytes again.
        CHECK(structure.Map(type) == bytes);
    }
}

TEST_CASE("RugnuxStructure_MapCoversAModelWiderThanTheCell", "[RugnuxLive]") {
    // A second atom 1.4 cells along a from the first: the map reaches along a over both, plus the margin.
    auto st = OneAtomModel();
    gemmi::Atom far = st->models[0].chains[0].residues[0].atoms[0];
    far.name = "CB";
    far.pos = st->cell.orthogonalize(gemmi::Fractional(ATOM_FRACT.x + 1.4, ATOM_FRACT.y, ATOM_FRACT.z));
    st->models[0].chains[0].residues[0].atoms.push_back(far);
    const RugnuxStructure structure(st, TEST_CELL, *gemmi::find_spacegroup_by_name("P 1"), PointCoefficients());
    CHECK(!structure.HasAnomalousMap());
    const std::string bytes = structure.Map(RugnuxMapType::TwoFoFc);
    gemmi::Ccp4<float> map;
    map.read_ccp4_from_memory(bytes.data(), bytes.size(), "test");
    const double a_span_A = TEST_CELL.a * map.grid.nu / map.header_i32(8);
    CHECK(a_span_A == Catch::Approx(1.4 * TEST_CELL.a + 2 * RugnuxStructure::MAP_MARGIN_A).margin(1.5));
    CHECK(map.grid.nv == map.header_i32(9));   // one cell along b and c
    CHECK(map.grid.nw == map.header_i32(10));
}

TEST_CASE("RugnuxStructure_ModelCif", "[RugnuxLive]") {
    const RugnuxStructure structure(OneAtomModel(), TEST_CELL, *gemmi::find_spacegroup_by_name("P 1"), {});
    CHECK(!structure.HasMaps());
    CHECK(!structure.HasAnomalousMap());
    CHECK(structure.Map(RugnuxMapType::Anomalous).empty());
    CHECK(structure.Map(RugnuxMapType::TwoFoFc).empty());
    const std::string cif = structure.ModelCif();
    CHECK(cif.find("_atom_site.") != std::string::npos);
    CHECK(cif.find("_cell.length_a") != std::string::npos);
    CHECK(cif.find("'P 1'") != std::string::npos);
}

TEST_CASE("RugnuxLiveStatus_StructureClearedAtStart", "[RugnuxLive]") {
    RugnuxLiveStatus status;
    status.Structure(std::make_shared<RugnuxStructure>(OneAtomModel(), TEST_CELL,
                                                       *gemmi::find_spacegroup_by_name("P 1"), std::nullopt));
    CHECK(status.Get().structure);
    status.Start(10, 1.0, std::chrono::steady_clock::now());
    CHECK(!status.Get().structure);
    status.Structure(std::make_shared<RugnuxStructure>(OneAtomModel(), TEST_CELL,
                                                       *gemmi::find_spacegroup_by_name("P 1"), std::nullopt));
    status.Clear();
    CHECK(!status.Get().structure);
}