// 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 <fstream>
#include <random>
#include <sstream>

#include <gemmi/read_cif.hpp>

#include "../image_analysis/geom_refinement/CellEsd.h"
#include "../image_analysis/WriteReflections.h"
#include "../image_analysis/scale_merge/Merge.h"
#include "../image_analysis/scale_merge/TwinningAnalysis.h"

namespace {
    Eigen::Matrix3d ReciprocalMetric(double a, double b, double c, double al, double be, double ga) {
        const gemmi::UnitCell cell(a, b, c, al, be, ga);
        Eigen::Matrix3d g;
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                g(i, j) = cell.frac.mat[i][0] * cell.frac.mat[j][0] + cell.frac.mat[i][1] * cell.frac.mat[j][1]
                          + cell.frac.mat[i][2] * cell.frac.mat[j][2];   // a*_i . a*_j, a*_i a row
        return g;
    }

    // The basis coordinates of G* (the basis is orthonormal over g11, g22, g33, g12, g13, g23), and
    // how far G* lies outside the basis' span.
    Eigen::VectorXd Coordinates(const std::vector<Eigen::Matrix3d> &basis, const Eigen::Matrix3d &g,
                                double &residual) {
        const auto vec = [](const Eigen::Matrix3d &m) {
            Eigen::Matrix<double, 6, 1> v;
            v << m(0, 0), m(1, 1), m(2, 2), m(0, 1), m(0, 2), m(1, 2);
            return v;
        };
        Eigen::VectorXd p(basis.size());
        Eigen::Matrix<double, 6, 1> rest = vec(g);
        for (size_t k = 0; k < basis.size(); k++) {
            p(k) = vec(basis[k]).dot(vec(g));
            rest -= p(k) * vec(basis[k]);
        }
        residual = rest.norm() / vec(g).norm();
        return p;
    }

    const gemmi::SpaceGroup &Sg(const char *name) { return *gemmi::find_spacegroup_by_name(name); }

    // The e.s.d.s of the cell under a covariance of the metric's coordinates, against the scatter of
    // the cells of metrics drawn from it.
    void CheckAgainstSampling(const char *sg, double a, double b, double c, double al, double be, double ga) {
        const auto basis = MetricBasis(Sg(sg));
        double residual = 0.0;
        const Eigen::VectorXd p0 = Coordinates(basis, ReciprocalMetric(a, b, c, al, be, ga), residual);
        REQUIRE(residual < 1e-12);
        const int n = static_cast<int>(basis.size());
        Eigen::MatrixXd cov = Eigen::MatrixXd::Zero(n, n);
        for (int k = 0; k < n; k++)
            cov(k, k) = std::pow(1e-5 * p0.norm() * (k + 1), 2);
        if (n > 1)
            cov(0, 1) = cov(1, 0) = 0.3 * std::sqrt(cov(0, 0) * cov(1, 1));
        const auto esd = CellEsdFromCovariance(basis, p0, cov);

        const Eigen::MatrixXd chol = cov.llt().matrixL();
        std::mt19937 rng(7);
        std::normal_distribution<double> normal;
        constexpr int SAMPLES = 20000;
        std::array<double, 6> sum{}, sum2{};
        for (int s = 0; s < SAMPLES; s++) {
            Eigen::VectorXd z(n);
            for (int k = 0; k < n; k++)
                z(k) = normal(rng);
            const Eigen::VectorXd p = p0 + chol * z;
            Eigen::Matrix3d g = Eigen::Matrix3d::Zero();
            for (int k = 0; k < n; k++)
                g += p(k) * basis[k];
            const auto cell = CellFromReciprocalMetric(g);
            for (int j = 0; j < 6; j++) {
                sum[j] += cell[j];
                sum2[j] += cell[j] * cell[j];
            }
        }
        for (int j = 0; j < 6; j++) {
            const double mean = sum[j] / SAMPLES;
            const double sd = std::sqrt(std::max(0.0, sum2[j] / SAMPLES - mean * mean));
            if (esd[j] < CELL_ESD_FIXED)
                CHECK(sd < 1e-7);
            else
                CHECK(sd == Catch::Approx(esd[j]).epsilon(0.03));
        }
    }
}

TEST_CASE("CellEsd: the metric basis has the free parameters of each crystal system", "[cell_esd]") {
    CHECK(MetricBasis(Sg("P 1")).size() == 6);
    CHECK(MetricBasis(Sg("P 1 21 1")).size() == 4);
    CHECK(MetricBasis(Sg("C 1 2/c 1")).size() == 4);
    CHECK(MetricBasis(Sg("P 21 21 21")).size() == 3);
    CHECK(MetricBasis(Sg("I 41/a")).size() == 2);
    CHECK(MetricBasis(Sg("P 61 2 2")).size() == 2);
    CHECK(MetricBasis(Sg("R 3 :H")).size() == 2);
    CHECK(MetricBasis(Sg("R 3 :R")).size() == 2);
    CHECK(MetricBasis(Sg("I a -3 d")).size() == 1);
}

TEST_CASE("CellEsd: the metric basis spans the reciprocal metric of a cell of the group", "[cell_esd]") {
    // R G* R^T, not R^T G* R: only the right one keeps a hexagonal G* (g12* = +a*^2/2) in the span.
    double residual = 1.0;
    Coordinates(MetricBasis(Sg("P 61 2 2")), ReciprocalMetric(5.5, 5.5, 56.0, 90, 90, 120), residual);
    CHECK(residual < 1e-12);
    Coordinates(MetricBasis(Sg("R 3 :R")), ReciprocalMetric(7.0, 7.0, 7.0, 80, 80, 80), residual);
    CHECK(residual < 1e-12);
    Coordinates(MetricBasis(Sg("P 1 21 1")), ReciprocalMetric(10.0, 7.0, 12.0, 90, 100, 90), residual);
    CHECK(residual < 1e-12);
    // and nothing outside it
    Coordinates(MetricBasis(Sg("P 1 21 1")), ReciprocalMetric(10.0, 7.0, 12.0, 91, 100, 90), residual);
    CHECK(residual > 1e-4);
}

TEST_CASE("CellEsd: e.s.d.s from a covariance follow the constraints and match sampling", "[cell_esd]") {
    CheckAgainstSampling("P 1", 7.1, 8.3, 9.2, 81.0, 76.0, 69.0);
    CheckAgainstSampling("P 1 21 1", 10.0, 7.0, 12.0, 90.0, 100.0, 90.0);
    CheckAgainstSampling("P 21 21 21", 8.0, 9.5, 27.0, 90.0, 90.0, 90.0);
    CheckAgainstSampling("P 61 2 2", 5.5, 5.5, 56.0, 90.0, 90.0, 120.0);
    CheckAgainstSampling("R 3 :R", 7.0, 7.0, 7.0, 80.0, 80.0, 80.0);

    // Tetragonal: a and b share one e.s.d., the angles have none.
    const auto basis = MetricBasis(Sg("I -4 2 d"));
    double residual = 0.0;
    const Eigen::VectorXd p = Coordinates(basis, ReciprocalMetric(7.5, 7.5, 7.0, 90, 90, 90), residual);
    const auto esd = CellEsdFromCovariance(basis, p, Eigen::MatrixXd::Identity(2, 2) * 1e-12);
    CHECK(esd[0] > 0.0);
    CHECK(esd[0] == Catch::Approx(esd[1]).epsilon(1e-9));
    for (int j = 3; j < 6; j++)
        CHECK(esd[j] < CELL_ESD_FIXED);
}

TEST_CASE("CellEsd: the d-spacing fit recovers the cell, and its e.s.d. is its scatter", "[cell_esd]") {
    const double a = 10.0, b = 7.0, c = 12.0, beta = 100.0;   // synthetic
    const Eigen::Matrix3d gstar = ReciprocalMetric(a, b, c, 90, beta, 90);
    std::vector<CellObservation> clean;
    for (int h = -8; h <= 8; h++)
        for (int k = -5; k <= 5; k++)
            for (int l = -8; l <= 8; l++) {
                const Eigen::Vector3d v(h, k, l);
                const double q = std::sqrt(v.dot(gstar * v));
                if (q > 0.0 && q < 1.0 / 1.0)
                    clean.push_back(CellObservation{h, k, l, static_cast<float>(q),
                                                    static_cast<int32_t>(clean.size() * 37 % 100)});
            }
    REQUIRE(clean.size() > 2500);

    const auto &sg = Sg("P 1 21 1");
    std::mt19937 rng(11);
    std::normal_distribution<double> noise(0.0, 2e-4);   // e.s.d. on q, 1/A
    std::vector<double> fitted_a, fitted_beta;
    std::vector<CellEsd> fits;
    for (int trial = 0; trial < 40; trial++) {
        std::vector<CellObservation> obs = clean;
        for (auto &o : obs)
            o.q = static_cast<float>(o.q + noise(rng));
        // One gross outlier, which the 4-sigma rejection takes out.
        obs[17].q *= 1.05f;
        const auto e = EstimateCellEsd(obs, sg, 0.62);
        REQUIRE(e.has_value());
        // a 4-sigma tail of 3000 Gaussian residuals may lose one or two more
        CHECK(e->observations <= static_cast<int64_t>(obs.size()) - 1);
        CHECK(e->observations >= static_cast<int64_t>(obs.size()) - 4);
        fitted_a.push_back(e->fitted.a);
        fitted_beta.push_back(e->fitted.beta);
        fits.push_back(*e);
    }
    const auto sd = [](const std::vector<double> &v) {
        double m = 0, s = 0;
        for (double x : v) m += x / v.size();
        for (double x : v) s += (x - m) * (x - m) / (v.size() - 1);
        return std::sqrt(s);
    };
    const CellEsd &e = fits.front();
    CHECK(sd(fitted_a) == Catch::Approx(e.esd_lsq[0]).epsilon(0.3));
    CHECK(sd(fitted_beta) == Catch::Approx(e.esd_lsq[4]).epsilon(0.3));
    CHECK(std::abs(e.fitted.a - a) < 5 * e.esd_lsq[0] + 1e-4);   // float cell
    CHECK(e.esd_lsq[3] < CELL_ESD_FIXED);
    CHECK(e.esd_lsq[5] < CELL_ESD_FIXED);
    // Pure centroid noise: the blocks agree with the least squares, and the larger of the two is written.
    REQUIRE(e.esd_blocks.has_value());
    CHECK(e.blocks == 8);
    CHECK((*e.esd_blocks)[0] == Catch::Approx(e.esd_lsq[0]).epsilon(0.8));
    for (int j = 0; j < 6; j++)
        CHECK(e.esd[j] == std::max(e.esd_lsq[j], (*e.esd_blocks)[j]));
    CHECK(e.theta_max_deg == Catch::Approx(std::asin(0.62 / 2.0) * 180.0 / M_PI).margin(0.1));
}

TEST_CASE("CellEsd: the SHELX .ins carries CELL, ZERR, LATT and SYMM", "[cell_esd][write_reflections]") {
    CellEsd esd;
    esd.esd = {0.0014, 0.0013, 0.0012, 0.0, 0.0034, 0.0};
    const UnitCell cell{10.0f, 7.0f, 12.0f, 90.0f, 100.0f, 90.0f};   // synthetic
    const auto path = (std::filesystem::temp_directory_path() / "rugnux_cell_esd_test.ins").string();
    const auto read_ins = [&](int sg_number, const CellEsd *e) {
        DiffractionExperiment x;
        x.IncidentEnergy_keV(20.0f);
        x.SpaceGroupNumber(sg_number);
        WriteShelxIns(cell, e, x, "p", path);
        std::ifstream in(path);
        std::stringstream ss;
        ss << in.rdbuf();
        std::filesystem::remove(path);
        return ss.str();
    };

    // P 21/c: centrosymmetric with -1 at the origin, so LATT +1 and one operator of each pair.
    std::string ins = read_ins(14, &esd);
    CHECK(ins.find("TITL p in P 1 21/c 1\n") == 0);
    CHECK(ins.find(" 10.00000 7.00000 12.00000 90.0000 100.0000 90.0000\n") != std::string::npos);
    CHECK(ins.find("\nCELL 0.6199") != std::string::npos);
    CHECK(ins.find("\nZERR 4 0.00140 0.00130 0.00120 0.0000 0.0034 0.0000\n") != std::string::npos);
    CHECK(ins.find("\nLATT 1\n") != std::string::npos);
    CHECK(ins.find("\nSYMM -X,Y+1/2,-Z+1/2\n") != std::string::npos);
    CHECK(ins.find("\nSYMM X,-Y+1/2,Z+1/2\n") == std::string::npos);
    CHECK(ins.find("\nHKLF 4\nEND\n") != std::string::npos);

    // P 21 21 21: no inversion, LATT -1 and three operators.
    ins = read_ins(19, nullptr);
    CHECK(ins.find("\nLATT -1\n") != std::string::npos);
    CHECK(ins.find("\nZERR") == std::string::npos);
    size_t symm = 0;
    for (size_t pos = ins.find("SYMM"); pos != std::string::npos; pos = ins.find("SYMM", pos + 1))
        symm++;
    CHECK(symm == 3);

    // C 2/c: the C centring, centrosymmetric, Z = 8.
    ins = read_ins(15, &esd);
    CHECK(ins.find("\nLATT 7\n") != std::string::npos);
    CHECK(ins.find("\nZERR 8 ") != std::string::npos);
}

TEST_CASE("CellEsd: the mmCIF carries the e.s.d.s, '.' where the symmetry fixes the parameter",
          "[cell_esd][write_reflections]") {
    DiffractionExperiment x;
    x.IncidentEnergy_keV(20.0f);
    x.SpaceGroupNumber(4);
    const UnitCell cell{10.0f, 7.0f, 12.0f, 90.0f, 100.0f, 90.0f};   // synthetic
    x.SetUnitCell(cell);
    CellEsd esd;
    esd.esd = {0.0014, 0.0013, 0.0012, 0.0, 0.0034, 0.0};
    esd.reflections = 1234;
    esd.theta_min_deg = 1.5;
    esd.theta_max_deg = 18.0;
    std::vector<MergedReflection> refl(1);
    refl[0].h = 1; refl[0].I = 10.0f; refl[0].sigma = 1.0f; refl[0].d = 5.0f;
    const auto path = (std::filesystem::temp_directory_path() / "rugnux_cell_esd_test.cif").string();
    WriteMmcifReflections(refl, cell, x, MergeStatistics{}, ErrorModelReport{}, TwinningAnalysisResult{},
                          OutputProvenance{}, path, 1, &esd);
    gemmi::cif::Document doc = gemmi::read_cif_gz(path);
    std::filesystem::remove(path);
    gemmi::cif::Block &block = doc.blocks[0];
    CHECK(*block.find_value("_cell.length_a") == "10.0000");
    CHECK(*block.find_value("_cell.length_a_esd") == "0.00140");
    CHECK(*block.find_value("_cell.angle_beta_esd") == "0.0034");
    CHECK(*block.find_value("_cell.angle_alpha_esd") == ".");
    CHECK(*block.find_value("_cell_measurement.reflns_used") == "1234");
    CHECK(*block.find_value("_cell_measurement.theta_max") == "18.00");
}