tests/test_token_efficiency.py

from pathlib import Path
import pandas as pd
import plotly.graph_objects as go

from src.charts import (
    create_performance_vs_resource_plot,
    create_score_vs_cost_plot,
    create_score_vs_tokens_plot,
)
from src.leaderboard import (
    ANALYSIS_COLUMNS,
    EFFICIENCY_RESOURCE_METRICS,
    TOKEN_EFFICIENCY_TABLE_COLUMNS,
    get_analysis_df,
    get_efficiency_resource_column,
    get_efficiency_df,
    get_token_efficiency_table_df,
    get_resource_pareto_frontier_df,
)
from src.models import Benchmark, Environment, Harness, Metrics, Model, Result


def make_result(
    *,
    benchmark: str = "Benchmark A",
    model: str = "model-a",
    harness: str = "harness-a",
    score: float = 0.5,
    n_tasks: int | None = 10,
    total_tokens: int | None = 100,
    input_tokens: int | None = 60,
    cache_tokens: int | None = 10,
    output_tokens: int | None = 30,
    cost_per_task: float | None = 0.25,
    agent_time_per_task: int | None = 10,
    cost_usd: float | None = None,
    agent_time_seconds: int | None = None,
) -> Result:
    return Result(
        benchmark=Benchmark(
            name=benchmark,
            repo="repo",
            num_tasks=10,
            url="https://example.com/benchmark",
        ),
        harness=Harness(
            name=harness,
            skills=[],
            is_oss=True,
            url="https://example.com/harness",
        ),
        model=Model(
            name=model,
            repo=None,
            is_oss=True,
            num_params=1,
            precision="fp16",
            url="https://example.com/model",
        ),
        environment=Environment(name="env", url="https://example.com/env"),
        metrics=Metrics(
            score=score,
            n_tasks=n_tasks,
            n_errors=1,
            mean_input_tokens_per_task=input_tokens,
            mean_cache_tokens_per_task=cache_tokens,
            mean_output_tokens_per_task=output_tokens,
            mean_tokens_per_task=total_tokens,
            mean_cost_usd_per_task=cost_per_task,
            mean_total_time_seconds_per_task=12,
            mean_agent_time_seconds_per_task=agent_time_per_task,
            cost_usd=cost_usd,
            agent_time_seconds=agent_time_seconds,
        ),
    )


def test_cost_vs_performance_tab_removed_and_navigation_is_single_layer():
    app_source = Path("app.py").read_text()

    assert 'gr.Tab("💰 Cost vs Performance")' not in app_source
    assert "cost_benchmark" not in app_source
    assert "cost_controls" not in app_source
    assert "render_score_vs_cost_plot" not in app_source
    assert 'with gr.Tab("Overview")' not in app_source
    assert app_source.count("with gr.Tabs():") == 1
    assert 'with gr.Tab("Rankings")' in app_source
    assert 'with gr.Tab("Trade-offs")' in app_source
    assert 'with gr.Tab("Matrices")' in app_source

def test_analysis_df_columns_and_derived_metrics():
    dataframe = get_analysis_df(
        [make_result(score=0.25, total_tokens=200, cost_per_task=0.125, agent_time_per_task=7)]
    )

    assert set(ANALYSIS_COLUMNS).issubset(dataframe.columns)
    assert dataframe.loc[0, "Score (%)"] == 25
    assert dataframe.loc[0, "Tokens Per Solved Task"] == 800
    assert dataframe.loc[0, "Cost Per Task"] == 0.125
    assert dataframe.loc[0, "Agent Time Per Task"] == 7
    assert bool(dataframe.loc[0, "Token Data Available"]) is True


def test_missing_zero_negative_resource_values_are_unavailable():
    dataframe = get_analysis_df(
        [
            make_result(model="missing", total_tokens=None, cost_per_task=None, agent_time_per_task=None),
            make_result(model="zero", total_tokens=0, cost_per_task=0, agent_time_per_task=0),
            make_result(model="negative", total_tokens=-10, cost_per_task=-1, agent_time_per_task=-3),
        ]
    )

    assert dataframe["Token Data Available"].tolist() == [False, False, False]
    assert dataframe["Tokens Per Solved Task"].isna().all()
    assert dataframe["Cost Per Task"].isna().all()
    assert dataframe["Agent Time Per Task"].isna().all()


def test_zero_score_does_not_divide_by_zero():
    dataframe = get_analysis_df([make_result(score=0, total_tokens=100)])

    assert pd.isna(dataframe.loc[0, "Tokens Per Solved Task"])
    assert bool(dataframe.loc[0, "Token Data Available"]) is True


def test_invalid_task_denominator_does_not_trigger_total_metric_fallback():
    dataframe = get_analysis_df(
        [
            make_result(
                n_tasks=0,
                cost_per_task=None,
                agent_time_per_task=None,
                cost_usd=1.5,
                agent_time_seconds=30,
            )
        ]
    )

    assert pd.isna(dataframe.loc[0, "Cost Per Task"])
    assert pd.isna(dataframe.loc[0, "Agent Time Per Task"])


def test_efficiency_resource_metric_choices_are_exact():
    assert list(EFFICIENCY_RESOURCE_METRICS) == [
        "Total tokens",
        "Cost per task",
        "Agent time per task",
    ]
    assert "Tokens Per Solved Task" not in EFFICIENCY_RESOURCE_METRICS
    assert get_efficiency_resource_column("Total tokens") == "Total Tokens Per Task"
    assert get_efficiency_resource_column("Cost per task") == "Cost Per Task"
    assert get_efficiency_resource_column("Agent time per task") == "Agent Time Per Task"


def test_efficiency_filtering_is_benchmark_specific_and_counts_exclusions():
    analysis_df = get_analysis_df(
        [
            make_result(benchmark="Benchmark A", model="valid", total_tokens=100),
            make_result(benchmark="Benchmark A", model="zero", total_tokens=0),
            make_result(benchmark="Benchmark A", model="missing", total_tokens=None),
            make_result(benchmark="Benchmark B", model="other", total_tokens=100),
        ]
    )

    filtered = get_efficiency_df(
        benchmark_name="Benchmark A",
        resource_metric="Total tokens",
        analysis_df=analysis_df,
    )

    assert filtered["Model"].tolist() == ["valid"]
    assert filtered["Benchmark"].unique().tolist() == ["Benchmark A"]
    assert filtered.attrs["exclusion_count"] == 2
    assert get_efficiency_df(
        "All benchmarks", "Total tokens", analysis_df
    ).empty


def test_cost_and_agent_time_filtering_use_positive_values_only():
    analysis_df = get_analysis_df(
        [
            make_result(model="valid", cost_per_task=0.2, agent_time_per_task=9),
            make_result(model="invalid", cost_per_task=0, agent_time_per_task=-1),
        ]
    )

    cost = get_efficiency_df("Benchmark A", "Cost per task", analysis_df)
    agent_time = get_efficiency_df("Benchmark A", "Agent time per task", analysis_df)

    assert cost["Model"].tolist() == ["valid"]
    assert agent_time["Model"].tolist() == ["valid"]
    assert cost.attrs["exclusion_count"] == 1
    assert agent_time.attrs["exclusion_count"] == 1


def test_efficiency_table_keeps_tokens_per_solved_task_and_expected_order():
    analysis_df = get_analysis_df(
        [make_result(score=0.333333333333, total_tokens=120, cost_per_task=0.123456)]
    )

    table = get_token_efficiency_table_df("Benchmark A", analysis_df)

    assert list(table.columns) == TOKEN_EFFICIENCY_TABLE_COLUMNS
    assert list(table.columns[:3]) == ["Model", "Harness", "Benchmark"]
    assert "Tokens Per Solved Task" in table.columns
    assert table.loc[0, "Score (%)"] == 33.3
    assert table.loc[0, "Cost Per Task"] == 0.1235


def test_pareto_frontier_for_all_resource_metrics():
    dataframe = pd.DataFrame(
        {
            "Run Label": ["a", "b", "c", "d"],
            "Model": ["a", "b", "c", "d"],
            "Total Tokens Per Task": [100, 200, 300, 400],
            "Cost Per Task": [0.1, 0.2, 0.3, 0.4],
            "Agent Time Per Task": [10, 20, 30, 40],
            "Score (%)": [50, 60, 55, 80],
        }
    )

    for metric in ("Total Tokens Per Task", "Cost Per Task", "Agent Time Per Task"):
        frontier = get_resource_pareto_frontier_df(dataframe, metric)
        assert frontier["Run Label"].tolist() == ["a", "b", "d"]


def test_pareto_equal_x_equal_score_ties_are_preserved():
    dataframe = pd.DataFrame(
        {
            "Run Label": ["z", "a", "dominated", "higher"],
            "Model": ["z", "a", "d", "h"],
            "Cost Per Task": [0.1, 0.1, 0.1, 0.2],
            "Score (%)": [50, 50, 40, 60],
        }
    )

    frontier = get_resource_pareto_frontier_df(dataframe, "Cost Per Task")

    assert frontier["Run Label"].tolist() == ["a", "z", "higher"]
    assert "dominated" not in frontier["Run Label"].tolist()


def test_pareto_excludes_missing_zero_and_negative_resources():
    dataframe = pd.DataFrame(
        {
            "Run Label": ["valid", "missing", "zero", "negative"],
            "Agent Time Per Task": [10, None, 0, -1],
            "Score (%)": [50, 100, 100, 100],
        }
    )

    frontier = get_resource_pareto_frontier_df(dataframe, "Agent Time Per Task")

    assert frontier["Run Label"].tolist() == ["valid"]


def test_performance_resource_charts_construct_with_metric_specific_axes():
    dataframe = get_analysis_df(
        [
            make_result(),
            make_result(model="model-b", score=0.7, total_tokens=200, cost_per_task=0.4, agent_time_per_task=20),
        ]
    )

    expected_titles = {
        "Total tokens": "Total tokens per task",
        "Cost per task": "Cost per task (USD)",
        "Agent time per task": "Agent time per task (seconds)",
    }
    for metric, title in expected_titles.items():
        figure = create_performance_vs_resource_plot(dataframe, resource_metric=metric)
        assert isinstance(figure, go.Figure)
        assert figure.layout.xaxis.type == "log"
        assert figure.layout.xaxis.title.text == title
        assert any(trace.name == "Pareto frontier" for trace in figure.data)

    compatibility = create_score_vs_tokens_plot(dataframe, token_metric="Total tokens")
    assert isinstance(compatibility, go.Figure)



def test_performance_resource_chart_rejects_multiple_benchmarks():
    dataframe = get_analysis_df(
        [
            make_result(benchmark="Benchmark A"),
            make_result(benchmark="Benchmark B", model="model-b"),
        ]
    )
    figure = create_performance_vs_resource_plot(dataframe, resource_metric="Total tokens")

    assert len(figure.layout.annotations) == 1
    assert "Select one benchmark" in figure.layout.annotations[0].text

def test_linear_scale_and_point_labels_still_work():
    dataframe = get_analysis_df([make_result()])
    figure = create_performance_vs_resource_plot(
        dataframe,
        resource_metric="Cost per task",
        x_scale="Linear",
        show_labels=True,
        show_pareto_frontier=False,
    )

    assert figure.layout.xaxis.type == "linear"
    assert figure.data[0].mode == "markers+text"
    assert list(figure.data[0].text) == ["model-a / harness-a"]


def test_color_by_benchmark_is_not_supported_in_efficiency_chart():
    dataframe = get_analysis_df([make_result()])
    figure = create_performance_vs_resource_plot(
        dataframe,
        resource_metric="Total tokens",
        color_by="Benchmark",  # type: ignore[arg-type]
    )

    assert len(figure.layout.annotations) == 1
    assert "Color dimension not available" in figure.layout.annotations[0].text


def test_empty_and_fully_invalid_resource_chart_data_are_graceful():
    empty = create_performance_vs_resource_plot(pd.DataFrame(), resource_metric="Total tokens")
    invalid_df = get_analysis_df([make_result(total_tokens=0)])
    invalid = create_performance_vs_resource_plot(invalid_df, resource_metric="Total tokens")

    assert isinstance(empty, go.Figure)
    assert isinstance(invalid, go.Figure)
    assert len(empty.layout.annotations) == 1
    assert len(invalid.layout.annotations) == 1


def test_palette_lookup_new_palettes_fallback_and_copy():
    from src.charts import COLOR_PALETTES, get_color_palette

    for palette_name in ("Grayscale", "Viridis", "Plasma", "Cividis"):
        assert get_color_palette(palette_name) == COLOR_PALETTES[palette_name]
        assert get_color_palette(palette_name) is not COLOR_PALETTES[palette_name]

    fallback = get_color_palette("unknown")
    assert fallback == COLOR_PALETTES["Citrus"]
    fallback.append("#000000")
    assert "#000000" not in COLOR_PALETTES["Citrus"]


def test_cost_and_efficiency_scatter_labels_toggle_consistently():
    cost_df = pd.DataFrame(
        {
            "Benchmark": ["Benchmark A"],
            "Model": ["model-a"],
            "Harness": ["harness-a"],
            "Score": [50.0],
            "Cost Per Task (USD)": [0.25],
            "Label": ["model-a<br>harness-a"],
        }
    )
    efficiency_df = get_analysis_df([make_result()])

    cost_without = create_score_vs_cost_plot(cost_df, "Benchmark A", show_labels=False)
    cost_with = create_score_vs_cost_plot(cost_df, "Benchmark A", show_labels=True)
    efficiency_without = create_performance_vs_resource_plot(
        efficiency_df,
        show_labels=False,
        show_pareto_frontier=False,
    )
    efficiency_with = create_performance_vs_resource_plot(
        efficiency_df,
        show_labels=True,
        show_pareto_frontier=False,
    )

    assert cost_without.data[0].mode == "markers"
    assert cost_with.data[0].mode == "markers+text"
    assert efficiency_without.data[0].mode == "markers"
    assert efficiency_with.data[0].mode == "markers+text"
    assert cost_without.data[0].hovertemplate == cost_with.data[0].hovertemplate
    assert efficiency_without.data[0].hovertemplate == efficiency_with.data[0].hovertemplate

def test_efficiency_figure_uses_responsive_autosizing_without_fixed_width():
    dataframe = get_analysis_df([make_result()])
    figure = create_performance_vs_resource_plot(
        dataframe,
        resource_metric="Total tokens",
    )

    assert figure.layout.autosize is True
    assert figure.layout.width is None
    assert figure.layout.height is None


def test_shared_plot_container_has_minimum_height_and_scrollable_tables():
    app_source = (Path(__file__).parents[1] / "app.py").read_text()

    assert "RESPONSIVE_PLOT_MIN_HEIGHT_PX = 420" in app_source
    assert "min-height: {RESPONSIVE_PLOT_MIN_HEIGHT_PX}px" in app_source
    assert "TABLE_MAX_HEIGHT_PX = 720" in app_source
    assert "max_height=TABLE_MAX_HEIGHT_PX" in app_source


def test_page_tables_are_single_and_below_visualizations_and_plots_resize():
    app_source = Path("app.py").read_text()

    rankings = app_source[app_source.index('with gr.Tab("Rankings")'):app_source.index('with gr.Tab("Trade-offs")')]
    tradeoffs = app_source[app_source.index('with gr.Tab("Trade-offs")'):app_source.index('with gr.Tab("Matrices")')]

    assert rankings.count("gr.Dataframe(") == 4
    assert tradeoffs.count("gr.Dataframe(") == 1
    assert rankings.rindex("gr.Dataframe(") > rankings.rindex("gr.Plot(")
    assert tradeoffs.rindex("gr.Dataframe(") > tradeoffs.rindex("gr.Plot(")
    assert '"Best first"' in rankings
    assert '"Best last"' in rankings
    assert '"Alphabetical (A–Z)"' in rankings
    assert "ResizeObserver" in app_source
    assert "window.Plotly.Plots.resize" in app_source