// 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 "../writer/HDF5Objects.h"
#include "../image_analysis/indexing/IndexerFactory.h"

TEST_CASE("CrystalLattice") {
    CrystalLattice l(50,60,80, 90, 90, 90);
    REQUIRE(l.Vec0().Length() == Catch::Approx(50));
    REQUIRE(l.Vec1().Length() == Catch::Approx(60));
    REQUIRE(l.Vec2().Length() == Catch::Approx(80));
    REQUIRE(angle_deg(l.Vec0(), l.Vec2()) == 90);
    REQUIRE(angle_deg(l.Vec0(), l.Vec1()) == 90);
    REQUIRE(angle_deg(l.Vec1(), l.Vec2()) == 90);

    auto uc0 = l.GetUnitCell();
    REQUIRE(uc0.a == Catch::Approx(50));
    REQUIRE(uc0.b == Catch::Approx(60));
    REQUIRE(uc0.c == Catch::Approx(80));
    REQUIRE(uc0.alpha == Catch::Approx(90));
    REQUIRE(uc0.beta == Catch::Approx(90));
    REQUIRE(uc0.gamma == Catch::Approx(90));

    l = CrystalLattice(30, 40, 70, 90, 95, 90);
    REQUIRE(l.Vec0().Length() == Catch::Approx(30));
    REQUIRE(l.Vec1().Length() == Catch::Approx(40));
    REQUIRE(l.Vec2().Length() == Catch::Approx(70));
    REQUIRE(angle_deg(l.Vec0(), l.Vec2()) == 95);
    REQUIRE(angle_deg(l.Vec0(), l.Vec1()) == 90);
    REQUIRE(angle_deg(l.Vec1(), l.Vec2()) == 90);

    auto uc1 = l.GetUnitCell();
    REQUIRE(uc1.a == Catch::Approx(30));
    REQUIRE(uc1.b == Catch::Approx(40));
    REQUIRE(uc1.c == Catch::Approx(70));
    REQUIRE(uc1.alpha == Catch::Approx(90));
    REQUIRE(uc1.beta == Catch::Approx(95));
    REQUIRE(uc1.gamma == Catch::Approx(90));

    l = CrystalLattice(45, 45, 70, 90, 90, 120);
    REQUIRE(l.Vec0().Length() == Catch::Approx(45));
    REQUIRE(l.Vec1().Length() == Catch::Approx(45));
    REQUIRE(l.Vec2().Length() == Catch::Approx(70));
    REQUIRE(angle_deg(l.Vec0(), l.Vec2()) == Catch::Approx(90));
    REQUIRE(angle_deg(l.Vec0(), l.Vec1()) == Catch::Approx(120));
    REQUIRE(angle_deg(l.Vec1(), l.Vec2()) == Catch::Approx(90));

    auto uc2 = l.GetUnitCell();
    REQUIRE(uc2.a == Catch::Approx(45));
    REQUIRE(uc2.b == Catch::Approx(45));
    REQUIRE(uc2.c == Catch::Approx(70));
    REQUIRE(uc2.alpha == Catch::Approx(90));
    REQUIRE(uc2.beta == Catch::Approx(90));
    REQUIRE(uc2.gamma == Catch::Approx(120));
}

TEST_CASE("CrystalLattice_Sort") {
    CrystalLattice l(80,60,50, 120, 90, 90);
    REQUIRE(l.Vec0().Length() == Catch::Approx(50));
    REQUIRE(l.Vec1().Length() == Catch::Approx(60));
    REQUIRE(l.Vec2().Length() == Catch::Approx(80));
    REQUIRE(angle_deg(l.Vec0(), l.Vec2()) == Catch::Approx(90));
    REQUIRE(angle_deg(l.Vec0(), l.Vec1()) == Catch::Approx(120));
    REQUIRE(angle_deg(l.Vec1(), l.Vec2()) == Catch::Approx(90));
}

TEST_CASE("CrystalLattice_Handedness") {
    CrystalLattice l(Coord(1,0,0), Coord(0,1,0), Coord(0,0,-1));
    REQUIRE(l.Vec0().x == Catch::Approx(1));
    REQUIRE(l.Vec1().y == Catch::Approx(1));
    REQUIRE(l.Vec2().z == Catch::Approx(1));
}

TEST_CASE("CrystalLattice_Volume") {
    CrystalLattice l(50, 60, 80, 90, 90, 90);
    REQUIRE(l.CalcVolume() == 50 * 60 * 80);

    CrystalLattice l2(50,60,80, 90, 120, 90);
    float sin120 = std::sqrt(3) / 2;
    REQUIRE(l2.CalcVolume() == Catch::Approx(50 * 60 * 80 * sin120));
}

TEST_CASE("CrystalLattice_Recip") {
    CrystalLattice l(50,60,80, 90, 90, 90);
    REQUIRE(l.Astar().Length() == Catch::Approx(1/50.0));
    REQUIRE(l.Astar().x == Catch::Approx(1/50.0));
    REQUIRE(l.Bstar().Length() == Catch::Approx(1/60.0));
    REQUIRE(l.Bstar().y == Catch::Approx(1/60.0));
    REQUIRE(l.Cstar().Length() == Catch::Approx(1/80.0));
    REQUIRE(l.Cstar().z == Catch::Approx(1/80.0));
}

inline double round_err(double x) {
    return std::abs(x - std::round(x));
}

#ifdef JFJOCH_USE_CUDA

#include <Eigen/Dense>

TEST_CASE("FastFeedbackIndexer","[Indexing]") {
    std::vector<Coord> hkl;
    for (int i = 1; i < 7; i++)
        for (int j = 1; j<6; j++)
            for (int k = 1; k < 4; k++)
                hkl.emplace_back(i,j,k);

    std::vector<UnitCell> cells;
    cells.emplace_back(30,40,50,90,90,90);
    cells.emplace_back(80,80,90,90,90,120);
    cells.emplace_back(40,45,80,90,82.5,90);

    DiffractionExperiment experiment;
    experiment.SetUnitCell(cells[0]);
    experiment.IndexingAlgorithm(IndexingAlgorithmEnum::FFBIDX);

    REQUIRE(experiment.GetIndexingAlgorithm() == IndexingAlgorithmEnum::FFBIDX);

    std::unique_ptr<Indexer> indexer = CreateIndexer(experiment);

    for (auto &c: cells) {
        CrystalLattice l(c);

        Eigen::Matrix3f m;
        m << l.Vec0().x, l.Vec0().y, l.Vec0().z,
                l.Vec1().x, l.Vec1().y, l.Vec1().z,
                l.Vec2().x, l.Vec2().y, l.Vec2().z;
        auto m1 = m.transpose().inverse();

        CrystalLattice recip_l(Coord(m1(0,0), m1(0,1), m1(0,2)),
                               Coord(m1(1,0), m1(1,1), m1(1,2)),
                               Coord(m1(2,0), m1(2,1), m1(2,2)));

        std::vector<Coord> recip;
        recip.reserve(hkl.size());
        for (const auto &i: hkl)
            recip.emplace_back(i.x * recip_l.Vec0() + i.y * recip_l.Vec1() + i.z * recip_l.Vec2());

        experiment.SetUnitCell(c);

        indexer->Setup(experiment);
        auto ret = indexer->Run(recip, recip.size());
        REQUIRE(!ret.empty());

        //auto uc = ret[0].GetUnitCell();
        //REQUIRE(c.a == Catch::Approx(uc.a));
        //REQUIRE(c.b == Catch::Approx(uc.b));
        //REQUIRE(c.c == Catch::Approx(uc.c));

        double err[3] = {0.0, 0.0, 0.0};
        for (const auto &iter: recip) {
            err[0] += round_err(ret[0].Vec0() * iter);
            err[1] += round_err(ret[0].Vec1() * iter);
            err[2] += round_err(ret[0].Vec2() * iter);
        }
        REQUIRE (err[0] < 0.001 * recip.size());
        REQUIRE (err[1] < 0.001 * recip.size());
        REQUIRE (err[2] < 0.001 * recip.size());
    }
}

#endif