"""Per-residue interface deviations behind the CDK2/CKS1B iRMSD.

Reproduces the iRMSD that DockQ 2.1.3 reports for AF-0000000210768563 against
PDB 1BUH (3.256 A) and lists the interface residues that deviate most after the
interface backbone superposition:

    python interface_deviation.py

Needs the DockQ 2.1.3 environment (pip install DockQ==2.1.3), which provides
Biopython and NumPy. Interface residues follow DockQ: residues with any heavy
atom within 10 A of the partner chain in the reference, backbone atoms N, CA,
C and O, sequence-aligned between model and reference.
"""

from pathlib import Path

import numpy as np
from Bio import pairwise2
from Bio.SeqUtils import seq1
from Bio.SVDSuperimposer import SVDSuperimposer
from DockQ.DockQ import load_PDB

HERE = Path(__file__).resolve().parent
BACKBONE = ("N", "CA", "C", "O")
INTERFACE_CUTOFF = 10.0


def residues(chain):
    return [r for r in chain if r.id[0] == " " and all(a in r for a in BACKBONE)]


def aligned_pairs(model_chain, native_chain):
    model_res, native_res = residues(model_chain), residues(native_chain)
    alignment = pairwise2.align.globalxx(
        "".join(seq1(r.get_resname()) for r in model_res),
        "".join(seq1(r.get_resname()) for r in native_res),
        one_alignment_only=True,
    )[0]
    pairs, i, j = [], 0, 0
    for a, b in zip(alignment.seqA, alignment.seqB):
        if a != "-" and b != "-":
            pairs.append((model_res[i], native_res[j]))
        i += a != "-"
        j += b != "-"
    return pairs


def min_distance(res1, res2):
    a = np.array([atom.coord for atom in res1 if atom.element != "H"])
    b = np.array([atom.coord for atom in res2 if atom.element != "H"])
    return np.sqrt(((a[:, None] - b[None]) ** 2).sum(-1)).min()


def main():
    model = load_PDB(str(HERE / "AF-0000000210768563-model_v1.cif"))
    native = load_PDB(str(HERE / "1BUH.cif"))
    cdk2 = aligned_pairs(model["A"], native["A"])
    cks1 = aligned_pairs(model["B"], native["B"])

    def interface(pairs, partner):
        return [
            (m, n) for m, n in pairs
            if any(min_distance(n, other) < INTERFACE_CUTOFF for _, other in partner)
        ]

    interface_pairs = [("CDK2", p) for p in interface(cdk2, cks1)] + [
        ("CKS1B", p) for p in interface(cks1, cdk2)
    ]
    model_xyz = np.array([m[a].coord for _, (m, _) in interface_pairs for a in BACKBONE])
    native_xyz = np.array([n[a].coord for _, (_, n) in interface_pairs for a in BACKBONE])

    sup = SVDSuperimposer()
    sup.set(native_xyz, model_xyz)
    sup.run()
    print(f"iRMSD {sup.get_rms():.3f} A over {len(interface_pairs)} interface residues")

    rot, tran = sup.get_rotran()
    moved = model_xyz @ rot + tran
    per_residue = np.sqrt(((moved - native_xyz) ** 2).sum(-1)).reshape(-1, 4).mean(1)
    ranked = sorted(zip(interface_pairs, per_residue), key=lambda item: -item[1])
    for (chain, (_, native_res)), deviation in ranked[:10]:
        print(f"{chain} {native_res.get_resname()}{native_res.id[1]} {deviation:.1f} A")


if __name__ == "__main__":
    main()