// 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 "../receiver/RugnuxLiveStatus.h"
#include "../receiver/JFJochReceiverRugnux.h"
#include "../broker/OpenAPIConvert.h"
#include "../image_puller/TestImagePuller.h"

using namespace std::chrono_literals;

TEST_CASE("RugnuxLiveStatus_TimesFromFirstImage", "[RugnuxLive]") {
    // 11 images at 100 Hz, the first 3 s after START, reported out of order as several forward threads do
    const auto start = std::chrono::steady_clock::now() - 10s;
    const auto first = start + 3s;

    RugnuxLiveStatus status;
    status.Start(11, 0.11, start);
    for (int i: {1, 0, 2, 3, 4, 5, 6, 7, 8, 10, 9})
        status.ImageReceived(first + i * 10ms);

    auto s = status.Get();
    CHECK(s.expected_collection_s.value() == Catch::Approx(0.11));
    REQUIRE(s.images_received == 11);
    REQUIRE(s.start_to_first_image_s.value() == Catch::Approx(3.0));
    CHECK(!s.last_image_s);
    // "now" is 10 s after START, 7 s after the first image
    CHECK(s.elapsed_s == Catch::Approx(7.0).margin(0.5));

    REQUIRE(status.LastImage() == first + 100ms);
    s = status.Get();
    REQUIRE(s.last_image_s.value() == Catch::Approx(0.1));
    CHECK((s.images_received - 1) / s.last_image_s.value() == Catch::Approx(100.0));
}

TEST_CASE("RugnuxLiveStatus_BeforeFirstImage", "[RugnuxLive]") {
    RugnuxLiveStatus status;
    status.Start(10, std::nullopt, std::chrono::steady_clock::now() - 2s);
    status.Step("waiting");

    const auto s = status.Get();
    CHECK(!s.start_to_first_image_s);
    CHECK(!s.expected_collection_s);
    CHECK(s.elapsed_s == 0.0);
    REQUIRE(s.steps.size() == 1);
    CHECK(s.steps[0].started_s <= 0.0);
    CHECK(!status.LastImage());
}

TEST_CASE("RugnuxLiveStatus_Clear", "[RugnuxLive]") {
    RugnuxLiveStatus status;
    status.Start(10, 1.0, std::chrono::steady_clock::now() - 2s);
    status.ImageReceived(std::chrono::steady_clock::now() - 1s);
    status.Step("indexing");
    status.Probe("lattice found");
    status.Result("OK", "fine", {{"Space group", "P 1"}}, {"a warning"});
    status.Finish(RugnuxLiveState::Done);
    REQUIRE(status.LastImage());

    status.Clear();
    const auto s = status.Get();
    CHECK(s.state == RugnuxLiveState::Idle);
    CHECK(s.images_expected == 0);
    CHECK(s.images_received == 0);
    CHECK(s.steps.empty());
    CHECK(s.probe.empty());
    CHECK(s.verdict.empty());
    CHECK(s.summary.empty());
    CHECK(s.warnings.empty());
    CHECK(!s.start_to_first_image_s);
    CHECK(!s.last_image_s);
    CHECK(s.elapsed_s == 0.0);
    CHECK(!status.LastImage());
}

TEST_CASE("TestImagePuller_SetsReceiveTime", "[RugnuxLive]") {
    TestImagePuller puller;
    const auto before = std::chrono::steady_clock::now();
    puller.Put(ImagePullerOutput{});
    const auto given = before - 5s;
    puller.Put(ImagePullerOutput{.received = given});

    auto out = puller.PollImage();
    REQUIRE(out);
    CHECK(out->received >= before);
    out = puller.PollImage();
    REQUIRE(out);
    CHECK(out->received == given);
}

TEST_CASE("RugnuxLiveModel_MolecularReplacement", "[RugnuxLive]") {
    ProcessResult r;
    CHECK(!LiveModel(r));

    MolecularReplacementResult mr;
    mr.ok = true;
    mr.model_name = "model.cif";
    mr.predicted_model = true;
    mr.decision = "AGREES";
    mr.d_min_search = 2.5;
    MRCandidate c;
    c.space_group = *gemmi::find_spacegroup_by_name("P 1 21 1");
    c.copies = 2;
    c.tf_z = 12.5;
    c.rf_z = 6.0;
    c.null_sigma = 20.0;
    mr.candidates = {MRCandidate{}, c};
    mr.chosen = 1;
    r.model_prior = mr;
    ModelValidationResult v;
    v.ok = true;
    v.r_work = 0.38;
    v.r_free = 0.40;
    v.cc_model_overall = 0.71;
    v.mean_atom_density_sigma = 1.6;
    v.anomalous_sites = {{"SD MET A1", 5.5}};
    r.model_validation = v;
    r.model_placement = "MR";

    const auto m = LiveModel(r);
    REQUIRE(m);
    CHECK(m->placement == "MR");
    CHECK(m->predicted);
    CHECK(m->mr_space_group == "P 1 21 1");
    CHECK(m->copies.value() == 2);
    CHECK(m->tf_z.value() == Catch::Approx(12.5));
    CHECK(!m->llg_z);
    CHECK(m->fit_sigma.value() == Catch::Approx(20.0));   // MR's null: the rigid body's was not tested
    CHECK(m->r_free.value() == Catch::Approx(0.40));
    CHECK(m->cc_model.value() == Catch::Approx(0.71));
    CHECK(!m->rigid_body_shift_A);
    CHECK(m->anomalous_site == "SD MET A1");

    RugnuxLiveStatus status;
    status.Start(10, 1.0, std::chrono::steady_clock::now());
    status.Result("OK", "fine", {}, {}, true);
    status.Model(*m);
    nlohmann::json j = Convert(status.Get());
    CHECK(j["anomalous_signal"] == true);
    CHECK(j["model"]["placement"] == "MR");
    CHECK(j["model"]["copies"] == 2);
    CHECK(j["model"]["r_free"].get<double>() == Catch::Approx(0.40));
    CHECK(!j["model"].contains("llg_z"));
    CHECK(!j["model"].contains("rigid_body_angle_deg"));

    status.Clear();
    CHECK(!status.Get().model);
    CHECK(!status.Get().anomalous_signal);
}

TEST_CASE("RugnuxLiveStatus_Headline", "[RugnuxLive]") {
    RugnuxLiveStatus status;
    status.Start(10, 1.0, std::chrono::steady_clock::now());
    nlohmann::json j = Convert(status.Get());
    CHECK(!j.contains("space_group"));
    CHECK(!j.contains("d_min"));
    CHECK(!j.contains("isa"));

    status.Headline("P 1 21 1", 1.25, 24.2);
    j = Convert(status.Get());
    CHECK(j["space_group"] == "P 1 21 1");
    CHECK(j["d_min"].get<double>() == Catch::Approx(1.25));
    CHECK(j["isa"].get<double>() == Catch::Approx(24.2));

    // ISa undetermined
    status.Headline("P 1 21 1", 1.25, std::nullopt);
    j = Convert(status.Get());
    CHECK(!j.contains("isa"));

    status.Clear();
    CHECK(status.Get().space_group.empty());
    CHECK(!status.Get().d_min);
}

TEST_CASE("RugnuxLiveModel_Refinement", "[RugnuxLive]") {
    ProcessResult r;
    ModelValidationResult v;
    v.ok = true;
    v.r_work = 0.38;
    v.r_free = 0.40;
    r.model_validation = v;
    r.model_placement = "RIGID_BODY";
    CHECK(!LiveModel(r)->refinement);

    RefinementResult rf;
    rf.ok = true;
    rf.on_gpu = true;
    rf.seconds = 3.5;
    rf.r_work_placement = 0.30;
    rf.r_free_placement = 0.33;
    rf.restrained = true;
    rf.r_work = 0.19;
    rf.r_free = 0.23;
    rf.geometry.rmsd_bonds = 0.011;
    rf.geometry.rmsd_angles = 1.6;
    DensityBlob b;
    b.kind = DensityBlob::Kind::Ligand;
    b.centroid = gemmi::Position(1.0, 2.0, 3.0);
    b.peak_sigma = 7.5;
    b.volume = 40.0;
    b.nearest_atom = "OG1 THR A42";
    b.nearest_atom_distance = 3.1;
    rf.blobs.assign(RugnuxLiveRefinement::MAX_BLOBS + 5, b);
    r.refinement = rf;

    const auto m = LiveModel(r);
    REQUIRE(m->refinement);
    const RugnuxLiveRefinement &lr = *m->refinement;
    CHECK(lr.performed);
    CHECK(lr.r_free.value() == Catch::Approx(0.23));
    CHECK(lr.rmsd_bonds_A.value() == Catch::Approx(0.011));
    CHECK(lr.blob_count == static_cast<int64_t>(RugnuxLiveRefinement::MAX_BLOBS + 5));
    CHECK(lr.blobs.size() == RugnuxLiveRefinement::MAX_BLOBS);
    // the placed model's R-factors are untouched
    CHECK(m->r_free.value() == Catch::Approx(0.40));

    RugnuxLiveStatus status;
    status.Start(10, 1.0, std::chrono::steady_clock::now());
    status.Model(*m);
    nlohmann::json j = Convert(status.Get());
    CHECK(j["refined_model"] == false);
    CHECK(j["model"]["r_free"].get<double>() == Catch::Approx(0.40));
    CHECK(j["model"]["refinement"]["r_free"].get<double>() == Catch::Approx(0.23));
    CHECK(j["model"]["refinement"]["on_gpu"] == true);
    CHECK(j["model"]["refinement"]["blob_count"] == RugnuxLiveRefinement::MAX_BLOBS + 5);
    CHECK(j["model"]["refinement"]["blobs"].size() == RugnuxLiveRefinement::MAX_BLOBS);
    CHECK(j["model"]["refinement"]["blobs"][0]["kind"] == "LIGAND");
    CHECK(j["model"]["refinement"]["blobs"][0]["z"].get<double>() == Catch::Approx(3.0));

    // Only the placement stage ran: its R-factors, no geometry
    rf.restrained = false;
    r.refinement = rf;
    const auto p = LiveModel(r);
    CHECK(p->refinement->r_free.value() == Catch::Approx(0.33));
    CHECK(!p->refinement->rmsd_bonds_A);

    // Failed
    RefinementResult failed;
    failed.failure_reason = "no dictionary";
    r.refinement = failed;
    status.Model(*LiveModel(r));
    j = Convert(status.Get());
    CHECK(j["model"]["refinement"]["performed"] == false);
    CHECK(j["model"]["refinement"]["reason"] == "no dictionary");
    CHECK(!j["model"]["refinement"].contains("r_free"));
}

TEST_CASE("RugnuxLiveModel_RefinementCancelled", "[RugnuxLive]") {
    ProcessResult r;
    ModelValidationResult v;
    v.ok = true;
    r.model_validation = v;
    r.model_placement = "MR";
    RefinementResult rf;
    rf.ok = true;
    rf.cancelled = true;
    r.refinement = rf;
    const auto m = LiveModel(r);
    CHECK(m->refinement->cancelled);
    RugnuxLiveStatus status;
    status.Start(10, 1.0, std::chrono::steady_clock::now());
    status.Model(*m);
    const nlohmann::json j = Convert(status.Get());
    CHECK(j["model"]["refinement"]["cancelled"] == true);
}

// A score that is not a number is left out, not written as null: the page reads every present one as a number.
TEST_CASE("RugnuxLiveModel_NonFiniteScores", "[RugnuxLive]") {
    ProcessResult r;
    MolecularReplacementResult mr;
    mr.ok = true;
    mr.d_min_search = NAN;
    MRCandidate c;
    c.space_group = *gemmi::find_spacegroup_by_name("P 1 21 1");
    c.tf_z = NAN;
    c.rf_z = INFINITY;
    mr.candidates = {c};
    mr.chosen = 0;
    r.model_prior = mr;
    ModelValidationResult v;
    v.ok = true;
    v.r_work = NAN;
    v.r_free = NAN;
    v.mean_atom_density_sigma = NAN;
    r.model_validation = v;
    r.model_placement = "MR";
    RefinementResult rf;
    rf.ok = true;
    rf.restrained = true;
    rf.r_work = NAN;
    rf.r_free = 0.25;
    rf.geometry.rmsd_bonds = NAN;
    DensityBlob b;
    b.nearest_atom_distance = INFINITY;
    rf.blobs = {b};
    r.refinement = rf;

    const auto m = LiveModel(r);
    REQUIRE(m);
    CHECK(!m->tf_z);
    CHECK(!m->rf_z);
    CHECK(!m->mr_d_min);
    CHECK(!m->r_free);
    CHECK(!m->map_atom_sigma);
    CHECK(!m->refinement->r_work);
    CHECK(m->refinement->r_free.value() == Catch::Approx(0.25));
    CHECK(!m->refinement->rmsd_bonds_A);
    CHECK(!m->refinement->blobs.at(0).nearest_atom_distance_A);

    RugnuxLiveStatus status;
    status.Start(10, 1.0, std::chrono::steady_clock::now());
    status.Model(*m);
    const nlohmann::json j = Convert(status.Get());
    CHECK(!j["model"].contains("tf_z"));
    CHECK(!j["model"].contains("rf_z"));
    CHECK(!j["model"].contains("r_free"));
    CHECK(!j["model"]["refinement"].contains("r_work"));
    CHECK(!j["model"]["refinement"]["blobs"][0].contains("nearest_atom_distance_A"));
    CHECK(j.dump().find("null") == std::string::npos);
}