"""Generate the BananaMind 2 Micro Base Bench efficiency chart."""

from dataclasses import dataclass
from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator


RANDOM_BASELINE = 25.0
OUTPUT_PATH = Path(__file__).with_name("parameter_efficiency.png")


@dataclass(frozen=True)
class ModelResult:
    name: str
    parameters: int
    accuracy: float
    highlighted: bool = False

    @property
    def excess_accuracy(self) -> float:
        return self.accuracy - RANDOM_BASELINE

    @property
    def efficiency(self) -> float:
        return self.excess_accuracy / (self.parameters / 100_000)

    @property
    def parameter_label(self) -> str:
        return f"{self.parameters / 1_000_000:.2f}M"


# Accuracy values are raw public BananaMind Base Bench 1.1 accuracy.
# Exact parameter counts and scores are preserved so the chart can be rebuilt.
MODELS = (
    ModelResult("BananaMind 2 Micro", 2_933_193, 34.57, highlighted=True),
    ModelResult("GPT-S-5M", 5_158_464, 37.14),
    ModelResult("GPT-S2-5M", 5_384_258, 35.71),
    ModelResult("Syn-2.6M", 2_604_210, 32.57),
    ModelResult("Ant-5M", 4_713_344, 25.43),
    ModelResult("Supra-Mini-v5-8M", 7_867_584, 36.29),
    ModelResult("cma-8M", 7_849_161, 40.86),
)


def build_chart(output_path: Path = OUTPUT_PATH) -> Path:
    ranked = sorted(MODELS, key=lambda model: model.efficiency, reverse=True)

    background = "#f4f7fb"
    ink = "#172033"
    muted = "#667085"
    grid = "#d8dee9"
    peer = "#5d7898"
    banana = "#f3b61f"
    banana_edge = "#d99a00"

    fig = plt.figure(figsize=(16, 9), dpi=120, facecolor=background)
    ax = fig.add_axes([0.255, 0.18, 0.49, 0.61], facecolor=background)

    positions = list(range(len(ranked)))
    colors = [banana if model.highlighted else peer for model in ranked]
    edges = [banana_edge if model.highlighted else peer for model in ranked]
    bars = ax.barh(
        positions,
        [model.efficiency for model in ranked],
        height=0.56,
        color=colors,
        edgecolor=edges,
        linewidth=1.2,
        zorder=3,
    )

    ax.invert_yaxis()
    ax.set_yticks(positions, [model.name for model in ranked])
    ax.tick_params(axis="y", length=0, pad=14, labelsize=14, colors=ink)
    ax.tick_params(axis="x", length=0, pad=8, labelsize=11, colors=muted)

    max_efficiency = max(model.efficiency for model in ranked)
    ax.set_xlim(0, max_efficiency * 1.18)
    ax.xaxis.set_major_locator(MultipleLocator(0.05))
    ax.grid(axis="x", color=grid, linewidth=1, zorder=0)
    ax.set_axisbelow(True)
    for spine in ax.spines.values():
        spine.set_visible(False)

    ax.set_xlabel(
        "Accuracy points above random per 100K parameters",
        fontsize=12,
        color=muted,
        labelpad=16,
    )

    for tick, model in zip(ax.get_yticklabels(), ranked):
        tick.set_fontweight("bold" if model.highlighted else "normal")
        tick.set_color("#9a6800" if model.highlighted else ink)

    for bar, model in zip(bars, ranked):
        ax.text(
            bar.get_width() + 0.006,
            bar.get_y() + bar.get_height() / 2,
            f"{model.efficiency:.3f}",
            va="center",
            ha="left",
            fontsize=12,
            fontweight="bold",
            color=ink,
        )

    column_transform = ax.get_yaxis_transform()
    ax.text(
        1.12,
        -0.82,
        "BASE BENCH\nACCURACY",
        transform=column_transform,
        ha="center",
        va="bottom",
        fontsize=9,
        fontweight="bold",
        color=muted,
        clip_on=False,
    )
    ax.text(
        1.34,
        -0.82,
        "PARAMETERS",
        transform=column_transform,
        ha="center",
        va="bottom",
        fontsize=9,
        fontweight="bold",
        color=muted,
        clip_on=False,
    )
    for position, model in zip(positions, ranked):
        text_color = "#9a6800" if model.highlighted else ink
        fontweight = "bold" if model.highlighted else "normal"
        ax.text(
            1.12,
            position,
            f"{model.accuracy:.2f}%",
            transform=column_transform,
            ha="center",
            va="center",
            fontsize=12,
            fontweight=fontweight,
            color=text_color,
            clip_on=False,
        )
        ax.text(
            1.34,
            position,
            model.parameter_label,
            transform=column_transform,
            ha="center",
            va="center",
            fontsize=12,
            fontweight=fontweight,
            color=text_color,
            clip_on=False,
        )

    fig.text(
        0.06,
        0.925,
        "BananaMind 2 Micro",
        fontsize=34,
        fontweight="bold",
        color=ink,
    )
    fig.text(
        0.06,
        0.872,
        "Base Bench parameter efficiency",
        fontsize=21,
        fontweight="bold",
        color=ink,
    )
    fig.text(
        0.06,
        0.835,
        "Seven sub-10M models ranked by useful accuracy per parameter",
        fontsize=13,
        color=muted,
    )

    fig.text(
        0.06,
        0.075,
        "Formula: (raw accuracy - 25% random baseline) / (parameters / 100,000)",
        fontsize=11,
        color=muted,
    )
    fig.text(
        0.94,
        0.075,
        "BananaMind Base Bench 1.1 | 350 questions",
        fontsize=11,
        color=muted,
        ha="right",
    )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output_path, facecolor=background)
    plt.close(fig)
    return output_path


if __name__ == "__main__":
    print(build_chart())