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

from dataclasses import dataclass
from pathlib import Path

import matplotlib.pyplot as plt
import pandas as pd


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


@dataclass(frozen=True)
class SaveSpec:
    stem: str
    title: str
    ylabel: str
    ylim: tuple[float, float] | None = (0.0, 1.05)


METHOD_ORDER = ["MChatbot (ours)", "GPT-5.3", "Grok-4", "Claude-Haiku"]
OURS_COLOR = "#1f77b4"
BASELINE_COLOR = "#9aa0a6"


def _style_axes(ax: plt.Axes, title: str, ylabel: str, ylim: tuple[float, float] | None) -> None:
    ax.set_title(title)
    ax.set_ylabel(ylabel)
    if ylim is not None:
        ax.set_ylim(*ylim)
    ax.grid(axis="y", alpha=0.25)
    ax.set_axisbelow(True)


def _bar_with_labels(ax: plt.Axes, x: list[str], y: list[float], colors: list[str]) -> None:
    bars = ax.bar(x, y, color=colors)
    for b, v in zip(bars, y):
        ax.text(
            b.get_x() + b.get_width() / 2,
            b.get_height() + 0.015,
            f"{v:.3f}".rstrip("0").rstrip("."),
            ha="center",
            va="bottom",
            fontsize=9,
        )


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 plot_qna_accuracy() -> None:
    df = pd.DataFrame(
        {
            "Method": METHOD_ORDER,
            "QnA Accuracy": [0.875, 0.750, 0.775, 0.800],
        }
    )
    colors = [OURS_COLOR] + [BASELINE_COLOR] * (len(df) - 1)
    fig, ax = plt.subplots(figsize=(6.6, 3.6))
    _bar_with_labels(ax, df["Method"].tolist(), df["QnA Accuracy"].tolist(), colors)
    _style_axes(ax, "EA: QnA Accuracy (40 queries)", "Accuracy", (0.0, 1.0))
    ax.tick_params(axis="x", rotation=10)
    _save(fig, "ea_qna_accuracy")


def plot_model_building() -> None:
    df = pd.DataFrame(
        {
            "Method": METHOD_ORDER,
            "MCR": [1.00, 0.30, 0.90, 1.00],
            "MAC": [0.82, 0.12, 0.56, 0.68],
        }
    )
    fig, axes = plt.subplots(1, 2, figsize=(9.2, 3.6), sharey=True)

    for ax, metric, title in zip(
        axes,
        ["MCR", "MAC"],
        ["EB: Model Compilation Rate (MCR)", "EB: Model Answer Correctness (MAC)"],
    ):
        colors = [OURS_COLOR] + [BASELINE_COLOR] * (len(df) - 1)
        _bar_with_labels(ax, df["Method"].tolist(), df[metric].tolist(), colors)
        _style_axes(ax, title, metric, (0.0, 1.0))
        ax.tick_params(axis="x", rotation=15)

    _save(fig, "eb_model_building_mcr_mac")


def plot_query_translation() -> None:
    df = pd.DataFrame(
        {
            "Method": METHOD_ORDER,
            "QCR": [1.000, 0.075, 0.075, 0.525],
            "QAR": [0.725, 0.075, 0.075, 0.450],
            "EM": [0.450, 0.050, 0.025, 0.425],
            "Token F1": [0.743, 0.641, 0.586, 0.725],
        }
    )

    fig, axes = plt.subplots(2, 2, figsize=(10.2, 6.6), sharex=True)
    metrics = [("QCR", "QCR"), ("QAR", "QAR"), ("EM", "Exact Match"), ("Token F1", "Token F1")]

    for ax, (col, ylabel) in zip(axes.flatten(), metrics):
        colors = [OURS_COLOR] + [BASELINE_COLOR] * (len(df) - 1)
        _bar_with_labels(ax, df["Method"].tolist(), df[col].tolist(), colors)
        _style_axes(ax, f"EB.2: {col}", ylabel, (0.0, 1.0))
        ax.tick_params(axis="x", rotation=15)

    _save(fig, "eb2_query_translation_qcr_qar_em_f1")


def plot_ablation() -> None:
    df = pd.DataFrame(
        {
            "Configuration": ["Full MChatbot", "No IR (direct generation)", "No repair loop"],
            "MCR": [1.00, 0.60, 0.80],
            "MAC": [0.84, 0.32, 0.58],
        }
    )

    fig, axes = plt.subplots(1, 2, figsize=(9.2, 3.6), sharey=True)

    for ax, metric in zip(axes, ["MCR", "MAC"]):
        colors = [OURS_COLOR, BASELINE_COLOR, BASELINE_COLOR]
        _bar_with_labels(ax, df["Configuration"].tolist(), df[metric].tolist(), colors)
        _style_axes(ax, f"RQ3 Ablation: {metric}", metric, (0.0, 1.0))
        ax.tick_params(axis="x", rotation=10)

    _save(fig, "rq3_ablation_mcr_mac")


def main() -> int:
    plt.rcParams.update(
        {
            "font.size": 11,
            "axes.titlesize": 12,
            "axes.labelsize": 11,
            "figure.dpi": 120,
        }
    )

    plot_qna_accuracy()
    plot_model_building()
    plot_query_translation()
    plot_ablation()

    print(f"Wrote plots to: {PLOTS_DIR.resolve()}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())