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

#include "SearchSpaceGroup.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <iomanip>
#include <limits>
#include <map>
#include <sstream>
#include <tuple>
#include <unordered_map>
#include <vector>

namespace {
    // A merged reflection, folded onto the +/- Friedel-equivalent it represents, used as a
    // hash key to match symmetry-related reflections.
    struct HKLKey {
        int h = 0, k = 0, l = 0;
        bool operator==(const HKLKey& o) const noexcept { return h == o.h && k == o.k && l == o.l; }
    };

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

    HKLKey Canonicalize(int h, int k, int l, bool merge_friedel) {
        if (merge_friedel && std::make_tuple(-h, -k, -l) < std::make_tuple(h, k, l))
            return {-h, -k, -l};
        return {h, k, l};
    }

    double PearsonCC(const std::vector<double>& x, const std::vector<double>& y) {
        if (x.size() < 2)
            return std::numeric_limits<double>::quiet_NaN();

        double sx = 0, sy = 0, sxx = 0, syy = 0, sxy = 0;
        for (size_t i = 0; i < x.size(); ++i) {
            sx += x[i]; sy += y[i];
            sxx += x[i] * x[i]; syy += y[i] * y[i]; sxy += x[i] * y[i];
        }
        const double n = static_cast<double>(x.size());
        const double vx = sxx - sx * sx / n;
        const double vy = syy - sy * sy / n;
        if (vx <= 0 || vy <= 0)
            return std::numeric_limits<double>::quiet_NaN();
        return (sxy - sx * sy / n) / std::sqrt(vx * vy);
    }

    // A reflection is extinct from lattice centering alone (independent of any screw/glide) when a
    // centering translation makes its structure factor cancel. Mirrors the centering half of
    // gemmi::GroupOps::is_systematically_absent, so screw absences can be judged separately.
    bool CenteringAbsent(const gemmi::GroupOps& gops, const gemmi::Op::Miller& hkl) {
        for (size_t i = 1; i < gops.cen_ops.size(); ++i) {
            const auto& t = gops.cen_ops[i];
            if ((t[0] * hkl[0] + t[1] * hkl[1] + t[2] * hkl[2]) % gemmi::Op::DEN != 0)
                return true;
        }
        return false;
    }

    std::array<int, 9> RotKey(const gemmi::Op& op) {
        std::array<int, 9> out{};
        for (int i = 0; i < 3; ++i)
            for (int j = 0; j < 3; ++j)
                out[i * 3 + j] = op.rot[i][j];
        return out;
    }

    // The rotation part of a space group in the reference setting (identity included), as a
    // sorted list of matrices - the key that groups space groups into a candidate point group.
    // It must be the rotation SET, not gemmi's PointGroup enum: P321 and P312 are both "32" yet
    // have their 2-folds along different directions, and only the matrices tell them apart.
    using RotationSet = std::vector<std::array<int, 9>>;

    RotationSet RotationSetOf(const gemmi::SpaceGroup& sg) {
        RotationSet out;
        for (const auto& op : sg.operations().derive_symmorphic().sym_ops)
            out.push_back(RotKey(op));
        std::sort(out.begin(), out.end());
        return out;
    }

    // Proper rotations of a crystal system's holohedry (the highest lattice symmetry it can host),
    // in the reference setting. Any candidate point group must be a subgroup of this.
    RotationSet HolohedryRotationSet(gemmi::CrystalSystem system) {
        int number = 0;
        switch (system) {
            case gemmi::CrystalSystem::Triclinic:    number = 1;   break;  // P1
            case gemmi::CrystalSystem::Monoclinic:   number = 3;   break;  // P2  (unique axis b)
            case gemmi::CrystalSystem::Orthorhombic: number = 16;  break;  // P222
            case gemmi::CrystalSystem::Tetragonal:   number = 89;  break;  // P422
            case gemmi::CrystalSystem::Trigonal:     number = 155; break;  // R32
            case gemmi::CrystalSystem::Hexagonal:    number = 177; break;  // P622
            case gemmi::CrystalSystem::Cubic:        number = 207; break;  // P432
        }
        const auto* sg = gemmi::find_spacegroup_by_number(number);
        return sg ? RotationSetOf(*sg) : RotationSet{};
    }

    // A candidate point group: its proper rotations (reference setting) and a representative
    // symmorphic space group (used when only the point group is wanted, or for display).
    struct PointGroupInfo {
        RotationSet rotation_set;
        std::vector<gemmi::Op> rotations;       // non-identity proper rotations
        const gemmi::SpaceGroup* representative = nullptr;
    };

    // Enumerate candidate point groups. When a holohedry is given (from the lattice metric), keep
    // only its subgroups - this both skips operators the lattice forbids and avoids accepting a
    // coincidental higher symmetry; all subgroups down to P1 are still candidates.
    std::vector<PointGroupInfo> EnumeratePointGroups(const std::optional<RotationSet>& holohedry) {
        std::vector<PointGroupInfo> out;
        std::map<RotationSet, size_t> index;

        for (const auto& sg : gemmi::spacegroup_tables::main) {
            if (!sg.is_sohncke() || !sg.is_reference_setting())
                continue;

            RotationSet rs = RotationSetOf(sg);
            if (holohedry.has_value() &&
                !std::includes(holohedry->begin(), holohedry->end(), rs.begin(), rs.end()))
                continue;

            auto it = index.find(rs);
            size_t pos;
            if (it == index.end()) {
                PointGroupInfo info;
                for (const auto& op : sg.operations().derive_symmorphic().sym_ops) {
                    if (op.rot == gemmi::Op::identity().rot)
                        continue;
                    info.rotations.push_back(gemmi::Op{op.rot, {0, 0, 0}, op.notation});
                }
                info.rotation_set = rs;
                pos = out.size();
                index[rs] = pos;
                out.push_back(std::move(info));
            } else {
                pos = it->second;
            }
            // Prefer a symmorphic representative (the plain point-group setting).
            auto& info = out[pos];
            if (info.representative == nullptr ||
                (!info.representative->is_symmorphic() && sg.is_symmorphic()))
                info.representative = &sg;
        }
        return out;
    }
}

SearchSpaceGroupResult SearchSpaceGroup(
    const std::vector<MergedReflection>& merged,
    const SearchSpaceGroupOptions& opt) {

    SearchSpaceGroupResult result;
    if (merged.empty())
        return result;

    const size_t n = merged.size();

    // Flatten the reflections and mark which ones each stage may use. The correlation stage drops
    // weak reflections; the absence stage must keep them - that is where the screw-axis signal is.
    std::vector<int> H(n), K(n), L(n);
    std::vector<double> I(n), Sigma(n), IoverSigma(n);
    std::vector<HKLKey> key(n);
    std::vector<char> pass_absence(n, 0), pass_cc(n, 0);

    for (size_t i = 0; i < n; ++i) {
        const auto& r = merged[i];
        H[i] = r.h; K[i] = r.k; L[i] = r.l;
        I[i] = r.I;
        Sigma[i] = std::isfinite(r.sigma) && r.sigma > 0 ? r.sigma : 0.0;
        key[i] = Canonicalize(r.h, r.k, r.l, opt.merge_friedel);

        const bool finite = std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0 &&
                            std::isfinite(r.d) && r.d > 0;
        const bool in_range = finite && (opt.d_min_limit_A <= 0 || r.d >= opt.d_min_limit_A);
        IoverSigma[i] = finite ? r.I / r.sigma : 0.0;
        pass_absence[i] = in_range;
        // The correlation stage uses only genuinely-present reflections. Near-zero (systematically
        // absent) reflections would otherwise form a second cluster at the origin and fake a high
        // correlation for false operators - fatal on centered lattices, where half the reflections
        // are extinct.
        pass_cc[i] = in_range && IoverSigma[i] >= opt.present_i_over_sigma &&
                     (opt.min_i_over_sigma <= 0 || IoverSigma[i] >= opt.min_i_over_sigma);
    }

    std::unordered_map<HKLKey, int, HKLKeyHash> key_to_index;
    key_to_index.reserve(n * 2);
    for (size_t i = 0; i < n; ++i)
        if (pass_absence[i])
            key_to_index.emplace(key[i], static_cast<int>(i));

    // --- Stage A: score each distinct rotation operator once ---
    std::vector<uint32_t> visited(n, 0);
    uint32_t epoch = 0;

    auto score_operator = [&](const gemmi::Op& op) -> SpaceGroupOperatorScore {
        ++epoch;
        std::vector<double> x, y;
        for (size_t i = 0; i < n; ++i) {
            if (!pass_cc[i] || visited[i] == epoch)
                continue;
            const auto m2 = op.apply_to_hkl(gemmi::Op::Miller{{H[i], K[i], L[i]}});
            const HKLKey k2 = Canonicalize(m2[0], m2[1], m2[2], opt.merge_friedel);
            if (k2 == key[i])
                continue;                       // reflection lies on this rotation axis
            const auto it = key_to_index.find(k2);
            if (it == key_to_index.end())
                continue;
            const int j = it->second;
            if (!pass_cc[j])
                continue;
            x.push_back(I[i]);
            y.push_back(I[j]);
            visited[i] = epoch;
            visited[j] = epoch;
        }
        SpaceGroupOperatorScore s;
        s.op_triplet_hkl = op.as_hkl().triplet('h');
        s.n_pairs = static_cast<int>(x.size());
        s.cc = PearsonCC(x, y);
        s.present = s.n_pairs >= opt.min_pairs_per_operator && std::isfinite(s.cc) &&
                    s.cc >= opt.min_operator_cc;
        return s;
    };

    std::map<std::array<int, 9>, SpaceGroupOperatorScore> op_cache;
    auto operator_score = [&](const gemmi::Op& op) -> const SpaceGroupOperatorScore& {
        const auto rk = RotKey(op);
        auto it = op_cache.find(rk);
        if (it != op_cache.end())
            return it->second;
        return op_cache.emplace(rk, score_operator(op)).first->second;
    };

    std::optional<RotationSet> holohedry;
    if (opt.lattice_system.has_value())
        holohedry = HolohedryRotationSet(opt.lattice_system.value());
    const auto point_groups = EnumeratePointGroups(holohedry);

    // Reduced chi^2 of the intensities merged under a point group's rotations - how well its symmetry
    // equivalents agree RELATIVE TO THEIR ERRORS. A real point group gives ~1; a false operator forces
    // non-equivalent reflections together, so they disagree by many sigma and chi^2 blows up. This is
    // more sensitive than R-meas to a strong pseudo-symmetry (where the intensities still correlate well
    // - high operator CC - but not within their errors). Inverse-variance weighted mean per orbit, over
    // the present (pass_cc) reflections.
    auto chi2_under = [&](const std::vector<gemmi::Op>& rotations) -> double {
        struct Acc { double sw = 0.0, swI = 0.0; int n = 0; };
        std::unordered_map<HKLKey, Acc, HKLKeyHash> grp;
        std::vector<HKLKey> rep(n);
        for (size_t i = 0; i < n; ++i) {
            if (!pass_cc[i] || !(Sigma[i] > 0.0))
                continue;
            HKLKey best = key[i];
            for (const auto& op : rotations) {
                const auto m = op.apply_to_hkl(gemmi::Op::Miller{{H[i], K[i], L[i]}});
                const HKLKey k2 = Canonicalize(m[0], m[1], m[2], opt.merge_friedel);
                if (std::make_tuple(k2.h, k2.k, k2.l) < std::make_tuple(best.h, best.k, best.l))
                    best = k2;
            }
            rep[i] = best;
            auto& g = grp[best];
            const double w = 1.0 / (Sigma[i] * Sigma[i]);
            g.sw += w; g.swI += w * I[i]; g.n += 1;
        }
        double chi2 = 0.0;
        long dof = 0;
        for (size_t i = 0; i < n; ++i) {
            if (!pass_cc[i] || !(Sigma[i] > 0.0))
                continue;
            const auto& g = grp[rep[i]];
            if (g.n < 2)
                continue;
            const double mean = g.swI / g.sw, dev = I[i] - mean;
            chi2 += dev * dev / (Sigma[i] * Sigma[i]);
        }
        for (const auto& [k, g] : grp)
            if (g.n >= 2)
                dof += g.n - 1;
        return dof > 0 ? chi2 / static_cast<double>(dof) : std::numeric_limits<double>::quiet_NaN();
    };

    // Operator-CC-confirmed candidates, each with its merge chi^2; chi2_ref = the most consistent.
    struct PGCand { const PointGroupInfo* pg; int order; double min_cc; double chi2; };
    std::vector<PGCand> pg_cands;
    double chi2_ref = std::numeric_limits<double>::infinity();
    for (const auto& pg : point_groups) {
        bool all_present = true;
        double min_cc = pg.rotations.empty() ? 1.0 : std::numeric_limits<double>::infinity();
        for (const auto& op : pg.rotations) {
            const auto& s = operator_score(op);
            all_present = all_present && s.present;
            min_cc = std::min(min_cc, s.cc);
        }
        if (!all_present)
            continue;
        const double ch = pg.rotations.empty() ? std::numeric_limits<double>::quiet_NaN()
                                               : chi2_under(pg.rotations);
        pg_cands.push_back({&pg, static_cast<int>(pg.rotations.size()) + 1, min_cc, ch});
        if (!pg.rotations.empty() && std::isfinite(ch))
            chi2_ref = std::min(chi2_ref, ch);
    }

    // Choose the largest point group that is both operator-confirmed AND self-consistent (its merge
    // chi^2 is not inflated past max_merge_chi2_ratio x the most-consistent candidate; ties -> higher
    // min CC). Identity (no operators) is always consistent, so it stays the P1 fallback.
    const PointGroupInfo* best_pg = nullptr;
    int best_pg_order = 0;
    double best_pg_min_cc = -2.0;
    for (const auto& c : pg_cands) {
        const bool consistent = c.pg->rotations.empty() || !std::isfinite(c.chi2) ||
                                !std::isfinite(chi2_ref) || c.chi2 <= chi2_ref * opt.max_merge_chi2_ratio;
        if (!consistent)
            continue;
        if (c.order > best_pg_order || (c.order == best_pg_order && c.min_cc > best_pg_min_cc)) {
            best_pg = c.pg;
            best_pg_order = c.order;
            best_pg_min_cc = c.min_cc;
        }
    }

    for (const auto& [rk, s] : op_cache)
        result.operator_scores.push_back(s);
    std::sort(result.operator_scores.begin(), result.operator_scores.end(),
              [](const auto& a, const auto& b) { return a.cc > b.cc; });

    if (best_pg == nullptr)                          // should not happen (C1 always qualifies)
        return result;

    if (best_pg->representative)
        result.point_group_hm = best_pg->representative->point_group_hm();

    // --- Stage B: pick the space group within the point group ---
    // Without screw/centering determination, return the symmorphic representative.
    if (!opt.determine_space_group || best_pg->rotations.empty()) {
        if (best_pg->representative)
            result.best_space_group = *best_pg->representative;
        return result;
    }

    for (const auto& sg : gemmi::spacegroup_tables::main) {
        if (!sg.is_sohncke() || !sg.is_reference_setting() || RotationSetOf(sg) != best_pg->rotation_set)
            continue;

        const gemmi::GroupOps gops = sg.operations();
        SpaceGroupCandidateScore s{.space_group = sg};
        double absent_sum = 0, present_sum = 0;
        int present_n = 0;

        // Judge centering and screw/glide absences on separate reflection sets. Lumping them lets
        // a large, correct centering-absent set hide a few strong screw violations and over-claim
        // screw axes (e.g. I4_132 on I432 data).
        int centering_absent = 0, centering_violations = 0;
        int screw_absent = 0, screw_violations = 0;

        for (size_t i = 0; i < n; ++i) {
            if (!pass_absence[i])
                continue;
            const gemmi::Op::Miller hkl{{H[i], K[i], L[i]}};
            const bool present = IoverSigma[i] > opt.present_i_over_sigma;

            if (CenteringAbsent(gops, hkl)) {
                s.absent_observed += 1;
                absent_sum += IoverSigma[i];
                centering_absent += 1;
                if (present) { s.absent_violations += 1; centering_violations += 1; }
            } else if (gops.is_systematically_absent(hkl)) {
                s.absent_observed += 1;
                absent_sum += IoverSigma[i];
                screw_absent += 1;
                if (present) { s.absent_violations += 1; screw_violations += 1; }
            } else {
                present_n += 1;
                present_sum += IoverSigma[i];
            }
        }

        if (s.absent_observed > 0)
            s.absent_mean_i_over_sigma = absent_sum / s.absent_observed;
        if (present_n > 0)
            s.present_mean_i_over_sigma = present_sum / present_n;

        const bool centering_ok = centering_absent == 0 ||
            centering_violations <= opt.max_absent_violation_fraction * centering_absent;
        const bool screw_ok = screw_absent == 0 ||
            screw_violations <= opt.max_absent_violation_fraction * screw_absent;
        s.consistent = centering_ok && screw_ok;
        result.candidates.push_back(std::move(s));
    }

    // A candidate is eligible when its absences are confirmed and there are enough of them to
    // trust (the symmorphic group, with no absences, is always eligible as the fallback). Rank
    // eligible candidates by how many absences they explain - the screw/centering content that is
    // both real and maximal wins, instead of defaulting to the symmorphic group.
    auto eligible = [&](const SpaceGroupCandidateScore& s) {
        return s.consistent && (s.absent_observed == 0 || s.absent_observed >= opt.min_absent_observed);
    };
    std::sort(result.candidates.begin(), result.candidates.end(),
              [&](const SpaceGroupCandidateScore& a, const SpaceGroupCandidateScore& b) {
                  if (eligible(a) != eligible(b))
                      return eligible(a);
                  if (a.absent_observed != b.absent_observed)
                      return a.absent_observed > b.absent_observed;
                  // Tie (e.g. I23 vs I2_13, indistinguishable by absences): lower space-group number.
                  return a.space_group.number < b.space_group.number;
              });

    if (!result.candidates.empty() && eligible(result.candidates.front())) {
        const int best_absent = result.candidates.front().absent_observed;
        for (auto& s : result.candidates) {
            if (!eligible(s) || s.absent_observed != best_absent)
                continue;
            s.selected = true;
            if (!result.best_space_group.has_value())
                result.best_space_group = s.space_group;   // representative (lowest number)
            else
                result.alternatives.push_back(s.space_group);
        }
    }

    return result;
}

std::string SearchSpaceGroupResultToText(const SearchSpaceGroupResult& result,
                                         size_t max_candidates_to_print) {
    std::ostringstream os;

    os << "Point group: " << (result.point_group_hm.empty() ? "?" : result.point_group_hm)
       << " (from intensity correlations)\n";

    os << "  " << std::setw(14) << std::left << "operator" << std::right
       << std::setw(9) << "CC" << std::setw(10) << "pairs" << std::setw(9) << "symm" << "\n";
    for (const auto& s : result.operator_scores) {
        os << "  " << std::setw(14) << std::left << s.op_triplet_hkl << std::right
           << std::setw(9) << std::fixed << std::setprecision(3) << s.cc
           << std::setw(10) << s.n_pairs
           << std::setw(9) << (s.present ? "yes" : "no") << "\n";
    }

    os << "\nSpace-group candidates\n";
    os << "  " << std::setw(10) << std::left << "SG" << std::right
       << std::setw(9) << "absent" << std::setw(7) << "viol"
       << std::setw(11) << "<I/s>abs" << std::setw(11) << "<I/s>pres"
       << std::setw(6) << "OK" << "\n";

    const size_t count = std::min(max_candidates_to_print, result.candidates.size());
    for (size_t i = 0; i < count; ++i) {
        const auto& c = result.candidates[i];
        os << (c.selected ? "* " : "  ")
           << std::setw(10) << std::left << c.space_group.short_name() << std::right
           << std::setw(9) << c.absent_observed << std::setw(7) << c.absent_violations
           << std::setw(11) << std::fixed << std::setprecision(2) << c.absent_mean_i_over_sigma
           << std::setw(11) << std::fixed << std::setprecision(2) << c.present_mean_i_over_sigma
           << std::setw(6) << (c.consistent ? "yes" : "no") << "\n";
    }

    if (result.best_space_group.has_value()) {
        os << "Best space group: " << result.best_space_group->short_name();
        for (const auto& alt : result.alternatives)
            os << " or " << alt.short_name();
        if (!result.alternatives.empty())
            os << " (indistinguishable from these data)";
        os << "\n";
    } else {
        os << "Best space group: none determined\n";
    }

    return os.str();
}