Jungfraujoch/tests/XtalOptimizerTest.cpp

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_all.hpp>
#include <iostream>

#include "../image_analysis/geom_refinement/XtalOptimizer.h"
#include "../image_analysis/bragg_prediction/BraggPrediction.h"

TEST_CASE("XtalOptimizer") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
            .BeamX_pxl(1000)
            .BeamY_pxl(1000)
            .PoniRot1_rad(0.01)
            .PoniRot2_rad(0.02)
            .DetectorDistance_mm(200);

    CrystalLattice latt_i(40,40,80,90,90,90);

    BraggPredictionSettings prediction_settings{
    .high_res_A = 1.5,
    .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt;
    xtal_opt.latt = CrystalLattice(40.2,39.4,80.2, 90,91, 89);
    xtal_opt.geom.BeamX_pxl(1010).BeamY_pxl(995).DetectorDistance_mm(200)
    .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Triclinic;

    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.05);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.05);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabsf(uc_i.alpha - uc_o.alpha) < 0.1);
    CHECK(fabsf(uc_i.beta - uc_o.beta) < 0.1);
    CHECK(fabsf(uc_i.gamma - uc_o.gamma) < 0.1);
}


TEST_CASE("XtalOptimizer_NoBeamCenter") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
            .BeamX_pxl(1000)
            .BeamY_pxl(1000)
            .PoniRot1_rad(0.01)
            .PoniRot2_rad(0.02)
            .DetectorDistance_mm(200);

    CrystalLattice latt_i(40,50,80,90,95,90);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt;
    xtal_opt.latt = CrystalLattice(40.2,49.4,80.2, 90,94, 89);
    xtal_opt.geom.BeamX_pxl(999.8).BeamY_pxl(1000.2).DetectorDistance_mm(200)
    .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Monoclinic;
    xtal_opt.refine_beam_center = false;

    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - 999.8) < 0.01);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - 1000.2) < 0.01);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabsf(uc_i.alpha - uc_o.alpha) < 0.1);
    CHECK(fabsf(uc_i.beta - uc_o.beta) < 0.1);
    CHECK(fabsf(uc_i.gamma - uc_o.gamma) < 0.1);
}

TEST_CASE("XtalOptimizer_orthorombic") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
            .BeamX_pxl(1000)
            .BeamY_pxl(1000)
            .PoniRot1_rad(0.01)
            .PoniRot2_rad(0.02)
            .DetectorDistance_mm(200);

    CrystalLattice latt_i(40,50,80,90,90,90);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }
    XtalOptimizerData xtal_opt;
    xtal_opt.latt = CrystalLattice(40.2,49.6,80.3, 90,91, 89);
    xtal_opt.geom.BeamX_pxl(1005).BeamY_pxl(997).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Orthorhombic;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.1);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.1);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabs(uc_o.alpha - 90) < 0.02);
    CHECK(fabs(uc_o.beta - 90) < 0.02);
    CHECK(fabs(uc_o.gamma - 90) < 0.02);
}

TEST_CASE("XtalOptimizer_triclinic") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
            .BeamX_pxl(1000)
            .BeamY_pxl(1000)
            .PoniRot1_rad(0.01)
            .PoniRot2_rad(0.02)
            .DetectorDistance_mm(200);

    CrystalLattice latt_i(40,55,120,95,97,100);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001,
    };
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }
    XtalOptimizerData xtal_opt;
    xtal_opt.latt = CrystalLattice(40.1,54.9,121, 95,97, 99.5);
    xtal_opt.geom.BeamX_pxl(997).BeamY_pxl(1005).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Triclinic;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.2);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.2);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.5);
    CHECK(fabsf(uc_i.alpha - uc_o.alpha) < 0.1);
    CHECK(fabsf(uc_i.beta - uc_o.beta) < 0.1);
    CHECK(fabsf(uc_i.gamma - uc_o.gamma) < 0.1);
}


TEST_CASE("XtalOptimizer_tetragonal") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
            .BeamX_pxl(1000)
            .BeamY_pxl(1000)
            .PoniRot1_rad(0.01)
            .PoniRot2_rad(0.02)
            .DetectorDistance_mm(200);

    CrystalLattice latt_i(40,40,80,90,90,90);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt;
    xtal_opt.latt = CrystalLattice(40.6,39.3,80.5, 90,91, 89);
    xtal_opt.geom.BeamX_pxl(1010).BeamY_pxl(995).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Tetragonal;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.1);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.1);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.5);
    CHECK(fabs(uc_o.alpha - 90) < 0.02);
    CHECK(fabs(uc_o.beta - 90) < 0.02);
    CHECK(fabs(uc_o.gamma - 90) < 0.02);
}

TEST_CASE("XtalOptimizer_hexagonal") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
            .BeamX_pxl(1000)
            .BeamY_pxl(1000)
            .PoniRot1_rad(0.01)
            .PoniRot2_rad(0.02)
            .DetectorDistance_mm(200);

    CrystalLattice latt_i(40,40,70,90,90,120);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt;
    xtal_opt.latt = CrystalLattice(39.5,39.8,70.1, 90,90, 119.5);
    xtal_opt.geom.BeamX_pxl(1007).BeamY_pxl(990).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Hexagonal;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << " " << uc_o.alpha << " " << uc_o.beta
              << " " << uc_o.gamma << std::endl;

    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.1);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.1);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabs(uc_o.alpha - 90) < 0.02);
    CHECK(fabs(uc_o.beta - 90) < 0.01);
    CHECK(fabs(uc_o.gamma - 120) < 0.01);
}

TEST_CASE("XtalOptimizer_hexagonal_unconstrained") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
            .BeamX_pxl(1000)
            .BeamY_pxl(1000)
            .PoniRot1_rad(0.01)
            .PoniRot2_rad(0.02)
            .DetectorDistance_mm(200);

    CrystalLattice latt_i(40,40,70,90,90,120);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt;
    xtal_opt.latt = CrystalLattice(39.9,39.8,70.1, 90,90, 120);
    xtal_opt.geom.BeamX_pxl(1002).BeamY_pxl(998).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Triclinic;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << " " << uc_o.alpha << " " << uc_o.beta
              << " " << uc_o.gamma << std::endl;

    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.3);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.3);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabs(uc_o.alpha - 90) < 0.1);
    CHECK(fabs(uc_o.beta - 90) < 0.1);
    CHECK(fabs(uc_o.gamma - 120) < 0.1);
}

TEST_CASE("XtalOptimizer_cubic") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
            .BeamX_pxl(1000)
            .BeamY_pxl(1000)
            .PoniRot1_rad(0.01)
            .PoniRot2_rad(0.02)
            .DetectorDistance_mm(200);

    CrystalLattice latt_i(Coord(40,0,0),
                          Coord(0, 40 / sqrt(2), -40 / sqrt(2)),
                          Coord(0, 40 / sqrt(2), 40 / sqrt(2)));
    auto uc_i = latt_i.GetUnitCell();

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt;
    xtal_opt.latt = CrystalLattice(Coord(39,0,0),
                                   Coord(0, 39.5 / sqrt(2), -40.5 / sqrt(2)),
                                   Coord(0, 39.2 / sqrt(2), 39.7 / sqrt(2)));
    xtal_opt.geom.BeamX_pxl(1007).BeamY_pxl(990).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Cubic;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.1);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.1);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabs(uc_o.alpha - 90) < 0.02);
    CHECK(fabs(uc_o.beta - 90) < 0.02);
    CHECK(fabs(uc_o.gamma - 90) < 0.02);
}

TEST_CASE("XtalOptimizer_monoclinic") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
            .BeamX_pxl(1000)
            .BeamY_pxl(1000)
            .PoniRot1_rad(0.01)
            .PoniRot2_rad(0.02)
            .DetectorDistance_mm(200);

    CrystalLattice latt_i(50,60,70,90,115,90);

    auto uc_i = latt_i.GetUnitCell();

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001
    };
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt;
    xtal_opt.latt = CrystalLattice(49.5, 60.5, 69.8, 90, 113.5, 90);
    xtal_opt.geom.BeamX_pxl(1007).BeamY_pxl(990).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Monoclinic;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << " " << uc_o.alpha << " " << uc_o.beta << " " << uc_o.gamma << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.2);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.2);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabs(uc_o.alpha - 90) < 0.05);
    CHECK(fabs(uc_o.beta - uc_i.beta) < 0.05);
    CHECK(fabs(uc_o.gamma - 90) < 0.05);
}

TEST_CASE("LatticeToRodrigues") {
    double rod[3];
    double lengths[3];

    CrystalLattice latt_i(40,50,80,90,90,90);

    LatticeToRodriguesAndLengths_GS(latt_i, rod, lengths);
    CHECK(lengths[0] == Catch::Approx(40.0));
    CHECK(lengths[1] == Catch::Approx(50.0));
    CHECK(lengths[2] == Catch::Approx(80.0));
    CHECK(fabs(rod[0]) < 0.001);
    CHECK(fabs(rod[1]) < 0.001);
    CHECK(fabs(rod[2]) < 0.001);

    auto latt_o = AngleAxisAndCellToLattice(rod, lengths, M_PI/2, M_PI/2, M_PI/2);

    CHECK(latt_o.Vec0().Length() == Catch::Approx(40.0));
    CHECK(latt_o.Vec1().Length() == Catch::Approx(50.0));
    CHECK(latt_o.Vec2().Length() == Catch::Approx(80.0));
}

TEST_CASE("LatticeToRodrigues_irregular") {
    double rod[3];
    double lengths[3];

    CrystalLattice latt_i(Coord(40,0,0),
                          Coord(0, 50 / sqrt(2), -50 / sqrt(2)),
                          Coord(0, 80 / sqrt(2), 80 / sqrt(2)));

    LatticeToRodriguesAndLengths_GS(latt_i, rod, lengths);
    CHECK(lengths[0] == Catch::Approx(40.0));
    CHECK(lengths[1] == Catch::Approx(50.0));
    CHECK(lengths[2] == Catch::Approx(80.0));

    auto latt_o = AngleAxisAndCellToLattice(rod, lengths, M_PI/2, M_PI/2, M_PI/2);

    CHECK(latt_o.Vec0().Length() == Catch::Approx(40.0));
    CHECK(latt_o.Vec1().Length() == Catch::Approx(50.0));
    CHECK(latt_o.Vec2().Length() == Catch::Approx(80.0));
}

TEST_CASE("LatticeToRodrigues_Hex") {
    double rod[3];
    double lengths[3];

    Coord a = Coord(40,0,0);
    Coord b = Coord(40 / 2, 40 * sqrt(3)/ 2.0, 0);
    Coord c = Coord(0, 0, 70);

    RotMatrix R(1.0, Coord(0,1,1));
    CrystalLattice latt_i(R*a,R*b,R*c);

    LatticeToRodriguesAndLengths_Hex(latt_i, rod, lengths);
    CHECK(lengths[0] == Catch::Approx(40.0));
    CHECK(lengths[1] == Catch::Approx(40.0));
    CHECK(lengths[2] == Catch::Approx(70.0));

    auto latt_o = AngleAxisAndCellToLattice(rod, lengths,M_PI / 2.0, M_PI / 2.0, 2.0 * M_PI / 3.0);
    auto uc_o = latt_o.GetUnitCell();
    CHECK(uc_o.a == Catch::Approx(40.0));
    CHECK(uc_o.b == Catch::Approx(40.0));
    CHECK(uc_o.c == Catch::Approx(70.0));
    CHECK(uc_o.alpha == Catch::Approx(90.0));
    CHECK(uc_o.beta == Catch::Approx(90.0));
    CHECK(uc_o.gamma == Catch::Approx(120.0));
}

TEST_CASE("XtalOptimizer_rotation") {
    // Geometry
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
         .BeamX_pxl(1000)
         .BeamY_pxl(1000)
         .PoniRot1_rad(0.01)
         .PoniRot2_rad(0.02)
         .DetectorDistance_mm(200);

    // Base lattice (non-pathological)
    CrystalLattice latt_base(40, 50, 80, 90, 95, 90);
    auto uc_ref = latt_base.GetUnitCell();

    // Rotation axis: around X with 3 deg per image
    GoniometerAxis axis("omega", 0.0f, 3.0f, Coord(1,0,0), std::nullopt);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.002
    };

    std::vector<SpotToSave> spots;
    BraggPrediction prediction;

    // Predict reflections for images at 0-30 deg.
    for (int img = 0; img < 10; ++img) {
        // For a rotated image, per-image lattice is obtained as Multiply(rot.transpose())
        const float angle_deg = axis.GetAngle_deg(img) + axis.GetWedge_deg() / 2.0f;
        const RotMatrix rot = axis.GetTransformationAngle(angle_deg);
        const CrystalLattice latt_img = latt_base.Multiply(rot.transpose());

        const auto n = prediction.Calc(exp_i, latt_img, prediction_settings);
        for (int i = 0; i < n; ++i) {
            const auto& r = prediction.GetReflections().at(i);
            SpotToSave s{};
            s.x = r.predicted_x;
            s.y = r.predicted_y;
            s.image = img;       // provide image index for rotation-aware refinement
            s.phi = angle_deg;
            s.intensity = 1.0f;  // minimal positive value
            s.ice_ring = false;
            s.indexed = true;
            spots.push_back(s);
        }
    }

    // Seed slightly perturbed geometry and lattice; provide rotation axis for refinement
    XtalOptimizerData xtal_opt;
    xtal_opt.latt = CrystalLattice(39.7f, 50.6f, 79.6f, 90.0f, 94.5f, 90.5f);
    xtal_opt.geom.BeamX_pxl(1003).BeamY_pxl(997).DetectorDistance_mm(203.0)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Monoclinic;
    xtal_opt.axis = axis;
    xtal_opt.min_spots = 200;
    xtal_opt.refine_beam_center = true;
    xtal_opt.refine_distance_mm = true;
    xtal_opt.refine_detector_angles = false;

    auto t0 = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto t1 = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer (rotation 4 images) took "
              << std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0).count()
              << " microseconds" << std::endl;

    const auto uc_out = xtal_opt.latt.GetUnitCell();

    // Geometry checks
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.2f);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.2f);

    // 700 micron is a lot, but for now it is important to see that refinement improves in the right direction
    CHECK(fabsf(xtal_opt.geom.GetDetectorDistance_mm() - exp_i.GetDetectorDistance_mm()) < 0.7f);

    // Lattice checks
    CHECK(fabsf(uc_ref.a - uc_out.a) < 0.2f);
    CHECK(fabsf(uc_ref.b - uc_out.b) < 0.2f);
    CHECK(fabsf(uc_ref.c - uc_out.c) < 0.4f);
    CHECK(fabsf(uc_ref.alpha - uc_out.alpha) < 0.2f);
    CHECK(fabsf(uc_ref.beta - uc_out.beta) < 0.2f);
    CHECK(fabsf(uc_ref.gamma - uc_out.gamma) < 0.2f);
}

TEST_CASE("XtalOptimizer_refine_rotation_axis") {
    // Geometry
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
         .BeamX_pxl(1000)
         .BeamY_pxl(1000)
         .PoniRot1_rad(0.01)
         .PoniRot2_rad(0.02)
         .DetectorDistance_mm(200);

    // Base lattice (non-pathological)
    CrystalLattice latt_base(40, 50, 80, 90, 95, 90);
    auto uc_ref = latt_base.GetUnitCell();

    // Rotation axis: around X with 3 deg per image
    GoniometerAxis axis("omega", 0.0f, 3.0f, Coord(1,0,0), std::nullopt);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.002
    };

    std::vector<SpotToSave> spots;
    BraggPrediction prediction;

    // Predict reflections for images at 0-30 deg.
    for (int img = 0; img < 10; ++img) {
        // For a rotated image, per-image lattice is obtained as Multiply(rot.transpose())
        const float angle_deg = axis.GetAngle_deg(img) + axis.GetWedge_deg() / 2.0f;
        const RotMatrix rot = axis.GetTransformationAngle(angle_deg);
        const CrystalLattice latt_img = latt_base.Multiply(rot.transpose());

        const auto n = prediction.Calc(exp_i, latt_img, prediction_settings);
        for (int i = 0; i < n; ++i) {
            const auto& r = prediction.GetReflections().at(i);
            SpotToSave s{};
            s.x = r.predicted_x;
            s.y = r.predicted_y;
            s.image = img;       // provide image index for rotation-aware refinement
            s.intensity = 1.0f;  // minimal positive value
            s.phi = angle_deg;
            s.ice_ring = false;
            s.indexed = true;
            spots.push_back(s);
        }
    }

    // Seed slightly perturbed geometry and lattice; provide rotation axis for refinement
    XtalOptimizerData xtal_opt;
    xtal_opt.latt = CrystalLattice(39.7f, 50.6f, 79.6f, 90.0f, 94.5f, 90.5f);
    xtal_opt.geom.BeamX_pxl(1003).BeamY_pxl(997).DetectorDistance_mm(200.0)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Monoclinic;
    xtal_opt.axis = GoniometerAxis("omega", 0.0f, 3.0f,
        Coord(0.8, 0.05, 0.05).Normalize(),
        std::nullopt);
    xtal_opt.min_spots = 200;
    xtal_opt.refine_beam_center = true;
    xtal_opt.refine_distance_mm = false;
    xtal_opt.refine_detector_angles = false;
    xtal_opt.refine_rotation_axis = true;

    auto t0 = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto t1 = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer (rotation 4 images) took "
              << std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0).count()
              << " microseconds" << std::endl;

    const auto uc_out = xtal_opt.latt.GetUnitCell();

    // Geometry checks
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.2f);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.2f);

    // Lattice checks
    CHECK(fabsf(uc_ref.a - uc_out.a) < 0.2f);
    CHECK(fabsf(uc_ref.b - uc_out.b) < 0.2f);
    CHECK(fabsf(uc_ref.c - uc_out.c) < 0.3f);
    CHECK(fabsf(uc_ref.alpha - uc_out.alpha) < 0.2f);
    CHECK(fabsf(uc_ref.beta - uc_out.beta) < 0.2f);
    CHECK(fabsf(uc_ref.gamma - uc_out.gamma) < 0.2f);
    CHECK(fabsf(xtal_opt.axis->GetAxis().x - 1.0) < 0.01f);
    CHECK(fabsf(xtal_opt.axis->GetAxis().y) < 0.01f);
    CHECK(fabsf(xtal_opt.axis->GetAxis().z) < 0.01f);
}

// --- helpers for lattice sanity tests ---
#include <Eigen/Dense>

namespace {
Eigen::Vector3d to_eigen(const Coord& v) {
    return {v[0], v[1], v[2]};
}

double angle_rad(const Eigen::Vector3d& a, const Eigen::Vector3d& b) {
    const double na = a.norm();
    const double nb = b.norm();
    if (na == 0.0 || nb == 0.0)
        return 0.0;
    double c = a.dot(b) / (na * nb);
    c = std::max(-1.0, std::min(1.0, c));
    return std::acos(c);
}

// Compare two lattices up to a global rotation: compare Gram matrices G = L^T L (rotation-invariant).
Eigen::Matrix3d gram(const CrystalLattice& latt) {
    const Eigen::Vector3d A = to_eigen(latt.Vec0());
    const Eigen::Vector3d B = to_eigen(latt.Vec1());
    const Eigen::Vector3d C = to_eigen(latt.Vec2());
    Eigen::Matrix3d G;
    G(0,0) = A.dot(A); G(0,1) = A.dot(B); G(0,2) = A.dot(C);
    G(1,0) = B.dot(A); G(1,1) = B.dot(B); G(1,2) = B.dot(C);
    G(2,0) = C.dot(A); G(2,1) = C.dot(B); G(2,2) = C.dot(C);
    return G;
}

    void check_gram_close(const CrystalLattice& a,
                          const CrystalLattice& b,
                          double abs_eps,
                          double rel_eps) {
    const Eigen::Matrix3d Ga = gram(a);
    const Eigen::Matrix3d Gb = gram(b);

    for (int r = 0; r < 3; ++r) {
        for (int c = 0; c < 3; ++c) {
            const double va = Ga(r, c);
            const double vb = Gb(r, c);

            // Scale for relative error; avoid blowing up around zero.
            const double scale = std::max({1.0, std::abs(va), std::abs(vb)});
            const double tol = std::max(abs_eps, rel_eps * scale);

            INFO("G(" << r << "," << c << ") va=" << va << " vb=" << vb
                      << " scale=" << scale << " tol=" << tol);

            CHECK(va == Catch::Approx(vb).margin(tol));
        }
    }
}
} // namespace

TEST_CASE("XtalOptimizer Lattice param roundtrip (GS) preserves Gram matrix") {
    // Non-orthogonal, irregular basis, but still a valid lattice
    CrystalLattice latt_i(Coord(40,  1,  2),
                          Coord( 3, 50, -4),
                          Coord(-5,  6, 80));

    double rod[3]{}, lengths[3]{};
    LatticeToRodriguesAndLengths_GS(latt_i, rod, lengths);
    auto latt_o = AngleAxisAndCellToLattice(rod, lengths, M_PI/2, M_PI/2, M_PI/2);

    // This parametrization only keeps "lengths + rotation", i.e. it cannot reproduce shear.
    // So we *do not* compare Gram matrices here.
    // Instead, sanity-check: reconstructed vectors have the requested lengths.
    CHECK(latt_o.Vec0().Length() == Catch::Approx(lengths[0]).margin(1e-9));
    CHECK(latt_o.Vec1().Length() == Catch::Approx(lengths[1]).margin(1e-9));
    CHECK(latt_o.Vec2().Length() == Catch::Approx(lengths[2]).margin(1e-9));
}

TEST_CASE("XtalOptimizer Lattice param roundtrip (Hex) preserves unit cell") {
    Coord a = Coord(40, 0, 0);
    Coord b = Coord(40 / 2.0, 40 * std::sqrt(3) / 2.0, 0);
    Coord c = Coord(0, 0, 70);

    // Apply an arbitrary rotation to ensure the rod extraction is meaningful
    RotMatrix R(1.0, Coord(0, 1, 1));
    CrystalLattice latt_i(R * a, R * b, R * c);

    double rod[3]{}, ac[3]{};
    LatticeToRodriguesAndLengths_Hex(latt_i, rod, ac);
    auto latt_o = AngleAxisAndCellToLattice(rod, ac, M_PI/2, M_PI/2, 2*M_PI/3);

    auto uc_o = latt_o.GetUnitCell();
    CHECK(uc_o.a == Catch::Approx(40.0).margin(1e-6));
    CHECK(uc_o.b == Catch::Approx(40.0).margin(1e-6));
    CHECK(uc_o.c == Catch::Approx(70.0).margin(1e-6));
    CHECK(uc_o.alpha == Catch::Approx(90.0).margin(1e-6));
    CHECK(uc_o.beta  == Catch::Approx(90.0).margin(1e-6));
    CHECK(uc_o.gamma == Catch::Approx(120.0).margin(1e-6));
}

TEST_CASE("XtalOptimizer Monoclinic param roundtrip preserves Gram matrix (beta far from 90)") {
    struct Case { double beta_deg; };
    const std::vector<Case> cases = {
        {60.0}, {75.0}, {115.0}, {130.0}
    };

    for (const auto& cs : cases) {
        INFO("beta_deg=" << cs.beta_deg);

        // Start from a clean monoclinic cell in its conventional setting (unique axis b).
        CrystalLattice latt0(50, 60, 70, 90, cs.beta_deg, 90);

        // Now apply a TRUE global rotation: rotate each basis vector (left-multiply).
        RotMatrix R(0.7, Coord(0.3, 0.9, 0.1));
        CrystalLattice latt_i(R * latt0.Vec0(),
                              R * latt0.Vec1(),
                              R * latt0.Vec2());

        double rod[3]{}, lengths[3]{}, beta_rad = 0.0;
        LatticeToRodriguesLengthsBeta_Mono(latt_i, rod, lengths, beta_rad);

        // Basic sanity
        CHECK(lengths[0] == Catch::Approx(50.0).margin(1e-6));
        CHECK(lengths[1] == Catch::Approx(60.0).margin(1e-6));
        CHECK(lengths[2] == Catch::Approx(70.0).margin(1e-6));
        CHECK(beta_rad * 180.0 / M_PI == Catch::Approx(cs.beta_deg).margin(1e-6));

        auto latt_o = AngleAxisAndCellToLattice(rod, lengths, M_PI/2, beta_rad, M_PI/2);

        // Rotation-invariant check: Gram matrices match.
        check_gram_close(latt_i, latt_o, /*abs_eps=*/5e-4, /*rel_eps=*/1e-10);

        // Also check the unit-cell angles we expect for monoclinic(unique b).
        auto uc_o = latt_o.GetUnitCell();
        CHECK(uc_o.alpha == Catch::Approx(90.0).margin(1e-4));
        CHECK(uc_o.gamma == Catch::Approx(90.0).margin(1e-4));
        CHECK(uc_o.beta  == Catch::Approx(cs.beta_deg).margin(1e-4));
    }
}

TEST_CASE("XtalOptimizer Monoclinic beta geometry: extracted beta equals angle(a,c)") {
    // This isolates only the beta definition, independent of other choices.
    CrystalLattice latt_i(50, 60, 70, 90, 130, 90);

    const Eigen::Vector3d A = to_eigen(latt_i.Vec0());
    const Eigen::Vector3d C = to_eigen(latt_i.Vec2());
    const double beta_geom = angle_rad(A, C);

    double rod[3]{}, lengths[3]{}, beta_rad = 0.0;
    LatticeToRodriguesLengthsBeta_Mono(latt_i, rod, lengths, beta_rad);

    CHECK(beta_rad == Catch::Approx(beta_geom).margin(1e-12));
}