// 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 <iostream>
#include "../common/DiffractionGeometry.h"
#include "../common/DiffractionExperiment.h"
TEST_CASE("RecipToDetector_1", "[LinearAlgebra][Coord]") {
    DiffractionExperiment x(DetJF(8, 2));
    x.BeamX_pxl(1024).BeamY_pxl(1024).DetectorDistance_mm(120);
    DiffractionGeometry geom = x.GetDiffractionGeometry();

    float pos_x = 512, pos_y = 512;

    auto recip = geom.DetectorToRecip(pos_x, pos_y);
    auto [proj_x, proj_y] = geom.RecipToDector(recip);
    REQUIRE(proj_x == Catch::Approx(pos_x));
    REQUIRE(proj_y == Catch::Approx(pos_y));
    REQUIRE((recip - geom.DetectorToRecip(proj_x, proj_y)).Length() < 0.00000001f);
    REQUIRE(geom.DistFromEwaldSphere(recip) < 4e-4);
}

TEST_CASE("RecipToDetector_2", "[LinearAlgebra][Coord]") {
    DiffractionExperiment x(DetJF(8, 2));
    x.BeamX_pxl(1024).BeamY_pxl(1024).DetectorDistance_mm(120);
    float pos_x = 1023, pos_y = 1023;
    DiffractionGeometry geom = x.GetDiffractionGeometry();

    auto recip = geom.DetectorToRecip(pos_x, pos_y);
    auto [proj_x, proj_y] = geom.RecipToDector(recip);
    REQUIRE(proj_x == Catch::Approx(pos_x));
    REQUIRE(proj_y == Catch::Approx(pos_y));
    REQUIRE((recip - geom.DetectorToRecip(proj_x, proj_y)).Length() < 0.00000001f);
    REQUIRE(geom.DistFromEwaldSphere(recip) < 4e-4);
}

TEST_CASE("RecipToDetector_3", "[LinearAlgebra][Coord]") {
    DiffractionExperiment x(DetJF(8, 2));
    x.BeamX_pxl(1024).BeamY_pxl(1024).DetectorDistance_mm(120);
    float pos_x = 30, pos_y = 30;
    DiffractionGeometry geom = x.GetDiffractionGeometry();

    auto recip = geom.DetectorToRecip(pos_x, pos_y);
    auto [proj_x, proj_y] = geom.RecipToDector(recip);
    REQUIRE(proj_x == Catch::Approx(pos_x));
    REQUIRE(proj_y == Catch::Approx(pos_y));
    REQUIRE((recip - geom.DetectorToRecip(proj_x, proj_y)).Length() < 0.00000001f);
    REQUIRE(geom.DistFromEwaldSphere(recip) < 4e-4);
}

TEST_CASE("DiffractionGeometry_Phi","") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(75).IncidentEnergy_keV(WVL_1A_IN_KEV);
    x.BeamX_pxl(1000).BeamY_pxl(1000);
    DiffractionGeometry geom = x.GetDiffractionGeometry();

    CHECK(geom.Phi(2000, 1000) * (180.0 / M_PI) == Catch::Approx(0.0));
    CHECK(geom.Phi(2000, 0) * (180.0 / M_PI) == Catch::Approx(45.0f));
    CHECK(geom.Phi(1000, 0) * (180.0 / M_PI) == Catch::Approx(90.0f));
    CHECK(geom.Phi(0, 0) * (180.0 / M_PI) == Catch::Approx(135.0f));
    CHECK(geom.Phi(0, 1000) * (180.0 / M_PI) == Catch::Approx(180.0f));
    CHECK(geom.Phi(1000, 2000) * (180.0 / M_PI) == Catch::Approx(270.f));
    CHECK(geom.Phi(2000, 2000) * (180.0 / M_PI) == Catch::Approx(315.0f));

}

TEST_CASE("DiffractionGeometry_Cos2Theta","") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(75).IncidentEnergy_keV(WVL_1A_IN_KEV);
    x.BeamX_pxl(1000).BeamY_pxl(1000);
    DiffractionGeometry geom = x.GetDiffractionGeometry();

    // det distance == 1000 pixel
    // theta = 30 deg
    // tan(2 * theta) = sqrt(3)
    REQUIRE(cosf(geom.TwoTheta(1000, 1000 * (1.0 + sqrt(3)))) == Catch::Approx(0.5f));
}

TEST_CASE("DiffractionGeometry_PxlToRes","") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(75).IncidentEnergy_keV(WVL_1A_IN_KEV);
    DiffractionGeometry geom = x.GetDiffractionGeometry();

    // sin(theta) = 1/2
    // theta = 30 deg
    // tan(2 * theta) = sqrt(3)
    REQUIRE(geom.PxlToRes( 0, 1000 * sqrt(3)) == Catch::Approx(1.0));

    // sin(theta) = 1/4
    // theta = 14.47 deg
    // tan(2 * theta) = 0.55328333517

    REQUIRE(geom.PxlToRes(1000 * 0.55328333517 * cosf(1), 1000 * 0.55328333517 * sinf(1)) == Catch::Approx(2.0));
}


TEST_CASE("DiffractionGeometry_ResToPxl","") {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(75).IncidentEnergy_keV(WVL_1A_IN_KEV);
    DiffractionGeometry geom = x.GetDiffractionGeometry();

    // sin(theta) = 1/2
    // theta = 30 deg
    // tan(2 * theta) = sqrt(3)
    REQUIRE(geom.ResToPxl(1.0) == Catch::Approx(1000 * sqrt(3)));

    // sin(theta) = 1/4
    // theta = 14.47 deg
    // tan(2 * theta) = 0.55328333517
    REQUIRE(geom.ResToPxl(2.0) == Catch::Approx(1000 * 0.55328333517));
}

TEST_CASE("DiffractionGeometry_SolidAngleCorrection","") {
    DiffractionExperiment x;
    x.IncidentEnergy_keV(WVL_1A_IN_KEV);
    x.BeamX_pxl(1000).BeamY_pxl(1000).DetectorDistance_mm(75);
    DiffractionGeometry geom = x.GetDiffractionGeometry();

    REQUIRE(geom.CalcAzIntSolidAngleCorr(0.0) == 1.0f);
    REQUIRE(geom.CalcAzIntSolidAngleCorr(2 * M_PI) == Catch::Approx(0.5f * 0.5f * 0.5f));

    // theta = 30 deg
    // cos (2 * theta) = 1/2
    REQUIRE(geom.CalcAzIntSolidAngleCorr(1000, 1000) == 1.0f);
    REQUIRE(geom.CalcAzIntSolidAngleCorr(1000 * (1.0 + sqrt(3)), 1000) == Catch::Approx(0.5f * 0.5f * 0.5f));
}

TEST_CASE("DiffractionGeometry_PolarizationCorrection","") {
    DiffractionExperiment x;
    x.IncidentEnergy_keV(WVL_1A_IN_KEV);
    x.BeamX_pxl(1000).BeamY_pxl(1000).DetectorDistance_mm(75);

    DiffractionGeometry geom = x.GetDiffractionGeometry();

    // Circular polarization 0.5*(1+cos(2theta)^2)
    x.PolarizationFactor(0);
    REQUIRE(geom.CalcAzIntPolarizationCorr(1000 * (1.0 + sqrt(3)), 1000, 0) == Catch::Approx(0.5f * (1 + 0.5f * 0.5f)));
    REQUIRE(geom.CalcAzIntPolarizationCorr(1000, 1000 * (1.0 + sqrt(3)), 0) == Catch::Approx(0.5f * (1 + 0.5f * 0.5f)));

    // Horizontal polarization
    x.PolarizationFactor(1);

    // No correction in vertical direction
    REQUIRE(geom.CalcAzIntPolarizationCorr(1000, 1000 * (1.0 + sqrt(3)), 1) == Catch::Approx(1.0f));
    REQUIRE(geom.CalcAzIntPolarizationCorr(1000, 1000 * (1.0 - sqrt(3)), 1) == Catch::Approx(1.0f));

    // cos(2*theta)^2 in horizontal direction
    REQUIRE(geom.CalcAzIntPolarizationCorr(1000 * (1.0 + sqrt(3)), 1000, 1) == Catch::Approx(0.5f * 0.5f));
    REQUIRE(geom.CalcAzIntPolarizationCorr(1000 * (1.0 - sqrt(3)), 1000, 1) == Catch::Approx(0.5f * 0.5f));
}

TEST_CASE("DiffractionGeometry_AngleFromEwaldSphere") {
    DiffractionGeometry geom;
    geom.Wavelength_A(1.0);

    // Center of Ewald sphere == (0,0,-1)

    // Points on Ewald sphere
    REQUIRE(geom.AngleFromEwaldSphere(Coord(1,0,-1)) == 0.0f);
    REQUIRE(geom.AngleFromEwaldSphere(Coord(1.0f/sqrtf(2.0f),1.0f/sqrtf(2.0f),-1)) == 0.0f);

    REQUIRE(geom.AngleFromEwaldSphere(Coord(1,0,1)) == Catch::Approx(90.0f));
    REQUIRE(geom.AngleFromEwaldSphere(Coord(-sqrtf(2.0f),0,0)) == Catch::Approx(45.0f));
    REQUIRE(geom.AngleFromEwaldSphere(Coord(-sqrtf(3.0f),0,0)) == Catch::Approx(60.0f));

    float cos_1deg = cosf(1.0f * M_PI / 180.0f);
    float sin_1deg = sinf(1.0f * M_PI / 180.0f);

    REQUIRE(fabsf(geom.AngleFromEwaldSphere((Coord(cos_1deg - sin_1deg, 0, -(cos_1deg + sin_1deg)))) - 1.0f) < 0.0005);

    // Cannot be rotated to fit into the Ewald sphere
    REQUIRE(isnanf(geom.AngleFromEwaldSphere(Coord(0,0,1))));
}

TEST_CASE("DiffractionGeometry_DirectBeam") {
    DiffractionGeometry geom;
    geom.Wavelength_A(1.0);
    geom.BeamX_pxl(1230).BeamY_pxl(1450);

    auto [x, y] = geom.GetDirectBeam_pxl();
    REQUIRE(x == Catch::Approx(1230.0f));
    REQUIRE(y == Catch::Approx(1450.0f));
}

TEST_CASE("DiffractionGeometry_DirectBeam_RotZ") {
    DiffractionGeometry geom;
    geom.Wavelength_A(1.0);
    geom.BeamX_pxl(1230).BeamY_pxl(1450);
    geom.PoniRot3_rad(M_PI_2);

    auto [x, y] = geom.GetDirectBeam_pxl();
    REQUIRE(x == Catch::Approx(1230.0f));
    REQUIRE(y == Catch::Approx(1450.0f));
}

TEST_CASE("DiffractionGeometry_DirectBeam_RotY") {
    DiffractionGeometry geom;
    geom.Wavelength_A(1.0);
    geom.DetectorDistance_mm(100);
    geom.PixelSize_mm(1.0);
    geom.BeamX_pxl(1230).BeamY_pxl(1450);
    geom.PoniRot2_rad(M_PI_4); // 45 deg rotation

    auto [x, y] = geom.GetDirectBeam_pxl();
    CHECK(x == Catch::Approx(1230.0f)); // no Change for X
    CHECK(y >1450.0f);
}