// 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/geom_refinement/AssignSpotsToRings.h"
#include "../common/Definitions.h"

TEST_CASE("DetGeomCalib_FindCircleCenter") {
     std::vector<SpotToSave> spots;

    // Make two colinear circles
    for (int i = 0; i < 7; i++) {
        float angle = i * 2 * M_PI / 10;
        float x = 200.0f + 100.0f * cosf(angle);
        float y = 200.0f + 100.0f * sinf(angle);
        spots.push_back(SpotToSave(x, y, 1000));
    }

    for (int i = 0; i < 7; i++) {
        float angle = i * 2 * M_PI / 10;
        float x = 200.0f + 50.0f * cosf(angle);
        float y = 200.0f + 50.0f * sinf(angle);
        spots.push_back(SpotToSave(x, y, 1000));
    }

    // Add some outliers
    spots.push_back(SpotToSave(1000.0f, 1000.0f, 1000));
    spots.push_back(SpotToSave(0.0f, 0.0f, 1000));

    auto ret = FindCircleCenter(spots);

    REQUIRE(ret.x == Catch::Approx(200.0f));
    REQUIRE(ret.y == Catch::Approx(200.0f));
}

TEST_CASE("DetGeomCalib_FindCircleCenter_250") {
    std::vector<SpotToSave> spots;

    for (int i = 0; i < 250; i++) {
        float angle = i * 2 * M_PI / 250;
        float x = 200.0f + 100.0f * cosf(angle);
        float y = 200.0f + 100.0f * sinf(angle);
        spots.push_back(SpotToSave(x, y, 1000));
    }

    auto ret = FindCircleCenter(spots);

    REQUIRE(ret.x == Catch::Approx(200.0f));
    REQUIRE(ret.y == Catch::Approx(200.0f));
    REQUIRE(ret.total_votes == 250 * 249 * 248 / 6); // N * (N-1) * (N-2) / 6
}

TEST_CASE("DetGeomCalib_dbscan") {
    std::vector<float> spots_r;
    for (int i = 0; i < 15; i++)
        spots_r.push_back(100.0);

    for (int i = 0; i < 15; i++)
        spots_r.push_back(50.0);

    for (int i = 0; i < 9; i++)
        spots_r.push_back(30.0);
    spots_r.push_back(70.0);
    spots_r.push_back(11.0);

    auto ret = ClusterSpotsIntoRings(spots_r, 0.1, 10);

    REQUIRE(ret.size() == 2);

    REQUIRE(ret[0][0] == 0);
    REQUIRE(ret[0][1] == 1);
    REQUIRE(ret[1][0] == 15);
}

TEST_CASE("DetGeomCalib_AnalyzeClusters") {
    std::vector<std::vector<int>> dbscan_result;

    std::vector<float> spots_r;

    std::vector<int> tmp_1;
    for (int i = 0; i < 15; i++) {
        spots_r.push_back(100.0);
        tmp_1.push_back(i);
    }

    std::vector<int> tmp_2;
    for (int i = 0; i < 15; i++) {
        spots_r.push_back(50.0);
        tmp_2.push_back(15+i);
    }
    dbscan_result.push_back(tmp_1);
    dbscan_result.push_back(tmp_2);

    auto ret = AnalyzeClusters(spots_r, dbscan_result);

    REQUIRE(ret.size() == 2);
    REQUIRE(ret[0].R_obs == Catch::Approx(50.0f));
    REQUIRE(ret[1].R_obs == Catch::Approx(100.0f));
}

TEST_CASE("DetGeomCalib_CalculateXtalRings_cubic") {
    auto ret = CalculateXtalRings(UnitCell(2.0, 2.0, 2.0, 90, 90, 90));
    CHECK(ret[0] == Catch::Approx(2.0 * M_PI * 1.0 / 2.0));
    CHECK(ret[1] == Catch::Approx(2.0 * M_PI * sqrt( 2.0 )/ 2.0));
    CHECK(ret[2] == Catch::Approx(2.0 * M_PI * sqrt( 3.0 )/ 2.0));
    CHECK(ret[3] == Catch::Approx(2.0 * M_PI * 2.0/ 2.0));
    // 7 cannot be obtained by h^2 + k^2 + l^2, while 8 can
    CHECK(ret[6] == Catch::Approx(2.0 * M_PI * sqrt( 8.0 )/ 2.0));
}

TEST_CASE("DetGeomCalib_CalculateXtalRings_one_long_axis") {
    auto ret = CalculateXtalRings(UnitCell(50.0, 2.0, 2.0, 90, 90, 90));
    CHECK(ret[0] == Catch::Approx(2.0 * M_PI * 1.0 / 50.0));
    CHECK(ret[1] == Catch::Approx(2.0 * M_PI * 2.0 / 50.0));
    CHECK(ret[2] == Catch::Approx(2.0 * M_PI * 3.0 / 50.0));
    CHECK(ret[3] == Catch::Approx(2.0 * M_PI * 4.0 / 50.0));
    CHECK(ret[4] == Catch::Approx(2.0 * M_PI * 5.0 / 50.0));



}

TEST_CASE("DetGeomCalib_GuessDetectorDistance") {
    std::vector<SpotToSave> spots;

    DiffractionGeometry geom;
    geom.Wavelength_A(1.0).BeamX_pxl(100.0).BeamY_pxl(200.0)
    .DetectorDistance_mm(1000);

    float lab6_a = 4.156468;

    float ring_radius_pxl = geom.ResToPxl(lab6_a);

    REQUIRE(GuessDetectorDistance(geom, ring_radius_pxl, lab6_a) == Catch::Approx(1000.0));
}

TEST_CASE("DetGeomCalib_GuessInitialGeometry") {
    std::vector<SpotToSave> spots;

    DiffractionGeometry geom;
    geom.Wavelength_A(1.0).BeamX_pxl(100.0).BeamY_pxl(200.0)
            .DetectorDistance_mm(1000);

    float lab6_a = LAB6_CELL_A;

    for (int i = 0; i < 30; i++) {
        auto [x, y] = geom.ResPhiToPxl(lab6_a, i * M_PI * 2.0 / 30.0);
        spots.push_back(SpotToSave{.x = x, .y = y});
    }

    DiffractionGeometry geom_out;
    geom_out.Wavelength_A(1.0);

    GuessInitialGeometry(geom_out, spots, LAB6_CELL_A);

    REQUIRE(geom_out.GetBeamX_pxl() == geom.GetBeamX_pxl());
    REQUIRE(geom_out.GetBeamY_pxl() == geom.GetBeamY_pxl());
    REQUIRE(geom_out.GetDetectorDistance_mm() == Catch::Approx(geom.GetDetectorDistance_mm()));
}

TEST_CASE("DetGeomCalib_GuessGeometry") {
    std::vector<SpotToSave> spots;

    DiffractionGeometry geom;
    geom.Wavelength_A(1.0).BeamX_pxl(100.0).BeamY_pxl(200.0)
            .DetectorDistance_mm(100);

    float lab6_a = LAB6_CELL_A;

    for (int i = 0; i < 30; i++) {
        auto [x, y] = geom.ResPhiToPxl(lab6_a, i * M_PI * 2.0 / 30.0);
        spots.push_back(SpotToSave{.x = x, .y = y});
    }

    for (int i = 0; i < 30; i++) {
        auto [x, y] = geom.ResPhiToPxl(lab6_a / sqrt(2), i * M_PI * 2.0 / 30.0);
        spots.push_back(SpotToSave{.x = x, .y = y});
    }

    for (int i = 0; i < 30; i++) {
        auto [x, y] = geom.ResPhiToPxl(lab6_a / sqrt(3.0f), i * M_PI * 2.0 / 30.0);
        spots.push_back(SpotToSave{.x = x, .y = y});
    }

    for (int i = 0; i < 30; i++) {
        auto [x, y] = geom.ResPhiToPxl(lab6_a / 2.0f, i * M_PI * 2.0 / 30.0);
        spots.push_back(SpotToSave{.x = x, .y = y});
    }

    DiffractionGeometry geom_out;
    geom_out.Wavelength_A(1.0).DetectorDistance_mm(200.0);

    GuessGeometry(geom_out, spots, UnitCell(LAB6_CELL_A, LAB6_CELL_A, LAB6_CELL_A, 90,90,90));

    CHECK(fabsf(geom_out.GetBeamX_pxl() - geom.GetBeamX_pxl()) < 0.001f);
    CHECK(fabsf(geom_out.GetBeamY_pxl() - geom.GetBeamY_pxl()) < 0.001f);
    // This is wrong!!!!
    CHECK(fabsf(geom_out.GetDetectorDistance_mm() - geom.GetDetectorDistance_mm()) < 0.01f);
}

TEST_CASE("DetGeomCalib_RingOptimizer") {
    std::vector<RingOptimizerInput> spots;

    DiffractionGeometry geom;
    geom.Wavelength_A(1.0).BeamX_pxl(100.0).BeamY_pxl(200.0)
            .DetectorDistance_mm(100).PoniRot1_rad(0.1).PoniRot2_rad(0.05);

    float lab6_a = LAB6_CELL_A;

    for (int i = 0; i < 30; i++) {
        auto [x, y] = geom.ResPhiToPxl(lab6_a, i * M_PI * 2.0 / 30.0);
        spots.push_back(RingOptimizerInput{
            .x = x, .y = y,
            .q_expected = M_PI * 2.0 / lab6_a
        });
    }

    for (int i = 0; i < 30; i++) {
        auto [x, y] = geom.ResPhiToPxl(lab6_a / sqrt(2), i * M_PI * 2.0 / 30.0);
        spots.push_back(RingOptimizerInput{
            .x = x, .y = y,
            .q_expected = sqrtf(2.0f) * M_PI * 2.0 / lab6_a
        });
    }

    DiffractionGeometry geom_i;
    geom_i.Wavelength_A(1.0).BeamX_pxl(105).BeamY_pxl(195).DetectorDistance_mm(110);

    RingOptimizer optimizer(geom_i);

    DiffractionGeometry geom_o = optimizer.Run(spots);

    CHECK(fabs(geom_o.GetBeamX_pxl() - geom.GetBeamX_pxl()) < 0.001f);
    CHECK(fabs(geom_o.GetBeamY_pxl() - geom.GetBeamY_pxl()) < 0.001f);
    CHECK(fabs(geom_o.GetDetectorDistance_mm() - geom.GetDetectorDistance_mm()) < 0.001f);
    CHECK(fabs(geom_o.GetPoniRot1_rad() - geom.GetPoniRot1_rad()) < 0.001f);
    CHECK(fabs(geom_o.GetPoniRot2_rad() - geom.GetPoniRot2_rad()) < 0.001f);
}

TEST_CASE("DetGeomCalib_OptimizeGeometry") {

    DiffractionGeometry geom;
    geom.Wavelength_A(1.0).BeamX_pxl(1000.0).BeamY_pxl(1275.0)
            .DetectorDistance_mm(100).PoniRot1_rad(0.1).PoniRot2_rad(0.05);

    float lab6_a = LAB6_CELL_A;

    std::vector<SpotToSave> spots;

    for (int d = 1; d < 7; d++) {
        for (int i = 0; i < 30; i++) {
            auto [x, y] = geom.ResPhiToPxl(lab6_a / sqrt(d), i * M_PI * 2.0 / 30.0);
            spots.push_back(SpotToSave{.x = x, .y = y});
        }
    }

    DiffractionGeometry geom_i;
    geom_i.Wavelength_A(1.0).BeamX_pxl(995.0).BeamY_pxl(1277.0)
            .DetectorDistance_mm(98).PoniRot1_rad(0.0975).PoniRot2_rad(0.055);

    OptimizeGeometry(geom_i, spots, UnitCell(LAB6_CELL_A, LAB6_CELL_A, LAB6_CELL_A, 90,90,90));

    CHECK(geom_i.GetBeamX_pxl() == Catch::Approx(geom.GetBeamX_pxl()));
    CHECK(geom_i.GetBeamY_pxl() == Catch::Approx(geom.GetBeamY_pxl()));
    CHECK(geom_i.GetDetectorDistance_mm() == Catch::Approx(geom.GetDetectorDistance_mm()));
    CHECK(geom_i.GetPoniRot1_rad() == Catch::Approx(geom.GetPoniRot1_rad()));
    CHECK(geom_i.GetPoniRot2_rad() == Catch::Approx(geom.GetPoniRot2_rad()));
}