#!/usr/bin/env python
"""Ball-and-stick renders of two PROTAC conformations from a replica trajectory:
the most extended early frame beside the most collapsed late frame. Both frames
are also written out as PDBs for PyMOL / VMD / ChimeraX.

Usage:
    python render_3d.py --rep egfrvhl_rank4/chloroform/rep1
"""
import argparse, glob, os
import numpy as np
import mdtraj as md
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Line3DCollection

# CPK element colors and relative sizes
CPK = {"C": ("#4d4d4d", 55), "N": ("#3050f8", 70), "O": ("#ff2d2d", 70),
       "S": ("#e6c200", 95), "Cl": ("#1fd01f", 95), "F": ("#6fd0c0", 60),
       "H": ("#dddddd", 20)}

def draw(ax, xyz, syms, bonds, title):
    xyz = xyz - xyz.mean(0)
    segs = [[xyz[i], xyz[j]] for i, j in bonds]
    ax.add_collection3d(Line3DCollection(segs, colors="#888888", linewidths=1.6))
    for el in set(syms):
        m = [k for k, s in enumerate(syms) if s == el]
        c, sz = CPK.get(el, ("#cc44cc", 60))
        ax.scatter(xyz[m, 0], xyz[m, 1], xyz[m, 2], c=c, s=sz,
                   edgecolors="white", linewidths=0.4, depthshade=True)
    ax.set_title(title, fontsize=12, weight="bold", color="#222", pad=0)
    r = np.abs(xyz).max() * 1.05
    ax.set_xlim(-r, r); ax.set_ylim(-r, r); ax.set_zlim(-r, r)
    ax.set_box_aspect((1, 1, 1)); ax.axis("off")

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--rep", required=True, help="replica directory, e.g. <name>/<solvent>/rep1")
    p.add_argument("--name", default=None, help="run name (default: the run directory's name)")
    p.add_argument("--report-ps", type=float, default=10.0, help="DCD frame interval")
    a = p.parse_args()
    rep = a.rep.rstrip("/")
    name = a.name or os.path.basename(os.path.dirname(os.path.dirname(rep)))

    solvated = glob.glob(os.path.join(rep, os.pardir, "*_solvated.pdb"))[0]
    traj = md.load(os.path.join(rep, f"{name}_prod.dcd"), top=solvated)
    top = traj.topology
    # the PROTAC is the one non-water residue with many atoms; water and ions are tiny
    lig = [at.index for res in top.residues
           if not res.is_water and res.n_atoms > 5 for at in res.atoms]
    heavy = [i for i in lig if top.atom(i).element.symbol != "H"]
    lset = set(heavy)
    remap = {g: k for k, g in enumerate(heavy)}
    bonds = [(remap[b0.index], remap[b1.index]) for b0, b1 in top.bonds
             if b0.index in lset and b1.index in lset]
    syms = [top.atom(i).element.symbol for i in heavy]

    sub = traj.atom_slice(heavy)
    rg = md.compute_rg(sub)  # nm
    n = len(rg); early = slice(0, n // 3); late = slice(2 * n // 3, n)
    fe = np.arange(n)[early][np.argmax(rg[early])]     # extended
    fc = np.arange(n)[late][np.argmin(rg[late])]       # collapsed
    xyz_A = sub.xyz * 10.0  # nm -> Angstrom
    ns = a.report_ps / 1000.0

    fig = plt.figure(figsize=(13, 6.5), dpi=150); fig.patch.set_facecolor("white")
    ax1 = fig.add_subplot(121, projection="3d")
    draw(ax1, xyz_A[fe], syms, bonds, f"Extended  (t={fe*ns:.1f} ns,  Rg={rg[fe]:.2f} nm)")
    ax2 = fig.add_subplot(122, projection="3d")
    draw(ax2, xyz_A[fc], syms, bonds, f"Collapsed  (t={fc*ns:.1f} ns,  Rg={rg[fc]:.2f} nm)")
    fig.suptitle(f"Free-PROTAC conformations from MD - {name}  (heavy atoms, H hidden)",
                 fontsize=13, weight="bold", color="#222", y=0.97)
    out = os.path.join(rep, f"{name}_3d.png")
    fig.savefig(out, facecolor="white", bbox_inches="tight"); print("wrote", out)

    sub[fe].save_pdb(os.path.join(rep, f"{name}_extended.pdb"))
    sub[fc].save_pdb(os.path.join(rep, f"{name}_collapsed.pdb"))
    print(f"extended frame {fe} (Rg {rg[fe]:.2f}), collapsed frame {fc} (Rg {rg[fc]:.2f})")

if __name__ == "__main__":
    main()