"""
Manuscript figures. Vector PDF for submission, PNG for preview.

Fig 1  R across major clades              -- the phenomenon, before any model
Fig 2  beta_kappa per clade with 95% CI   -- consistency of the effect
Fig 3  out-of-clade error decomposition   -- why prediction is not just phylogeny
Fig 4  similarity vs kappa                -- the negative result
"""

import os as _os
HERE = _os.path.dirname(_os.path.abspath(__file__))
ROOT_DIR = _os.path.dirname(HERE)


import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

OUT = ROOT_DIR + "/figures/"

plt.rcParams.update({
    "font.family": "DejaVu Sans",
    "font.size": 9,
    "axes.linewidth": 0.8,
    "axes.spines.top": False,
    "axes.spines.right": False,
    "xtick.major.width": 0.8,
    "ytick.major.width": 0.8,
    "legend.frameon": False,
    "figure.dpi": 150,
    # Springer wants embedded outline fonts in vector figures; Matplotlib's
    # PostScript default is Type 3, which some pipelines reject.
    "ps.fonttype": 42,
    "pdf.fonttype": 42,
})

# Springer's column grid: 84 mm single, 174 mm double. Figures are generated
# at one of the two so that nothing is rescaled — rescaling a 139 mm figure
# to 84 mm would take its smallest label below 4 pt.
MM = 1 / 25.4
SINGLE, DOUBLE = 84 * MM, 174 * MM


def grid(width_mm, aspect):
    """Figure size in inches at a column width, given a height:width ratio."""
    w = (SINGLE if width_mm == 84 else DOUBLE)
    return (w, w * aspect)

INK = "#1a1a1a"
ACCENT = "#B3472A"
MUTED = "#7a7a7a"
LIGHT = "#c9c9c9"


FIGNUM = {"figure1_R_by_clade": 1,
          "figure3_beta_kappa_by_clade": 3,
          "figure4_error_decomposition": 4,
          "figure5_optimisation_similarity": 5}


def save(fig, name):
    """Write the figure under both its descriptive name and its Fig<n> alias.

    Figure 2 of the manuscript is produced by code/divergence.py, so the
    manuscript numbering runs 1, 3, 4, 5 here and the numbers in the names
    below are the manuscript's, not the order in this file.
    """
    """Springer asks for >=300 dpi halftone, 600 for combination art, and
    prefers vector. We write both, named Fig<n> as the guidelines require."""
    n = FIGNUM[name]
    fig.savefig(OUT + f"Fig{n}.pdf", bbox_inches="tight")
    fig.savefig(OUT + f"Fig{n}.eps", bbox_inches="tight", format="eps")
    fig.savefig(OUT + f"Fig{n}.png", bbox_inches="tight", dpi=600)
    plt.close(fig)


# ------------------------------------------------------------------ data
conf = pd.read_csv(ROOT_DIR + "/data/confirmatory_aligned.csv")
comp = pd.read_csv(ROOT_DIR + "/data/spectrum_components.csv")
raw = pd.read_excel(_os.environ.get("CORAL_XLSX","media-1.xlsx"),
                    sheet_name="CORALMutationsDB")
al = conf.merge(comp, on="species")
al["clade"] = al["species"].map(raw.groupby("taxa1")["Clade"].first())
BIG = ["Fish", "Plants", "Birds", "Mammals", "Fungi", "Insects", "Reptiles"]
al["fold"] = np.where(al["clade"].isin(BIG), al["clade"], "Other")

slopes = pd.read_csv(ROOT_DIR + "/data/per_clade_slopes_phylo.csv")
rmse = pd.read_csv(ROOT_DIR + "/data/rmse_decomposition.csv")


# ------------------------------------------------------- Fig 1: the phenomenon
order = al.groupby("fold")["R"].median().sort_values().index.tolist()
fig, ax = plt.subplots(figsize=grid(174, 0.62))
for i, f in enumerate(order):
    v = al.loc[al["fold"] == f, "R"].values
    x = np.random.default_rng(i).normal(i, 0.075, len(v))
    ax.scatter(x, v, s=2.2, alpha=0.16, color=MUTED, linewidths=0, rasterized=True)
    ax.plot([i - 0.3, i + 0.3], [np.median(v)] * 2, color=ACCENT, lw=2.0,
            solid_capstyle="butt", zorder=3)
ax.set_xticks(range(len(order)))
ax.set_xticklabels(order, rotation=30, ha="right")
ax.set_ylabel("spectrum-conditioned advantage  $R$")
ax.set_xlabel("")
save(fig, "figure1_R_by_clade")  # -> Fig1


# ---------------------------------------------------- Fig 2: coefficient forest
s = slopes.sort_values("beta")
fig, ax = plt.subplots(figsize=grid(174, 0.64))
y = np.arange(len(s))
ax.hlines(y, s["lo"], s["hi"], color=INK, lw=1.4)
ax.plot(s["beta"], y, "o", ms=5, color=ACCENT, zorder=3)
ax.axvline(0, color=LIGHT, lw=1.0, zorder=0)
ax.set_yticks(y)
ax.set_yticklabels([f"{c}  (n={n:,})" for c, n in zip(s["clade"], s["n"])])
ax.set_xlabel(r"effect of transition bias, $\beta_\kappa$")
save(fig, "figure3_beta_kappa_by_clade")  # -> Fig3


# ------------------------------------------------- Fig 3: error decomposition
d = rmse.sort_values("rmse1", ascending=False)
fig, ax = plt.subplots(figsize=grid(174, 0.59))
y = np.arange(len(d))
h = 0.34
ax.barh(y + h / 2, d["rmse0"], height=h, color=LIGHT,
        label="phylogeny only")
ax.barh(y - h / 2, np.abs(d["bias1"]), height=h, color=ACCENT,
        label="spectrum model: intercept bias")
ax.barh(y - h / 2, d["disp1"], height=h, left=np.abs(d["bias1"]),
        color=MUTED, label="spectrum model: dispersion")
ax.set_yticks(y)
ax.set_yticklabels([f"{f}  (n={n:,})" for f, n in zip(d["fold"], d["n"])])
ax.set_xlabel("out-of-clade prediction error (RMSE)")
ax.legend(loc="lower right", fontsize=7.5)
save(fig, "figure4_error_decomposition")  # -> Fig4


# --------------------------------------------------- Fig 4: negative result
kap = np.array([1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 8.0])
sim = np.array([12.4, 7.9, 3.1, 4.1, 3.8, 2.1, 4.5])
fig, ax = plt.subplots(figsize=grid(174, 0.67))
ax.axhline(5.0, color=LIGHT, lw=1.2, ls="--", zorder=0)
ax.text(8.3, 5.0, "chance", va="center", fontsize=7.5, color=MUTED)
ax.plot(kap, sim, "-o", color=ACCENT, ms=5, lw=1.4)
ax.set_xscale("log")
ax.set_xticks(kap)
ax.set_xticklabels([str(k).rstrip("0").rstrip(".") for k in kap])
ax.set_xlabel(r"transition / transversion ratio  $\kappa$")
ax.set_ylabel("similarity to standard code (%)")
ax.set_ylim(0, 14)
save(fig, "figure5_optimisation_similarity")  # -> Fig5

print("four figures written to", OUT)
