// 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 <filesystem>
#include <string>

#include <gemmi/mtz.hpp>

#include "../common/DiffractionExperiment.h"
#include "../image_analysis/WriteReflections.h"
#include "../image_analysis/IntegrationOutcome.h"
#include <cmath>
#include "SyntheticMergedReflections.h"

namespace {
    // The wavelength CCP4's mtzlib substitutes when the data columns belong to dataset 0, which it
    // takes for the reserved HKL_base: Cu K-alpha, and the wrong edge for anything reading f'/f''
    // out of the file.
    constexpr double CU_KALPHA_A = 1.54187;

    constexpr UnitCell TETRAGONAL_CELL{47.0f, 47.0f, 63.0f, 90.0f, 90.0f, 90.0f};

    DiffractionExperiment TestExperiment() {
        DiffractionExperiment x;
        x.IncidentEnergy_keV(12.7f);         // ~0.976 A, nowhere near the Cu K-alpha default
        x.SpaceGroupNumber(96);              // P 43 21 2
        x.SetUnitCell(TETRAGONAL_CELL);
        return x;
    }
}

TEST_CASE("Merged MTZ: the data dataset is id 1 and carries the wavelength", "[write_reflections]") {
    jfjoch_test::SyntheticMergeParams params;
    params.true_space_group = "P 43 21 2";
    params.twin_supergroup = "P 43 21 2";
    params.d_min_A = 5.0;
    const auto reflections = jfjoch_test::GenerateSyntheticMerged(params);
    REQUIRE(!reflections.empty());

    const auto experiment = TestExperiment();
    const auto path = (std::filesystem::temp_directory_path() / "rugnux_merged_wavelength.mtz").string();
    WriteMtzReflections(reflections, TETRAGONAL_CELL, experiment, path);
    const gemmi::Mtz mtz = gemmi::read_mtz_file(path);
    std::filesystem::remove(path);

    // HKL_base at id 0, the data at id 1. A data dataset written at id 0 occupies the id MTZ
    // reserves for the base, and mtzlib then reports CU_KALPHA_A instead of the real wavelength.
    REQUIRE(mtz.datasets.size() == 2);
    CHECK(mtz.datasets[0].id == 0);
    CHECK(mtz.datasets[0].dataset_name == "HKL_base");
    CHECK(mtz.datasets[1].id == 1);
    CHECK(mtz.datasets[1].wavelength == Catch::Approx(experiment.GetWavelength_A()).epsilon(1e-5));
    CHECK(mtz.datasets[1].wavelength != Catch::Approx(CU_KALPHA_A).epsilon(1e-3));

    // The wavelength is read off the dataset the data columns belong to, so they have to be on the
    // data dataset and not on the base.
    for (const char *label : {"IMEAN", "SIGIMEAN", "F", "SIGF", "FreeR_flag"}) {
        const gemmi::Mtz::Column *col = mtz.column_with_label(label);
        REQUIRE(col != nullptr);
        CHECK(col->dataset_id == 1);
    }

    // Cell and space group travel in the same header.
    CHECK(mtz.spacegroup != nullptr);
    CHECK(mtz.spacegroup->number == 96);
    CHECK(mtz.cell.a == Catch::Approx(TETRAGONAL_CELL.a).epsilon(1e-5));
    CHECK(mtz.cell.c == Catch::Approx(TETRAGONAL_CELL.c).epsilon(1e-5));
    CHECK(mtz.cell.gamma == Catch::Approx(90.0).epsilon(1e-5));
    CHECK(mtz.datasets[1].cell.a == Catch::Approx(TETRAGONAL_CELL.a).epsilon(1e-5));

    CHECK(mtz.nreflections == static_cast<int>(reflections.size()));
}

TEST_CASE("Unmerged MTZ: LP is Lorentz-polarization and QE carries the sensor efficiency",
          "[write_reflections]") {
    // The whole point of the split: LP must mean what XDS and DIALS mean by it, and the raw count
    // sum must still be recoverable from the file alone, as I / LP * QE.
    auto experiment = TestExperiment();
    experiment.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {}));

    IntegrationOutcome outcome;
    const float raw[3] = {1000.0f, 250.0f, 40.0f};
    const float lp[3]  = {1.75f, 2.50f, 0.90f};    // Lorentz x polarization, and nothing else
    const float qe[3]  = {0.9375f, 0.8125f, 1.0f};    // 1.0 = the sensor said nothing to correct
    for (int i = 0; i < 3; ++i) {
        Reflection r{};
        r.h = 4 + i; r.k = 2; r.l = 6;
        r.image_number = static_cast<float>(i);
        r.d = 5.0f + i;
        r.I = raw[i];                                 // the writer is what applies the factor
        r.sigma = std::sqrt(raw[i]);
        r.prescaling_corr = lp[i];
        r.qe_corr = qe[i];
        r.partiality = 1.0f;
        r.predicted_x = 100.0f + i; r.predicted_y = 200.0f + i;
        r.observed_x = NAN; r.observed_y = NAN;
        outcome.reflections.push_back(r);
    }

    const auto path = (std::filesystem::temp_directory_path() / "rugnux_unmerged_qe.mtz").string();
    WriteUnmergedMtzReflections({outcome}, TETRAGONAL_CELL, experiment, false, path);
    const gemmi::Mtz mtz = gemmi::read_mtz_file(path);
    std::filesystem::remove(path);

    const gemmi::Mtz::Column *c_I = mtz.column_with_label("I");
    const gemmi::Mtz::Column *c_lp = mtz.column_with_label("LP");
    const gemmi::Mtz::Column *c_qe = mtz.column_with_label("QE");
    REQUIRE(c_I != nullptr);
    REQUIRE(c_lp != nullptr);
    REQUIRE(c_qe != nullptr);          // DIALS writes this column even when there is nothing in it
    REQUIRE(mtz.nreflections == 3);

    for (int i = 0; i < 3; ++i) {
        const float I = mtz.data[i * mtz.columns.size() + c_I->idx];
        const float LP = mtz.data[i * mtz.columns.size() + c_lp->idx];
        const float QE = mtz.data[i * mtz.columns.size() + c_qe->idx];
        INFO("row " << i);
        // LP holds Lorentz x polarization alone: the sensor term was never inside it.
        CHECK(LP == Catch::Approx(lp[i]).epsilon(1e-5));
        // QE is a divisor normalised to 1 at normal incidence, so it never drops below 1.
        CHECK(QE == Catch::Approx(1.0f / qe[i]).epsilon(1e-5));
        CHECK(QE >= 1.0f);
        // ... and the two together put the raw counts back.
        CHECK(I / LP * QE == Catch::Approx(raw[i]).epsilon(1e-4));
        // The intensity itself is the fully corrected value - both halves applied.
        CHECK(I == Catch::Approx(raw[i] * lp[i] * qe[i]).epsilon(1e-5));
    }
}