// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#pragma once

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <optional>
#include <random>
#include <string>
#include <tuple>
#include <vector>

#include "../common/Reflection.h"
#include "gemmi/symmetry.hpp"
#include "gemmi/unitcell.hpp"

// Synthetic merged intensities for the space-group / point-group search tests.
//
// The set produced here is what SearchSpaceGroup is fed in production: a P1 merge (one entry per
// Friedel-canonical hkl) with intensities, sigmas and half-set intensities. Unlike the noise-free
// set in SearchSpaceGroupTest.cpp it models the three things the point-group decision actually
// depends on:
//
//   * a merohedral TWIN - the crystal's true point group is a subgroup of index 2 of the metric
//     (lattice) point group, and the twin law is the operator that separates them:
//         I_obs(h) = (1-alpha) I_true(h) + alpha I_true(twin h)
//     I_true is a function of the TRUE (sub)group's asu, so the subgroup symmetry is exact and only
//     the extra supergroup operator is broken - by (1-2 alpha) (I_true(h) - I_true(twin h)). At
//     alpha = 0.5 the two are identical and the twin is indistinguishable from real symmetry.
//     Setting the true group to the SUPERgroup instead gives the untwinned high-symmetry control.
//
//   * the MERGE MULTIPLICITY, modelled the way the real merge behaves (Merge.h,
//     SigmaWithSystematicFloor): the random part of the merged sigma averages down as
//     1/sqrt(multiplicity) while the systematic part (b*I - absorption, partiality, beam flicker;
//     correlated across a reflection's repeats) does not, so the merged sigma is
//     max(sigma_statistical, b*|I|). Multiplicity therefore changes the sigmas but NOT the physics,
//     and no symmetry decision may depend on it.
//
//   * an ERROR-MODEL MISCALIBRATION - real merged sigmas come out under-estimated (~1.7x), which is
//     what pushes the merge's reduced chi^2 to ~3 and switches SearchSpaceGroup between its
//     chi^2-ratio and systematic-b regimes.
//
// Intensities follow a Wilson (exponential) distribution with a resolution fall-off, so they span a
// realistic dynamic range; the "structure factor" is a hash of the asu index, not a draw from the
// RNG stream, so the same crystal is reproduced bit-for-bit whatever the multiplicity or the twin
// fraction and two runs differ only in the knob under test.

namespace jfjoch_test {

    struct SyntheticMergeParams {
        // Symmetry the structure factors actually have: the generated intensities are exactly
        // invariant under it, whatever the twin fraction.
        std::string true_space_group = "R 3 :H";

        // Supergroup of index 2 over true_space_group; its extra operator is the twin law. Set it
        // equal to true_space_group to model a crystal that GENUINELY has the higher symmetry -
        // there is then no extra operator, twinning by a real symmetry operator is a no-op, and the
        // twin fraction has no effect.
        std::string twin_supergroup = "R 32 :H";

        // Merohedral twin fraction alpha in [0, 0.5]. 0 = untwinned, 0.5 = perfect twin (whose
        // intensities are exactly invariant under the twin law, hence indistinguishable from a
        // crystal that really has the supergroup symmetry).
        double twin_fraction = 0.0;

        // Number of observations merged into each reflection; at least 2, so both half-sets exist.
        int multiplicity = 6;

        // Error model. error_model_b is the b the merge FITTED and floors its merged sigmas with
        // (ISa = 1/b); true_systematic_b is the intensity-proportional systematic scatter actually
        // present in the data - unset means the two agree, i.e. a perfectly fitted error model.
        // sigma_miscalibration is how many times too SMALL the merged sigmas come out overall
        // (> 1 = under-estimated, the usual case; < 1 = over-estimated).
        double error_model_b = 0.05;
        std::optional<double> true_systematic_b;
        double sigma_miscalibration = 1.7;

        // Intensity distribution: mean intensity at infinite resolution, Wilson B fall-off, and the
        // background variance that keeps sigma finite for near-zero (systematically absent)
        // reflections.
        double mean_intensity = 8000.0;
        double wilson_b_A2 = 25.0;
        double background_variance = 400.0;

        double d_min_A = 2.5;

        uint32_t seed = 20260727;
    };

    namespace detail {
        inline uint64_t Mix64(uint64_t x) {
            x ^= x >> 33; x *= 0xff51afd7ed558ccdULL;
            x ^= x >> 33; x *= 0xc4ceb9fe1a85ec53ULL;
            x ^= x >> 33; return x;
        }

        // Uniform in (0, 1), a pure function of the Miller index.
        inline double UniformFromHkl(const gemmi::Op::Miller& hkl) {
            const uint64_t x = Mix64(static_cast<uint64_t>(hkl[0] + 512) * 0x9e3779b97f4a7c15ULL ^
                                     Mix64(static_cast<uint64_t>(hkl[1] + 512)) ^
                                     (Mix64(static_cast<uint64_t>(hkl[2] + 512)) << 1));
            return (static_cast<double>(x >> 11) + 0.5) * (1.0 / 9007199254740992.0);
        }

        inline std::vector<std::array<int, 9>> RotationSetOf(const gemmi::SpaceGroup& sg) {
            std::vector<std::array<int, 9>> out;
            for (const auto& op : sg.operations().derive_symmorphic().sym_ops) {
                std::array<int, 9> rot{};
                for (int i = 0; i < 3; ++i)
                    for (int j = 0; j < 3; ++j)
                        rot[i * 3 + j] = op.rot[i][j];
                out.push_back(rot);
            }
            std::sort(out.begin(), out.end());
            return out;
        }
    }

    // The twin law: a rotation of the supergroup that is not in the subgroup. Any of them gives the
    // same twinned intensities (an index-2 subgroup is normal, so the coset members differ by a
    // subgroup operator, which leaves I_true unchanged), so the first one found is used.
    inline gemmi::Op TwinLaw(const gemmi::SpaceGroup& sub, const gemmi::SpaceGroup& super) {
        const auto sub_rots = detail::RotationSetOf(sub);
        for (const auto& op : super.operations().derive_symmorphic().sym_ops) {
            std::array<int, 9> rot{};
            for (int i = 0; i < 3; ++i)
                for (int j = 0; j < 3; ++j)
                    rot[i * 3 + j] = op.rot[i][j];
            if (!std::binary_search(sub_rots.begin(), sub_rots.end(), rot))
                return gemmi::Op{op.rot, {0, 0, 0}, op.notation};
        }
        return gemmi::Op::identity();
    }

    // A cell consistent with the space group's crystal system. Synthetic throughout - the tests must
    // not carry the cell of any real sample.
    inline gemmi::UnitCell SyntheticCellFor(const gemmi::SpaceGroup& sg) {
        switch (sg.crystal_system()) {
            case gemmi::CrystalSystem::Triclinic:    return {33, 37, 41, 85, 95, 105};
            case gemmi::CrystalSystem::Monoclinic:   return {37, 43, 51, 90, 101, 90};
            case gemmi::CrystalSystem::Orthorhombic: return {37, 43, 51, 90, 90, 90};
            case gemmi::CrystalSystem::Tetragonal:   return {47, 47, 63, 90, 90, 90};
            case gemmi::CrystalSystem::Trigonal:
            case gemmi::CrystalSystem::Hexagonal:    return {51, 51, 71, 90, 90, 120};
            case gemmi::CrystalSystem::Cubic:        return {57, 57, 57, 90, 90, 90};
        }
        return {50, 50, 50, 90, 90, 90};
    }

    inline std::vector<MergedReflection> GenerateSyntheticMerged(const SyntheticMergeParams& p) {
        const gemmi::SpaceGroup& sub = gemmi::get_spacegroup_by_name(p.true_space_group);
        const gemmi::SpaceGroup& super = gemmi::get_spacegroup_by_name(p.twin_supergroup);
        const gemmi::Op twin = TwinLaw(sub, super);
        const gemmi::UnitCell cell = SyntheticCellFor(sub);
        const gemmi::GroupOps gops = sub.operations();
        const gemmi::ReciprocalAsu rasu(&sub);

        // True (untwinned) intensity: Wilson-distributed |F|^2 of the subgroup asu, with a
        // resolution fall-off. Systematically absent reflections (here: the lattice centering) carry
        // no intensity - they are what Stage B confirms the centering from.
        auto true_intensity = [&](const gemmi::Op::Miller& hkl) -> double {
            if (gops.is_systematically_absent(hkl))
                return 0.0;
            const auto asu = rasu.to_asu_sign(hkl, gops).first;
            const double e_squared = -std::log(detail::UniformFromHkl(asu));   // mean 1, exponential
            const double d = cell.calculate_d(hkl);
            return p.mean_intensity * e_squared * std::exp(-p.wilson_b_A2 / (2.0 * d * d));
        };

        // Half-set split of the multiplicity (n0 >= n1); both halves see the same systematic error.
        const int n_obs = std::max(2, p.multiplicity);
        const int n_half[2] = {(n_obs + 1) / 2, n_obs / 2};

        std::mt19937 rng(p.seed);
        std::normal_distribution<double> gauss(0.0, 1.0);

        const int hmax = static_cast<int>(std::ceil(cell.a / p.d_min_A)) + 1;
        const int kmax = static_cast<int>(std::ceil(cell.b / p.d_min_A)) + 1;
        const int lmax = static_cast<int>(std::ceil(cell.c / p.d_min_A)) + 1;

        std::vector<MergedReflection> merged;

        for (int h = -hmax; h <= hmax; ++h)
            for (int k = -kmax; k <= kmax; ++k)
                for (int l = -lmax; l <= lmax; ++l) {
                    // One entry per Friedel pair, matching the Friedel-merged P1 set the search gets.
                    if (std::make_tuple(h, k, l) <= std::make_tuple(-h, -k, -l))
                        continue;
                    const gemmi::Op::Miller hkl{{h, k, l}};
                    const double d = cell.calculate_d(hkl);
                    if (!(d >= p.d_min_A))
                        continue;

                    const auto twinned = twin.apply_to_hkl(hkl);
                    const double i_obs = (1.0 - p.twin_fraction) * true_intensity(hkl) +
                                         p.twin_fraction * true_intensity(twinned);

                    // Statistical error of one observation, and of the merge of n of them.
                    const double sigma_one = std::sqrt(i_obs + p.background_variance);
                    // Systematic error: a property of the reflection, identical in every observation
                    // of it, so it survives the merge - this is what the b*|I| sigma floor models.
                    const double systematic =
                        p.true_systematic_b.value_or(p.error_model_b) * i_obs * gauss(rng);

                    MergedReflection r;
                    r.h = h;
                    r.k = k;
                    r.l = l;
                    r.d = static_cast<float>(d);

                    double sum_n_i = 0.0;
                    for (int half = 0; half < 2; ++half) {
                        const double sigma_stat_half =
                            sigma_one / std::sqrt(static_cast<double>(n_half[half]));
                        const double i_half = i_obs + systematic + sigma_stat_half * gauss(rng);
                        r.I_half[half] = static_cast<float>(i_half);
                        r.sigma_half[half] = static_cast<float>(
                            std::max(sigma_stat_half, p.error_model_b * std::abs(i_half)) /
                            p.sigma_miscalibration);
                        sum_n_i += n_half[half] * i_half;
                    }

                    const double i_merged = sum_n_i / n_obs;
                    const double sigma_stat = sigma_one / std::sqrt(static_cast<double>(n_obs));
                    r.I = static_cast<float>(i_merged);
                    // Merge.h SigmaWithSystematicFloor, then thrown off by the error-model
                    // miscalibration.
                    r.sigma = static_cast<float>(
                        std::max(sigma_stat, p.error_model_b * std::abs(i_merged)) / p.sigma_miscalibration);

                    merged.push_back(r);
                }

        return merged;
    }
}