#!/usr/bin/env python
import argparse
import json
from pathlib import Path

import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.ticker import LogLocator

ROOT = Path(__file__).resolve().parents[2]
CACHE = ROOT / "data" / "p53" / "fig_data"


FIGSIZE = (7.2, 5.4)
RC = {
    "font.size": 8,
    "axes.labelsize": 8.5,
    "axes.titlesize": 8.5,
    "xtick.labelsize": 7.5,
    "ytick.labelsize": 7.5,
    "legend.fontsize": 7,
    "axes.spines.top": False,
    "axes.spines.right": False,
    "pdf.fonttype": 42,
    "ps.fonttype": 42,
}
GREY = "#b8b8b8"
ACC = "#2f6fb5"
SEED_C = "#e8a33d"
HI = "#c8412b"
HL_COLORS = {56: HI, 0: ACC}
APO_C = "#4a4a4a"
STAR_SIZE, DOT_SIZE, HL_DOT = 70, 16, 42
HIST_BINS = 70


ANNOT_C = {56: (5, -9), 0: (4, 6)}
ANNOT_A = (11, -4)
ANNOT_B = (0.06, 0.06)
TITLES = {
    "a": "a   Ligandability is not equilibrium weight",
    "b": "b   Rare, and sampled by the reference runs",
    "c": "c   Holo-like conformers exist in the ensemble",
    "d": "d   Compared with the unbound crystal",
}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=Path, default=ROOT / "figures" / "p53_discovery.pdf")
    ap.add_argument("--png", action="store_true", help="also write a PNG alongside")
    a = ap.parse_args()

    npz = CACHE / "discovery.npz"
    if not npz.exists():
        raise SystemExit(
            f"missing {npz}\nrun paper_figures/data/p53_discovery_data.py first"
        )
    d = np.load(npz)
    S = json.loads((CACHE / "discovery.json").read_text())

    mids = list(d["mids"])
    drug, truth, pi = d["drug"], d["truth"], d["pi"]
    order, v = d["order"], d["frame_scores"]
    v_ref, v_seed = d["frame_scores_ref"], d["frame_scores_seed"]
    highlight = list(d["highlight"])
    profiles = d["highlight_rmsd_profiles"]
    n_cx = int(d["n_complexes"])
    rank = {int(m): S["highlight"][str(m)]["rank"] for m in highlight}
    i56 = mids.index(56)

    plt.rcParams.update(RC)
    fig, ax = plt.subplots(2, 2, figsize=FIGSIZE)

    p = ax[0, 0]
    p.scatter(pi, drug, s=DOT_SIZE, c=GREY, edgecolor="none", zorder=2)
    p.scatter(
        [pi[i56]],
        [drug[i56]],
        s=STAR_SIZE,
        c=HI,
        marker="*",
        edgecolor="k",
        linewidth=0.4,
        zorder=4,
    )
    p.annotate(
        f"microstate 56\n$\\pi$={pi[i56]:.4f}, score {drug[i56]:.3f}",
        (pi[i56], drug[i56]),
        textcoords="offset points",
        xytext=ANNOT_A,
        ha="left",
        va="top",
        fontsize=6.5,
        color=HI,
    )
    p.set_xscale("log")
    p.xaxis.set_major_locator(LogLocator(numticks=5))
    p.set_xlabel("stationary probability $\\pi$")
    p.set_ylabel("representative-conformer\npocket druggability")
    p.set_title(TITLES["a"], loc="left", fontweight="bold")

    p = ax[0, 1]
    bins = np.linspace(0.0, 1.0, HIST_BINS + 1)
    p.hist(
        v_ref,
        bins=bins,
        color=ACC,
        alpha=0.80,
        edgecolor="none",
        label=f"reference (n={len(v_ref):,})",
    )
    p.hist(
        v_seed,
        bins=bins,
        color=SEED_C,
        alpha=0.60,
        edgecolor="none",
        label=f"seeded (n={len(v_seed):,})",
    )
    p.set_yscale("log")
    p.axvline(drug[i56], color=HI, ls="--", lw=1.2)
    sh = S["high_druggability_share"]["ge_0.8"]

    p.text(
        ANNOT_B[0],
        ANNOT_B[1],
        f"micro-56 rep.  {drug[i56]:.3f}\n"
        f"{S['n_frames_ge_thresh']}/{S['n_frames']} frames "
        f"$\\geq${S['rare_threshold']:g}\n"
        f"{100 * sh['reference'] / sh['n']:.0f}% of frames $\\geq$0.8 are reference",
        transform=p.transAxes,
        ha="left",
        va="bottom",
        fontsize=6.3,
        color=HI,
    )
    p.set_xlabel("fpocket druggability score, per frame")
    p.set_ylabel("frames (log)")
    p.legend(frameon=False, loc="upper right", fontsize=6.5)
    p.set_title(TITLES["b"], loc="left", fontweight="bold")

    p = ax[1, 0]
    xs = np.arange(1, len(truth) + 1)
    p.plot(xs, truth[order], "-", color=GREY, lw=1.0, zorder=1)
    p.scatter(xs, truth[order], s=12, c=GREY, edgecolor="none", zorder=2)
    for mic in highlight:
        mic = int(mic)
        j = mids.index(mic)
        p.scatter(
            [rank[mic]],
            [truth[j]],
            s=HL_DOT,
            c=HL_COLORS[mic],
            edgecolor="k",
            linewidth=0.4,
            zorder=4,
        )
        p.annotate(
            f"{mic}",
            (rank[mic], truth[j]),
            textcoords="offset points",
            xytext=ANNOT_C[mic],
            fontsize=6.5,
            color=HL_COLORS[mic],
            fontweight="bold",
        )
    p.set_xlabel("microstate rank by median holo RMSD")
    p.set_ylabel(f"median pocket RMSD to\n{n_cx} Y220C complexes ($\\AA$)")
    p.set_title(TITLES["c"], loc="left", fontweight="bold")

    p = ax[1, 1]
    for k, mic in enumerate(highlight):
        mic = int(mic)
        p.plot(
            profiles[k],
            np.arange(1, n_cx + 1),
            "-",
            color=HL_COLORS[mic],
            lw=1.3,
            label=f"micro {mic} (rank {rank[mic]}/{len(mids)})",
        )
    p.plot(
        d["apo_rmsd_profile"],
        np.arange(1, n_cx + 1),
        "--",
        color=APO_C,
        lw=1.3,
        label=f"apo crystal (median {S['apo_to_holo_median']:.2f} $\\AA$)",
    )
    p.text(
        0.035,
        0.97,
        f"ensemble closer than apo\nfor {S['n_ensemble_closer_than_apo']}/{n_cx} complexes",
        transform=p.transAxes,
        ha="left",
        va="top",
        fontsize=6.3,
        color=APO_C,
    )
    p.set_xlabel("pocket RMSD to an individual complex ($\\AA$)")
    p.set_ylabel("crystal complexes (cumulative)")
    p.legend(frameon=False, loc="lower right")
    p.set_title(TITLES["d"], loc="left", fontweight="bold")

    fig.tight_layout(pad=0.9)
    a.out.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(a.out, format="pdf", bbox_inches="tight")
    print(f"-> {a.out}")
    if a.png:
        fig.savefig(a.out.with_suffix(".png"), dpi=200, bbox_inches="tight")
        print(f"-> {a.out.with_suffix('.png')}")


if __name__ == "__main__":
    main()