#!/usr/bin/env python
"""
Conformational analysis of free-PROTAC MD replicas.

Reads all rep*/ trajectories under a solvent directory and computes, per replica
and pooled: radius of gyration (Rg), RMSD to the first frame (convergence),
solvent-accessible 3D polar surface area (3D PSA) and intramolecular H-bond
count, then PCA + K-means clustering into conformational states with populations
and saved centroid structures. The dashboard PNG carries the literature
yardsticks (7 A Rg permeability cutoff, 3D PSA ~200-210 A^2).

Time axes assume 10 ps per frame, protac_md.py's default --report-ps; a run made
with a different --report-ps gets ns labels scaled by that ratio.

Usage:
    python analyze_conformations.py --dir egfrvhl_rank4/chloroform --name egfrvhl_rank4
"""
import argparse, glob, os
import numpy as np
import mdtraj as md
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.cluster.vq import kmeans2

RG_CUT = 0.70          # nm, passive-permeability cutoff (Atilaw 2021)
PSA_LO, PSA_HI = 200.0, 210.0   # A^2; low-moderate permeability band
COLORS = ["#0072B2", "#D55E00", "#009E73", "#CC79A7", "#E69F00"]
INK, MUTED, GRID = "#222", "#666", "#DDD"


def ligand_indices(top):
    return [at.index for res in top.residues
            if not res.is_water and res.n_atoms > 5 for at in res.atoms]


def polar_slice_indices(top_lig):
    """Indices (within the ligand-only topology) of polar atoms: N, O, and H on N/O."""
    nbr = {}
    for b in top_lig.bonds:
        nbr.setdefault(b[0].index, []).append(b[1].index)
        nbr.setdefault(b[1].index, []).append(b[0].index)
    polar = []
    for at in top_lig.atoms:
        s = at.element.symbol
        if s in ("N", "O"):
            polar.append(at.index)
        elif s == "H" and any(top_lig.atom(j).element.symbol in ("N", "O")
                               for j in nbr.get(at.index, [])):
            polar.append(at.index)
    return np.array(sorted(set(polar)))


def load_replicas(solvent_dir, name):
    top_pdbs = glob.glob(os.path.join(solvent_dir, f"{name}_*_solvated.pdb"))
    if not top_pdbs:
        raise SystemExit(f"no *_solvated.pdb in {solvent_dir}")
    top = top_pdbs[0]
    reps = []
    for rd in sorted(glob.glob(os.path.join(solvent_dir, "rep*"))):
        dcd = os.path.join(rd, f"{name}_prod.dcd")
        if os.path.exists(dcd):
            reps.append((os.path.basename(rd), md.load(dcd, top=top)))
    if not reps:
        raise SystemExit(f"no rep*/{name}_prod.dcd trajectories in {solvent_dir}")
    return reps


def per_replica_metrics(traj):
    top = traj.topology
    lig = ligand_indices(top)
    ligtraj = traj.atom_slice(lig)
    ligtraj.center_coordinates()
    rg = md.compute_rg(ligtraj)                                   # nm
    rmsd = md.rmsd(ligtraj, ligtraj, 0) * 10.0                    # A, vs frame 0
    sasa = md.shrake_rupley(ligtraj, mode="atom")                # nm^2 per atom
    polar = polar_slice_indices(ligtraj.topology)
    psa = sasa[:, polar].sum(axis=1) * 100.0                      # A^2
    hb = [len(x) for x in md.wernet_nilsson(ligtraj)]            # intramolecular H-bonds/frame
    return dict(rg=rg, rmsd=rmsd, psa=psa, hb=np.array(hb, float), ligtraj=ligtraj)


def cluster(ligtrajs, k=5, seed=0):
    pooled = md.join(ligtrajs)
    pooled.superpose(pooled, 0)
    X = pooled.xyz.reshape(pooled.n_frames, -1)
    X = X - X.mean(0)
    # PCA via SVD, keeping components up to 90% variance (cap 10)
    U, S, Vt = np.linalg.svd(X, full_matrices=False)
    var = (S ** 2) / (S ** 2).sum()
    ncomp = min(10, max(2, int(np.searchsorted(np.cumsum(var), 0.90) + 1)))
    proj = U[:, :ncomp] * S[:ncomp]
    cent, labels = kmeans2(proj, k, minit="++", seed=seed)
    pops = np.array([(labels == i).mean() for i in range(k)])
    # global frame index closest to each cluster centroid, in PC space
    cframes = []
    for i in range(k):
        idx = np.where(labels == i)[0]
        cframes.append(int(idx[np.argmin(((proj[idx] - cent[i]) ** 2).sum(1))]) if len(idx) else -1)
    return pooled, labels, pops, cframes, proj


def style(ax, title, xlabel, ylabel):
    ax.set_title(title, fontsize=11, color=INK, loc="left", weight="bold", pad=6)
    ax.set_xlabel(xlabel, fontsize=9, color=MUTED); ax.set_ylabel(ylabel, fontsize=9, color=MUTED)
    ax.grid(True, color=GRID, lw=0.6, alpha=0.8); ax.set_axisbelow(True)
    for s in ("top", "right"): ax.spines[s].set_visible(False)
    for s in ("left", "bottom"): ax.spines[s].set_color(GRID)
    ax.tick_params(colors=MUTED, labelsize=8)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", required=True, help="solvent dir containing rep*/")
    ap.add_argument("--name", required=True)
    ap.add_argument("--clusters", type=int, default=5)
    a = ap.parse_args()

    reps = load_replicas(a.dir, a.name)
    solvent = os.path.basename(a.dir.rstrip("/"))
    print(f"[analyze] {a.name} / {solvent}: {len(reps)} replicas")

    M = {rn: per_replica_metrics(tr) for rn, tr in reps}
    ligtrajs = [M[rn]["ligtraj"] for rn, _ in reps]
    pooled, labels, pops, cframes, proj = cluster(ligtrajs, a.clusters)

    rg_all = np.concatenate([M[rn]["rg"] for rn, _ in reps])
    psa_all = np.concatenate([M[rn]["psa"] for rn, _ in reps])
    hb_all = np.concatenate([M[rn]["hb"] for rn, _ in reps])
    rep_means = [M[rn]["rg"].mean() for rn, _ in reps]

    # ---- dashboard ----
    fig, axes = plt.subplots(2, 3, figsize=(15, 8.5), dpi=140)
    fig.patch.set_facecolor("white")

    ax = axes[0, 0]
    for i, (rn, _) in enumerate(reps):
        t = np.arange(len(M[rn]["rg"])) * 0.01  # ns; assumes protac_md.py default --report-ps 10
        ax.plot(t, M[rn]["rg"], color=COLORS[i % 5], lw=0.8, alpha=0.85, label=rn)
    ax.axhline(RG_CUT, color="k", ls="--", lw=1); ax.text(0.99, RG_CUT+0.003, "7 A permeability cutoff",
             transform=ax.get_yaxis_transform(), ha="right", fontsize=7, color=MUTED)
    style(ax, "Radius of gyration", "time (ns)", "Rg (nm)"); ax.legend(fontsize=7, frameon=False)

    ax = axes[0, 1]
    for i, (rn, _) in enumerate(reps):
        t = np.arange(len(M[rn]["rmsd"])) * 0.01
        ax.plot(t, M[rn]["rmsd"], color=COLORS[i % 5], lw=0.8, alpha=0.85)
    style(ax, "RMSD to start (convergence)", "time (ns)", "RMSD (A)")

    ax = axes[0, 2]
    ax.hexbin(rg_all, psa_all, gridsize=35, cmap="Blues", mincnt=1)
    ax.axvline(RG_CUT, color="k", ls="--", lw=1)
    ax.axhspan(PSA_LO, PSA_HI, color="orange", alpha=0.15)
    ax.text(0.02, PSA_HI, " PSA ~200-210 A^2", transform=ax.get_yaxis_transform(),
            fontsize=7, color=MUTED, va="bottom")
    style(ax, "3D PSA vs Rg (folding->permeability)", "Rg (nm)", "3D PSA (A^2)")

    ax = axes[1, 0]
    order = np.argsort(-pops)
    ax.bar(range(len(pops)), pops[order] * 100, color=[COLORS[i % 5] for i in range(len(pops))])
    for r, i in enumerate(order):
        idx = labels == i
        if idx.any():
            ax.text(r, pops[i]*100 + 0.5, f"Rg {md.compute_rg(pooled[idx]).mean():.2f}",
                    ha="center", fontsize=7, color=MUTED)
    style(ax, f"Conformational clusters (k={a.clusters})", "cluster (by population)", "population (%)")

    ax = axes[1, 1]
    ax.hist(hb_all, bins=np.arange(0, hb_all.max()+2)-0.5, color="#0072B2", alpha=0.85, rwidth=0.9)
    style(ax, "Intramolecular H-bonds", "H-bonds per frame", "count")

    ax = axes[1, 2]; ax.axis("off")
    summary = [
        f"PROTAC : {a.name}",
        f"Solvent: {solvent}   Replicas: {len(reps)}",
        f"Frames : {pooled.n_frames} pooled",
        "",
        f"Rg      : {rg_all.mean():.3f} ± {rg_all.std():.3f} nm",
        f"  per-rep means: " + ", ".join(f"{x:.3f}" for x in rep_means),
        f"  replica spread (std of means): {np.std(rep_means):.3f} nm",
        f"  frac below 7 A cutoff: {100*np.mean(rg_all<RG_CUT):.0f}%",
        f"3D PSA  : {psa_all.mean():.0f} ± {psa_all.std():.0f} A^2",
        f"IMHBs   : {hb_all.mean():.1f} ± {hb_all.std():.1f} per frame",
        "",
        "Reference: compact minimum Rg ~5.4-5.6 A;",
        "3D PSA ~190-265 A^2 permeable vs 290-330 poor.",
    ]
    ax.text(0.0, 1.0, "\n".join(summary), va="top", family="monospace", fontsize=9, color=INK)

    fig.suptitle(f"Free-PROTAC conformational analysis - {a.name} in {solvent} "
                 f"({len(reps)}x{len(M[reps[0][0]]['rg'])*0.01:g} ns)",
                 fontsize=13, weight="bold", color=INK, x=0.06, ha="left")
    fig.tight_layout(rect=(0, 0, 1, 0.97))
    out = os.path.join(a.dir, f"{a.name}_analysis.png")
    fig.savefig(out, facecolor="white"); print(f"[analyze] wrote {out}")

    # centroid structures
    for r, i in enumerate(order):
        if cframes[i] >= 0:
            pooled[cframes[i]].save_pdb(os.path.join(a.dir, f"{a.name}_cluster{r+1}_pop{pops[i]*100:.0f}.pdb"))
    print(f"[analyze] Rg {rg_all.mean():.3f} nm | 3D PSA {psa_all.mean():.0f} A^2 | "
          f"IMHB {hb_all.mean():.1f} | replica-mean spread {np.std(rep_means):.3f} nm")


if __name__ == "__main__":
    main()