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

#include "../../common/JFJochMath.h"
#include "RotationIndexer.h"

#include "../geom_refinement/XtalOptimizer.h"
#include "../indexing/FFTIndexer.h"
#include "../lattice_search/LatticeSearch.h"
#include "../indexing/MultiLatticeSearch.h"

namespace {
    // Re-express a primitive hexagonal/trigonal lattice in the conventional hexagonal setting
    // (a = b, gamma = 120). The Niggli-reduced primitive cell carries the two equal-length axes
    // at gamma = 60; replacing b with b - a opens that angle to 120 without changing the lattice.
    CrystalLattice HexagonalConventional(CrystalLattice latt) {
        latt.ReorderABEqual();                       // put the equal-length pair in a, b
        Coord a = latt.Vec0(), b = latt.Vec1(), c = latt.Vec2();
        if (angle_deg(a, b) < 90.0f)
            b -= a;
        return CrystalLattice(a, b, c);              // constructor fixes handedness
    }

    bool IsHexagonalSystem(gemmi::CrystalSystem s) {
        return s == gemmi::CrystalSystem::Trigonal || s == gemmi::CrystalSystem::Hexagonal;
    }

    // The hexagonal lattice metric (two equal axes at 60/120 deg, both perpendicular to the third) is
    // also satisfied by its ortho-hexagonal C-centred supercell, so the geometry-keyed LatticeSearch can
    // land there. Detect the hexagonal metric on the reduced PRIMITIVE cell so the de-novo path (no space
    // group to key on) can re-express it in conventional hexagonal axes.
    bool IsMetricallyHexagonal(CrystalLattice latt, float rel_tol = 0.03f, float angle_tol_deg = 3.0f) {
        latt.ReorderABEqual();
        const Coord a = latt.Vec0(), b = latt.Vec1(), c = latt.Vec2();
        const float la = a.Length(), lb = b.Length();
        if (la <= 0.0f || lb <= 0.0f || std::fabs(la - lb) > rel_tol * std::max(la, lb))
            return false;
        const float gab = angle_deg(a, b);
        if (std::fabs(gab - 60.0f) > angle_tol_deg && std::fabs(gab - 120.0f) > angle_tol_deg)
            return false;
        return std::fabs(angle_deg(a, c) - 90.0f) <= angle_tol_deg &&
               std::fabs(angle_deg(b, c) - 90.0f) <= angle_tol_deg;
    }

    // Fraction of the accumulated reciprocal-space spots that a lattice indexes to near-integer
    // Miller indices within tol. Comparing a symmetry-constrained refinement against an
    // unconstrained (triclinic) one is a data-driven test for a false promotion: a wrong
    // higher-symmetry constraint snaps a pseudo cell onto ideal angles and misplaces most spots.
    float IndexedFraction(const CrystalLattice &latt, const std::vector<Coord> &coords, float tol) {
        if (coords.empty())
            return 0.0f;
        const Coord a = latt.Vec0(), b = latt.Vec1(), c = latt.Vec2();
        const float tol_sq = tol * tol;
        size_t indexed = 0;
        for (const Coord &s : coords) {
            const float dh = a * s - std::round(a * s);   // Coord operator* = dot product = Miller index
            const float dk = b * s - std::round(b * s);
            const float dl = c * s - std::round(c * s);
            if (dh * dh + dk * dk + dl * dl < tol_sq)
                ++indexed;
        }
        return static_cast<float>(indexed) / static_cast<float>(coords.size());
    }
}

RotationIndexer::RotationIndexer(const DiffractionExperiment &x, IndexerThreadPool &indexer)
    : experiment(x),
      index_ice_rings(x.GetIndexingSettings().GetIndexIceRings()),
      v_(experiment.GetImageNum()),
      angle_deg_(experiment.GetImageNum()),
      axis_(x.GetGoniometer()),
      geom_(x.GetDiffractionGeometry()),
      updated_geom_(geom_),
      indexer_(indexer) {
}

void RotationIndexer::RunIndexing() {
    std::unique_lock ul(m);
    if (!axis_)
        return;

    std::vector<Coord> coords;
    coords.reserve(max_spots_per_image * v_.size());

    for (int i = 0; i < v_.size(); i++) {
        const float angle_deg = angle_deg_[i].value_or(axis_->GetAngle_deg(i) + axis_->GetWedge_deg() / 2.0f);
        const auto rot = axis_->GetTransformationAngle(angle_deg);

        for (const auto &s: v_[i])
            coords.emplace_back(rot * s.ReciprocalCoord(geom_));
    }
    const auto indexer_result = indexer_.Run(experiment, coords);

    if (!indexer_result.lattice.empty() && indexer_result.lattice[0].CalcVolume() > 1.0) {
        auto sg = experiment.GetGemmiSpaceGroup();
        DiffractionExperiment experiment_copy(experiment);
        const float index_tol = experiment.GetIndexingSettings().GetTolerance();
        const auto orig_axis = axis_;

        // Map an FFT candidate cell to a (metric) space-group setting: the user-fixed SG's conventional
        // cell, or the de-novo Bravais lattice. Re-express a metrically-hexagonal cell in conventional
        // hexagonal axes (LatticeSearch can land on the ortho-hexagonal C setting) so the 3-fold is not
        // hidden from scaling.
        auto build_sr = [&](const CrystalLattice &cand) -> LatticeSearchResult {
            auto ls = LatticeSearch(cand);
            if (sg) {
                const auto is_hexagonal = [](gemmi::CrystalSystem s) {
                    return s == gemmi::CrystalSystem::Trigonal || s == gemmi::CrystalSystem::Hexagonal;
                };
                CrystalLattice conventional = ls.conventional;
                if (is_hexagonal(sg->crystal_system()) && !is_hexagonal(ls.system))
                    conventional = HexagonalConventional(ls.primitive_reduced);
                return LatticeSearchResult{
                    .niggli_class = ls.niggli_class,
                    .primitive_reduced = ls.primitive_reduced,
                    .conventional = conventional,
                    .system = sg->crystal_system(),
                    .centering = sg->centring_type(),
                    .reindex = ls.reindex,
                };
            }
            if (!IsHexagonalSystem(ls.system) && IsMetricallyHexagonal(ls.primitive_reduced)) {
                ls.conventional = HexagonalConventional(ls.primitive_reduced);
                ls.system = gemmi::CrystalSystem::Hexagonal;
                ls.centering = 'P';
            }
            return ls;
        };

        // Re-accumulate the reciprocal spots under a refined geometry/axis, to score a refined cell.
        auto accumulate = [&](const DiffractionGeometry &g, const std::optional<GoniometerAxis> &ax) {
            std::vector<Coord> c;
            c.reserve(max_spots_per_image * v_.size());
            for (int i = 0; i < v_.size(); i++) {
                const float a = angle_deg_[i].value_or(ax->GetAngle_deg(i) + ax->GetWedge_deg() / 2.0f);
                const auto rot = ax->GetTransformationAngle(a);
                for (const auto &s : v_[i])
                    c.emplace_back(rot * s.ReciprocalCoord(g));
            }
            return c;
        };

        // The FFT offers a few candidate cells (its best reduction plus, for large/elongated cells, a
        // widened alternative). Fully refine each and keep the one that indexes the most spots AFTER
        // geometry refinement - the pre-refinement fraction is not a reliable discriminator (an
        // incorrect larger cell can fit more of the un-refined accumulated spots than the correct one).
        const size_t n_try = std::min<size_t>(indexer_result.lattice.size(), 4);
        float best_frac = -1.0f;
        bool have_best = false;
        size_t best_ci = 0;
        XtalOptimizerData best_data;
        LatticeSearchResult best_sr;

        for (size_t ci = 0; ci < n_try; ci++) {
            const CrystalLattice &cand = indexer_result.lattice[ci];
            if (cand.CalcVolume() <= 1.0)
                continue;

            LatticeSearchResult sr = build_sr(cand);
            const auto conv_uc = sr.conventional.GetUnitCell();
            const float length_bound_A = 1.2f * static_cast<float>(std::max({conv_uc.a, conv_uc.b, conv_uc.c}));

            // Bound the axis lengths just above the found cell so a free (triclinic) refine cannot drift
            // onto a pseudo-translation / modulation supercell (a modulated crystal whose satellites
            // define a ~4x period would otherwise inflate one axis to the max-length clamp).
            auto make_data = [&](const CrystalLattice &latt, gemmi::CrystalSystem sys) {
                XtalOptimizerData d{
                    .geom = experiment_copy.GetDiffractionGeometry(),
                    .latt = latt,
                    .crystal_system = sys,
                    .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
                    .max_length_A = length_bound_A,
                    .refine_beam_center = true,
                    .refine_distance_mm =  false,
                    .refine_detector_angles = true,
                    .refine_rotation_axis = true,
                    .index_ice_rings = experiment.GetIndexingSettings().GetIndexIceRings(),
                    .axis = orig_axis
                };
                if (d.crystal_system == gemmi::CrystalSystem::Trigonal)
                    d.crystal_system = gemmi::CrystalSystem::Hexagonal;
                if (d.crystal_system == gemmi::CrystalSystem::Monoclinic)
                    d.latt.ReorderMonoclinic();
                return d;
            };

            // Refine under the metric crystal system, and score by the fraction indexed under the
            // refined geometry (the reliable discriminator - the raw pre-refinement fraction is not).
            XtalOptimizerData data = make_data(sr.conventional, sr.system);
            bool ok = XtalOptimizer(data, v_);
            float frac = ok ? IndexedFraction(data.latt, accumulate(data.geom, data.axis), index_tol) : 0.0f;

            // Pseudo-symmetry guard (de-novo only - never override a user-fixed space group): also refine
            // unconstrained (triclinic) on the primitive cell. Adopt it only if it indexes CLEARLY more
            // than the constrained cell - a false promotion (a near-90 pseudo cell forced to ideal angles +
            // a bogus centering) misplaces most reflections, whereas genuine higher symmetry ties.
            // Preferring the constrained cell on a near-tie keeps the real symmetry/centering; the
            // intensities settle the final space group.
            if (!sg && sr.system != gemmi::CrystalSystem::Triclinic) {
                XtalOptimizerData data_tri = make_data(sr.primitive_reduced, gemmi::CrystalSystem::Triclinic);
                if (XtalOptimizer(data_tri, v_)) {
                    const float frac_tri = IndexedFraction(data_tri.latt, accumulate(data_tri.geom, data_tri.axis), index_tol);
                    // Demote only when the constrained cell indexes less than HALF of what the free
                    // primitive does: a false promotion misplaces most reflections (CQ066 ratio ~0.1),
                    // whereas genuine higher symmetry (incl. R-centred) indexes comparably (ratio ~0.7).
                    if (frac_tri > 0.3f && frac < 0.5f * frac_tri) {
                        data = data_tri;
                        ok = true;
                        frac = frac_tri;
                        sr.system = gemmi::CrystalSystem::Triclinic;
                        sr.centering = 'P';
                        sr.conventional = sr.primitive_reduced;
                        sr.reindex = gemmi::Mat33(1, 0, 0, 0, 1, 0, 0, 0, 1);
                    }
                }
            }

            if (!ok)
                continue;

            // Prefer the indexer's earlier (primary) candidate; adopt a later one only if it indexes
            // clearly more AND indexes reasonably well in absolute terms. The absolute floor stops a
            // marginally-higher alternative from displacing the primary when both index poorly (e.g. a
            // twin, where the accumulated-spot fraction is a noisy proxy) - only a decisively better
            // cell (a superstructure's true cell vs its sublattice) takes over.
            if (!have_best || (frac > best_frac + 0.05f && frac > 0.15f)) {
                best_frac = frac;
                have_best = true;
                best_data = data;
                best_sr = sr;
                best_ci = ci;
            }
        }

        if (have_best) {
            search_result_ = best_sr;
            indexed_lattice = best_data.latt;
            updated_geom_ = best_data.geom;
            axis_ = best_data.axis;
        }

        // Extra (twin) lattices: MultiLatticeSearch derives each rotation by relating the FFT's primary
        // lattice[0] to its near-copies, so only apply it when the chosen cell IS that primary. If a
        // widened alternative won (a superstructure/large cell), lattice[0] is a different (sublattice)
        // metric and its rotations would misorient the chosen cell.
        if (have_best && best_ci == 0 && indexer_result.lattice.size() > 1) {
            auto ml_latt = MultiLatticeSearch(indexer_result.lattice);
            for (auto &l : ml_latt) {
                if (extra_lattices_.size() >= experiment.GetIndexingSettings().GetMaxExtraLattices())
                    break;

                // Ignore lattices oriented by less than 3.0 degree
                if (l.rotation_vector.Length() < 3.0 * PI / 180.0)
                    continue;

                RotMatrix rot(l.rotation_vector.Length(), l.rotation_vector.Normalize());

                XtalOptimizerData data_multi{
                    .geom = experiment_copy.GetDiffractionGeometry(),
                    .latt = indexed_lattice->Multiply(rot),
                    .crystal_system = search_result_.system,
                    .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
                    .refine_beam_center = false,
                    .refine_distance_mm =  false,
                    .refine_detector_angles = false,
                    .refine_unit_cell = false,
                    .refine_rotation_axis = false,
                    .index_ice_rings = experiment.GetIndexingSettings().GetIndexIceRings(),
                    .axis = axis_
                };

                // Quick refinement: orientation only. Cell size/angles, beam center,
                // detector angles and rotation axis are all kept from the first lattice.
                // XtalOptimizer always refines orientation; everything else is frozen above.
                XtalOptimizer(data_multi, v_);

                extra_lattices_.push_back(data_multi.latt);
            }
        }
    }
}

void RotationIndexer::ProcessImage(int64_t image, const std::vector<SpotToSave> &spots,
                                   std::optional<float> angle_deg) {
    std::unique_lock ul(m);

    // For non-rotation just ignore the whole procedure
    if (!axis_)
        return;

    // Guard: `image` is a slot in [0, image count); a bad index (e.g. a global number for a subset
    // run) must not corrupt memory.
    if (image < 0 || image >= static_cast<int64_t>(v_.size()))
        return;

    if (accumulated_spots >= max_spots)
        return;

    if (indexed_lattice)
        return;

    angle_deg_[image] = angle_deg;
    v_[image].reserve(spots.size());

    for (const auto &s: spots) {
        if (index_ice_rings || !s.ice_ring)
            v_[image].emplace_back(s);
    }

    // truncate spots, so we don't get above max_spots (total) and max_spots_per_image (for this image)
    size_t max_spots_limit = std::min(max_spots_per_image, max_spots - accumulated_spots);

    if (v_[image].size() > max_spots_limit) {
        std::ranges::nth_element(v_[image], v_[image].begin() + max_spots_limit,
                                 [](const SpotToSave &a, const SpotToSave &b) {
                                     return a.intensity > b.intensity;
                                 }
        );

        v_[image].resize(max_spots_limit);
    }

    accumulated_spots += v_[image].size();
}

std::optional<RotationIndexerResult> RotationIndexer::GetLattice() const {
    std::unique_lock ul(m);

    if (!indexed_lattice)
        return {};
    return RotationIndexerResult{
        .lattice = indexed_lattice.value(),
        .extra_lattices = extra_lattices_,
        .search_result = search_result_,
        .geom = updated_geom_,
        .axis = axis_,
    };
}

void RotationIndexer::ForceResult(const RotationIndexerResult &result) {
    std::unique_lock ul(m);
    indexed_lattice = result.lattice;
    extra_lattices_ = result.extra_lattices;
    search_result_ = result.search_result;
    updated_geom_ = result.geom;
    axis_ = result.axis;
}

bool RotationIndexer::AccumulationFull() const {
    std::unique_lock ul(m);
    return accumulated_spots >= max_spots;
}

void RotationIndexer::ForceLattice(const CrystalLattice &lattice) {
    indexed_lattice = lattice;
    auto sg_num = experiment.GetSpaceGroupNumber().value_or(1);
    auto sg = gemmi::find_spacegroup_by_number(sg_num);
    if (sg != nullptr) {
        search_result_ = LatticeSearchResult{
            .niggli_class = 0, // Since Niggli class was not searched for, we don't know which one
            .conventional = lattice, // If lattice provided, it is for now primitive == conventional
            .system = sg->crystal_system(),
            .centering = sg->centring_type(),
        };
    } else
        search_result_ = LatticeSearchResult{
            .niggli_class = 0, // Since Niggli class was not searched for, we don't know which one
            .conventional = lattice, // If lattice provided, it is for now primitive == conventional
            .system = gemmi::CrystalSystem::Triclinic,
            .centering = 'P',
        };
}