#!/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

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


FIGSIZE = (7.2, 6.0)
HEIGHT_RATIOS, HSPACE, WSPACE = [1.15, 1.0], 0.42, 0.32
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,
}
ACC = "#2f6fb5"
HI = "#c8412b"
NULL_GREY = "#c9c9c9"
CMAP = "RdBu_r"
CLIP_PCT = 98
BAR_W, BAR_DX = 0.32, 0.17
YLIM_B = (-0.12, 1.12)
TITLES = {
    "a": "a   Ligand $\\times$ conformer interaction $\\gamma$\n"
    "     (ligand, cavity-volume and receptor main effects removed)",
    "b": "b   Preference survives removing\n     between-stratum contrasts",
    "c": "c   Preference lives inside states",
}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=Path, default=ROOT / "figures" / "p53_preference.pdf")
    ap.add_argument("--png", action="store_true")
    a = ap.parse_args()

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

    gamma, strat = d["gamma"], d["strat"]
    lig_order, rec_order = d["lig_order"], d["rec_order"]
    n_lig, n_rec = gamma.shape
    labs = np.unique(strat)
    Gp = gamma[np.ix_(lig_order, rec_order)]

    plt.rcParams.update(RC)
    fig = plt.figure(figsize=FIGSIZE)
    gs = fig.add_gridspec(
        2, 2, height_ratios=HEIGHT_RATIOS, hspace=HSPACE, wspace=WSPACE
    )

    p = fig.add_subplot(gs[0, :])
    lim = np.percentile(np.abs(Gp), CLIP_PCT)
    im = p.imshow(
        Gp, aspect="auto", cmap=CMAP, vmin=-lim, vmax=lim, interpolation="nearest"
    )
    b0, centres = 0, []
    for l in labs:
        n = int((strat == l).sum())
        centres.append(b0 + n / 2 - 0.5)
        b0 += n
        if b0 < n_rec:
            p.axvline(b0 - 0.5, color="k", lw=0.8)
    p.set_xticks(centres)
    p.set_xticklabels([f"stratum {int(l) + 1}" for l in labs])
    p.tick_params(axis="x", length=0)
    p.set_xlabel(f"{n_rec} receptor conformers, grouped by structural stratum")
    p.set_ylabel(f"{n_lig} ligands (clustered)")
    p.set_yticks([])
    p.set_title(TITLES["a"], loc="left", fontweight="bold", pad=6)
    cb = fig.colorbar(im, ax=p, pad=0.012, fraction=0.021)
    cb.set_label("$\\gamma$  (blue = better than expected)", fontsize=6.5)
    cb.ax.tick_params(labelsize=6)

    p = fig.add_subplot(gs[1, 0])
    series = [
        ("all\nreceptors", d["obs_all"], d["nul_all"]),
        ("within\nstratum only", d["obs_within"], d["nul_within"]),
    ]
    for i, (nm, o, nl) in enumerate(series):
        p.bar(
            i - BAR_DX,
            o.mean(),
            BAR_W,
            yerr=o.std(),
            color=ACC,
            capsize=2.5,
            label="observed" if i == 0 else None,
        )
        p.bar(
            i + BAR_DX,
            nl.mean(),
            BAR_W,
            yerr=nl.std(),
            color=NULL_GREY,
            capsize=2.5,
            label="permuted" if i == 0 else None,
        )
        p.text(
            i - BAR_DX,
            o.mean() / 2,
            f"{o.mean():.3f}",
            ha="center",
            va="center",
            fontsize=7.5,
            color="w",
            fontweight="bold",
        )
    p.axhline(0, color="k", lw=0.7)
    p.set_xticks(range(len(series)))
    p.set_xticklabels([s[0] for s in series])
    p.set_ylabel("split-half Spearman of\nligand-preference structure")
    p.set_ylim(*YLIM_B)
    p.legend(frameon=False, loc="upper center", ncol=2, columnspacing=1.0)
    p.set_title(TITLES["b"], loc="left", fontweight="bold")

    p = fig.add_subplot(gs[1, 1])
    fw, fb = float(d["frac_within"]), float(d["frac_between"])
    p.bar([0], [100 * fw], 0.5, color=ACC, label="within strata")
    p.bar([0], [100 * fb], 0.5, bottom=[100 * fw], color=HI, label="between strata")
    p.text(
        0,
        100 * fw / 2,
        f"{100 * fw:.1f}%",
        ha="center",
        va="center",
        color="w",
        fontweight="bold",
        fontsize=9,
    )
    p.text(
        0,
        100 * fw + 100 * fb / 2 + 3,
        f"{100 * fb:.1f}%",
        ha="center",
        va="center",
        color=HI,
        fontweight="bold",
        fontsize=8,
    )
    p.set_xlim(-0.6, 1.4)
    p.set_xticks([])
    p.set_ylabel("share of $\\mathrm{var}(\\gamma)$  (%)")
    p.legend(frameon=False, loc="center right")
    p.set_title(TITLES["c"], loc="left", fontweight="bold")

    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')}")
    print(
        f"   (cache: {S['n_ligands']} ligands, {S['n_receptors']} receptors, "
        f"{S['n_splits']} splits, within {100 * S['frac_within']:.1f}%)"
    )


if __name__ == "__main__":
    main()