Spaces:

InstaDeepAI
/

ntv3_benchmark

Running

Replace the 2025 benchmark with the revised NTv3 benchmark, keeping the 2025 leaderboard as a page
326d57c verified 11 days ago
27.1 kB
import os

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import streamlit as st
from plotly.subplots import make_subplots

# ---------------------------------------------------------------------
# Page config (must be the first Streamlit command)
# ---------------------------------------------------------------------
st.set_page_config(
    page_title="NTv3 Benchmark",
    layout="wide",
)

import initial_release  # noqa: E402  (reads its CSV on import, so after the page config)

# ---------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------
COLORS = {
    'blue_0': '#004697',
    'blue_1': '#3973fc',
    'blue_2': '#7ea4fc',
    'blue_3': '#c3d5fc',
    'red_1': '#ff554d',
    'green_1': '#00b050',
    'green_3': '#c6e0b4',
    'gold_1': '#fdb932',
    'orange_1': '#ff975e',
    'purple_1': '#9a6ce4',
    'purple_2': '#bb9aef',
    'gray_1': '#808080',
    'gray_2': '#b3b3b3',
    'cyan_1': '#0096b4',
    'magenta_1': '#b428a0',
}

ASSAY_TYPE_MAPPING = {
    'ATAC-seq': 'chromatin accessibility',
    'DNase-seq': 'chromatin accessibility',
    'Histone ChIP-seq': 'histone modifications',
    'TF ChIP-seq': 'TF binding',
    'PRO-cap': 'transcription initiation',
    'eCLIP': 'RNA binding sites',
    'RNA-seq': 'gene expression',
    'ribo-seq': 'mRNA translation',
}

ASSAY_COLORS = {
    'chromatin accessibility': '#004697',
    'histone modifications': '#cc0000',
    'TF binding': '#00b050',
    'transcription initiation': '#ff9900',
    'RNA binding sites': '#9933cc',
    'gene expression': '#009900',
    'mRNA translation': '#ff6699',
    'other': '#808080',
}

CLADE_COLORS = {
    "Animals": COLORS['blue_1'],
    "Plants": COLORS['green_1'],
    "Fungi": COLORS['gold_1'],
    "Protists": COLORS['purple_1'],
}

# CSV model name -> the label the paper figure uses. Other models in the CSV (v1 NTv3
# checkpoints, small BPNet / Residual CNN) are kept in the data but not shown.
MODEL_LABELS = {
    "BPNet arch. 6M": "BPNet",
    "Residual CNN 44M": "Residual CNN",
    "HyenaDNA 7M": "HyenaDNA",
    "Caduceus 7M": "Caduceus",
    "NTv2 500M": "NTv2",
    "Evo2 1B": "Evo2",
    "PlantCAD2 88M": "PlantCAD2",
    "Borzoi": "Borzoi",
    "AlphaGenome": "AlphaGenome",
    "NTv3 700M (tracks)": "NTv3",
}
NTV3 = "NTv3"

# In the paper figure's architecture blocks, with its colours.
MODEL_COLORS = {
    "BPNet": COLORS['cyan_1'],
    "Residual CNN": COLORS['magenta_1'],
    "HyenaDNA": COLORS['gold_1'],
    "Caduceus": COLORS['purple_2'],
    "NTv2": COLORS['gray_1'],
    "Evo2": COLORS['green_3'],
    "PlantCAD2": COLORS['purple_1'],
    "Borzoi": COLORS['gray_2'],
    "AlphaGenome": COLORS['orange_1'],
    NTV3: COLORS['blue_0'],
}

MODEL_TRAINING_STATUS = {
    "BPNet": "SCRATCH",
    "Residual CNN": "SCRATCH",
    "HyenaDNA": "PRE",
    "Caduceus": "PRE",
    "NTv2": "PRE",
    "Evo2": "PRE",
    "PlantCAD2": "PRE",
    "Borzoi": "PRE",
    "AlphaGenome": "PRE",
    NTV3: "POS",
}

MODEL_NAMES = list(MODEL_COLORS.keys())

# CSV species key -> label, grouped by clade, in the paper's order.
SPECIES_GROUPS = {
    "Animals": {"human": "Human", "cattle": "Cattle", "rat": "Rat"},
    "Plants": {
        "arabidopsis": "A. thaliana",
        "rice": "Rice",
        "maize": "Maize",
        "tomato": "Tomato",
    },
    "Fungi": {"yeast": "S. cerevisiae", "fission yeast": "S. pombe"},
    "Protists": {"plasmodium": "P. falciparum", "dictyostelium": "D. discoideum"},
}
SPECIES_LABEL = {k: v for group in SPECIES_GROUPS.values() for k, v in group.items()}
SPECIES_ORDER = list(SPECIES_LABEL)
SPECIES_CLADE = {s: name for name, group in SPECIES_GROUPS.items() for s in group}
SPECIES_ICON = {
    "human": "🧍",
    "cattle": "🐄",
    "rat": "🐀",
    "arabidopsis": "🌱",
    "rice": "🌾",
    "maize": "🌽",
    "tomato": "🍅",
    "yeast": "🍞",
    "fission yeast": "🍺",
    "plasmodium": "🦟",
    "dictyostelium": "🦠",
}
SPECIES_COMMON_NAME = {
    "arabidopsis": "thale cress",
    "yeast": "baker's yeast",
    "fission yeast": "fission yeast",
    "plasmodium": "malaria parasite",
    "dictyostelium": "social amoeba",
}
METRIC_LABEL = "Pearson correlation"

_LAST_UPDATED = "Sep 24, 2026"
_INTRO = """
The **NTv3 Benchmark** is 155 long-range functional genomics tracks for evaluating models
under realistic 32 kb input, single-base-pair output settings. The tracks span chromatin accessibility (ATAC-seq),
histone modifications and transcription factor binding (ChIP-seq), transcription initiation (PRO-cap),
RNA binding (eCLIP), gene expression (RNA-seq) and translation (Ribo-seq).
The benchmark covers 11 phylogenetically diverse species:
**animals** (human, cattle, rat), **plants** (*A. thaliana*, rice, maize, tomato),
**fungi** (*S. cerevisiae*, *S. pombe*) and **protists** (*P. falciparum*, *D. discoideum*).
The metric is the **Pearson correlation** on the test set, per track, averaged over three
seeds where available. Not every model was run on every species, so scores are never
averaged across species: pick a species in the sidebar for its own leaderboard, where every
model shown was evaluated on all of that species' tracks. AlphaGenome is fine-tuned from its
fold-1 checkpoint on human, whose test regions the all-folds checkpoint saw, and from its
all-folds checkpoint on every other species.
"""

HERE = os.path.dirname(os.path.abspath(__file__))  # /app/src
PROJECT_ROOT = os.path.dirname(HERE)               # /app
DATA_DIR = os.path.join(PROJECT_ROOT, "data")
SINGLE_TABLE_PATH = os.path.join(DATA_DIR, "ntv3_benchmark_results.csv")

# ---------------------------------------------------------------------
# Data loading & preprocessing
# ---------------------------------------------------------------------


@st.cache_data
def load_raw_data() -> pd.DataFrame:
    df = pd.read_csv(SINGLE_TABLE_PATH)
    df.columns = [c.strip() for c in df.columns]
    return df


@st.cache_data
def load_track_data() -> pd.DataFrame:
    """One row per (species, track, model): the score averaged over seeds."""
    df = load_raw_data().rename(columns={"Metric": "Score", "model_name": "Model"})
    df["Model"] = df["Model"].map(MODEL_LABELS)
    df = df[df["Model"].notna() & df["species"].isin(SPECIES_ORDER)]
    df["Score"] = pd.to_numeric(df["Score"], errors="coerce")
    keys = ["species", "datasets", "Model", "assay_type", "track_name_clean"]
    out = df.groupby(keys, as_index=False, dropna=False).agg(
        Score=("Score", "mean"), Seeds=("Score", "size")
    )
    out["assay_type"] = out["assay_type"].map(ASSAY_TYPE_MAPPING).fillna("other")
    return out


@st.cache_data
def load_run_data() -> pd.DataFrame:
    """One row per training run: its best checkpoint's step and GPU hours."""
    df = load_raw_data().rename(columns={"model_name": "Model"})
    df["Model"] = df["Model"].map(MODEL_LABELS)
    df = df[df["Model"].notna() & df["species"].isin(SPECIES_ORDER)]
    run_cols = ["species", "Model", "best_step", "best_step_time", "running_time"]
    runs = df.drop_duplicates(run_cols).copy()
    # Training steps only: excludes job setup and validation, whose cadence differs
    # between runs of the benchmark.
    # Parsed one by one: pandas 3 misparses a column mixing fraction lengths.
    hours = runs["best_step_train_time"].map(
        lambda v: pd.Timedelta(v).total_seconds() / 3600 if pd.notna(v) else np.nan
    )
    runs["GPU hours"] = hours * pd.to_numeric(runs["num_gpus"], errors="coerce")
    runs["Tokens (billions)"] = pd.to_numeric(runs["best_step"], errors="coerce") / 1e9
    return runs


_TRACKS = load_track_data()
_RUNS = load_run_data()


def models_in(df: pd.DataFrame) -> list[str]:
    """Models present in `df`, in the paper's block order."""
    present = set(df["Model"])
    return [m for m in MODEL_NAMES if m in present]


# ---------------------------------------------------------------------
# Computation helpers
# ---------------------------------------------------------------------


def build_species_scores(df: pd.DataFrame) -> pd.DataFrame:
    """Mean over tracks per (species, model): the paper figure's points."""
    return df.groupby(["species", "Model"], as_index=False).agg(
        Score=("Score", "mean"),
        Tracks=("datasets", "nunique"),
        Seeds=("Seeds", "max"),
    )


def build_leaderboard(df: pd.DataFrame) -> pd.DataFrame:
    if df.empty:
        return pd.DataFrame(columns=["Model", "Training", "Num entries", "Mean score"])
    agg = df.groupby("Model").agg(
        **{
            "Num entries": ("Score", "count"),
            "Seeds": ("Seeds", "max"),
            "Mean score": ("Score", "mean"),
        }
    ).reset_index()
    agg["Mean score"] = agg["Mean score"].round(3)
    agg["Training"] = agg["Model"].map(MODEL_TRAINING_STATUS).fillna("UNKNOWN")
    agg = agg.sort_values("Mean score", ascending=False).reset_index(drop=True)
    agg = agg[["Model", "Training", "Num entries", "Seeds", "Mean score"]]
    agg.index += 1
    return agg


def build_category_model_df(df: pd.DataFrame) -> pd.DataFrame:
    """Mean score per (assay category, model)."""
    if df.empty:
        return pd.DataFrame(columns=["Category", "Model", "Mean score"])
    out = (
        df.groupby(["assay_type", "Model"], as_index=False)["Score"].mean()
        .rename(columns={"assay_type": "Category", "Score": "Mean score"})
    )
    out["Mean score"] = out["Mean score"].round(3)
    return out


def build_pairwise_scatter_df(df: pd.DataFrame, model_a: str, model_b: str) -> pd.DataFrame:
    """Scores of two models on the tracks both were evaluated on."""
    sub = df[df["Model"].isin([model_a, model_b])]
    wide = sub.pivot_table(
        index=["species", "assay_type", "datasets", "track_name_clean"],
        columns="Model",
        values="Score",
    ).reset_index()
    if model_a not in wide.columns or model_b not in wide.columns:
        return pd.DataFrame()
    wide = wide.dropna(subset=[model_a, model_b])
    wide["Species"] = wide["species"].map(SPECIES_LABEL)
    wide["Track"] = wide["Species"] + " / " + wide["track_name_clean"].astype(str)
    return wide.rename(columns={model_a: "Model A", model_b: "Model B"})


def aggregate_pairs_by_species(pairs: pd.DataFrame) -> pd.DataFrame:
    """One point per species: both models' means over the tracks they share there."""
    clade = {s: name for name, group in SPECIES_GROUPS.items() for s in group}
    out = pairs.groupby(["species", "Species"], as_index=False).agg(
        **{
            "Model A": ("Model A", "mean"),
            "Model B": ("Model B", "mean"),
            "Shared tracks": ("datasets", "size"),
        }
    )
    out["Clade"] = out["species"].map(clade)
    out["Track"] = out["Species"]
    return out


# ---------------------------------------------------------------------
# Plots
# ---------------------------------------------------------------------


def plot_species_overview(scores: pd.DataFrame, models: list[str]) -> go.Figure:
    """The paper figure: one column per species, one point per model evaluated on it."""
    cols = [s for s in SPECIES_ORDER if s in set(scores["species"])]
    labels = [SPECIES_LABEL[s] for s in cols]
    fig = go.Figure()
    # Thin guide line per species, from the weakest to the strongest model.
    span = scores.groupby("species")["Score"].agg(["min", "max"]).reindex(cols)
    for label, (lo, hi) in zip(labels, span.itertuples(index=False)):
        fig.add_shape(
            type="line", x0=label, x1=label, y0=lo, y1=hi, layer="below",
            line=dict(color="rgba(128,128,128,0.35)", width=2),
        )
    # AlphaGenome is drawn last so it stays visible where it sits on NTv3's marker;
    # `legendrank` keeps the legend in the paper's block order.
    draw_order = [m for m in models if m != "AlphaGenome"] + (
        ["AlphaGenome"] if "AlphaGenome" in models else []
    )
    for model in draw_order:
        sub = scores[scores["Model"] == model].set_index("species").reindex(cols)
        is_ntv3 = model == NTV3
        fig.add_trace(
            go.Scatter(
                x=labels,
                y=sub["Score"],
                mode="markers",
                name=model,
                legendrank=models.index(model),
                marker=dict(
                    color=MODEL_COLORS[model],
                    size=16 if is_ntv3 else 11,
                    symbol="diamond" if is_ntv3 else "circle",
                    line=dict(color="rgba(0,0,0,0.5)", width=1),
                ),
                customdata=np.stack([sub["Tracks"], sub["Seeds"]], axis=-1),
                hovertemplate=(
                    f"<b>{model}</b><br>%{{x}}<br>{METRIC_LABEL}: %{{y:.3f}}"
                    "<br>%{customdata[0]} tracks, %{customdata[1]} seed(s)<extra></extra>"
                ),
            )
        )
    # Clade separators and clade names above the plot.
    x = 0
    for name, group in SPECIES_GROUPS.items():
        n = sum(1 for s in group if s in cols)
        if n == 0:
            continue
        fig.add_annotation(
            x=x + (n - 1) / 2, y=1.0, xref="x", yref="paper", yanchor="bottom",
            text=f"<b>{name}</b>", showarrow=False,
        )
        x += n
        if x < len(cols):
            fig.add_vline(x=x - 0.5, line=dict(color="rgba(128,128,128,0.4)", dash="dot"))
    fig.update_xaxes(categoryorder="array", categoryarray=labels, title_text="")
    fig.update_yaxes(
        title_text=f"{METRIC_LABEL} on test set",
        showgrid=True,
        gridcolor="rgba(128,128,128,0.15)",
        automargin=True,
        tickformat=".1f",
    )
    fig.update_layout(
        height=600,
        plot_bgcolor="rgba(0,0,0,0)",
        paper_bgcolor="rgba(0,0,0,0)",
        legend=dict(orientation="h", yanchor="top", y=-0.1, x=0.5, xanchor="center"),
        margin=dict(t=40, l=80, r=10, b=40),
    )
    return fig


def plot_mean_bars(leaderboard: pd.DataFrame) -> go.Figure:
    bar_df = leaderboard.sort_values("Mean score", ascending=True)
    fig = px.bar(
        bar_df,
        x="Model",
        y="Mean score",
        color="Model",
        color_discrete_map=MODEL_COLORS,
        category_orders={"Model": bar_df["Model"].tolist()},
    )
    fig.update_layout(
        height=500,
        xaxis_title="",
        yaxis_title=METRIC_LABEL,
        plot_bgcolor="rgba(0,0,0,0)",
        paper_bgcolor="rgba(0,0,0,0)",
        bargap=0.08,
    )
    fig.update_xaxes(showticklabels=False)
    return fig


def plot_breakdown_facets_sorted_models(breakdown_df: pd.DataFrame, height: int = 300):
    categories = sorted(breakdown_df["Category"].dropna().unique())
    if not categories:
        return None
    y_min = breakdown_df["Mean score"].min()
    y_max = breakdown_df["Mean score"].max()
    pad = 0.05 * (y_max - y_min if y_max > y_min else 1.0)
    y_range = [min(0, y_min - pad), y_max + pad]

    fig = make_subplots(
        rows=1,
        cols=len(categories),
        subplot_titles=categories,
        shared_yaxes=True,
        horizontal_spacing=0.04,
    )
    for i, cat in enumerate(categories):
        sub = breakdown_df[breakdown_df["Category"] == cat].sort_values("Mean score")
        fig.add_trace(
            go.Bar(
                x=sub["Model"],
                y=sub["Mean score"],
                marker_color=[MODEL_COLORS.get(m, "#808080") for m in sub["Model"]],
                showlegend=False,
                hovertemplate="%{x}: %{y:.3f}<extra></extra>",
            ),
            row=1,
            col=i + 1,
        )
        fig.update_xaxes(showticklabels=False, title_text="", row=1, col=i + 1)
        fig.update_yaxes(range=y_range, title_text="", row=1, col=i + 1)
    fig.update_layout(
        height=height,
        plot_bgcolor="rgba(0,0,0,0)",
        paper_bgcolor="rgba(0,0,0,0)",
        margin=dict(t=60, l=10, r=10, b=10),
    )
    fig.update_yaxes(title_text=METRIC_LABEL, row=1, col=1)
    return fig


def show_head_to_head(
    df: pd.DataFrame, key: str, per_species: bool = False, title: str | None = None
) -> None:
    """Two models on their shared tracks only, with the number of tracks compared.
    Per track by default, one point per track coloured by assay; `per_species` shows
    one point per species instead, each the mean over that species' shared tracks.
    """
    st.markdown(title or "#### ⚖️ Head-to-head (per track)")
    options = models_in(df)
    if len(options) < 2:
        st.info("Fewer than two models were evaluated here.")
        return
    default_a = options.index(NTV3) if NTV3 in options else 0
    c_a, c_b = st.columns([1, 1])
    with c_a:
        model_a = st.selectbox(
            "Model A (y-axis)", options=options, index=default_a, key=f"a_{key}"
        )
    with c_b:
        b_options = [m for m in options if m != model_a]
        default_b = b_options.index("AlphaGenome") if "AlphaGenome" in b_options else 0
        model_b = st.selectbox(
            "Model B (x-axis)", options=b_options, index=default_b, key=f"b_{key}"
        )

    scatter_df = build_pairwise_scatter_df(df, model_a, model_b)
    if scatter_df.empty:
        st.info(f"{model_a} and {model_b} share no evaluated track.")
        return

    n_tracks = len(scatter_df)
    n_species = scatter_df["species"].nunique()
    if per_species:
        scatter_df = aggregate_pairs_by_species(scatter_df)
        wins = int((scatter_df["Model A"] > scatter_df["Model B"]).sum())
        st.caption(
            f"Compared on **n = {n_species} shared species** ({n_tracks} shared "
            f"tracks). {model_a} is higher on {wins} of them."
        )
        color_col, hover = "Clade", {"Shared tracks": True}
    else:
        wins = int((scatter_df["Model A"] > scatter_df["Model B"]).sum())
        st.caption(
            f"Compared on **n = {n_tracks} shared tracks** over {n_species} "
            f"species. {model_a} is higher on {wins} of them."
        )
        color_col, hover = "assay_type", {"Species": True, "datasets": True}

    min_v = float(min(scatter_df["Model A"].min(), scatter_df["Model B"].min()))
    max_v = float(max(scatter_df["Model A"].max(), scatter_df["Model B"].max()))
    pad = 0.05 * (max_v - min_v if max_v > min_v else 1.0)
    axis_range = [min_v - pad, max_v + pad]

    color_map = ASSAY_COLORS if color_col == "assay_type" else CLADE_COLORS
    fig = px.scatter(
        scatter_df,
        x="Model B",
        y="Model A",
        color=color_col,
        color_discrete_map=color_map,
        category_orders={"Clade": list(SPECIES_GROUPS)},
        hover_name="Track",
        hover_data={**hover, "Model A": ":.3f", "Model B": ":.3f"},
        text="Species" if per_species else None,
    )
    if per_species:
        fig.update_traces(marker=dict(size=13), textposition="top center")
    fig.add_shape(
        type="line",
        x0=axis_range[0], y0=axis_range[0], x1=axis_range[1], y1=axis_range[1],
        xref="x", yref="y",
        line=dict(color="red", dash="dot", width=2),
    )
    fig.update_layout(
        height=620,
        margin=dict(l=60, r=20, t=20, b=110),
        xaxis=dict(
            title=f"{model_b} — {METRIC_LABEL}",
            range=axis_range, dtick=0.1, tickformat=".1f", constrain="domain",
        ),
        yaxis=dict(
            title=f"{model_a} — {METRIC_LABEL}",
            range=axis_range, dtick=0.1, tickformat=".1f",
            scaleanchor="x", scaleratio=1, constrain="domain",
        ),
        plot_bgcolor="rgba(0,0,0,0)",
        paper_bgcolor="rgba(0,0,0,0)",
        legend=dict(
            title="",
            orientation="h",
            x=0.5, xanchor="center", y=-0.15, yanchor="top",
        ),
    )
    st.plotly_chart(fig, width="stretch")


def build_convergence_df(tracks: pd.DataFrame, runs: pd.DataFrame, x_col: str) -> pd.DataFrame:
    """Per model: mean score and the mean cost of reaching its best checkpoint."""
    perf = tracks.groupby("Model", as_index=False)["Score"].mean()
    cost = runs.groupby("Model", as_index=False).agg(X=(x_col, "mean"), Runs=(x_col, "count"))
    out = perf.merge(cost, on="Model").dropna(subset=["X"])
    out["Score"] = out["Score"].round(3)
    return out


def show_convergence(tracks: pd.DataFrame, runs: pd.DataFrame, key: str) -> None:
    st.markdown("#### ⏱️ Time to convergence")
    x_mode = st.selectbox(
        "X-axis",
        options=["GPU hours", "Tokens (billions)"],
        index=0,
        key=f"conv_{key}",
    )
    conv_df = build_convergence_df(tracks, runs, x_mode)
    if conv_df.empty:
        st.info("No convergence data for this species.")
        return
    if x_mode == "GPU hours":
        caption = (
            "GPU hours spent in training steps until the best validation checkpoint "
            "(the one scored on the test set), × GPUs used, averaged over seeds. Job "
            "setup and validation are excluded. All runs on H100s."
        )
    else:
        caption = (
            "Training tokens seen at the best validation checkpoint, averaged over "
            "seeds. Validation ran every 1.05B tokens on most runs, so this is "
            "resolved to about 1B tokens."
        )
    st.caption(caption)
    fig = px.scatter(
        conv_df,
        x="X",
        y="Score",
        text="Model",
        color="Model",
        color_discrete_map=MODEL_COLORS,
        hover_data={"Model": True, "X": ":.2f", "Score": ":.3f", "Runs": True},
    )
    fig.update_traces(marker=dict(size=14), textposition="middle right")
    fig.update_layout(
        height=620,
        xaxis_title=x_mode + (" (log scale)" if x_mode == "GPU hours" else ""),
        yaxis_title=METRIC_LABEL,
        plot_bgcolor="rgba(0,0,0,0)",
        paper_bgcolor="rgba(0,0,0,0)",
        showlegend=False,
    )
    if x_mode == "GPU hours":
        lo, hi = np.log10(conv_df["X"].min()), np.log10(conv_df["X"].max())
        fig.update_xaxes(
            type="log", dtick=1, range=[lo - 0.15, hi + 0.45],
            minor=dict(ticks="", showgrid=False),
        )
    else:
        lo, hi = conv_df["X"].min(), conv_df["X"].max()
        fig.update_xaxes(range=[lo - 0.05 * (hi - lo), hi + 0.3 * (hi - lo)])
    st.plotly_chart(fig, width="stretch")


def plot_violin(df: pd.DataFrame) -> go.Figure:
    order = df.groupby("Model")["Score"].median().sort_values().index.tolist()
    fig = px.violin(
        df,
        x="Model",
        y="Score",
        color="Model",
        color_discrete_map=MODEL_COLORS,
        box=True,
        points=False,
        category_orders={"Model": order},
        hover_data=["track_name_clean"],
    )
    fig.update_layout(
        height=650,
        xaxis_title="",
        yaxis_title=METRIC_LABEL,
        plot_bgcolor="rgba(0,0,0,0)",
        paper_bgcolor="rgba(0,0,0,0)",
        showlegend=False,
    )
    fig.update_traces(box_visible=True, meanline_visible=False)
    return fig


# ---------------------------------------------------------------------
# Pages
# ---------------------------------------------------------------------


SIDEBAR_CSS = """
<style>
[data-testid="stSidebarNav"] a, [data-testid="stSidebarNavLink"] {
    padding-top: 0.35rem; padding-bottom: 0.35rem; border-radius: 0.5rem;
}
[data-testid="stSidebarNav"] a span, [data-testid="stSidebarNavLink"] span {
    font-size: 1.08rem;
}
[data-testid="stNavSectionHeader"], [data-testid="stSidebarNav"] header {
    font-size: 0.85rem; font-weight: 700; letter-spacing: 0.06em;
    text-transform: uppercase; opacity: 0.7; margin-top: 0.8rem;
}
</style>
"""


def overview_page() -> None:
    df = _TRACKS
    st.title("🧬 NTv3 Benchmark")
    st.markdown(_INTRO)
    st.markdown(f"_Last updated: **{_LAST_UPDATED}**_")
    st.divider()

    st.subheader("📊 Every model on every species")
    st.caption(
        "One point per model and species: the mean over that species' tracks. "
        "A model with no point on a species was not evaluated on it. Hover for values."
    )
    scores = build_species_scores(df)
    st.plotly_chart(plot_species_overview(scores, models_in(df)), width="stretch")

    left, right = st.columns([1, 1], gap="large")
    with left:
        show_head_to_head(df, key="overview_track")
    with right:
        show_head_to_head(
            df,
            key="overview_species",
            per_species=True,
            title="#### ⚖️ Head-to-head (per species)",
        )


def species_page(species: str) -> None:
    sub = _TRACKS[_TRACKS["species"] == species]
    runs = _RUNS[(_RUNS["species"] == species) & _RUNS["Model"].isin(set(sub["Model"]))]
    n_tracks = sub["datasets"].nunique()
    common = SPECIES_COMMON_NAME.get(species)
    st.title(f"{SPECIES_ICON[species]} {SPECIES_LABEL[species]}")
    st.caption(
        f"{SPECIES_CLADE[species]}{f' · {common}' if common else ''} · {n_tracks} tracks · "
        f"{sub['Model'].nunique()} models, each evaluated on all {n_tracks} tracks."
    )

    leaderboard = build_leaderboard(sub)
    col1, col2 = st.columns([1, 1])
    with col1:
        st.subheader("🏅 Leaderboard")
        st.write("\n")
        st.write("\n")
        st.write("\n")
        st.dataframe(leaderboard, width="stretch")
    with col2:
        st.subheader("📈 Mean score per model")
        st.plotly_chart(plot_mean_bars(leaderboard), width="stretch")

    breakdown_df = build_category_model_df(sub)
    if breakdown_df["Category"].nunique() > 1:
        st.subheader("🧪 Mean score by assay type")
        st.plotly_chart(plot_breakdown_facets_sorted_models(breakdown_df), width="stretch")

    left, right = st.columns([1, 1], gap="large")
    with left:
        show_head_to_head(sub, key=species)
    with right:
        show_convergence(sub, runs, key=species)

    st.subheader("🎻 Performance across tracks")
    st.plotly_chart(plot_violin(sub), width="stretch")


def make_species_page(species: str) -> st.Page:
    def render() -> None:
        species_page(species)

    render.__name__ = f"page_{species.replace(' ', '_')}"
    return st.Page(
        render,
        title=SPECIES_LABEL[species],
        icon=SPECIES_ICON[species],
        url_path=species.replace(" ", "-"),
    )


# ---------------------------------------------------------------------
# Streamlit UI
# ---------------------------------------------------------------------
def main():
    st.markdown(SIDEBAR_CSS, unsafe_allow_html=True)
    available = set(_TRACKS["species"])
    sections = {
        "": [st.Page(overview_page, title="Overview", icon="🌍", default=True)],
    }
    for clade, group in SPECIES_GROUPS.items():
        pages = [make_species_page(s) for s in group if s in available]
        if pages:
            sections[clade] = pages
    sections["Initial release (2025)"] = [
        st.Page(initial_release.main, title="2025 leaderboard", icon="🗄️", url_path="initial-release")
    ]
    st.navigation(sections, position="sidebar").run()


if __name__ == "__main__":
    main()