#!/usr/bin/env python3
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import matplotlib.pyplot as plt


ROOT = Path(__file__).resolve().parents[1]
PLOTS_DIR = ROOT / "artifacts" / "plots"

OURS_COLOR = "#1f77b4"
BASELINE_COLOR = "#9aa0a6"


@dataclass(frozen=True)
class Baseline:
    name: str
    seconds: float


def _save(fig: plt.Figure, stem: str) -> None:
    PLOTS_DIR.mkdir(parents=True, exist_ok=True)
    fig.tight_layout()
    fig.savefig(PLOTS_DIR / f"{stem}.png", dpi=300, bbox_inches="tight")
    fig.savefig(PLOTS_DIR / f"{stem}.pdf", bbox_inches="tight")
    plt.close(fig)


def main() -> int:
    # ---- Reported averages over 40 NL queries (seconds/query) ----
    # MChatbot pipeline includes: schema extraction + compilation + repair loop + TCTL verification.
    mchatbot_total_s = 90.0
    breakdown = [
        ("JSON schema extraction (LLM)", 35.0),
        ("UPPAAL XML compilation", 5.0),
        ("Repair loop (≤3 iters)", 40.0),
        ("TCTL translation + verification", 10.0),
    ]
    if abs(sum(v for _, v in breakdown) - mchatbot_total_s) > 1e-9:
        raise SystemExit("Breakdown does not sum to MChatbot total runtime.")

    baselines = [
        Baseline("GPT-5.3", 2.0),
        Baseline("Grok-4", 3.0),
        Baseline("Claude-Haiku", 4.0),
    ]

    fig, (ax_total, ax_break) = plt.subplots(
        1,
        2,
        figsize=(11.2, 4.2),
        gridspec_kw={"width_ratios": [1.0, 1.4]},
    )

    # ---- Panel (a): total runtime comparison ----
    total_names = ["MChatbot (ours)"] + [b.name for b in baselines]
    total_vals = [mchatbot_total_s] + [b.seconds for b in baselines]
    total_colors = [OURS_COLOR] + [BASELINE_COLOR] * len(baselines)

    bars = ax_total.bar(total_names, total_vals, color=total_colors)
    ax_total.set_title("Total runtime (avg / query)")
    ax_total.set_ylabel("Seconds")
    ax_total.grid(axis="y", alpha=0.25)
    ax_total.set_axisbelow(True)
    ax_total.tick_params(axis="x", rotation=10)

    for b, v in zip(bars, total_vals):
        ax_total.text(
            b.get_x() + b.get_width() / 2,
            b.get_height() + max(0.8, 0.02 * max(total_vals)),
            f"{v:.0f}s",
            ha="center",
            va="bottom",
            fontsize=10,
        )

    # ---- Panel (b): MChatbot breakdown (stacked) ----
    labels = [k for k, _ in breakdown]
    vals = [v for _, v in breakdown]
    stack_colors = ["#4c78a8", "#72b7b2", "#f58518", "#54a24b"]

    bottom = 0.0
    for (label, v), c in zip(breakdown, stack_colors, strict=True):
        ax_break.bar(["MChatbot (ours)"], [v], bottom=bottom, color=c, label=label)
        ax_break.text(
            0,
            bottom + v / 2,
            f"{v:.0f}s",
            ha="center",
            va="center",
            fontsize=10,
            color="white" if c in {"#4c78a8", "#f58518"} else "black",
            fontweight="bold",
        )
        bottom += v

    ax_break.set_title("MChatbot breakdown (avg / query)")
    ax_break.set_ylabel("Seconds")
    ax_break.set_ylim(0, mchatbot_total_s * 1.15)
    ax_break.grid(axis="y", alpha=0.25)
    ax_break.set_axisbelow(True)
    ax_break.legend(loc="upper left", frameon=True, fontsize=9)

    fig.suptitle("Time analysis: MChatbot vs direct-LLM baselines", fontsize=13, fontweight="bold")
    _save(fig, "time_analysis_runtime_breakdown")
    print(f"Wrote plots to: {PLOTS_DIR.resolve()}")
    return 0


if __name__ == "__main__":
    plt.rcParams.update(
        {
            "font.size": 11,
            "axes.titlesize": 12,
            "axes.labelsize": 11,
            "figure.dpi": 120,
        }
    )
    raise SystemExit(main())