#!/usr/bin/env python3
"""Figures for the MSQ study. Reads results/metrics.json and results/raw/*.jsonl."""
from __future__ import annotations

import json
import pathlib
import sys

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

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
import msq  # noqa: E402

ROOT = pathlib.Path(__file__).resolve().parent.parent
RES = ROOT / "results"
FIG = ROOT / "figures"

INK = "#1f2933"
ACCENT = "#1d6fd0"
WARM = "#e85d4c"
TEAL = "#2a9d8f"
AMBER = "#e09f3e"
GREY = "#9aa5b1"


def style(ax):
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_color(GREY)
    ax.spines["bottom"].set_color(GREY)
    ax.tick_params(colors=INK, labelsize=9)
    ax.grid(axis="y", color="#e4e7eb", lw=0.8)
    ax.set_axisbelow(True)


def load_metrics():
    p = RES / "metrics.json"
    return json.loads(p.read_text()) if p.exists() else {}


def load_raw(model):
    p = RES / "raw" / f"{model}.jsonl"
    if not p.exists():
        return []
    return [json.loads(l) for l in p.open(encoding="utf-8") if l.strip()]


def fig_prr_msq(metrics, primary):
    """PRR and MSQ by class, for every model: the three-symptom structure."""
    models = [m for m in metrics if metrics[m].get("mcq")]
    if not models:
        return
    classes = ["A", "B", "C", "E"]
    labels = ["A\nfiqh\n(madhhab)", "B\nlive\ncontroversy", "C\njurisdiction",
              "E\nPython\nversion"]
    fig, axes = plt.subplots(1, 3, figsize=(13.0, 4.3))
    width = 0.8 / max(len(models), 1)

    def draw(ax, section, key, title, ylabel):
        for k, m in enumerate(models):
            vals = []
            for c in classes:
                v = metrics[m]["mcq"].get(section, {}).get(c, {}).get(key)
                vals.append((v or 0) * 100)
            xs = [i + (k - (len(models) - 1) / 2) * width for i in range(len(classes))]
            ax.bar(xs, vals, width=width * 0.9,
                   color=ACCENT if m == primary else TEAL,
                   label=m if k == 0 else None)
        ax.set_xticks(range(len(classes)))
        ax.set_xticklabels(labels, fontsize=8)
        ax.set_title(title, fontsize=10, color=INK, pad=10, loc="left")
        ax.set_ylabel(ylabel)
        ax.set_ylim(0, 108)
        style(ax)

    draw(axes[0], "mcq_neutral", "plurality_recognition_rate",
         "Does it NOTICE the question is conditional?", "% choosing 'it depends'")
    draw(axes[1], "mcq_specified", "school_adherence_rate",
         "Does it USE the context it was given?", "% following the named context")
    ax = axes[2]
    for k, m in enumerate(models):
        cr = metrics[m].get("context_response", {}).get("ALL", {})
        v = (cr.get("mean_context_response_tv") or 0)
        ax.bar([k], [v], width=0.5, color=ACCENT if m == primary else TEAL)
        ax.text(k, v + 0.012, f"{v:.3f}", ha="center", fontsize=9, color=INK)
        u = cr.get("argmax_unchanged_rate")
        if u is not None:
            ax.text(k, 0.012, f"{u*100:.0f}% unchanged", ha="center", fontsize=8,
                    color="white", rotation=90, va="bottom")
    ax.set_xticks(range(len(models)))
    ax.set_xticklabels([m.replace("-", "\n") for m in models], fontsize=8)
    ax.set_title("How far does naming the context move the answer?",
                 fontsize=10, color=INK, pad=10, loc="left")
    ax.set_ylabel("total-variation distance between the two conditions")
    ax.set_ylim(0, 0.55)
    style(ax)
    handles = [plt.Rectangle((0, 0), 1, 1, color=ACCENT),
               plt.Rectangle((0, 0), 1, 1, color=TEAL)]
    fig.legend(handles, ["primary", "comparison models"], loc="upper right",
               frameon=False, fontsize=9, ncol=2)
    fig.tight_layout()
    fig.savefig(FIG / "fig1_prr_msq.png", dpi=170, bbox_inches="tight")
    plt.close(fig)


def fig_disclosure(metrics, primary):
    if primary not in metrics:
        return
    ft = metrics[primary]["free_text"]
    pc = metrics[primary].get("positions_contained", {})
    probes = ["ft_default", "ft_repair", "ft_technical", "ft_two_schools", "ft_pair"]
    labels = ["plain\nquestion", "one-line\nrepair", "technical\nregister",
              "asked to\nenumerate", "asked for two\nnamed schools"]
    series = [
        ("disclosure_rate", "signals that the answer is context-dependent", ACCENT, None),
        ("plural_attribution_rate", "names two or more competing schools", WARM, None),
        ("mean_positions_contained", "positions actually carried", TEAL, "pc"),
    ]
    fig, axes = plt.subplots(1, 3, figsize=(12.4, 4.0), sharey=False)
    for ax, (key, title, col, kind) in zip(axes, series):
        vals = []
        for p in probes:
            if kind == "pc":
                v = pc.get(p, {}).get("ALL", {}).get("mean_positions_contained")
            else:
                v = ft.get(p, {}).get("ALL", {}).get(key)
            vals.append(v or 0)
        if kind == "pc":
            vals = [v * 2 for v in vals]  # scale to a 0-2 "positions" axis
            top = 2.1
        else:
            vals = [v * 100 for v in vals]
            top = 105
        bars = ax.bar(range(len(probes)), vals, color=col, width=0.62)
        for b, v in zip(bars, vals):
            ax.text(b.get_x() + b.get_width() / 2, v + top * 0.02,
                    f"{v:.2f}" if kind == "pc" else f"{v:.0f}%",
                    ha="center", fontsize=8.5, color=INK)
        ax.set_xticks(range(len(probes)))
        ax.set_xticklabels(labels, fontsize=7.6)
        ax.set_title(title, fontsize=9.5, color=INK, pad=10, loc="left")
        ax.set_ylim(0, top)
        if kind == "pc":
            ax.set_yticks([0, 1, 2])
            ax.set_yticklabels(["0", "1", "2"])
        style(ax)
    axes[0].set_ylabel("% of answers")
    fig.suptitle(f"Disclosure behaviour on contested questions  ·  {primary}",
                 fontsize=11.5, color=INK, x=0.005, ha="left", y=1.02)
    fig.tight_layout()
    fig.savefig(FIG / "fig2_disclosure.png", dpi=170, bbox_inches="tight")
    plt.close(fig)


def fig_confidence(metrics, primary):
    if primary not in metrics:
        return
    mn = metrics[primary]["mcq"]["mcq_neutral"]
    so = metrics[primary].get("specified_outcome", {}).get("mcq_specified", {})
    cr = metrics[primary].get("context_response", {})
    classes = ["A", "B", "C", "E"]
    labels = ["A\nfiqh", "B\ncontroversy", "C\njurisdiction", "E\nPython\nversion"]
    fig, axes = plt.subplots(1, 3, figsize=(13.4, 4.2))

    # panel 1 -- what it picks once the context is named
    ax = axes[0]
    keys = [("named_school", "the context it was given", ACCENT),
            ("divided_anyway", "'it depends' anyway", AMBER),
            ("other_position", "some other position", GREY)]
    bottoms = [0.0] * len(classes)
    for key, name, col in keys:
        vals = []
        for c in classes:
            v = so.get(c, {}).get(key)
            vals.append((v or 0) * 100)
        ax.bar(range(len(classes)), vals, bottom=bottoms, color=col, width=0.6,
               label=name)
        bottoms = [b + v for b, v in zip(bottoms, vals)]
    ax.set_xticks(range(len(classes)))
    ax.set_xticklabels(labels, fontsize=8.5)
    ax.set_ylabel("% of items")
    ax.set_title("Once the context IS named, what does it pick?", fontsize=10,
                 color=INK, pad=10, loc="left")
    ax.set_ylim(0, 100)
    ax.legend(fontsize=8, frameon=False, loc="upper center",
              bbox_to_anchor=(0.5, -0.18), ncol=1)
    style(ax)

    # panel 2 -- how far the instruction moved the answer
    ax = axes[1]
    for c, col in zip(classes, [WARM, AMBER, TEAL, "#8e7dbe"]):
        d = cr.get(c, {}).get("median_delta_on_named_position")
        if d is None:
            continue
        ax.bar([classes.index(c)], [d], width=0.55, color=col)
        ax.text(classes.index(c), d + 0.004, f"{d:.3f}", ha="center", fontsize=9,
                color=INK)
    ax.set_xticks(range(len(classes)))
    ax.set_xticklabels(labels, fontsize=8.5)
    ax.set_ylabel("median change in probability on the correct answer")
    ax.set_title("Naming the context barely moves the answer", fontsize=10,
                 color=INK, pad=10, loc="left")
    ax.set_ylim(0, 0.42)
    style(ax)

    # panel 3 -- confidence
    ax = axes[2]
    allc = classes + ["D"]
    vals = [(mn.get(c, {}).get("mean_max_prob") or 0) * 100 for c in allc]
    cols = [WARM, AMBER, TEAL, "#8e7dbe", GREY]
    bars = ax.bar(range(len(allc)), vals, color=cols, width=0.6)
    for b, v in zip(bars, vals):
        ax.text(b.get_x() + b.get_width() / 2, v + 1.5, f"{v:.0f}%", ha="center",
                fontsize=9.5, color=INK)
    ax.set_xticks(range(len(allc)))
    ax.set_xticklabels(["A\nfiqh", "B\ncontroversy", "C\njurisdiction",
                         "E\nPython", "D\ncontrol"], fontsize=8.5)
    ax.set_ylabel("mean confidence in its own chosen option")
    ax.set_title("Confidence never drops", fontsize=10, color=INK, pad=10, loc="left")
    ax.set_ylim(0, 112)
    style(ax)

    fig.tight_layout()
    fig.savefig(FIG / "fig3_confidence_prr.png", dpi=170, bbox_inches="tight")
    plt.close(fig)


def fig_per_item(metrics, primary):
    recs = load_raw(primary)
    if not recs:
        return
    items = {i["id"]: i for i in msq.load_items()}
    rows = [r for r in recs if r["probe"] == "mcq_neutral" and r["cls"] in ("A", "B")]
    if not rows:
        return
    rows.sort(key=lambda r: r["item_id"])
    fig, ax = plt.subplots(figsize=(9.6, 5.2))
    xs, ys, cs = [], [], []
    for r in rows:
        it = items[r["item_id"]]
        divided = max((i for i, o in enumerate(r["options"])
                       if o.lower().startswith(("the scholars are divided",
                                                "this is a contested matter",
                                                "it depends on the jurisdiction"))),
                      default=None)
        if divided is None:
            continue
        xs.append(r["item_id"])
        ys.append(r["probs"][divided] * 100)
        arg = r["argmax_index"]
        cs.append("#2a9d8f" if arg == divided else "#e85d4c")
    ax.bar(range(len(xs)), ys, color=cs, width=0.7)
    ax.set_xticks(range(len(xs)))
    ax.set_xticklabels(xs, rotation=90, fontsize=7)
    ax.set_ylabel("probability assigned to the 'divided' option")
    ax.set_title("Per-item probability of the plurality answer (teal = the model "
                 "chose it; red = it chose one school)",
                 fontsize=10.5, color=INK, pad=10, loc="left")
    style(ax)
    fig.tight_layout()
    fig.savefig(FIG / "fig4_per_item.png", dpi=170)
    plt.close(fig)


def fig_tokenizer():
    p = RES / "repr_probe.json"
    if not p.exists():
        return
    rep = json.loads(p.read_text())
    fert = rep.get("fertility_by_script", {})
    names = {"en": "English", "zh": "Chinese", "ar_vocalized": "Arabic\n(vocalised)",
             "ar_plain": "Arabic\n(plain)", "ur": "Urdu", "fa": "Persian",
             "hi": "Hindi", "arabizi": "Romanised\nArabic"}
    fig, axes = plt.subplots(1, 2, figsize=(10.8, 4.2))
    keys = [k for k in names if k in fert]
    ax = axes[0]
    vals = [fert[k]["tokens_per_char"] * 100 for k in keys]
    colors = []
    for k in keys:
        if k.startswith("ar"):
            colors.append(WARM)
        elif k in ("ur", "fa", "hi"):
            colors.append(TEAL)
        else:
            colors.append(GREY)
    bars = ax.bar(range(len(keys)), vals, color=colors, width=0.62)
    for b, v in zip(bars, vals):
        ax.text(b.get_x() + b.get_width() / 2, v + 0.6, f"{v:.1f}", ha="center",
                fontsize=8.5, color=INK)
    ax.set_xticks(range(len(keys)))
    ax.set_xticklabels([names[k] for k in keys], fontsize=8)
    ax.set_ylabel("tokens per 100 characters")
    ax.set_title("Token cost by script", fontsize=10.5, color=INK, pad=10, loc="left")
    style(ax)

    ax = axes[1]
    keys2 = [k for k in keys]
    vals2 = [fert[k]["byte_level_rate"] * 100 for k in keys2]
    bars = ax.bar(range(len(keys2)), vals2, color=colors, width=0.62)
    for b, v in zip(bars, vals2):
        ax.text(b.get_x() + b.get_width() / 2, v + 1.2, f"{v:.0f}%", ha="center",
                fontsize=8.5, color=INK)
    ax.set_xticks(range(len(keys2)))
    ax.set_xticklabels([names[k] for k in keys2], fontsize=8)
    ax.set_ylabel("% of tokens that are raw byte fragments")
    ax.set_title("Where the vocabulary gives up", fontsize=10.5, color=INK, pad=10,
                 loc="left")
    ax.set_ylim(0, 68)
    style(ax)
    fig.tight_layout()
    fig.savefig(FIG / "fig5_tokenizer.png", dpi=170)
    plt.close(fig)


def main():
    FIG.mkdir(parents=True, exist_ok=True)
    metrics = load_metrics()
    primary = "minicpm5-2b" if "minicpm5-2b" in metrics else next(iter(metrics), None)
    if metrics:
        fig_prr_msq(metrics, primary)
        fig_disclosure(metrics, primary)
        fig_confidence(metrics, primary)
        fig_per_item(metrics, primary)
    fig_tokenizer()
    print("figures:", sorted(p.name for p in FIG.glob("*.png")))


if __name__ == "__main__":
    main()