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

#include <catch2/catch_all.hpp>

#include "../image_analysis/bragg_prediction/BraggPrediction.h"
#include <iostream>

TEST_CASE("BraggPrediction_11keV") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
    .IncidentEnergy_keV(11.0);

    DiffractionGeometry geom = experiment.GetDiffractionGeometry();

    CrystalLattice lattice(Coord{20, 10, 0}, Coord{-20, 40, 0},
                           Coord{0, 0, 200});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.001,
        .max_hkl = 40
    };
    BraggPrediction prediction;

    int count = prediction.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);

    for (int i = 0; i < count; i++) {
        auto r = prediction.GetReflections().at(i);
        auto recip = r.h * lattice.Astar() + r.k * lattice.Bstar() + r.l * lattice.Cstar();

        REQUIRE(std::abs(r.h) < settings.max_hkl );
        REQUIRE(std::abs(r.k) < settings.max_hkl );
        REQUIRE(std::abs(r.l) < settings.max_hkl );
        REQUIRE(r.d >= settings.high_res_A);
        REQUIRE(r.d == Catch::Approx(1/std::sqrt(recip * recip)).margin(0.01f));
        REQUIRE(r.dist_ewald == Catch::Approx(std::abs(geom.DistFromEwaldSphere(recip))).epsilon(1e-5));
        auto [x,y] = geom.RecipToDector(recip);
        REQUIRE(r.predicted_x == Catch::Approx(x).margin(0.01));
        REQUIRE(r.predicted_y == Catch::Approx(y).margin(0.01));
    }
}

TEST_CASE("BraggPrediction_15keV") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
    .IncidentEnergy_keV(15.0);

    DiffractionGeometry geom = experiment.GetDiffractionGeometry();

    CrystalLattice lattice(Coord{20, 10, 0}, Coord{-20, 40, 0},
                           Coord{0, 0, 200});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.001,
        .max_hkl = 40
    };
    BraggPrediction prediction;

    int count = prediction.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);

    for (int i = 0; i < count; i++) {
        auto r = prediction.GetReflections().at(i);
        auto recip = r.h * lattice.Astar() + r.k * lattice.Bstar() + r.l * lattice.Cstar();

        REQUIRE(std::abs(r.h) < settings.max_hkl );
        REQUIRE(std::abs(r.k) < settings.max_hkl );
        REQUIRE(std::abs(r.l) < settings.max_hkl );
        REQUIRE(r.d >= settings.high_res_A);
        REQUIRE(r.d == Catch::Approx(1/std::sqrt(recip * recip)).margin(0.01f));
        REQUIRE(r.dist_ewald == Catch::Approx(std::abs(geom.DistFromEwaldSphere(recip))).epsilon(1e-5));
        auto [x,y] = geom.RecipToDector(recip);
        REQUIRE(r.predicted_x == Catch::Approx(x).margin(0.01));
        REQUIRE(r.predicted_y == Catch::Approx(y).margin(0.01));
    }
}

TEST_CASE("BraggPrediction_Rot1_Rot2") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
    .PoniRot1_rad(2.0/180.0 * M_PI).PoniRot2_rad(3.0/180.0 * M_PI)
    .IncidentEnergy_keV(11.0);

    DiffractionGeometry geom = experiment.GetDiffractionGeometry();

    CrystalLattice lattice(Coord{20, 10, 0}, Coord{-20, 40, 0},
                           Coord{0, 0, 200});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.001,
        .max_hkl = 40
    };
    BraggPrediction prediction;

    int count = prediction.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);

    for (int i = 0; i < count; i++) {
        auto r = prediction.GetReflections().at(i);
        auto recip = r.h * lattice.Astar() + r.k * lattice.Bstar() + r.l * lattice.Cstar();

        REQUIRE(std::abs(r.h) < settings.max_hkl );
        REQUIRE(std::abs(r.k) < settings.max_hkl );
        REQUIRE(std::abs(r.l) < settings.max_hkl );
        REQUIRE(r.d >= settings.high_res_A);
        REQUIRE(r.d == Catch::Approx(1/std::sqrt(recip * recip)).margin(0.01f));
        REQUIRE(r.dist_ewald == Catch::Approx(std::abs(geom.DistFromEwaldSphere(recip))).epsilon(1e-5));
        auto [x,y] = geom.RecipToDector(recip);
        REQUIRE(r.predicted_x == Catch::Approx(x).margin(0.01));
        REQUIRE(r.predicted_y == Catch::Approx(y).margin(0.01));
    }
}

TEST_CASE("BraggPrediction_backscattering") {
    DiffractionExperiment experiment(DetJF9M());
    experiment.DetectorDistance_mm(120.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
        .IncidentEnergy_keV(3.0/WVL_1A_IN_KEV);

    // Orthogonal basis, sufficient to test parity rules
    CrystalLattice lattice(
        Coord{40, 0, 0},
        Coord{0, 50, 0},
        Coord{0, 0, 60}
    );

    BraggPredictionSettings settings{
        .high_res_A = 3.0f,
        .ewald_dist_cutoff = 0.1f, // Very large cutoff, to be able to see as many reflections as possible
        .max_hkl = 50
    };

    BraggPrediction pred;
    int count = pred.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);
    for (const auto& r : pred.GetReflections()) {
        if (r.d == 0) break;
        REQUIRE(r.d > 3.0 / sqrt(2.0));
    }
}

TEST_CASE("BraggPrediction_systematic_absences") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(120.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
        .IncidentEnergy_keV(12.0);

    // Orthogonal basis, sufficient to test parity rules
    CrystalLattice lattice(
        Coord{40, 0, 0},
        Coord{0, 50, 0},
        Coord{0, 0, 60}
    );

    BraggPredictionSettings settings{
        .high_res_A = 3.0f,
        .ewald_dist_cutoff = 0.1f, // Very large cutoff, to be able to see as many reflections as possible
        .max_hkl = 50
    };

    BraggPrediction pred;

    SECTION("I centering") {
        // 1) Body-centered I: reflections with h+k+l odd must be absent
        settings.centering = 'I';
        int count_I = pred.Calc(experiment, lattice, settings);
        REQUIRE(count_I > 0);
        for (const auto& r : pred.GetReflections()) {
            if (r.d == 0) break; // ignore unfilled tail if any
            REQUIRE(((r.h + r.k + r.l) % 2) == 0);
        }
    }

    SECTION ("F centering") {
        // 2) Face-centered F: h,k,l all even or all odd
        settings.centering = 'F';

        int count_F = pred.Calc(experiment, lattice, settings);
        REQUIRE(count_F > 0);
        for (const auto& r : pred.GetReflections()) {
            if (r.d == 0) break;
            const bool he = (r.h & 1) == 0, ke = (r.k & 1) == 0, le = (r.l & 1) == 0;
            const bool all_even = he && ke && le;
            const bool all_odd = (!he) && (!ke) && (!le);
            REQUIRE((all_even || all_odd));
        }
    }

    SECTION("R centering") {
        settings.centering = 'R';
        int count_R = pred.Calc(experiment, lattice, settings);
        REQUIRE(count_R > 0);

        // R (hexagonal setting): -h + k + l = 3n
        for (const auto& r : pred.GetReflections()) {
            if (r.d == 0) break;
            int cond = (-r.h + r.k + r.l) % 3;
            if (cond < 0) cond += 3;
            REQUIRE(cond == 0);
        }
    }
}

#ifdef JFJOCH_USE_CUDA
#include "../image_analysis/bragg_prediction/BraggPredictionGPU.h"

TEST_CASE("BraggPredictionGPU") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
    .IncidentEnergy_keV(13.0);

    DiffractionGeometry geom = experiment.GetDiffractionGeometry();

    CrystalLattice lattice(Coord{20, 10, 0}, Coord{-20, 40, 0},
                           Coord{0, 0, 200});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.001,
        .max_hkl = 40
    };
    BraggPredictionGPU prediction;

    int count = prediction.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);

    for (int i = 0; i < count; i++) {
        auto r = prediction.GetReflections().at(i);
        auto recip = r.h * lattice.Astar() + r.k * lattice.Bstar() + r.l * lattice.Cstar();

        REQUIRE(std::abs(r.h) < settings.max_hkl );
        REQUIRE(std::abs(r.k) < settings.max_hkl );
        REQUIRE(std::abs(r.l) < settings.max_hkl );
        REQUIRE(r.d >= settings.high_res_A);
        REQUIRE(r.d == Catch::Approx(1/std::sqrt(recip * recip)).margin(0.01f));
        REQUIRE(r.dist_ewald == Catch::Approx(std::abs(geom.DistFromEwaldSphere(recip))).epsilon(2e-2));
        auto [x,y] = geom.RecipToDector(recip);
        REQUIRE(r.predicted_x == Catch::Approx(x).margin(0.05));
        REQUIRE(r.predicted_y == Catch::Approx(y).margin(0.05));
    }
}

TEST_CASE("BraggPredictionGPU_Rot1_Rot2") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
    .PoniRot1_rad(2.0/180.0 * M_PI).PoniRot2_rad(3.0/180.0 * M_PI)
    .IncidentEnergy_keV(11.0);

    DiffractionGeometry geom = experiment.GetDiffractionGeometry();

    CrystalLattice lattice(Coord{20, 10, 0}, Coord{-20, 40, 0},
                           Coord{0, 0, 200});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.001,
        .max_hkl = 40
    };
    BraggPredictionGPU prediction;

    int count = prediction.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);

    for (int i = 0; i < count; i++) {
        auto r = prediction.GetReflections().at(i);
        auto recip = r.h * lattice.Astar() + r.k * lattice.Bstar() + r.l * lattice.Cstar();

        REQUIRE(std::abs(r.h) < settings.max_hkl );
        REQUIRE(std::abs(r.k) < settings.max_hkl );
        REQUIRE(std::abs(r.l) < settings.max_hkl );
        REQUIRE(r.d >= settings.high_res_A);
        REQUIRE(r.d == Catch::Approx(1/std::sqrt(recip * recip)).margin(0.01f));
        REQUIRE(r.dist_ewald == Catch::Approx(std::abs(geom.DistFromEwaldSphere(recip))).epsilon(1e-3));
        auto [x,y] = geom.RecipToDector(recip);
        REQUIRE(r.predicted_x == Catch::Approx(x).margin(0.05));
        REQUIRE(r.predicted_y == Catch::Approx(y).margin(0.05));
    }
}

TEST_CASE("BraggPredictionGPU_systematic_absences") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(120.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
        .IncidentEnergy_keV(12.0);

    CrystalLattice lattice(
        Coord{40, 0, 0},
        Coord{0, 50, 0},
        Coord{0, 0, 60}
    );

    BraggPredictionSettings settings{
        .high_res_A = 3.0f,
        .ewald_dist_cutoff = 0.1f,
        .max_hkl = 50
    };

    BraggPredictionGPU pred;
    SECTION ("I centering") {
        // 1) Body-centered I
        settings.centering = 'I';

        int count_I = pred.Calc(experiment, lattice, settings);
        REQUIRE(count_I > 0);
        for (const auto& r : pred.GetReflections()) {
            if (r.d == 0) break;
            REQUIRE(((r.h + r.k + r.l) % 2) == 0);
        }
    }

    SECTION ("F centering") {
        // 2) Face-centered F
        settings.centering = 'F';
        int count_F = pred.Calc(experiment, lattice, settings);
        REQUIRE(count_F > 0);
        for (const auto& r : pred.GetReflections()) {
            if (r.d == 0) break;
            const bool he = (r.h & 1) == 0, ke = (r.k & 1) == 0, le = (r.l & 1) == 0;
            const bool all_even = he && ke && le;
            const bool all_odd = (!he) && (!ke) && (!le);
            REQUIRE((all_even || all_odd));
        }
    }

    SECTION("R centering") {
        settings.centering = 'R';
        int count_R = pred.Calc(experiment, lattice, settings);
        REQUIRE(count_R > 0);

        // R (hexagonal setting): -h + k + l = 3n
        for (const auto& r : pred.GetReflections()) {
            if (r.d == 0) break;
            int cond = (-r.h + r.k + r.l) % 3;
            if (cond < 0) cond += 3;
            REQUIRE(cond == 0);
        }
    }
}

TEST_CASE("BraggPredictionGPU_backscattering") {
    DiffractionExperiment experiment(DetJF9M());
    experiment.DetectorDistance_mm(120.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
        .IncidentEnergy_keV(3.0/WVL_1A_IN_KEV);

    // Orthogonal basis, sufficient to test parity rules
    CrystalLattice lattice(
        Coord{40, 0, 0},
        Coord{0, 50, 0},
        Coord{0, 0, 60}
    );

    BraggPredictionSettings settings{
        .high_res_A = 3.0f,
        .ewald_dist_cutoff = 0.1f, // Very large cutoff, to be able to see as many reflections as possible
        .max_hkl = 50
    };

    BraggPredictionGPU pred;
    int count = pred.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);
    for (const auto& r : pred.GetReflections()) {
        if (r.d == 0) break;
        REQUIRE(r.d > 3.0 / sqrt(2.0));
    }
}

TEST_CASE("BraggPrediction_CPU_GPU_consistency_tilted") {
    // Verify CPU and GPU implementations produce identical results with tilted detector
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
        .PoniRot1_rad(0.04).PoniRot2_rad(-0.025)
        .IncidentEnergy_keV(12.0);

    CrystalLattice lattice(Coord{30, 10, 0}, Coord{-15, 45, 0}, Coord{0, 0, 150});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.0015,
        .max_hkl = 30
    };

    BraggPrediction cpu_pred;
    BraggPredictionGPU gpu_pred;

    int cpu_count = cpu_pred.Calc(experiment, lattice, settings);
    int gpu_count = gpu_pred.Calc(experiment, lattice, settings);

    REQUIRE(cpu_count > 0);
    REQUIRE(gpu_count > 0);

    // Build map of GPU reflections by hkl
    std::map<std::tuple<int,int,int>, const Reflection*> gpu_refl_map;
    for (int i = 0; i < gpu_count; ++i) {
        const auto& r = gpu_pred.GetReflections().at(i);
        gpu_refl_map[{r.h, r.k, r.l}] = &r;
    }

    // Check that each CPU reflection has a matching GPU reflection
    int matched = 0;
    for (int i = 0; i < cpu_count; ++i) {
        const auto& cpu_r = cpu_pred.GetReflections().at(i);
        auto key = std::make_tuple(cpu_r.h, cpu_r.k, cpu_r.l);
        auto it = gpu_refl_map.find(key);
        if (it != gpu_refl_map.end()) {
            const auto& gpu_r = *it->second;
            CHECK(cpu_r.predicted_x == Catch::Approx(gpu_r.predicted_x).margin(0.1));
            CHECK(cpu_r.predicted_y == Catch::Approx(gpu_r.predicted_y).margin(0.1));
            CHECK(cpu_r.d == Catch::Approx(gpu_r.d).margin(0.01));
            matched++;
        }
    }

    // Most reflections should match (allow for some numerical differences at boundaries)
    CHECK(matched > cpu_count * 0.95);
}
#endif