"""Shared matplotlib styling and the figures used by train_nn.py, run_demo.py and run_experiments.py."""
from __future__ import annotations

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

# validated categorical slots (first three pass all-pairs CVD checks) + neutral inks
SERIES = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#6250d6"]
VAN_COLORS = SERIES[:3]
SURFACE, INK, INK2, MUTED, GRID = "#fcfcfb", "#0b0b0b", "#52514e", "#8a8984", "#e4e3df"
LATE = "#e34948"


def style():
    plt.rcParams.update({
        "figure.facecolor": SURFACE, "axes.facecolor": SURFACE, "savefig.facecolor": SURFACE,
        "axes.edgecolor": MUTED, "axes.labelcolor": INK2, "xtick.color": INK2, "ytick.color": INK2,
        "text.color": INK, "axes.titlecolor": INK, "axes.titleweight": "bold", "axes.titlesize": 12,
        "axes.spines.top": False, "axes.spines.right": False, "axes.grid": True,
        "grid.color": GRID, "grid.linewidth": 0.8, "axes.axisbelow": True,
        "lines.linewidth": 2.0, "font.size": 10, "legend.frameon": False, "figure.dpi": 110,
    })


style()


def save(fig, path):
    fig.tight_layout()
    fig.savefig(path, dpi=130, bbox_inches="tight")
    plt.close(fig)
    return path


def hm(t):
    t = int(round(t))
    return f"{t // 60:02d}:{t % 60:02d}"


# --------------------------------------------------------------------------- NN figures
def learning_curve(history, path):
    h = np.asarray(history)
    fig, ax = plt.subplots(figsize=(6, 3.4))
    ax.plot(h[:, 0], h[:, 2], color=SERIES[0], marker="o", ms=4, label="validation MAE (min)")
    best = int(np.argmin(h[:, 2]))
    ax.scatter([h[best, 0]], [h[best, 2]], s=80, facecolor="none", edgecolor=INK, zorder=5)
    ax.annotate(f"best epoch {int(h[best, 0])}: {h[best, 2]:.2f} min", (h[best, 0], h[best, 2]),
                xytext=(8, 10), textcoords="offset points", color=INK2)
    ax.set_xlabel("epoch"); ax.set_ylabel("validation MAE (minutes)")
    ax.set_title("NN training: early stopping on validation MAE")
    return save(fig, path)


def mae_by_hour(test, preds: dict, path):
    fig, ax = plt.subplots(figsize=(7.5, 3.6))
    hours = np.arange(24)
    for k, (name, p) in enumerate(preds.items()):
        err = np.abs(np.asarray(p) - test.duration_min.values)
        m = [err[test.hour.values == h].mean() if (test.hour.values == h).any() else np.nan for h in hours]
        ax.plot(hours, m, color=SERIES[k], marker="o", ms=3.5, label=name)
    cnt = np.bincount(test.hour.values, minlength=24)
    ax2 = ax.inset_axes([0.0, -0.42, 1.0, 0.22], sharex=ax)
    ax2.bar(hours, cnt, color=MUTED, width=0.8)
    ax2.set_ylabel("test trips", fontsize=8); ax2.tick_params(labelsize=8); ax2.grid(False)
    ax.set_xlim(-0.5, 23.5); ax.set_xticks(range(0, 24, 3))
    ax.set_ylabel("MAE (minutes)"); ax.set_title("Test MAE by pickup hour")
    ax.legend(loc="upper left", ncol=len(preds))
    return save(fig, path)


def pred_vs_actual(y, p, path, n=4000, seed=0):
    rng = np.random.default_rng(seed)
    idx = rng.choice(len(y), size=min(n, len(y)), replace=False)
    fig, ax = plt.subplots(figsize=(4.6, 4.4))
    ax.scatter(np.asarray(y)[idx], np.asarray(p)[idx], s=8, alpha=0.25, color=SERIES[0], edgecolor="none")
    ax.plot([0, 60], [0, 60], color=INK2, lw=1, ls="--")
    ax.set_xlim(0, 60); ax.set_ylim(0, 60)
    ax.set_xlabel("actual duration (min)"); ax.set_ylabel("NN predicted (min)")
    ax.set_title("NN on held-out week (4k sample)")
    return save(fig, path)


def matrix_heatmaps(M, zones, path, hours=(8, 18)):
    labels = [z[:24] for z in zones.sort_values("idx").zone]
    fig, axes = plt.subplots(1, len(hours) + 1, figsize=(16, 5.6),
                             gridspec_kw={"width_ratios": [1] * len(hours) + [0.9], "wspace": 0.12})
    vmax = np.percentile(M[list(hours)], 98)
    for k, (ax, h) in enumerate(zip(axes, hours)):
        im = ax.imshow(M[h], cmap="Blues", vmin=0, vmax=vmax)
        ax.set_title(f"NN predicted taxi minutes, {h:02d}:00"); ax.grid(False)
        ax.set_xticks(range(len(labels))); ax.set_yticks(range(len(labels)))
        ax.set_xticklabels(labels, rotation=90, fontsize=6.5)
        ax.set_yticklabels(labels if k == 0 else [], fontsize=6.5)
        ax.set_xlabel("drop-off zone")
        if k == 0:
            ax.set_ylabel("pickup zone")
    fig.colorbar(im, ax=axes[:len(hours)], shrink=0.7, label="minutes")
    ax = axes[-1]
    off = ~np.eye(M.shape[1], dtype=bool)
    ax.plot(range(24), [M[h][off].mean() for h in range(24)], color=SERIES[0], marker="o", ms=3)
    ax.set_title("Mean over zone pairs by hour"); ax.set_xlabel("departure hour"); ax.set_ylabel("minutes")
    ax.set_xticks(range(0, 24, 3))
    fig.savefig(path, dpi=130, bbox_inches="tight"); plt.close(fig)
    return path