// 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_build_u") {
    auto ret = CalculateCubicXtalRings();
    REQUIRE(ret[0] == Catch::Approx(1.0));
    REQUIRE(ret[1] == Catch::Approx(sqrtf(2.0)));
    REQUIRE(ret[2] == Catch::Approx(sqrtf(3.0)));
    REQUIRE(ret[3] == Catch::Approx(2.0));
    // 7 cannot be obtained by h^2 + k^2 + l^2, while 8 can
    REQUIRE(ret[6] == Catch::Approx(sqrtf(8.0)));
}

TEST_CASE("DetGeomCalib_AssignRings") {

    std::vector<RingClusters> obs_sorted;
    float S_val = 100.0f;

    obs_sorted.push_back(RingClusters{{}, S_val * 1.0f,-1 });
    obs_sorted.push_back(RingClusters{{}, S_val * 1.2f,-1 });
    obs_sorted.push_back(RingClusters{{}, S_val * sqrtf(3.0f),-1 });
    obs_sorted.push_back(RingClusters{{}, S_val * sqrtf(5.0f),-1 });
    obs_sorted.push_back(RingClusters{{}, S_val * 3.100f,-1 });

    auto res = AssignRings(obs_sorted);
    CHECK(res.S == S_val);
    CHECK(res.num_rings == 3);

    REQUIRE(obs_sorted[0].ring_idx == 0);
    REQUIRE(obs_sorted[2].ring_idx == 2);
    REQUIRE(obs_sorted[3].ring_idx == 4);
}

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, LAB6_CELL_A);

    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, .d_expected = 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, .d_expected = lab6_a / sqrtf(2.0f)});
    }


    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);
}