#!/usr/bin/env python
"""Warhead-to-E3 span distribution from the water ensembles.

Per frame, the distance between the centroid of the warhead core (the group that
occupies the target pocket) and the centroid of the E3 ligand core. In a
productive ternary complex that span has to bridge the two pockets, so the
distribution shows which separations a linker can hold. Water, not chloroform:
the ternary complex forms in the cytosol.

This is a description of reach, not a degradation score; scoring it needs a
target ternary geometry from docking. Replica spread is printed alongside so
poorly converged runs are visible.

Default SMARTS are for the EGFR-VHL series (anilinoquinazoline warhead,
4-hydroxyproline E3 anchor); override with --warhead-smarts / --e3-smarts.

Usage:
    python ternary_span.py egfrvhl_rank4 ms39_ref --out span.png
"""
import argparse, glob, os
import numpy as np
import mdtraj as md
from rdkit import Chem
from scipy.stats import gaussian_kde
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt

WH_SMARTS = "c1cc2ncncc2cc1"            # quinazoline bicyclic, target pocket
E3_SMARTS = "[#6]1[#6][#6](O)[#6][#7]1"  # 4-hydroxypyrrolidine, VHL pocket


def anchors(sdf, wh_smarts, e3_smarts):
    m = Chem.MolFromMolFile(sdf, removeHs=False)
    wh = list(m.GetSubstructMatch(Chem.MolFromSmarts(wh_smarts)))
    e3 = list(m.GetSubstructMatch(Chem.MolFromSmarts(e3_smarts)))
    syms = [a.GetSymbol() for a in m.GetAtoms()]
    return wh, e3, syms


def spans_for_run(run, solvent, wh_smarts, e3_smarts):
    """Per-replica span arrays (A), or None if the run can't be used."""
    name = os.path.basename(run.rstrip("/"))
    tops = glob.glob(os.path.join(run, solvent, "*_solvated.pdb"))
    dcds = sorted(glob.glob(os.path.join(run, solvent, "rep*", f"{name}_prod.dcd")))
    if not tops or not dcds:
        print(f"{name:14s}  no {solvent} data")
        return None
    wh, e3, sdf_syms = anchors(os.path.join(run, f"{name}.sdf"), wh_smarts, e3_smarts)
    if not wh or not e3:
        print(f"{name:14s}  SMARTS match failed (wh={len(wh)} e3={len(e3)})")
        return None
    out = []
    for d in dcds:
        tr = md.load(d, top=tops[0])
        lig = [a.index for r in tr.topology.residues
               if not r.is_water and r.n_atoms > 5 for a in r.atoms]
        sub = tr.atom_slice(lig)
        if [a.element.symbol for a in sub.topology.atoms] != sdf_syms:
            print(f"{name:14s}  SDF/MD atom-order mismatch, skipping")
            return None
        xyz = sub.xyz * 10.0                       # nm -> A
        out.append(np.linalg.norm(xyz[:, wh, :].mean(1) - xyz[:, e3, :].mean(1), axis=1))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("runs", nargs="+", help="run directories, one per compound")
    ap.add_argument("--solvent", default="water")
    ap.add_argument("--warhead-smarts", default=WH_SMARTS)
    ap.add_argument("--e3-smarts", default=E3_SMARTS)
    ap.add_argument("--out", default="ternary_span.png")
    a = ap.parse_args()

    fig, ax = plt.subplots(figsize=(11, 6.5), dpi=150)
    fig.patch.set_facecolor("white")
    cmap = plt.get_cmap("viridis")
    print(f"{'compound':14s}{'mean span':>10s}{'sd':>6s}{'min':>6s}{'max':>6s}"
          f"{'range':>7s}{'rep spread':>12s}")
    print("-" * 67)
    drawn = 0
    for i, run in enumerate(a.runs):
        label = os.path.basename(run.rstrip("/"))
        per_rep = spans_for_run(run, a.solvent, a.warhead_smarts, a.e3_smarts)
        if not per_rep:
            continue
        s = np.concatenate(per_rep)
        spread = float(np.std([r.mean() for r in per_rep]))
        print(f"{label:14s}{s.mean():10.1f}{s.std():6.1f}{s.min():6.1f}"
              f"{s.max():6.1f}{s.max()-s.min():7.1f}{spread:12.1f}")
        grid = np.linspace(3, 25, 300)
        ax.plot(grid, gaussian_kde(s, bw_method=0.15)(grid), lw=2.2,
                color=cmap(i / max(1, len(a.runs) - 1)),
                label=f"{label} ({s.mean():.0f} A)")
        drawn += 1
    if not drawn:
        return
    ax.set_xlabel("warhead to E3 pharmacophore separation (A)")
    ax.set_ylabel("probability density")
    ax.set_title("Warhead to E3 span\n"
                 f"centroid distance between the pocket-binding groups ({a.solvent})",
                 fontsize=12.5, weight="bold", loc="left")
    ax.legend(fontsize=8.5, frameon=False, ncol=2)
    ax.grid(alpha=0.25)
    fig.tight_layout()
    fig.savefig(a.out, facecolor="white")
    print(f"\nwrote {a.out}")


if __name__ == "__main__":
    main()