#!/usr/bin/env python
"""Rg free-energy profiles and basin populations for one or more finished runs.

Mean Rg does not separate permeable from non-permeable PROTACs in the J. Med.
Chem. 2022 study: their most permeable compound sits mainly just above 8 A,
overlapping their least permeable one at 7-9.5 A. What separates them is whether
a compact minimum near 5.5 A is populated at all. This script computes, per
compound and solvent,

    PMF(Rg) = -kT ln P(Rg)

locates its basins (local minima), and reports each basin's position and
population, plus the fraction of frames below 6 and 7 A.

Ligand selection and the Rg calculation match analyze_conformations.py exactly
(non-water residue with >5 atoms, md.compute_rg on the centred ligand-only
trajectory).

Usage:
    python rg_basins.py egfrvhl_rank4 ms39_ref --solvents chloroform water
"""
import argparse, glob, os
import numpy as np
import mdtraj as md
from scipy.stats import gaussian_kde
from scipy.signal import find_peaks
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt

KT = 0.593  # kcal/mol at 300 K


def pooled_rg(run_dir, solvent):
    """Rg in Angstrom, pooled over replicas."""
    name = os.path.basename(run_dir.rstrip("/"))
    tops = glob.glob(os.path.join(run_dir, solvent, f"{name}_*_solvated.pdb"))
    dcds = sorted(glob.glob(os.path.join(run_dir, solvent, "rep*", f"{name}_prod.dcd")))
    if not tops or not dcds:
        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]
        lt = tr.atom_slice(lig)
        lt.center_coordinates()
        out.append(md.compute_rg(lt) * 10.0)          # nm -> A
    return np.concatenate(out)


def basins(rg, prom_frac=0.05):
    """Basins = peaks of P(Rg). Returns (grid, pmf, [(rg_position, population)])."""
    grid = np.linspace(rg.min() - 0.3, rg.max() + 0.3, 400)
    dens = gaussian_kde(rg, bw_method=0.12)(grid)
    pmf = -KT * np.log(np.clip(dens, 1e-12, None))
    pmf -= pmf.min()
    pk, _ = find_peaks(dens, prominence=prom_frac * dens.max())
    if len(pk) == 0:
        pk = np.array([int(np.argmax(dens))])
    # valley boundaries between adjacent peaks assign frames to basins
    edges = [rg.min() - 1e-6]
    for a, b in zip(pk[:-1], pk[1:]):
        edges.append(grid[a + int(np.argmin(dens[a:b]))])
    edges.append(rg.max() + 1e-6)
    out = []
    for i, p in enumerate(pk):
        pop = float(((rg > edges[i]) & (rg <= edges[i + 1])).mean())
        out.append((float(grid[p]), pop))
    return grid, pmf, out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("runs", nargs="+", help="run directories, one per compound")
    ap.add_argument("--solvents", nargs="+", default=["chloroform", "water"])
    ap.add_argument("--plot-solvent", default="chloroform",
                    help="solvent whose PMFs go into the figure")
    ap.add_argument("--out", default="rg_basins.png")
    a = ap.parse_args()

    print(f"{'compound':22s}{'solv':6s}{'mean':>7s}{'basins (Rg A, pop%)':>44s}{'<6A':>7s}{'<7A':>7s}")
    print("-" * 96)
    results = {}
    for run in a.runs:
        label = os.path.basename(run.rstrip("/"))
        for sol in a.solvents:
            rg = pooled_rg(run, sol)
            if rg is None:
                continue
            grid, pmf, bs = basins(rg)
            results[(label, sol)] = (rg, grid, pmf, bs)
            btxt = "  ".join(f"{p:.2f}({100*w:.0f}%)" for p, w in bs)
            print(f"{label:22s}{sol[:5]:6s}{rg.mean():7.2f}{btxt:>44s}"
                  f"{100*(rg<6.0).mean():6.0f}%{100*(rg<7.0).mean():6.0f}%")

    have = [k for k in results if k[1] == a.plot_solvent]
    if not have:
        return
    fig, ax = plt.subplots(figsize=(11, 6.5), dpi=150)
    fig.patch.set_facecolor("white")
    cmap = plt.get_cmap("viridis")
    for i, key in enumerate(have):
        _, grid, pmf, _ = results[key]
        ax.plot(grid, pmf, lw=2.2, color=cmap(i / max(1, len(have) - 1)), label=key[0])
    ax.axvspan(5.4, 5.7, color="#009E73", alpha=0.12)
    ax.text(5.55, ax.get_ylim()[1] * 0.94, "reference compact\nminimum 5.5-5.6 A",
            ha="center", fontsize=8.5, color="#009E73", weight="bold")
    ax.axvspan(7.0, 9.5, color="#C0392B", alpha=0.07)
    ax.text(8.25, ax.get_ylim()[1] * 0.94, "reference least-permeable\ncompound 7-9.5 A",
            ha="center", fontsize=8.5, color="#C0392B", weight="bold")
    ax.set_xlabel("Radius of gyration (A)")
    ax.set_ylabel("PMF  -kT ln P(Rg)   (kcal/mol)")
    ax.set_title(f"Rg free-energy profiles in {a.plot_solvent}",
                 fontsize=13, weight="bold", loc="left")
    ax.set_ylim(0, 3.2); 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()