"""Research-grade benchmark visualizations for the OBLITERATUS dashboard.

Generates publication-quality figures from benchmark results, designed
for rapid visual comprehension by AI researchers.

Chart types:
  1. Pareto Frontier — refusal rate vs perplexity (the money chart)
  2. Method Radar — multi-axis capability profile per method
  3. Metric Bars — grouped bar chart of all metrics side-by-side
  4. Model Scaling — how a technique degrades across model sizes
  5. Timing Efficiency — time vs quality scatter (bang for buck)
  6. Layer Heatmap — strong layer distribution across methods/models

All functions return matplotlib Figure objects compatible with
Gradio's gr.Plot component.
"""

from __future__ import annotations


import re

import matplotlib
matplotlib.use("Agg")  # Non-interactive backend for server use
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import numpy as np


def _sanitize_label(text: str, max_len: int = 80) -> str:
    """Strip filesystem paths, tokens, and overly-long strings from labels."""
    text = re.sub(r"(/[a-zA-Z0-9_./-]{3,})", lambda m: m.group(0).rsplit("/", 1)[-1], text)
    text = re.sub(r"\bhf_[A-Za-z0-9]{6,}\b", "<TOKEN>", text)
    text = re.sub(r"\b[0-9a-fA-F]{32,}\b", "<REDACTED>", text)
    if len(text) > max_len:
        text = text[: max_len - 3] + "..."
    return text


# ── Color palette (research-paper friendly, colorblind-safe) ──────────────

PALETTE = {
    "basic": "#4C72B0",
    "advanced": "#55A868",
    "aggressive": "#C44E52",
    "surgical": "#8172B3",
    "optimized": "#CCB974",
    "inverted": "#64B5CD",
    "nuclear": "#E5583B",
}

MODEL_PALETTE = [
    "#4C72B0", "#55A868", "#C44E52", "#8172B3",
    "#CCB974", "#64B5CD", "#E5583B", "#917254",
]

def _style_ax(ax, title: str = "", xlabel: str = "", ylabel: str = ""):
    """Apply consistent styling to an axis."""
    ax.set_title(title, fontsize=13, fontweight="bold", pad=12)
    ax.set_xlabel(xlabel, fontsize=10)
    ax.set_ylabel(ylabel, fontsize=10)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.tick_params(labelsize=9)
    ax.grid(axis="y", alpha=0.3, linestyle="--")


def _get_color(method: str, idx: int = 0) -> str:
    return PALETTE.get(method, MODEL_PALETTE[idx % len(MODEL_PALETTE)])


# ── 1. Pareto Frontier ───────────────────────────────────────────────────

def plot_pareto_frontier(results: list[dict], title_suffix: str = "") -> plt.Figure:
    """Refusal rate vs perplexity Pareto frontier.

    The most important chart for abliteration research: shows the
    capability-safety tradeoff.  Points in the bottom-left are ideal
    (low refusal AND low perplexity).

    Pareto-optimal points are connected with a frontier line.
    """
    fig, ax = plt.subplots(figsize=(8, 6))

    valid = [r for r in results if r.get("perplexity") is not None and r.get("refusal_rate") is not None]
    if not valid:
        ax.text(0.5, 0.5, "No valid data points", ha="center", va="center", fontsize=14, transform=ax.transAxes)
        _style_ax(ax, f"Pareto Frontier{title_suffix}")
        fig.tight_layout()
        return fig

    xs = [r["refusal_rate"] * 100 for r in valid]  # percentage
    ys = [r["perplexity"] for r in valid]
    labels = [_sanitize_label(r.get("method", r.get("model_short", "?"))) for r in valid]
    colors = [_get_color(r.get("method", ""), i) for i, r in enumerate(valid)]

    # Scatter
    for i, (x, y, label, color) in enumerate(zip(xs, ys, labels, colors)):
        ax.scatter(x, y, c=color, s=120, zorder=5, edgecolors="white", linewidth=1.5)
        # Label offset to avoid overlap
        offset_x = 2 if i % 2 == 0 else -2
        offset_y = 0.3 if i % 3 == 0 else -0.3
        ax.annotate(
            label, (x, y),
            textcoords="offset points", xytext=(offset_x, 8 + offset_y),
            fontsize=8.5, fontweight="bold", color=color,
            ha="center",
        )

    # Pareto frontier line
    points = sorted(zip(xs, ys), key=lambda p: p[0])
    frontier_x, frontier_y = [], []
    best_y = float("inf")
    for px, py in points:
        if py <= best_y:
            frontier_x.append(px)
            frontier_y.append(py)
            best_y = py
    if len(frontier_x) > 1:
        ax.plot(frontier_x, frontier_y, "--", color="#888888", alpha=0.5, linewidth=1.5, label="Pareto frontier")

    # Ideal zone annotation
    ax.annotate(
        "IDEAL", xy=(0, min(ys) * 0.9),
        fontsize=9, color="#2ecc71", fontweight="bold", alpha=0.6,
    )
    ax.annotate(
        "", xy=(5, min(ys)), xytext=(max(xs) * 0.7, max(ys) * 0.8),
        arrowprops=dict(arrowstyle="->", color="#cccccc", lw=1.5),
    )

    _style_ax(ax, f"Capability-Safety Pareto Frontier{title_suffix}",
              "Refusal Rate (%)", "Perplexity (lower = better)")
    ax.set_xlim(left=-2)
    if len(frontier_x) > 1:
        ax.legend(loc="upper right", fontsize=8)
    # Fixed margins avoid renderer/version-dependent tight-layout warnings for
    # sparse (especially single-point) frontiers under strict warning policy.
    fig.subplots_adjust(top=0.88, bottom=0.14, left=0.12, right=0.96)
    return fig


# ── 2. Method Radar Chart ────────────────────────────────────────────────

def plot_method_radar(results: list[dict], title_suffix: str = "") -> plt.Figure:
    """Multi-axis radar chart comparing methods across all metrics.

    Axes: Refusal Removal, Coherence, Low Perplexity, Speed, Layer Coverage.
    Larger area = better overall profile.
    """
    valid = [r for r in results if r.get("perplexity") is not None]
    if not valid:
        fig, ax = plt.subplots(figsize=(7, 7))
        ax.text(0.5, 0.5, "No valid data points", ha="center", va="center", fontsize=14, transform=ax.transAxes)
        fig.tight_layout()
        return fig

    # Metrics to plot (name, how to extract, higher=better?)
    categories = [
        "Refusal\nRemoval",
        "Coherence",
        "Low\nPerplexity",
        "Speed",
        "Layer\nCoverage",
    ]
    n_cats = len(categories)

    # Normalize all metrics to 0-1 (higher = better)
    max_ppl = max((r.get("perplexity", 1) for r in valid), default=50)
    max_time = max((r.get("time_s", 1) for r in valid), default=300)
    max_layers = max((r.get("strong_layers", 1) for r in valid), default=20)

    fig, ax = plt.subplots(figsize=(7, 7), subplot_kw=dict(polar=True))
    angles = np.linspace(0, 2 * np.pi, n_cats, endpoint=False).tolist()
    angles += angles[:1]  # close the polygon

    for i, r in enumerate(valid):
        label = _sanitize_label(r.get("method", r.get("model_short", f"run-{i}")))
        color = _get_color(r.get("method", ""), i)

        refusal_removal = 1.0 - (r.get("refusal_rate") or 1.0)
        coherence = r.get("coherence") or 0.0
        low_ppl = 1.0 - min((r.get("perplexity", max_ppl) / max_ppl), 1.0) if max_ppl > 0 else 0
        speed = 1.0 - min((r.get("time_s", max_time) / max_time), 1.0) if max_time > 0 else 0
        layer_cov = min((r.get("strong_layers", 0) / max(max_layers, 1)), 1.0)

        values = [refusal_removal, coherence, low_ppl, speed, layer_cov]
        values += values[:1]

        ax.plot(angles, values, "o-", linewidth=2, label=label, color=color, markersize=5)
        ax.fill(angles, values, alpha=0.1, color=color)

    ax.set_xticks(angles[:-1])
    ax.set_xticklabels(categories, fontsize=9)
    ax.set_ylim(0, 1.05)
    ax.set_yticks([0.25, 0.5, 0.75, 1.0])
    ax.set_yticklabels(["25%", "50%", "75%", "100%"], fontsize=7, alpha=0.5)
    ax.legend(loc="upper right", bbox_to_anchor=(1.3, 1.1), fontsize=8)
    ax.set_title(f"Method Profile Radar{title_suffix}", fontsize=13, fontweight="bold", y=1.08)
    fig.tight_layout()
    return fig


# ── 3. Grouped Bar Chart ─────────────────────────────────────────────────

def plot_metric_bars(results: list[dict], title_suffix: str = "") -> plt.Figure:
    """Grouped bar chart of key metrics across methods/models.

    Shows refusal rate, coherence, and normalized perplexity side by side
    for quick visual comparison.
    """
    valid = [r for r in results if r.get("perplexity") is not None or r.get("refusal_rate") is not None]
    if not valid:
        fig, ax = plt.subplots(figsize=(10, 5))
        ax.text(0.5, 0.5, "No valid data points", ha="center", va="center", fontsize=14, transform=ax.transAxes)
        fig.tight_layout()
        return fig

    labels = [_sanitize_label(r.get("method", r.get("model_short", "?"))) for r in valid]
    n = len(valid)

    # Metrics
    refusal = [(1.0 - (r.get("refusal_rate") or 1.0)) * 100 for r in valid]
    coherence = [(r.get("coherence") or 0) * 100 for r in valid]
    max_ppl = max((r.get("perplexity", 1) for r in valid), default=50)
    norm_ppl = [(1.0 - min(r.get("perplexity", max_ppl) / max(max_ppl, 1), 1.0)) * 100 for r in valid]

    x = np.arange(n)
    width = 0.25

    fig, ax = plt.subplots(figsize=(max(10, n * 1.8), 5.5))

    bars1 = ax.bar(x - width, refusal, width, label="Refusal Removal %", color="#e74c3c", alpha=0.85)
    bars2 = ax.bar(x, coherence, width, label="Coherence %", color="#2ecc71", alpha=0.85)
    bars3 = ax.bar(x + width, norm_ppl, width, label="Capability (inv. PPL) %", color="#3498db", alpha=0.85)

    # Value labels on bars
    for bars in [bars1, bars2, bars3]:
        for bar in bars:
            h = bar.get_height()
            if h > 5:
                ax.text(bar.get_x() + bar.get_width() / 2, h + 1, f"{h:.0f}",
                        ha="center", va="bottom", fontsize=7.5, fontweight="bold")

    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=30 if n > 4 else 0, ha="right" if n > 4 else "center", fontsize=9)
    ax.set_ylim(0, 110)
    ax.yaxis.set_major_formatter(mticker.FuncFormatter(lambda v, _: f"{v:.0f}%"))
    ax.legend(loc="upper right", fontsize=8.5)
    _style_ax(ax, f"Metric Comparison{title_suffix}", "", "Score (%)")
    fig.tight_layout()
    return fig


# ── 4. Timing Efficiency Scatter ─────────────────────────────────────────

def plot_timing_efficiency(results: list[dict], title_suffix: str = "") -> plt.Figure:
    """Time vs quality scatter — bang for your compute buck.

    X-axis: wall-clock time.  Y-axis: composite quality score.
    Bubble size: number of strong layers touched.
    Points in the top-left are most efficient.
    """
    valid = [r for r in results if r.get("time_s") is not None and r.get("refusal_rate") is not None]
    if not valid:
        fig, ax = plt.subplots(figsize=(8, 5))
        ax.text(0.5, 0.5, "No valid data points", ha="center", va="center", fontsize=14, transform=ax.transAxes)
        fig.tight_layout()
        return fig

    fig, ax = plt.subplots(figsize=(8, 5.5))

    for i, r in enumerate(valid):
        label = _sanitize_label(r.get("method", r.get("model_short", "?")))
        color = _get_color(r.get("method", ""), i)
        x = r["time_s"]
        # Composite quality: weighted sum of refusal removal and coherence
        refusal_removal = 1.0 - (r.get("refusal_rate") or 1.0)
        coherence = r.get("coherence") or 0.0
        quality = 0.6 * refusal_removal + 0.4 * coherence
        size = max(40, (r.get("strong_layers", 5)) * 12)

        ax.scatter(x, quality * 100, s=size, c=color, alpha=0.8, edgecolors="white", linewidth=1.5, zorder=5)
        ax.annotate(label, (x, quality * 100), textcoords="offset points", xytext=(8, 5),
                    fontsize=8.5, fontweight="bold", color=color)

    # Efficiency frontier
    ax.annotate(
        "EFFICIENT", xy=(min(r["time_s"] for r in valid) * 0.8, 95),
        fontsize=9, color="#2ecc71", fontweight="bold", alpha=0.5,
    )

    _style_ax(ax, f"Compute Efficiency{title_suffix}",
              "Wall-Clock Time (seconds)", "Quality Score (% — higher = better)")
    ax.yaxis.set_major_formatter(mticker.FuncFormatter(lambda v, _: f"{v:.0f}%"))
    fig.tight_layout()
    return fig


# ── 5. MoE Metrics Panel ─────────────────────────────────────────────────

def plot_moe_metrics(results: list[dict], title_suffix: str = "") -> plt.Figure:
    """MoE-specific metrics: EGA directions, CoT preservation, expert coverage.

    Only meaningful for results that include MoE-aware techniques.
    Shows a stacked horizontal bar chart of MoE feature activation.
    """
    valid = [r for r in results if r.get("ega_expert_dirs", 0) > 0 or r.get("cot_preserved", 0) > 0]
    if not valid:
        # Show all results with a note about no MoE activity
        fig, ax = plt.subplots(figsize=(8, 4))
        all_labels = [_sanitize_label(r.get("method", r.get("model_short", "?"))) for r in results]
        if results:
            ax.barh(all_labels, [0] * len(results), color="#cccccc")
            ax.text(0.5, 0.5, "No MoE-specific features activated\n(use surgical/optimized/nuclear methods)",
                    ha="center", va="center", fontsize=11, color="#999999", transform=ax.transAxes)
        else:
            ax.text(0.5, 0.5, "No data", ha="center", va="center", fontsize=14, transform=ax.transAxes)
        _style_ax(ax, f"MoE Feature Activation{title_suffix}")
        fig.tight_layout()
        return fig

    labels = [_sanitize_label(r.get("method", r.get("model_short", "?"))) for r in valid]
    ega = [r.get("ega_expert_dirs", 0) for r in valid]
    cot = [r.get("cot_preserved", 0) for r in valid]
    safety = [r.get("expert_classified_layers", r.get("ega_safety_layers", 0)) for r in valid]

    fig, ax = plt.subplots(figsize=(9, max(3, len(valid) * 0.8 + 1)))

    y = np.arange(len(valid))
    h = 0.25

    ax.barh(y - h, ega, h, label="EGA Expert Directions", color="#8172B3", alpha=0.85)
    ax.barh(y, cot, h, label="CoT Preserved Layers", color="#CCB974", alpha=0.85)
    ax.barh(y + h, safety, h, label="Expert-Classified Layers", color="#55A868", alpha=0.85)

    # Value labels
    for vals, offset in [(ega, -h), (cot, 0), (safety, h)]:
        for j, v in enumerate(vals):
            if v > 0:
                ax.text(v + 0.3, j + offset, str(v), va="center", fontsize=8, fontweight="bold")

    ax.set_yticks(y)
    ax.set_yticklabels(labels, fontsize=9)
    ax.legend(loc="lower right", fontsize=8)
    _style_ax(ax, f"MoE Feature Activation{title_suffix}", "Count", "")
    fig.tight_layout()
    return fig


# ── 6. Multi-Model Scaling Chart ─────────────────────────────────────────

def plot_model_scaling(results: list[dict], title_suffix: str = "") -> plt.Figure:
    """How metrics change as model size grows.

    For multi-model benchmarks: plots refusal rate and perplexity
    against model name, showing scaling behavior.
    """
    valid = [r for r in results if r.get("perplexity") is not None]
    if not valid:
        fig, ax = plt.subplots(figsize=(10, 5))
        ax.text(0.5, 0.5, "No valid data points", ha="center", va="center", fontsize=14, transform=ax.transAxes)
        fig.tight_layout()
        return fig

    labels = [_sanitize_label(r.get("model_short", r.get("model", "?").split("/")[-1])) for r in valid]
    n = len(valid)
    x = np.arange(n)

    fig, ax1 = plt.subplots(figsize=(max(9, n * 2), 5.5))
    ax2 = ax1.twinx()

    refusal = [(r.get("refusal_rate") or 0) * 100 for r in valid]
    ppl = [r.get("perplexity", 0) for r in valid]

    line1 = ax1.plot(x, refusal, "o-", color="#e74c3c", linewidth=2.5, markersize=8, label="Refusal Rate (%)", zorder=5)
    line2 = ax2.plot(x, ppl, "s--", color="#3498db", linewidth=2.5, markersize=8, label="Perplexity", zorder=5)

    # Fill between for visual emphasis
    ax1.fill_between(x, refusal, alpha=0.08, color="#e74c3c")
    ax2.fill_between(x, ppl, alpha=0.08, color="#3498db")

    ax1.set_xticks(x)
    ax1.set_xticklabels(labels, rotation=20, ha="right", fontsize=9)
    ax1.set_ylabel("Refusal Rate (%)", color="#e74c3c", fontsize=10)
    ax1.tick_params(axis="y", labelcolor="#e74c3c")
    ax1.yaxis.set_major_formatter(mticker.FuncFormatter(lambda v, _: f"{v:.0f}%"))

    ax2.set_ylabel("Perplexity", color="#3498db", fontsize=10)
    ax2.tick_params(axis="y", labelcolor="#3498db")

    # Combined legend
    lines = line1 + line2
    labels_legend = [ln.get_label() for ln in lines]
    ax1.legend(lines, labels_legend, loc="upper left", fontsize=8.5)

    ax1.set_title(f"Cross-Model Scaling{title_suffix}", fontsize=13, fontweight="bold", pad=12)
    ax1.spines["top"].set_visible(False)
    ax2.spines["top"].set_visible(False)
    ax1.grid(axis="y", alpha=0.2, linestyle="--")
    fig.tight_layout()
    return fig


# ── Composite Dashboard ──────────────────────────────────────────────────

def generate_benchmark_dashboard(
    results: list[dict],
    mode: str = "multi_method",
    title_suffix: str = "",
) -> list[plt.Figure]:
    """Generate a full dashboard of benchmark visualizations.

    Args:
        results: List of benchmark result dicts.
        mode: "multi_method" (N methods x 1 model) or
              "multi_model" (1 method x N models).
        title_suffix: Optional suffix for chart titles.

    Returns:
        List of matplotlib Figure objects ready for Gradio gr.Plot.
    """
    if not results:
        return []

    figs = []

    if mode == "multi_method":
        # 1. Pareto frontier (the headline chart)
        figs.append(plot_pareto_frontier(results, title_suffix))
        # 2. Radar chart
        figs.append(plot_method_radar(results, title_suffix))
        # 3. Metric bars
        figs.append(plot_metric_bars(results, title_suffix))
        # 4. Timing efficiency
        figs.append(plot_timing_efficiency(results, title_suffix))
        # 5. MoE metrics (if any)
        if any(r.get("ega_expert_dirs", 0) > 0 for r in results):
            figs.append(plot_moe_metrics(results, title_suffix))
    elif mode == "multi_model":
        # 1. Scaling chart (headline)
        figs.append(plot_model_scaling(results, title_suffix))
        # 2. Metric bars
        figs.append(plot_metric_bars(results, title_suffix))
        # 3. Pareto frontier
        figs.append(plot_pareto_frontier(results, title_suffix))
        # 4. Timing
        figs.append(plot_timing_efficiency(results, title_suffix))
        # 5. MoE metrics
        if any(r.get("ega_expert_dirs", 0) > 0 for r in results):
            figs.append(plot_moe_metrics(results, title_suffix))

    return figs
