import json
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import streamlit as st
from datasets import load_dataset

CALCULATIONS_DATASET_ID = "hoololi/llm-calculations"
NEXT_PRIME_DATASET_ID = "hoololi/llm-next-prime"
MODEL_INFO_PATH = Path("model_info_snapshot.json")

MODE_ORDER = ["raw", "prompted", "tool"]
MODE_COLORS = {"raw": "#6EA8FE", "prompted": "#FFB86B", "tool": "#69DB7C"}

st.set_page_config(page_title="LLM & Arithmetic", page_icon="🔢", layout="wide")


@dataclass(frozen=True)
class DatasetConfig:
    key: str
    title: str
    dataset_id: str
    description: str
    default_stat_view: str
    category_label: str
    category_col: str
    task_label: str = "Task"


DATASETS = {
    "calculations": DatasetConfig(
        key="calculations",
        title="Calculations",
        dataset_id=CALCULATIONS_DATASET_ID,
        description=(
            "This experiment asks LLMs to solve arithmetic operations such as `45 × 23` "
            "under raw, prompted, and calculator-tool modes. Each row is one operation, "
            "one model, and one mode."
        ),
        default_stat_view="Accuracy by category",
        category_label="Operation categories",
        category_col="operation_category",
        task_label="Operation",
    ),
    "next_prime": DatasetConfig(
        key="next_prime",
        title="Next Prime",
        dataset_id=NEXT_PRIME_DATASET_ID,
        description=(
            "This experiment asks LLMs to answer questions of the form: "
            "`What is the smallest prime number that is strictly greater than n?` "
            "Inputs are sampled across magnitude ranges and tested in raw, prompted, "
            "and next-prime-tool modes."
        ),
        default_stat_view="Accuracy by magnitude",
        category_label="Magnitude groups",
        category_col="magnitude_group",
        task_label="Input number",
    ),
}


@st.cache_data(show_spinner=False)
def load_hf_dataset(dataset_id: str) -> pd.DataFrame:
    ds = load_dataset(dataset_id, split="train")
    df = ds.to_pandas()
    return normalize_frame(df)


@st.cache_data(show_spinner=False)
def load_model_info() -> Dict[str, Any]:
    if not MODEL_INFO_PATH.exists():
        return {}
    return json.loads(MODEL_INFO_PATH.read_text(encoding="utf-8"))


def normalize_frame(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    if "operation" not in df.columns and "question" in df.columns:
        df["operation"] = df["question"]
    if "operation_id" not in df.columns and "case_id" in df.columns:
        df["operation_id"] = df["case_id"]
    if "operation_category" not in df.columns and "category" in df.columns:
        df["operation_category"] = df["category"]
    if "correct_result_str" not in df.columns and "correct_result" in df.columns:
        df["correct_result_str"] = df["correct_result"].astype("string")
    if "extracted_answer_str" not in df.columns and "extracted_answer" in df.columns:
        df["extracted_answer_str"] = df["extracted_answer"].astype("string")

    for col in ["cost", "latency_seconds", "total_tokens", "input_tokens", "output_tokens"]:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")
    if "mode" in df.columns:
        df["mode"] = pd.Categorical(df["mode"].astype(str), categories=MODE_ORDER, ordered=True)
    return df


def sanitize_text(value: Any) -> str:
    if value is None or pd.isna(value):
        return ""
    text = str(value)
    text = re.sub(r"'user_id':\s*'[^']+'", "'user_id': '<redacted>'", text)
    text = re.sub(r'"user_id"\s*:\s*"[^"]+"', '"user_id": "<redacted>"', text)
    text = re.sub(r"user_[A-Za-z0-9]+", "user_<redacted>", text)
    return text


def bool_mean_pct(s: pd.Series) -> float:
    return float(s.eq(True).mean() * 100)


def selected_models(model_a: str, model_b: str) -> List[str]:
    models = [model_a]
    if model_b and model_b != model_a:
        models.append(model_b)
    return models


def sort_magnitude(values: List[Any]) -> List[str]:
    def key(v: Any) -> int:
        match = re.search(r"10\^(\d+)", str(v))
        return int(match.group(1)) if match else 999

    return sorted([str(v) for v in values if pd.notna(v)], key=key)


def summarize(df: pd.DataFrame, model_order: Optional[List[str]] = None) -> pd.DataFrame:
    if df.empty:
        return pd.DataFrame()
    summary = (
        df.groupby(["model", "mode"], observed=True)
        .agg(
            rows=("mode", "size"),
            accuracy_pct=("final_result_correct", bool_mean_pct),
            run_success_pct=("run_success", bool_mean_pct),
            avg_total_tokens=("total_tokens", "mean"),
            avg_cost=("cost", "mean"),
            total_cost=("cost", "sum"),
            avg_latency_seconds=("latency_seconds", "mean"),
            median_latency_seconds=("latency_seconds", "median"),
        )
        .reset_index()
    )
    summary["mode"] = pd.Categorical(summary["mode"].astype(str), categories=MODE_ORDER, ordered=True)
    if model_order:
        summary["model"] = pd.Categorical(summary["model"], categories=model_order, ordered=True)
    summary = summary.sort_values(["model", "mode"])
    summary["model"] = summary["model"].astype(str)
    summary["mode"] = summary["mode"].astype(str)
    return summary.round(4)


def metric_chart(summary: pd.DataFrame, metric: str, yaxis_title: str, percent: bool = False) -> go.Figure:
    fig = go.Figure()
    if summary.empty or metric not in summary.columns:
        return fig
    model_order = summary["model"].drop_duplicates().tolist()
    for mode in MODE_ORDER:
        part = summary[summary["mode"] == mode]
        fig.add_bar(
            x=part["model"],
            y=part[metric],
            name=mode,
            marker_color=MODE_COLORS[mode],
            hovertemplate=f"{yaxis_title}: %{{y:.4g}}<extra>{mode}</extra>",
        )
    fig.update_layout(
        template="plotly_dark",
        barmode="group",
        height=300,
        margin=dict(l=35, r=10, t=20, b=45),
        legend=dict(orientation="h", yanchor="top", y=0.99, xanchor="right", x=0.99),
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(255,255,255,0.035)",
        yaxis_title=yaxis_title,
        xaxis_title=None,
    )
    fig.update_xaxes(categoryorder="array", categoryarray=model_order, tickangle=0)
    fig.update_yaxes(gridcolor="rgba(255,255,255,0.12)")
    if percent:
        fig.update_yaxes(range=[0, 100], tickvals=[0, 20, 40, 60, 80, 100])
    return fig


def grouped_accuracy_chart(df: pd.DataFrame, group_col: str, title: str, group_order: Optional[List[str]] = None) -> go.Figure:
    if df.empty or group_col not in df.columns:
        return go.Figure()
    grouped = (
        df.groupby(["model", group_col, "mode"], observed=True)
        .agg(rows=("mode", "size"), accuracy_pct=("final_result_correct", bool_mean_pct))
        .reset_index()
    )
    grouped[group_col] = grouped[group_col].astype(str)
    model_count = grouped["model"].nunique()
    if group_order is None:
        group_order = sorted(grouped[group_col].dropna().astype(str).unique().tolist())

    fig = px.bar(
        grouped,
        x=group_col,
        y="accuracy_pct",
        color="mode",
        facet_col="model" if model_count > 1 else None,
        barmode="group",
        category_orders={group_col: group_order, "mode": MODE_ORDER},
        color_discrete_map=MODE_COLORS,
        hover_data={"rows": True, "accuracy_pct": ":.2f"},
        labels={"accuracy_pct": "Accuracy (%)", group_col: ""},
        title=title,
    )
    fig.update_layout(
        template="plotly_dark",
        height=420 if model_count > 1 else 380,
        margin=dict(l=35, r=10, t=45, b=55),
        legend=dict(orientation="h", yanchor="top", y=0.98, xanchor="right", x=0.99),
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(255,255,255,0.035)",
        yaxis_title="Accuracy (%)",
    )
    fig.update_xaxes(categoryorder="array", categoryarray=group_order)
    fig.update_yaxes(range=[0, 100], gridcolor="rgba(255,255,255,0.12)")
    fig.for_each_annotation(lambda a: a.update(text=a.text.replace("model=", "")))
    return fig


def tool_pipeline_summary(df: pd.DataFrame) -> pd.DataFrame:
    tool = df[df["mode"].astype(str) == "tool"].copy()
    if tool.empty:
        return pd.DataFrame()
    tool["final_correct_bool"] = tool["final_result_correct"].eq(True)
    tool["tool_correct_bool"] = tool.get("tool_result_correct", pd.Series(False, index=tool.index)).eq(True)
    tool["final_or_tool_correct"] = tool["final_correct_bool"] | tool["tool_correct_bool"]
    return (
        tool.groupby(["model_family", "model"], observed=True)
        .agg(
            rows=("mode", "size"),
            final_answer_accuracy_pct=("final_correct_bool", lambda s: s.mean() * 100),
            tool_output_accuracy_pct=("tool_correct_bool", lambda s: s.mean() * 100),
            final_or_tool_accuracy_pct=("final_or_tool_correct", lambda s: s.mean() * 100),
            run_success_pct=("run_success", bool_mean_pct),
        )
        .reset_index()
        .sort_values(["final_or_tool_accuracy_pct", "final_answer_accuracy_pct"], ascending=False)
        .round(3)
    )


def tradeoff_chart(df: pd.DataFrame) -> go.Figure:
    if df.empty:
        return go.Figure()
    tradeoff = (
        df.groupby(["model_family", "model", "mode"], observed=True)
        .agg(
            accuracy_pct=("final_result_correct", bool_mean_pct),
            success_rate_pct=("run_success", bool_mean_pct),
            median_latency=("latency_seconds", "median"),
            avg_cost=("cost", "mean"),
        )
        .reset_index()
        .dropna(subset=["avg_cost", "median_latency", "accuracy_pct"])
    )
    if tradeoff.empty:
        return go.Figure()
    tradeoff["viable"] = (tradeoff["accuracy_pct"] >= 95) & (tradeoff["success_rate_pct"] >= 95)
    tradeoff["label"] = tradeoff["model"].str.split("/").str[-1] + " (" + tradeoff["mode"].astype(str) + ")"
    min_positive_cost = tradeoff.loc[tradeoff["avg_cost"] > 0, "avg_cost"].min()
    zero_cost_display = (min_positive_cost / 10) if pd.notna(min_positive_cost) else 1e-9
    tradeoff["avg_cost_display"] = tradeoff["avg_cost"].where(tradeoff["avg_cost"] > 0, zero_cost_display)
    fig = px.scatter(
        tradeoff,
        x="avg_cost_display",
        y="median_latency",
        color="accuracy_pct",
        symbol="mode",
        size="success_rate_pct",
        hover_name="label",
        hover_data={"model_family": True, "avg_cost": ":.6f", "avg_cost_display": False, "median_latency": ":.3f", "accuracy_pct": ":.2f", "success_rate_pct": ":.2f"},
        color_continuous_scale=[[0, "#D73027"], [0.5, "#FEE08B"], [1, "#1A9850"]],
        range_color=[50, 100],
        symbol_map={"raw": "circle", "prompted": "x", "tool": "square"},
    )
    viable = tradeoff[tradeoff["viable"]]
    fig.add_trace(
        go.Scatter(
            x=viable["avg_cost_display"],
            y=viable["median_latency"],
            mode="markers",
            marker=dict(symbol="circle-open", size=18, color="white", line=dict(color="black", width=2)),
            name="viable ≥95% accuracy/success",
            hoverinfo="skip",
        )
    )
    fig.update_xaxes(type="log", title="Average cost per call, USD (zero-cost points shown near axis minimum)", gridcolor="rgba(255,255,255,0.12)")
    fig.update_yaxes(type="log", title="Median latency, seconds", gridcolor="rgba(255,255,255,0.12)")
    fig.update_layout(
        template="plotly_dark",
        height=520,
        margin=dict(l=35, r=10, t=35, b=45),
        title="Cost / latency trade-off, colored by accuracy",
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(255,255,255,0.035)",
        legend=dict(orientation="h", yanchor="bottom", y=-0.28, xanchor="center", x=0.5),
    )
    return fig


def price_per_million(value: Any) -> str:
    if value in (None, ""):
        return "n/a"
    try:
        return f"${float(value) * 1_000_000:.2f}/1M tokens"
    except (TypeError, ValueError):
        return str(value)


def model_card(model: str, model_info: Dict[str, Any]) -> None:
    info = model_info.get(model, {})
    if not info:
        st.caption(f"No OpenRouter snapshot information available for `{model}`.")
        return
    pricing = info.get("pricing", {}) if isinstance(info.get("pricing", {}), dict) else {}
    architecture = info.get("architecture", {}) if isinstance(info.get("architecture", {}), dict) else {}
    provider = info.get("top_provider", {}) if isinstance(info.get("top_provider", {}), dict) else {}
    supported = info.get("supported_parameters", []) or []
    description = (info.get("description") or "").strip().replace("\n", " ")
    if len(description) > 260:
        description = description[:260].rstrip() + "…"
    st.markdown(f"#### {info.get('name') or model}")
    st.code(model, language=None)
    st.markdown(
        f"""
- **Context length:** {info.get('context_length') or provider.get('context_length') or 'n/a'}
- **Max completion tokens:** {provider.get('max_completion_tokens') or 'n/a'}
- **Input price:** {price_per_million(pricing.get('prompt') or pricing.get('input'))}
- **Output price:** {price_per_million(pricing.get('completion') or pricing.get('output'))}
- **Tokenizer:** {architecture.get('tokenizer') or 'n/a'}
- **Supported parameters:** {', '.join(supported[:7]) if supported else 'n/a'}{'…' if len(supported) > 7 else ''}
"""
    )
    if description:
        st.caption(description)


def apply_global_filters(df: pd.DataFrame, config: DatasetConfig, key_prefix: str) -> tuple[pd.DataFrame, List[str]]:
    models = sorted(df["model"].dropna().unique().tolist())
    categories = sorted(df[config.category_col].dropna().astype(str).unique().tolist()) if config.category_col in df.columns else []
    if config.category_col == "magnitude_group":
        categories = sort_magnitude(categories)

    with st.container(border=True):
        c1, c2, c3, c4 = st.columns([1.15, 1.15, 1.45, 1])
        with c1:
            model_a = st.selectbox("Model A / primary", models, index=0, key=f"{key_prefix}_model_a")
        with c2:
            model_b = st.selectbox("Model B / compare optional", [""] + models, index=0, key=f"{key_prefix}_model_b")
        with c3:
            selected_categories = st.multiselect(
                config.category_label,
                categories,
                default=categories,
                key=f"{key_prefix}_categories",
            )
        with c4:
            run_error_policy = st.selectbox(
                "Run errors in metrics",
                ["Count as failures", "Exclude from metrics"],
                index=0,
                key=f"{key_prefix}_run_errors",
            )

    st.caption("Charts use the selected model names directly. If two models are selected, category/magnitude charts are split by model.")

    chosen = selected_models(model_a, model_b)
    view = df[df["model"].isin(chosen)].copy()
    if selected_categories and config.category_col in view.columns:
        view = view[view[config.category_col].astype(str).isin(selected_categories)]
    if run_error_policy == "Exclude from metrics" and "run_success" in view.columns:
        view = view[view["run_success"].eq(True)]
    return view, chosen


def render_dataset_intro(df: pd.DataFrame, config: DatasetConfig) -> None:
    st.markdown(f"### {config.title}")
    st.markdown(config.description)
    st.markdown(f"Dataset: [`{config.dataset_id}`](https://huggingface.co/datasets/{config.dataset_id})")
    c1, c2, c3, c4 = st.columns(4)
    c1.metric("Rows", f"{len(df):,}")
    c2.metric("Tasks", f"{df['operation_id'].nunique():,}" if "operation_id" in df.columns else "n/a")
    c3.metric("Models", f"{df['model'].nunique():,}" if "model" in df.columns else "n/a")
    c4.metric("Modes", f"{df['mode'].nunique():,}" if "mode" in df.columns else "n/a")


def render_statistics(df: pd.DataFrame, config: DatasetConfig, model_info: Dict[str, Any], key_prefix: str) -> None:
    st.caption("Explore aggregate statistics and compare models across raw, prompted, and tool modes.")
    view, chosen = apply_global_filters(df, config, key_prefix + "_stats")
    summary = summarize(view, chosen)

    stat_views = ["Accuracy by category", "Accuracy", "Latency", "Cost", "Tokens", "Tool pipeline", "Cost/latency trade-off"]
    if config.key == "next_prime":
        stat_views[0] = "Accuracy by magnitude"
    default_index = stat_views.index(config.default_stat_view) if config.default_stat_view in stat_views else 0
    stat_view = st.selectbox("Statistic view", stat_views, index=default_index, key=f"{key_prefix}_stat_view")

    if stat_view in ["Accuracy by category", "Accuracy by magnitude"]:
        order = None
        if config.category_col == "magnitude_group" and config.category_col in view.columns:
            order = sort_magnitude(view[config.category_col].dropna().unique().tolist())
        st.plotly_chart(
            grouped_accuracy_chart(view, config.category_col, stat_view, order),
            width="stretch",
        )
    elif stat_view == "Accuracy":
        st.plotly_chart(metric_chart(summary, "accuracy_pct", "Accuracy (%)", True), width="stretch")
    elif stat_view == "Latency":
        st.plotly_chart(metric_chart(summary, "avg_latency_seconds", "Average latency, seconds"), width="stretch")
    elif stat_view == "Cost":
        st.plotly_chart(metric_chart(summary, "avg_cost", "Average cost, USD"), width="stretch")
    elif stat_view == "Tokens":
        st.plotly_chart(metric_chart(summary, "avg_total_tokens", "Average total tokens"), width="stretch")
    elif stat_view == "Tool pipeline":
        st.dataframe(tool_pipeline_summary(view), width="stretch", hide_index=True)
    elif stat_view == "Cost/latency trade-off":
        st.plotly_chart(tradeoff_chart(view), width="stretch")

    st.markdown("#### Summary table")
    st.dataframe(summary, width="stretch", hide_index=True)

    with st.expander("Selected model details from OpenRouter snapshot", expanded=False):
        info_cols = st.columns(max(1, len(chosen)))
        for col, model in zip(info_cols, chosen):
            with col:
                model_card(model, model_info)


def quick_filter(df: pd.DataFrame, filter_name: str, config: DatasetConfig) -> pd.DataFrame:
    out = df.copy()
    if filter_name == "Wrong answers":
        out = out[out["final_result_correct"].ne(True)]
    elif filter_name == "Correct answers":
        out = out[out["final_result_correct"].eq(True)]
    elif filter_name == "Run failures":
        out = out[out["run_success"].ne(True)]
    elif filter_name == "Tool issues":
        out = out[
            (out["mode"].astype(str) == "tool")
            & (
                out.get("tool_called", pd.Series(False, index=out.index)).ne(True)
                | out.get("tool_result_correct", pd.Series(True, index=out.index)).ne(True)
                | out.get("tool_error_type", pd.Series("none", index=out.index)).fillna("none").ne("none")
            )
        ]
    elif filter_name == "Tool correct, final failed":
        out = out[(out["mode"].astype(str) == "tool") & out.get("tool_result_correct", pd.Series(False, index=out.index)).eq(True) & out["final_result_correct"].ne(True)]
    elif filter_name == "Long answers":
        lengths = out["answer_text"].fillna("").astype(str).str.len()
        out = out[lengths >= lengths.quantile(0.90)]
    elif filter_name == "Slowest answers":
        out = out[out["latency_seconds"] >= out["latency_seconds"].quantile(0.90)]
    elif filter_name == "Fastest answers":
        out = out[out["latency_seconds"] <= out["latency_seconds"].quantile(0.10)]
    elif filter_name == "Highest cost":
        if out["cost"].notna().any():
            out = out[out["cost"] >= out["cost"].quantile(0.90)]
    elif filter_name == "High magnitude" and config.key == "next_prime":
        out = out[out.get("magnitude_group", pd.Series("", index=out.index)).astype(str).isin(["10^4", "10^5"])]
    elif filter_name == "Large prime gap" and config.key == "next_prime":
        out = out[out.get("distance_to_next_prime", pd.Series(0, index=out.index)) >= out.get("distance_to_next_prime", pd.Series(0, index=out.index)).quantile(0.90)]
    return out


def record_label(row: pd.Series) -> str:
    status = "✅" if row.get("final_result_correct") == True else "❌"
    return f"{status} {row.get('model')} | {row.get('mode')}"


def display_result_panel(row: Optional[pd.Series], title: str) -> None:
    st.markdown(f"#### {title}")
    if row is None:
        st.info("No matching result.")
        return
    status = "✅ correct" if row.get("final_result_correct") == True else "❌ not correct"
    st.markdown(f"**{row.get('model')}** · `{row.get('mode')}` · {status}")
    m1, m2, m3 = st.columns(3)
    m1.metric("Extracted", str(row.get("extracted_answer_str", row.get("extracted_answer"))))
    m2.metric("Latency", f"{row.get('latency_seconds'):.3g}s" if pd.notna(row.get("latency_seconds")) else "n/a")
    m3.metric("Cost", f"${row.get('cost'):.6f}" if pd.notna(row.get("cost")) else "n/a")
    with st.expander("Diagnostics", expanded=False):
        st.json(
            {
                "run_success": bool(row.get("run_success")) if pd.notna(row.get("run_success")) else None,
                "run_error": sanitize_text(row.get("run_error")),
                "extraction_error_type": row.get("extraction_error_type"),
                "tool_called": row.get("tool_called"),
                "tool_result": row.get("tool_result_str", row.get("tool_result")),
                "tool_result_correct": row.get("tool_result_correct"),
                "tool_error_type": row.get("tool_error_type"),
            }
        )
    st.markdown("**Model answer**")
    st.code(sanitize_text(row.get("answer_text")), language=None)
    if str(row.get("mode")) == "tool":
        with st.expander("Tool details", expanded=False):
            st.json(
                {
                    "tool_expression": row.get("tool_expression"),
                    "tool_result": row.get("tool_result_str", row.get("tool_result")),
                    "tool_error": sanitize_text(row.get("tool_error")),
                    "tool_calls": row.get("tool_calls"),
                }
            )


def render_task_details(case_rows: pd.DataFrame, config: DatasetConfig, key_prefix: str) -> None:
    if case_rows.empty:
        st.info("No rows for this task.")
        return
    first = case_rows.iloc[0]
    with st.container(border=True):
        st.markdown(f"#### {config.task_label}")
        st.write(first.get("operation"))
        meta = {
            "operation_id": first.get("operation_id"),
            "category": first.get("operation_category"),
            "correct_result": first.get("correct_result_str", first.get("correct_result")),
        }
        for col in ["input_number", "magnitude_group", "distance_to_next_prime", "input_is_prime"]:
            if col in case_rows.columns:
                meta[col] = first.get(col)
        st.json(meta)

    available = case_rows.sort_values(["model", "mode"]).reset_index(drop=True)
    labels = available.apply(record_label, axis=1).tolist()
    task_id_for_key = str(first.get("operation_id"))
    c1, c2 = st.columns(2)
    with c1:
        left_label = st.selectbox("Left result", labels, index=0, key=f"{key_prefix}_{task_id_for_key}_left_result")
    with c2:
        right_index = min(1, len(labels) - 1)
        right_label = st.selectbox("Right result", labels, index=right_index, key=f"{key_prefix}_{task_id_for_key}_right_result")
    left_row = available.iloc[labels.index(left_label)] if labels else None
    right_row = available.iloc[labels.index(right_label)] if labels else None

    r1, r2 = st.columns(2)
    with r1:
        display_result_panel(left_row, "Result A")
    with r2:
        display_result_panel(right_row, "Result B")


def render_details(df: pd.DataFrame, config: DatasetConfig, key_prefix: str) -> None:
    st.caption("Pick a task, then compare two model × mode results side by side.")
    filters = [
        "All",
        "Wrong answers",
        "Correct answers",
        "Run failures",
        "Tool issues",
        "Tool correct, final failed",
        "Long answers",
        "Slowest answers",
        "Fastest answers",
        "Highest cost",
    ]
    if config.key == "next_prime":
        filters += ["High magnitude", "Large prime gap"]

    c1, c2, c3 = st.columns([1.1, 1.1, 1.8])
    previous_filter = st.session_state.get(f"{key_prefix}_active_filter")
    with c1:
        filter_name = st.selectbox("Quick filter", filters, index=0, key=f"{key_prefix}_quick_filter")
    selected_task_key = f"{key_prefix}_selected_task_id"
    if previous_filter != filter_name:
        st.session_state[f"{key_prefix}_active_filter"] = filter_name
        st.session_state.pop(selected_task_key, None)
    with c2:
        modes = st.multiselect("Modes included", MODE_ORDER, default=MODE_ORDER, key=f"{key_prefix}_detail_modes")
    with c3:
        search_text = st.text_input("Search task or answer", value="", key=f"{key_prefix}_search")

    filtered = quick_filter(df, filter_name, config)
    if modes:
        filtered = filtered[filtered["mode"].astype(str).isin(modes)]
    if search_text.strip():
        needle = search_text.strip().lower()
        filtered = filtered[
            filtered["operation"].fillna("").str.lower().str.contains(needle, regex=False)
            | filtered["answer_text"].fillna("").str.lower().str.contains(needle, regex=False)
            | filtered["model"].fillna("").str.lower().str.contains(needle, regex=False)
        ]

    st.caption(
        f"Quick filter matches {len(filtered):,} rows across "
        f"{filtered['operation_id'].nunique() if 'operation_id' in filtered.columns else 0:,} tasks. "
        "The comparison area below uses only rows matching the current filter/mode/search selection."
    )

    if filter_name == "Long answers":
        filtered = filtered.assign(_answer_length=filtered["answer_text"].fillna("").astype(str).str.len())
        task_ids = (
            filtered.groupby("operation_id", observed=True)["_answer_length"]
            .max()
            .sort_values(ascending=False)
            .index.astype(str)
            .tolist()
        )
    elif filter_name == "Slowest answers":
        task_ids = (
            filtered.groupby("operation_id", observed=True)["latency_seconds"]
            .max()
            .sort_values(ascending=False)
            .index.astype(str)
            .tolist()
        )
    elif filter_name == "Fastest answers":
        task_ids = (
            filtered.groupby("operation_id", observed=True)["latency_seconds"]
            .min()
            .sort_values(ascending=True)
            .index.astype(str)
            .tolist()
        )
    else:
        task_ids = filtered["operation_id"].dropna().astype(str).drop_duplicates().sort_values().tolist()
    if not task_ids:
        st.info("No tasks match the current filters.")
        return

    if st.session_state.get(selected_task_key) not in task_ids:
        st.session_state[selected_task_key] = task_ids[0]

    current_index = task_ids.index(st.session_state[selected_task_key])
    n1, n2, n3 = st.columns([0.7, 2.6, 0.7])
    with n1:
        if st.button("◀ Previous", key=f"{key_prefix}_prev", disabled=current_index == 0):
            st.session_state[selected_task_key] = task_ids[current_index - 1]
    with n3:
        if st.button("Next ▶", key=f"{key_prefix}_next", disabled=current_index == len(task_ids) - 1):
            st.session_state[selected_task_key] = task_ids[current_index + 1]

    current_index = task_ids.index(st.session_state[selected_task_key])
    with n2:
        selected_id = st.selectbox(
            "Task",
            task_ids,
            index=current_index,
            key=selected_task_key,
        )

    case_rows = filtered[filtered["operation_id"].astype(str) == selected_id].copy()
    render_task_details(case_rows, config, key_prefix)

    with st.expander("Rows matching current quick filter", expanded=False):
        display_cols = [
            "operation_id",
            "operation_category",
            "model",
            "mode",
            "operation",
            "correct_result_str",
            "extracted_answer_str",
            "final_result_correct",
            "run_success",
            "latency_seconds",
            "cost",
            "tool_error_type",
        ]
        display_cols = [c for c in display_cols if c in filtered.columns]
        st.dataframe(filtered[display_cols].head(500), width="stretch", hide_index=True, height=260)


def render_dataset_tab(df: pd.DataFrame, config: DatasetConfig, model_info: Dict[str, Any]) -> None:
    render_dataset_intro(df, config)
    st.divider()
    stats_tab, details_tab = st.tabs(["Statistics", "Detailed results"])
    with stats_tab:
        render_statistics(df, config, model_info, config.key)
    with details_tab:
        render_details(df, config, config.key)


st.title("🔢 LLM & Arithmetic")
st.markdown(
    "Explore two small, public Hugging Face datasets about LLM behavior on arithmetic-like tasks: "
    "standard calculations and next-prime search."
)

with st.spinner("Loading Hugging Face datasets..."):
    calculations_df = load_hf_dataset(CALCULATIONS_DATASET_ID)
    next_prime_df = load_hf_dataset(NEXT_PRIME_DATASET_ID)
model_info_snapshot = load_model_info()

calc_tab, prime_tab = st.tabs(["Calculations", "Next Prime"])
with calc_tab:
    render_dataset_tab(calculations_df, DATASETS["calculations"], model_info_snapshot)
with prime_tab:
    render_dataset_tab(next_prime_df, DATASETS["next_prime"], model_info_snapshot)