#!/usr/bin/env python
"""One-page summary figure for a single replica of a protac_md.py run: the
radius of gyration, potential energy and density traces of the production run,
with the PROTAC 2D structure beside them when --smiles is given. Writes
<rep>/<name>_summary.png.

Usage:
    python plot_results.py --rep egfrvhl_rank4/chloroform/rep1
    python plot_results.py --rep egfrvhl_rank4/chloroform/rep1 --smiles "<SMILES>"
"""
import argparse, csv, glob, io, os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import mdtraj as md
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem.Draw import rdMolDraw2D
from PIL import Image

# Okabe-Ito colorblind-safe hues
BLUE, VERM, GREEN = "#0072B2", "#D55E00", "#009E73"
INK, MUTED, GRID = "#222222", "#666666", "#DDDDDD"


def load_csv(path):
    with open(path) as fh:
        rows = list(csv.reader(fh))
    header = [h.strip().strip('#"').strip() for h in rows[0]]
    data = np.array([[float(x) for x in r] for r in rows[1:] if r])
    col = {h: data[:, i] for i, h in enumerate(header)}

    def find(key):
        return next(c for h, c in col.items() if key.lower() in h.lower())

    t = find("time") / 1000.0          # ps -> ns
    return t - t[0], find("potential"), find("temperature"), find("density")


def load_rg(rep_dir, name):
    """Rg of the PROTAC (nm), same selection as analyze_conformations.py."""
    tops = glob.glob(os.path.join(rep_dir, os.pardir, "*_solvated.pdb"))
    traj = md.load(os.path.join(rep_dir, f"{name}_prod.dcd"), top=tops[0])
    lig = [a.index for r in traj.topology.residues
           if not r.is_water and r.n_atoms > 5 for a in r.atoms]
    sub = traj.atom_slice(lig)
    sub.center_coordinates()
    return md.compute_rg(sub)


def mol_image(smiles, px=700):
    mol = Chem.MolFromSmiles(smiles)
    AllChem.Compute2DCoords(mol)
    d = rdMolDraw2D.MolDraw2DCairo(px, px)
    d.drawOptions().bondLineWidth = 2
    d.drawOptions().padding = 0.06
    rdMolDraw2D.PrepareAndDrawMolecule(d, mol)
    d.FinishDrawing()
    return Image.open(io.BytesIO(d.GetDrawingText()))


def style(ax, title, ylabel):
    ax.set_title(title, fontsize=11, color=INK, loc="left", pad=6, weight="bold")
    ax.set_ylabel(ylabel, fontsize=9, color=MUTED)
    ax.set_xlabel("production time (ns)", 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, length=3)


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("--smiles", default=None, help="draw this structure in the left panel")
    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)))
    solvent = os.path.basename(os.path.dirname(rep))

    t, pe, temp, dens = load_csv(os.path.join(rep, f"{name}_prod.csv"))
    rg = load_rg(rep, name)
    rg_t = np.arange(len(rg)) * a.report_ps / 1000.0       # frame index -> ns

    fig = plt.figure(figsize=(13, 7.5), dpi=150)
    fig.patch.set_facecolor("white")
    ncol = 3 if a.smiles else 2
    gs = GridSpec(2, ncol, figure=fig, hspace=0.42, wspace=0.30,
                  left=0.06, right=0.98, top=0.90, bottom=0.09)
    off = 0

    if a.smiles:
        axm = fig.add_subplot(gs[:, 0])
        axm.axis("off")
        axm.imshow(mol_image(a.smiles))
        axm.set_title(name, fontsize=11, color=INK, loc="left", weight="bold", pad=6)
        off = 1

    ax1 = fig.add_subplot(gs[0, off:])
    ax1.plot(rg_t, rg, color=BLUE, lw=1.2)
    m = rg.mean()
    ax1.axhline(m, color=MUTED, lw=1.0, ls="--")
    ax1.annotate(f"mean {m:.2f} nm", (rg_t[-1], m), color=MUTED, fontsize=8,
                 ha="right", va="bottom", xytext=(0, 3), textcoords="offset points")
    style(ax1, "Radius of gyration", "Rg (nm)")

    ax2 = fig.add_subplot(gs[1, off])
    ax2.plot(t, pe / 1000.0, color=VERM, lw=1.0)
    style(ax2, "Potential energy", "PE (10^3 kJ/mol)")

    ax3 = fig.add_subplot(gs[1, off + 1])
    ax3.plot(t, dens, color=GREEN, lw=1.0)
    style(ax3, "System density", "density (g/mL)")

    fig.suptitle(f"Free-PROTAC MD summary - {name} in {solvent} "
                 f"({t[-1]:.0f} ns, 300 K / 1 bar)",
                 fontsize=13, color=INK, x=0.06, ha="left", weight="bold")

    out = os.path.join(rep, f"{name}_summary.png")
    fig.savefig(out, facecolor="white")
    print("wrote", out)
    print(f"Rg {rg.mean():.3f}±{rg.std():.3f} nm | T {temp.mean():.1f} K | "
          f"rho {dens.mean():.3f} g/mL")


if __name__ == "__main__":
    main()