#!/usr/bin/env python3
"""Render public benchmark assets for TensorMind 1.5 Preview."""

from __future__ import annotations

import json
from pathlib import Path

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.offsetbox import AnnotationBbox, OffsetImage


ROOT = Path(__file__).resolve().parent
DATA = json.loads((ROOT / "benchmark-results.json").read_text())
LOGO = ROOT / "tensorplay-ai-logo.png"

BG = "#07111F"
PANEL = "#0C1A2C"
PANEL_ALT = "#10243A"
WHITE = "#F4F8FF"
MUTED = "#8EA5C3"
GRID = "#263A52"
CYAN = "#2ED7FF"
BLUE = "#267BFF"
ORANGE = "#FF9D42"

mpl.rcParams.update(
    {
        "font.family": "DejaVu Sans",
        "axes.facecolor": BG,
        "figure.facecolor": BG,
        "savefig.facecolor": BG,
        "text.color": WHITE,
        "axes.labelcolor": MUTED,
        "xtick.color": MUTED,
        "ytick.color": MUTED,
        "axes.edgecolor": GRID,
        "svg.fonttype": "none",
    }
)


def add_logo(fig: plt.Figure) -> None:
    rgba = plt.imread(LOGO)
    alpha = rgba[..., 3].copy() if rgba.shape[-1] == 4 else 1.0 - rgba[..., :3].mean(axis=-1)
    white_logo = np.ones((*alpha.shape, 4), dtype=float)
    white_logo[..., 3] = alpha
    fig.add_artist(
        AnnotationBbox(
            OffsetImage(white_logo, zoom=0.105),
            (0.84, 0.927),
            xycoords="figure fraction",
            frameon=False,
        )
    )


def header(fig: plt.Figure, title: str, subtitle: str) -> None:
    fig.text(0.055, 0.905, "TensorMind 1.5 Preview", fontsize=13, color=CYAN, weight="bold")
    fig.text(0.055, 0.843, title, fontsize=28, weight="bold")
    fig.text(0.055, 0.792, subtitle, fontsize=11, color=MUTED)
    add_logo(fig)


def footer(fig: plt.Figure) -> None:
    fig.text(
        0.055,
        0.055,
        "Protocol  lm-eval 0.4.12 · SGLang 0.5.14 · 0-shot · full datasets · batch 48 · fixed seeds",
        fontsize=8.5,
        color=MUTED,
    )
    fig.text(0.955, 0.055, "Accuracy, higher is better", ha="right", fontsize=8.5, color=MUTED)


def save(fig: plt.Figure, name: str) -> None:
    fig.savefig(ROOT / f"{name}.png", dpi=200)
    fig.savefig(ROOT / f"{name}.svg")
    plt.close(fig)


def render_suite() -> None:
    results = DATA["results"]
    metrics = [
        ("CMMLU", results["cmmlu"]),
        ("AGIEval-CN", results["agieval_cn"]),
        ("A-CLUE", results["a_clue"]),
        ("C-Eval", results["c_eval"]),
        ("TMMLU+", results["tmmlu_plus"]),
    ]

    fig = plt.figure(figsize=(14, 7.875), dpi=200)
    header(fig, "Chinese benchmark suite", "Five full-dataset evaluations under one matched zero-shot protocol")
    gs = fig.add_gridspec(1, 12, left=0.095, right=0.955, top=0.72, bottom=0.13, wspace=1.2)

    ax = fig.add_subplot(gs[0, :8])
    ax.set_facecolor(PANEL)
    for spine in ax.spines.values():
        spine.set_visible(False)
    names = [name for name, _ in metrics][::-1]
    values = [value for _, value in metrics][::-1]
    y = np.arange(len(metrics))
    ax.barh(y, values, height=0.46, color=[BLUE, CYAN, BLUE, CYAN, BLUE], alpha=0.95)
    ax.set_xlim(0, 35)
    ax.set_yticks(y, names, fontsize=10.5)
    ax.set_xticks([0, 10, 20, 30])
    ax.tick_params(axis="both", length=0, pad=10)
    ax.grid(axis="x", color=GRID, linewidth=0.8, alpha=0.75)
    ax.set_axisbelow(True)
    for yi, value in enumerate(values):
        ax.text(value + 0.45, yi, f"{value:.4f}", va="center", fontsize=10, color=WHITE, weight="bold")
    ax.set_xlabel("Accuracy (%)", loc="right", fontsize=9, labelpad=10)

    ax_card = fig.add_subplot(gs[0, 9:])
    ax_card.set_facecolor(PANEL_ALT)
    ax_card.set_xticks([])
    ax_card.set_yticks([])
    for spine in ax_card.spines.values():
        spine.set_visible(False)
    ax_card.text(0.10, 0.86, "FIVE-SUITE MACRO", fontsize=8.5, color=CYAN, weight="bold", transform=ax_card.transAxes)
    ax_card.text(0.10, 0.64, f"{results['five_suite_macro']:.4f}", fontsize=36, color=WHITE, weight="bold", transform=ax_card.transAxes)
    ax_card.text(0.10, 0.53, "full-dataset accuracy", fontsize=9.5, color=MUTED, transform=ax_card.transAxes)
    ax_card.plot([0.10, 0.90], [0.43, 0.43], color=GRID, linewidth=1.0, transform=ax_card.transAxes)
    ax_card.text(0.10, 0.32, "5", fontsize=19, color=CYAN, weight="bold", transform=ax_card.transAxes)
    ax_card.text(0.21, 0.33, "benchmark suites", fontsize=9.5, color=MUTED, transform=ax_card.transAxes)
    ax_card.text(0.10, 0.18, "0-shot", fontsize=19, color=ORANGE, weight="bold", transform=ax_card.transAxes)
    ax_card.text(0.52, 0.19, "matched protocol", fontsize=9.5, color=MUTED, transform=ax_card.transAxes)

    footer(fig)
    save(fig, "benchmark-suite")


def render_scorecard() -> None:
    results = DATA["results"]
    cards = [
        ("CMMLU", results["cmmlu"], BLUE),
        ("AGIEval-CN", results["agieval_cn"], CYAN),
        ("A-CLUE", results["a_clue"], BLUE),
        ("C-Eval", results["c_eval"], ORANGE),
        ("TMMLU+", results["tmmlu_plus"], CYAN),
        ("5-suite macro", results["five_suite_macro"], WHITE),
    ]

    fig = plt.figure(figsize=(14, 7.875), dpi=200)
    header(fig, "Benchmark scorecard", "TensorMind 1.5 Preview · full-dataset accuracy (%)")
    gs = fig.add_gridspec(2, 3, left=0.08, right=0.92, top=0.70, bottom=0.18, wspace=0.10, hspace=0.14)
    for idx, (name, value, color) in enumerate(cards):
        ax = fig.add_subplot(gs[idx // 3, idx % 3])
        ax.set_facecolor(PANEL_ALT if idx == 5 else PANEL)
        ax.set_xticks([])
        ax.set_yticks([])
        for spine in ax.spines.values():
            spine.set_visible(False)
        ax.add_patch(plt.Rectangle((0.0, 0.0), 0.018, 1.0, color=color, transform=ax.transAxes, lw=0))
        ax.text(0.09, 0.70, name.upper(), fontsize=9, color=MUTED, weight="bold", transform=ax.transAxes)
        ax.text(0.09, 0.29, f"{value:.4f}", fontsize=27, color=color, weight="bold", transform=ax.transAxes)
        ax.text(0.09, 0.12, "accuracy", fontsize=8.5, color=MUTED, transform=ax.transAxes)

    footer(fig)
    save(fig, "benchmark-matrix")


if __name__ == "__main__":
    render_suite()
    render_scorecard()