#include <catch2/catch_all.hpp>

#include "../image_analysis/scale_merge/SearchSpaceGroup.h"
#include "gemmi/symmetry.hpp"

#include <cmath>
#include <cstdint>
#include <string>
#include <tuple>
#include <unordered_set>
#include <vector>

namespace {
    struct HKL {
        int h = 0;
        int k = 0;
        int l = 0;

        bool operator==(const HKL& o) const noexcept {
            return h == o.h && k == o.k && l == o.l;
        }
    };

    struct HKLHash {
        size_t operator()(const HKL& x) const noexcept {
            auto mix = [](uint64_t v) {
                v ^= v >> 33;
                v *= 0xff51afd7ed558ccdULL;
                v ^= v >> 33;
                v *= 0xc4ceb9fe1a85ec53ULL;
                v ^= v >> 33;
                return v;
            };
            return static_cast<size_t>(
                mix(static_cast<uint64_t>(x.h)) ^
                (mix(static_cast<uint64_t>(x.k)) << 1) ^
                (mix(static_cast<uint64_t>(x.l)) << 2));
        }
    };

    double CalcSyntheticD(int h, int k, int l) {
        const double q2 = static_cast<double>(h * h + k * k + l * l);
        return 40.0 / std::sqrt(q2 + 1.0);
    }

    double SyntheticIntensityFromAsu(const gemmi::Op::Miller& asu) {
        uint64_t x = static_cast<uint64_t>((asu[0] + 31) * 73856093u) ^
                     static_cast<uint64_t>((asu[1] + 37) * 19349663u) ^
                     static_cast<uint64_t>((asu[2] + 41) * 83492791u);
        x ^= x >> 13;
        x *= 0x9e3779b97f4a7c15ULL;
        x ^= x >> 17;
        return 100.0 + static_cast<double>(x % 500);
    }

    std::vector<MergedReflection> GenerateMergedReflectionsForSpaceGroup(
        const gemmi::SpaceGroup& sg,
        int hmax = 8) {

        std::vector<MergedReflection> merged;
        std::unordered_set<HKL, HKLHash> added;

        const gemmi::GroupOps gops = sg.operations();
        const gemmi::ReciprocalAsu rasu(&sg);

        for (int h = -hmax; h <= hmax; ++h) {
            for (int k = -hmax; k <= hmax; ++k) {
                for (int l = -hmax; l <= hmax; ++l) {
                    if (h == 0 && k == 0 && l == 0)
                        continue;

                    gemmi::Op::Miller hkl{{h, k, l}};
                    if (gops.is_systematically_absent(hkl))
                        continue;

                    const auto [asu, sign_plus] = rasu.to_asu_sign(hkl, gops);
                    if (!sign_plus)
                        continue;

                    const HKL key{h, k, l};
                    if (added.find(key) != added.end())
                        continue;
                    added.insert(key);

                    merged.push_back(MergedReflection{
                        .h = h,
                        .k = k,
                        .l = l,
                        .I = SyntheticIntensityFromAsu(asu),
                        .sigma = 1.0,
                        .d = CalcSyntheticD(h, k, l)
                    });
                }
            }
        }

        return merged;
    }
}

TEST_CASE("SearchSpaceGroup detects synthetic space groups") {
    struct Case {
        std::string input_name;
        std::string expected_short_name;
    };

    const std::vector<Case> cases = {
        {"P 1", "P1"},
        {"P 1 2 1", "P2"},
        {"P 3 2 1", "P321"},
        {"P 4 2 2", "P422"},
        {"P 4 3 2", "P432"},
        {"P 6 2 2", "P622"},
        {"C 1 2 1", "C2"},
        {"C 2 2 2", "C222"},
        {"I 4 3 2", "I432"},
    };

    for (const auto& tc : cases) {
        DYNAMIC_SECTION(tc.expected_short_name) {
            const gemmi::SpaceGroup& sg = gemmi::get_spacegroup_by_name(tc.input_name);
            const auto merged = GenerateMergedReflectionsForSpaceGroup(sg);

            SearchSpaceGroupOptions opt;
            opt.crystal_system.reset();
            opt.centering = '\0';
            opt.merge_friedel = true;
            opt.d_min_limit_A = 2.0;
            opt.min_operator_cc = 0.99;
            opt.min_pairs_per_operator = 20;
            opt.min_total_compared = 80;
            opt.test_systematic_absences = true;
            opt.max_mean_absent_i_over_sigma = 0.1;
            opt.max_absent_violations = 0;

            const auto result = SearchSpaceGroup(merged, opt);

            INFO(SearchSpaceGroupResultToText(result));
            REQUIRE(result.best_space_group.has_value());
            CHECK(result.best_space_group->short_name() == tc.expected_short_name);
        }
    }
}