// 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 "../common/AzimuthalIntegrationProfile.h"
#include "../common/AzimuthalIntegration.h"

TEST_CASE("AzimuthalIntegrationMapping_Constructor","[AzimuthalIntegration]") {
    DiffractionExperiment x(DetJF4M());

    REQUIRE(x.GetPixelsNum() == 2164*2068);

    std::unique_ptr<AzimuthalIntegration> radial;

    x.QSpacingForAzimInt_recipA(0.1).QRangeForAzimInt_recipA(0.1, 5);

    PixelMask pixel_mask(x);
    REQUIRE_NOTHROW(radial = std::make_unique<AzimuthalIntegration>(x, pixel_mask));
}

TEST_CASE("AzimuthalIntegrationMapping_GetBinNumber","[AzimuthalIntegration]") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(50).BeamX_pxl(1000).BeamY_pxl(1000);
    x.QSpacingForAzimInt_recipA(0.1).QRangeForAzimInt_recipA(0.1, 4);

    PixelMask pixel_mask(x);
    AzimuthalIntegration mapping(x, pixel_mask);
    REQUIRE(mapping.GetBinNumber() == 39);
}

TEST_CASE("AzimuthalIntegrationMapping_GetBinNumber_mask","[AzimuthalIntegration]") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(50).BeamX_pxl(1000).BeamY_pxl(1000);
    x.QSpacingForAzimInt_recipA(0.1).QRangeForAzimInt_recipA(0.1, 9);

    std::vector<uint32_t> pixel_mask(x.GetPixelsNum(), 0);
    auto geom = x.GetDiffractionGeometry();
    for (int row = 0; row < x.GetYPixelsNum(); row++) {
        for (int col = 0; col < x.GetXPixelsNum(); col++) {
            float d = geom.PxlToRes(col, row);
            float q = 2 * M_PI / d;
            if (q >= 3.1)
                pixel_mask[row * x.GetXPixelsNum() + col] = 1;
        }
    }

    PixelMask pixel_mask_obj(x);
    pixel_mask_obj.LoadUserMask(x, pixel_mask);
    AzimuthalIntegration mapping(x, pixel_mask_obj);
    REQUIRE(mapping.GetBinNumber() == 89);
}

TEST_CASE("AzimuthalIntegrationMapping_GetBinNumber_DetectorLimit","[AzimuthalIntegration]") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(50).BeamX_pxl(1000).BeamY_pxl(1000);

    x.QSpacingForAzimInt_recipA(0.1).QRangeForAzimInt_recipA(0.1, 9.9);
    PixelMask pixel_mask(x);
    AzimuthalIntegration mapping(x, pixel_mask);

    REQUIRE(mapping.GetBinNumber() == 98);
}

TEST_CASE("AzimuthalIntegrationMapping_GetBinToQ","[AzimuthalIntegration]") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(50).BeamX_pxl(1000).BeamY_pxl(1000);
    x.QSpacingForAzimInt_recipA(0.1).QRangeForAzimInt_recipA(0.1, 4);

    PixelMask pixel_mask(x);
    AzimuthalIntegration mapping(x, pixel_mask);

    auto bin_to_q = mapping.GetBinToQ();

    REQUIRE(bin_to_q.size() == 39);
    CHECK(bin_to_q[0] == Catch::Approx(0.15));
    CHECK(bin_to_q[1] == Catch::Approx(0.25));
    CHECK(bin_to_q[15] == Catch::Approx(1.65));
    CHECK(bin_to_q[38] == Catch::Approx(3.95));
}

TEST_CASE("AzimuthalIntegrationMapping_GetBinToPhi","[AzimuthalIntegration]") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(50).BeamX_pxl(1000).BeamY_pxl(1000);

    AzimuthalIntegrationSettings settings;

    settings.QSpacing_recipA(0.1).QRange_recipA(0.1, 4.1).AzimuthalBinCount(4);
    x.ImportAzimuthalIntegrationSettings(settings);

    PixelMask pixel_mask(x);
    AzimuthalIntegration mapping(x, pixel_mask);

    auto &bin_to_q = mapping.GetBinToQ();

    REQUIRE(bin_to_q.size() == 40 * 4);
    CHECK(bin_to_q[0] == Catch::Approx(0.15));
    CHECK(bin_to_q[1] == Catch::Approx(0.25));
    CHECK(bin_to_q[15] == Catch::Approx(1.65));
    CHECK(bin_to_q[38] == Catch::Approx(3.95));

    auto &bin_to_phi = mapping.GetBinToPhi();

    REQUIRE(bin_to_phi.size() == 40 * 4);
    CHECK(bin_to_phi[0] == Catch::Approx(0));
    CHECK(bin_to_phi[1] == Catch::Approx(0));
    CHECK(bin_to_phi[38] == Catch::Approx(0));
    CHECK(bin_to_phi[40] == Catch::Approx(90));
}

TEST_CASE("AzimuthalIntegrationMapping_GetBin","[AzimuthalIntegration]") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(50).BeamX_pxl(1000).BeamY_pxl(1000);

    AzimuthalIntegrationSettings settings;

    settings.QSpacing_recipA(0.1).QRange_recipA(0.1, 4.1).AzimuthalBinCount(4);

    CHECK(settings.GetBin(0.11f, 0.0f) == 0);
    CHECK(settings.GetBin(0.11f, 95.0f) == 40);
    CHECK(settings.GetBin(0.11f, 185.0f) == 80);
    CHECK(settings.GetBin(4.02f, 280.0f) == 120 + 39);
}

#include "../preview/JFJochTIFF.h"

TEST_CASE("AzimuthalIntegrationMapping_GetMapping","[AzimuthalIntegration]") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(100).BeamX_pxl(1000).BeamY_pxl(1000);

    AzimuthalIntegrationSettings settings;

    settings.QSpacing_recipA(0.1).QRange_recipA(0.1, 4.1).AzimuthalBinCount(4);
    x.ImportAzimuthalIntegrationSettings(settings);

    PixelMask pixel_mask(x);
    AzimuthalIntegration mapping(x, pixel_mask);

    auto map = mapping.GetPixelToBin();

    CompressedImage image(map, x.GetXPixelsNum(), x.GetYPixelsNum());
    WriteTIFFToFile("test.tiff", image);

    CHECK(map[x.GetXPixelsNum() * 500 + 1500] / 40 == 3);
    CHECK(map[x.GetXPixelsNum() * 500 + 500] / 40 == 2);
    CHECK(map[x.GetXPixelsNum() * 1500 + 500] / 40 == 1);
    CHECK(map[x.GetXPixelsNum() * 1500 + 1500] / 40 == 0);
}

TEST_CASE("AzimuthalIntegrationMapping_QToBin","[AzimuthalIntegration]") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(50).BeamX_pxl(1000).BeamY_pxl(1000);
    x.QSpacingForAzimInt_recipA(0.1).QRangeForAzimInt_recipA(0.1, 4);

    PixelMask pixel_mask(x);
    AzimuthalIntegration mapping(x, pixel_mask);

    REQUIRE(mapping.QToBin(0.0) == 0);
    REQUIRE(std::floor(mapping.QToBin(0.200001)) == 1);
    REQUIRE(mapping.QToBin(0.6) == Catch::Approx(5));
    REQUIRE(mapping.QToBin(50.0) == Catch::Approx(38));
}

TEST_CASE("AzimuthalIntegrationProfile","[AzimuthalIntegration]") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(50).BeamX_pxl(1000).BeamY_pxl(1000);
    x.QSpacingForAzimInt_recipA(0.1).QRangeForAzimInt_recipA(0.1, 4);

    PixelMask pixel_mask(x);
    AzimuthalIntegration mapping(x, pixel_mask);

    AzimuthalIntegrationProfile profile(mapping);

    std::vector<float> sum(mapping.GetBinNumber());
    std::vector<uint32_t> count(mapping.GetBinNumber());

    for (int i = 0; i < mapping.GetBinNumber(); i++) {
        sum[i] = i * i * 4;
        count[i] = i;
    }
    REQUIRE_NOTHROW(profile.Add(sum, count));
    REQUIRE_NOTHROW(profile.Add(sum, count));

    std::vector<float> sum_wr(mapping.GetBinNumber() - 1);
    REQUIRE_THROWS(profile.Add(sum_wr, count));

    auto plot = profile.GetPlot();
    REQUIRE(plot.GetPlots().size() == 1);

    REQUIRE(plot.GetPlots()[0].x.size() == mapping.GetBinNumber());
    REQUIRE(plot.GetPlots()[0].y.size() == mapping.GetBinNumber());
    REQUIRE(plot.GetPlots()[0].x[0] == Catch::Approx(mapping.GetBinToQ()[0]));
    REQUIRE(std::isnan(plot.GetPlots()[0].y[0]));

    for (int i = 1; i < mapping.GetBinNumber(); i++) {
        REQUIRE(plot.GetPlots()[0].x[i] == Catch::Approx(mapping.GetBinToQ()[i]));
        REQUIRE(plot.GetPlots()[0].y[i] == Catch::Approx(i * 4));
    }
}

TEST_CASE("AzimuthalIntegrationProfile_operatorAdd","[AzimuthalIntegration]") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(50).BeamX_pxl(1000).BeamY_pxl(1000);
    x.QSpacingForAzimInt_recipA(0.1).QRangeForAzimInt_recipA(0.1, 4);

    PixelMask pixel_mask(x);
    AzimuthalIntegration mapping(x, pixel_mask);

    AzimuthalIntegrationProfile profile0(mapping), profile1(mapping);

    std::vector<float> sum(mapping.GetBinNumber());
    std::vector<uint32_t> count(mapping.GetBinNumber());

    for (int i = 0; i < mapping.GetBinNumber(); i++) {
        sum[i] = (i + 1) * i * 4;
        count[i] = i + 1;
    }
    REQUIRE_NOTHROW(profile0.Add(sum, count));
    REQUIRE_NOTHROW(profile1 += profile0);

    auto plot = profile1.GetPlot();

    REQUIRE(plot.GetPlots().size() == 1);
    REQUIRE(plot.GetPlots()[0].x.size() == mapping.GetBinNumber());
    REQUIRE(plot.GetPlots()[0].y.size() == mapping.GetBinNumber());
    for (int i = 0; i < mapping.GetBinNumber(); i++) {
        REQUIRE(plot.GetPlots()[0].x[i] == Catch::Approx(mapping.GetBinToQ()[i]));
        REQUIRE(plot.GetPlots()[0].y[i] == Catch::Approx(i * 4));
    }
}

TEST_CASE("AzimuthalIntegrationProfile_GetMeanValueOfBins","[AzimuthalIntegration]") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(50).BeamX_pxl(1000).BeamY_pxl(1000);
    x.QSpacingForAzimInt_recipA(0.1).QRangeForAzimInt_recipA(0.1, 4);

    PixelMask pixel_mask(x);
    AzimuthalIntegration mapping(x, pixel_mask);

    AzimuthalIntegrationProfile profile(mapping);

    std::vector<float> sum(mapping.GetBinNumber());
    std::vector<uint32_t> count(mapping.GetBinNumber());

    for (int i = 0; i < mapping.GetBinNumber(); i++) {
        sum[i] = i * i * 4;
        count[i] = i;
    }
    REQUIRE_NOTHROW(profile.Add(sum, count));

    REQUIRE(profile.GetMeanValueOfBins(0,2) == Catch::Approx((sum[0] + sum[1] + sum[2])  / double(count[0] + count[1] + count[2])));
    REQUIRE(profile.GetMeanValueOfBins(5,7) == Catch::Approx((sum[5] + sum[6] + sum[7])  / double (count[5] + count[6] + count[7])));

    x.BkgEstimateQRange_recipA(0.7, 0.8);

    REQUIRE(profile.GetBkgEstimate(x.GetAzimuthalIntegrationSettings()) == Catch::Approx((sum[5] + sum[6] + sum[7])  / double (count[5] + count[6] + count[7])));

    x.BkgEstimateQRange_recipA(0.01, 0.345);
    REQUIRE(profile.GetBkgEstimate(x.GetAzimuthalIntegrationSettings()) == Catch::Approx((sum[0] + sum[1] + sum[2])  / double(count[0] + count[1] + count[2])));
}