// 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 <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() {
        MtzTable mtz;
        mtz.unit_cell = TEST_CELL;
        mtz.space_group_name = "P 1";
        for (const char *label: {"H", "K", "L", "FWT", "PHWT", "DELFWT", "PHDELWT"})
            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);
                    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))};
                    for (size_t c = 0; c < mtz.columns.size(); c++)
                        mtz.columns[c].values.push_back(row[c]);
                }
        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) {
        gemmi::Ccp4<float> map;
        map.read_ccp4_from_memory(bytes.data(), bytes.size(), "test");
        points = map.grid.data.size();
        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());
    REQUIRE(structure.HasMaps());

    for (const auto type: {RugnuxMapType::TwoFoFc, RugnuxMapType::FoFc}) {
        const std::string bytes = structure.Map(type);
        REQUIRE(bytes.size() > 1024);
        CHECK(bytes.substr(208, 4) == "MAP ");
        size_t points = 0;
        const gemmi::Fractional peak = Extremum(bytes, type == RugnuxMapType::FoFc, 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));
        // A box of 2 x 5 A around the atom, not the 30-odd A cell.
        CHECK(points < 20 * 20 * 20);
        // Cached: the same bytes again.
        CHECK(structure.Map(type) == bytes);
    }
}

TEST_CASE("RugnuxStructure_ModelCif", "[RugnuxLive]") {
    const RugnuxStructure structure(OneAtomModel(), TEST_CELL, *gemmi::find_spacegroup_by_name("P 1"), {});
    CHECK(!structure.HasMaps());
    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);
}