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

#include "RefinementNcs.h"

#include <algorithm>
#include <cmath>
#include <limits>
#include <map>
#include <set>
#include <unordered_map>

#include "gemmi/polyheur.hpp"   // check_polymer_type
#include "gemmi/qcp.hpp"        // superpose_positions
#include "gemmi/seqalign.hpp"   // align_string_sequences

#include "RefinementMath.h"

namespace {

// The distances restrained: those shorter than this in either copy, the bonded and angle-related pairs left out.
// Following Nicholls, Long & Murshudov (2012) Acta Cryst. D68, 404-417
constexpr double NCS_DISTANCE_MAX = 4.2;   // A
// Where a pair's restraint lets go (NcsDistanceTerm): copies whose distance differs by this much are taken to differ.
constexpr float NCS_RELEASE = 0.5f;        // A
// Copies: more than this many residues aligned, at this identity at least (of the longer chain).
// Following Murshudov et al. (2011) Acta Cryst. D67, 355-367
constexpr int NCS_MIN_ALIGNED = 15;
constexpr double NCS_MIN_IDENTITY = 80.0;  // %
// and their C-alpha (P) atoms no further apart than this after superposition: the same sequence further away (a
// quasi-equivalent subunit of a virus capsid) is another molecule.
// Following Murshudov et al. (2011) Acta Cryst. D67, 355-367
constexpr double NCS_MAX_RMSD = 2.5;       // A

// One polymer chain: its residues and, per residue, its atoms' indices by name.
struct NcsChain {
    std::string name;
    gemmi::PolymerType type = gemmi::PolymerType::Unknown;
    std::vector<std::string> sequence;
    std::vector<int> residue;                           // numbered over the model
    std::vector<std::map<std::string, int>> atoms;      // name -> atom index
    std::vector<int> anchor;                            // CA or P, -1 for none
};

// Two copies' corresponding atoms (aligned residues of the same name, atoms of the same name), the residue of each in
// the first copy, and their anchors.
void Correspond(const NcsChain &a, const NcsChain &b, std::vector<std::array<int, 2>> &atoms,
                std::vector<int> &residue_of, std::vector<std::array<int, 2>> &anchors) {
    const gemmi::AlignmentResult al = gemmi::align_string_sequences(a.sequence, b.sequence, {},
                                                                    gemmi::AlignmentScoring::simple());
    size_t i = 0, j = 0;
    for (const gemmi::AlignmentResult::Item &item : al.cigar)
        for (uint32_t k = 0; k < item.len(); k++) {
            const char op = item.op();
            if (op == 'M' && a.sequence[i] == b.sequence[j]) {
                for (const auto &[name, ia] : a.atoms[i]) {
                    auto it = b.atoms[j].find(name);
                    if (it != b.atoms[j].end()) {
                        atoms.push_back({ia, it->second});
                        residue_of.push_back(a.residue[i]);
                    }
                }
                if (a.anchor[i] >= 0 && b.anchor[j] >= 0)
                    anchors.push_back({a.anchor[i], b.anchor[j]});
            }
            if (op == 'M' || op == 'I')
                i++;
            if (op == 'M' || op == 'D')
                j++;
        }
}

// The C-alpha (P) rmsd of two chains after superposition where their sequences make them copies of one molecule,
// infinity where they do not.
double CopyRmsd(const NcsChain &a, const NcsChain &b, const std::vector<gemmi::Position> &pos) {
    const double none = std::numeric_limits<double>::infinity();
    if (a.type != b.type || a.sequence.size() <= NCS_MIN_ALIGNED || b.sequence.size() <= NCS_MIN_ALIGNED)
        return none;
    const gemmi::AlignmentResult al = gemmi::align_string_sequences(a.sequence, b.sequence, {},
                                                                    gemmi::AlignmentScoring::simple());
    // identity over the longer chain: over the shorter one, a short RNA aligned with gaps into a long one finds most of
    // its four letters matched
    const size_t longer = std::max(a.sequence.size(), b.sequence.size());
    if (al.match_count <= NCS_MIN_ALIGNED || 100.0 * al.match_count < NCS_MIN_IDENTITY * double(longer))
        return none;
    std::vector<std::array<int, 2>> atoms, anchors;
    std::vector<int> residue_of;
    Correspond(a, b, atoms, residue_of, anchors);
    if (anchors.size() < 3)
        return none;
    std::vector<gemmi::Position> pa, pb;
    for (const std::array<int, 2> &p : anchors) {
        pa.push_back(pos[p[0]]);
        pb.push_back(pos[p[1]]);
    }
    return gemmi::superpose_positions(pa.data(), pb.data(), pa.size(), nullptr).rmsd;
}

// The pairs (i, j), i < j, of `atoms` closer than `reach` at `pos`, by a grid of cells of that size.
std::set<std::pair<int, int>> PairsWithin(const std::vector<int> &atoms, const std::vector<gemmi::Position> &pos,
                                          double reach) {
    std::map<std::array<int, 3>, std::vector<int>> cells;   // cell -> positions in `atoms`
    auto cell_of = [&](const gemmi::Position &p) {
        return std::array<int, 3>{static_cast<int>(std::floor(p.x / reach)), static_cast<int>(std::floor(p.y / reach)),
                                  static_cast<int>(std::floor(p.z / reach))};
    };
    for (size_t k = 0; k < atoms.size(); k++)
        cells[cell_of(pos[atoms[k]])].push_back(static_cast<int>(k));
    std::set<std::pair<int, int>> out;
    for (size_t k = 0; k < atoms.size(); k++) {
        const std::array<int, 3> c = cell_of(pos[atoms[k]]);
        for (int dx = -1; dx <= 1; dx++)
            for (int dy = -1; dy <= 1; dy++)
                for (int dz = -1; dz <= 1; dz++) {
                    auto it = cells.find({c[0] + dx, c[1] + dy, c[2] + dz});
                    if (it == cells.end())
                        continue;
                    for (int l : it->second)
                        if (atoms[l] > atoms[k] && pos[atoms[k]].dist(pos[atoms[l]]) < reach)
                            out.insert({atoms[k], atoms[l]});
                }
    }
    return out;
}

}  // namespace

NcsCopies FindNcsCopies(const gemmi::Model &model, const RefinementTopology &topo) {
    NcsCopies out;
    std::vector<gemmi::Position> pos;
    std::vector<NcsChain> chains;
    int atom = 0, residue = 0;
    for (const gemmi::Chain &ch : model.chains) {
        NcsChain c;
        c.name = ch.name;
        const gemmi::ConstResidueSpan polymer = ch.get_polymer();
        c.type = gemmi::check_polymer_type(polymer);
        const bool na = gemmi::is_polynucleotide(c.type);
        for (const gemmi::Residue &r : ch.residues) {
            const bool in_polymer = !polymer.empty() && &r >= &polymer.front() && &r <= &polymer.back();
            std::map<std::string, int> names;
            int anchor = -1;
            for (const gemmi::Atom &a : r.atoms) {
                pos.push_back(a.pos);
                if (in_polymer && !a.is_hydrogen() && a.altloc == '\0' && a.occ > 0 && topo.refinable[atom]) {
                    names[a.name] = atom;
                    if (a.name == (na ? "P" : "CA"))
                        anchor = atom;
                }
                atom++;
            }
            if (in_polymer) {
                c.sequence.push_back(r.name);
                c.residue.push_back(residue);
                c.atoms.push_back(std::move(names));
                c.anchor.push_back(anchor);
            }
            residue++;
        }
        if (gemmi::is_polypeptide(c.type) || gemmi::is_polynucleotide(c.type))
            chains.push_back(std::move(c));
    }

    // The molecules: each chain joins the first group whose first chain it is a copy of.
    std::vector<std::vector<int>> groups;
    for (size_t k = 0; k < chains.size(); k++) {
        bool placed = false;
        for (std::vector<int> &g : groups)
            if (CopyRmsd(chains[g[0]], chains[k], pos) <= NCS_MAX_RMSD) {
                g.push_back(static_cast<int>(k));
                placed = true;
                break;
            }
        if (!placed)
            groups.push_back({static_cast<int>(k)});
    }
    // The distances the dictionary fixes: bonded atoms and those bonded to a common atom.
    std::set<std::pair<int, int>> fixed;
    for (const RestraintRecord &r : topo.geometry.records)
        if (r.kind == static_cast<int>(RestraintKind::Bond))
            fixed.insert(std::minmax(r.atom[0], r.atom[1]));
        else if (r.kind == static_cast<int>(RestraintKind::Angle))
            fixed.insert(std::minmax(r.atom[0], r.atom[2]));

    for (const std::vector<int> &g : groups) {
        if (g.size() < 2)
            continue;
        out.groups++;
        out.copies += static_cast<int>(g.size());
        std::string names;
        for (int k : g)
            names += (names.empty() ? "" : "=") + chains[k].name;
        out.description += (out.description.empty() ? "" : " ") + names;
        // each copy with the next, and the last with the first where that is a pair not made already
        const size_t n_pairs = g.size() == 2 ? 1 : g.size();
        for (size_t p = 0; p < n_pairs; p++) {
            const NcsChain &a = chains[g[p]], &b = chains[g[(p + 1) % g.size()]];
            std::vector<std::array<int, 2>> corr, anchors;
            std::vector<int> residue_of;
            Correspond(a, b, corr, residue_of, anchors);
            out.anchors.push_back(std::move(anchors));
            std::unordered_map<int, int> to_b, res_of;
            std::vector<int> atoms_a, atoms_b;
            for (size_t k = 0; k < corr.size(); k++) {
                to_b[corr[k][0]] = corr[k][1];
                res_of[corr[k][0]] = residue_of[k];
                atoms_a.push_back(corr[k][0]);
                atoms_b.push_back(corr[k][1]);
            }
            std::set<std::pair<int, int>> pairs = PairsWithin(atoms_a, pos, NCS_DISTANCE_MAX);
            std::unordered_map<int, int> to_a;
            for (const auto &[ia, ib] : to_b)
                to_a[ib] = ia;
            for (const auto &[i, j] : PairsWithin(atoms_b, pos, NCS_DISTANCE_MAX))
                pairs.insert(std::minmax(to_a[i], to_a[j]));
            for (const auto &[i, j] : pairs) {
                if (fixed.count({i, j}) > 0)
                    continue;
                RestraintRecord r;
                r.kind = static_cast<int>(RestraintKind::NcsDistance);
                r.atom[0] = i;
                r.atom[1] = j;
                r.atom[2] = to_b[i];
                r.atom[3] = to_b[j];
                r.ideal = NCS_RELEASE;
                r.sigma = 1.0f;
                out.distances.Add(r, 4);
                out.record_residue.push_back(res_of[i]);
            }
        }
    }
    return out;
}

RestraintSet NcsRestraintsAt(const NcsCopies &ncs, double sigma) {
    RestraintSet set = ncs.distances;
    for (RestraintRecord &r : set.records)
        r.sigma = static_cast<float>(sigma);
    return set;
}

double NcsRmsd(const NcsCopies &ncs, const std::vector<std::array<double, 4>> &xyzb) {
    double sum = 0;
    int n = 0;
    for (const std::vector<std::array<int, 2>> &anchors : ncs.anchors) {
        if (anchors.size() < 3)
            continue;
        std::vector<gemmi::Position> pa, pb;
        for (const std::array<int, 2> &p : anchors) {
            pa.emplace_back(xyzb[p[0]][0], xyzb[p[0]][1], xyzb[p[0]][2]);
            pb.emplace_back(xyzb[p[1]][0], xyzb[p[1]][1], xyzb[p[1]][2]);
        }
        sum += gemmi::superpose_positions(pa.data(), pb.data(), pa.size(), nullptr).rmsd;
        n++;
    }
    return n > 0 ? sum / n : NAN;
}

int NcsReleasedResidues(const NcsCopies &ncs, const std::vector<std::array<double, 4>> &xyzb) {
    std::map<int, std::array<int, 2>> count;   // residue -> restraints released, all
    for (size_t k = 0; k < ncs.distances.records.size(); k++) {
        const RestraintRecord &r = ncs.distances.records[k];
        const double delta = RestraintValue(ncs.distances, r, xyzb);
        std::array<int, 2> &c = count[ncs.record_residue[k]];
        c[0] += std::fabs(delta) > r.ideal;
        c[1]++;
    }
    int released = 0;
    for (const auto &[res, c] : count)
        released += 2 * c[0] > c[1];
    return released;
}