app.py

import json
import os
import sys
from html import escape
from pathlib import Path

import gradio as gr

SPACE_DIR = Path(__file__).resolve().parent
if str(SPACE_DIR) not in sys.path:
    sys.path.insert(0, str(SPACE_DIR))

from src.about import (
    CITATION_BUTTON_LABEL,
    CITATION_BUTTON_TEXT,
    INTRODUCTION_TEXT,
    LLM_BENCHMARKS_TEXT,
    TITLE,
)
from src.display.css_html_js import custom_css
from src.utils import (
    build_leaderboard_summary_html,
    build_domain_status_md,
    dataset_section_title,
    get_grouped_dfs,
    group_datasets_by_domain,
    load_fixed_window_table,
    load_gift_aggregate_metadata_md,
    load_gift_aggregate_table,
    load_forecast_snapshots,
    make_domain_pie_chart,
    make_forecast_plot,
    prepare_ranks_table,
    prepare_values_table,
)

RESULTS_PATH = os.getenv("TSFM_RESULTS_PATH", "space/results" if Path("space/results").exists() else "results")
REFRESH_SECONDS = int(os.getenv("TSFM_REFRESH_SECONDS", "300"))
DOMAIN_OUTPUT_OFFSET = 14
MODEL_COLUMN_WIDTH = "180px"
RANK_MODEL_COLUMN_WIDTH = MODEL_COLUMN_WIDTH
RANK_INDEX_COLUMN_WIDTH = MODEL_COLUMN_WIDTH
METRIC_COLUMN_WIDTH = "118px"
MODEL_COLUMN_NAMES = {"model"}

def _fixed_window_outputs(window: str) -> tuple[str, str]:
    return (
        _rank_table_html(load_fixed_window_table(RESULTS_PATH, window, "rank")),
        _rank_table_html(load_fixed_window_table(RESULTS_PATH, window, "overall")),
    )


def _model_column_index(value) -> int | None:
    columns = getattr(value, "columns", None)
    if columns is None:
        return None
    for idx, column in enumerate(columns):
        if str(column).strip().lower() in MODEL_COLUMN_NAMES:
            return idx
    return None


def _model_cell_html(value) -> str:
    text = "" if value is None else str(value)
    return (
        '<div class="leaderboard-model-cell" title="'
        + escape(text, quote=True)
        + '">'
        + escape(text)
        + "</div>"
    )


def _rank_display_table(value):
    model_col_idx = _model_column_index(value)
    if model_col_idx is None:
        return value
    columns = getattr(value, "columns", None)
    if columns is None:
        return value
    table = value.copy()
    model_col = columns[model_col_idx]
    table[model_col] = table[model_col].map(_model_cell_html)
    return table


def _dataframe_datatypes(value, *, rank_table: bool = False) -> list[str] | None:
    columns = getattr(value, "columns", None)
    if columns is None:
        return None
    try:
        column_count = len(columns)
    except TypeError:
        return None
    if column_count <= 0:
        return None
    datatypes = ["str"] * column_count
    model_col_idx = _model_column_index(value)
    if rank_table and model_col_idx is not None:
        datatypes[model_col_idx] = "html"
    return datatypes


def _dataframe_column_widths(value, *, rank_table: bool = False) -> list[str] | None:
    columns = getattr(value, "columns", None)
    if columns is None:
        return None
    try:
        column_count = len(columns)
    except TypeError:
        return None
    if column_count <= 0:
        return None
    model_width = RANK_MODEL_COLUMN_WIDTH if rank_table else MODEL_COLUMN_WIDTH
    widths = [METRIC_COLUMN_WIDTH] * column_count
    model_col_idx = _model_column_index(value)
    if model_col_idx is None:
        model_col_idx = 0
    widths[model_col_idx] = model_width
    return widths


def _leaderboard_dataframe(value, *, rank_table: bool = False, **kwargs):
    classes = kwargs.pop("elem_classes", []) or []
    if isinstance(classes, str):
        classes = [classes]
    if "leaderboard-dataframe" not in classes:
        classes.append("leaderboard-dataframe")
    if rank_table and "rank-leaderboard-dataframe" not in classes:
        classes.append("rank-leaderboard-dataframe")
    model_col_idx = _model_column_index(value)
    if model_col_idx is not None:
        model_col_class = f"leaderboard-model-col-{model_col_idx}"
        if model_col_class not in classes:
            classes.append(model_col_class)
    kwargs["elem_classes"] = classes
    kwargs["interactive"] = False
    kwargs["wrap"] = False
    kwargs["line_breaks"] = False
    kwargs.setdefault("datatype", _dataframe_datatypes(value, rank_table=rank_table))
    kwargs.setdefault("column_widths", _dataframe_column_widths(value, rank_table=rank_table))
    return gr.Dataframe(value=value, **kwargs)


def _rank_table_html(table) -> str:
    columns = list(getattr(table, "columns", []))
    if not columns:
        return '<div class="rank-table-scroll"><table class="rank-html-table"></table></div>'

    model_col_idx = _model_column_index(table)
    if model_col_idx is None:
        fixed_col_idx = 0
        fixed_col_width = RANK_INDEX_COLUMN_WIDTH
    else:
        fixed_col_idx = model_col_idx
        fixed_col_width = RANK_MODEL_COLUMN_WIDTH

    colgroup = []
    for idx in range(len(columns)):
        width = fixed_col_width if idx == fixed_col_idx else METRIC_COLUMN_WIDTH
        colgroup.append(f'<col style="width: {width}; min-width: {width};">')

    header_cells = []
    for idx, column in enumerate(columns):
        class_name = ' class="rank-fixed-column"' if idx == fixed_col_idx else ""
        header_cells.append(f"<th{class_name}>{escape(str(column))}</th>")

    body_rows = []
    for _, row in table.iterrows():
        cells = []
        for idx, column in enumerate(columns):
            text = "" if row[column] is None else str(row[column])
            if idx == fixed_col_idx:
                cells.append(
                    '<td class="rank-fixed-column"><div class="rank-fixed-window" title="'
                    + escape(text, quote=True)
                    + '">'
                    + escape(text)
                    + "</div></td>"
                )
            else:
                cells.append(f"<td>{escape(text)}</td>")
        body_rows.append("<tr>" + "".join(cells) + "</tr>")

    return (
        '<div class="rank-table-scroll">'
        '<table class="rank-html-table">'
        "<colgroup>"
        + "".join(colgroup)
        + "</colgroup><thead><tr>"
        + "".join(header_cells)
        + "</tr></thead><tbody>"
        + "".join(body_rows)
        + "</tbody></table></div>"
    )


def _dataset_tables(grouped: dict, datasets: list[str]) -> tuple[list, list]:
    values_tables = [
        prepare_values_table(grouped["dataset_values"].get(dataset, None))
        for dataset in datasets
    ]
    rank_tables = [
        _rank_table_html(prepare_ranks_table(grouped["dataset_ranks"].get(dataset, None)))
        for dataset in datasets
    ]
    return values_tables, rank_tables


def _model_label(root: Path, model_slug: str) -> str:
    config_path = root / model_slug / "config.json"
    if not config_path.exists():
        return model_slug
    try:
        config = json.loads(config_path.read_text())
        return str(config.get("model") or model_slug)
    except Exception:
        return model_slug


def _forecast_choices(root_dir: str, datasets: list[str]) -> dict[str, dict[str, dict]]:
    dataset_keys = {dataset.split("/")[0] for dataset in datasets}
    choices: dict[str, dict[str, dict]] = {dataset_key: {} for dataset_key in dataset_keys}
    root = Path(root_dir)
    if not root.exists():
        return choices

    for model_dir in sorted(root.iterdir()):
        if not model_dir.is_dir():
            continue
        model_label = _model_label(root, model_dir.name)
        for snapshot_dataset, snapshot in load_forecast_snapshots(root_dir, model_dir.name).items():
            dataset_key = str(snapshot_dataset).split("/")[0]
            if dataset_key not in choices:
                continue
            choices[dataset_key][model_label] = snapshot

    return choices


def _available_forecast_models(dataset_key: str) -> list[str]:
    return sorted(FORECAST_CHOICES.get(dataset_key, {}).keys())


def _forecast_plot_for(dataset_key: str, model_label: str):
    if not model_label:
        return None
    return make_forecast_plot(FORECAST_CHOICES.get(dataset_key, {}).get(model_label))


def _make_forecast_plotter(dataset_key: str):
    def update_forecast_plot(model_label: str):
        return _forecast_plot_for(dataset_key, model_label)

    return update_forecast_plot


def refresh_leaderboard() -> tuple:
    grouped = get_grouped_dfs(RESULTS_PATH)
    datasets = grouped["datasets"]
    one_day_rank, one_day_metrics = _fixed_window_outputs("1d")
    seven_day_rank, seven_day_metrics = _fixed_window_outputs("7d")
    thirty_day_rank, thirty_day_metrics = _fixed_window_outputs("30d")

    outputs: list = [
        build_leaderboard_summary_html(RESULTS_PATH),       # [0]
        prepare_values_table(grouped["overall_values"]),    # [1]
        _rank_table_html(prepare_ranks_table(grouped["overall_ranks"])),  # [2]
        make_domain_pie_chart(RESULTS_PATH),                # [3]
        one_day_rank,                                       # [4]
        one_day_metrics,                                    # [5]
        seven_day_rank,                                     # [6]
        seven_day_metrics,                                  # [7]
        thirty_day_rank,                                    # [8]
        thirty_day_metrics,                                 # [9]
        load_gift_aggregate_metadata_md(RESULTS_PATH),       # [10]
        load_gift_aggregate_table(RESULTS_PATH, "prediction_length"),  # [11]
        load_gift_aggregate_table(RESULTS_PATH, "domain"),             # [12]
        load_gift_aggregate_table(RESULTS_PATH, "frequency"),          # [13]
    ]

    # [14 .. 14+D-1] per-domain status lines
    for domain, domain_datasets in DOMAIN_GROUPS.items():
        outputs.append(build_domain_status_md(domain, domain_datasets, RESULTS_PATH))

    value_tables, rank_tables = _dataset_tables(grouped, datasets)
    # [14+D .. 14+D+N-1] value tables
    outputs.extend(value_tables)
    # [14+D+N .. 14+D+2N-1] rank tables
    outputs.extend(rank_tables)
    return tuple(outputs)


grouped_initial = get_grouped_dfs(RESULTS_PATH)
DATASETS: list[str] = grouped_initial["datasets"]
FORECAST_CHOICES = _forecast_choices(RESULTS_PATH, DATASETS)
DOMAIN_GROUPS = group_datasets_by_domain(DATASETS, RESULTS_PATH)
N_DOMAINS = len(DOMAIN_GROUPS)
INITIAL_OUTPUTS = refresh_leaderboard()

demo = gr.Blocks(css=custom_css)
with demo:
    gr.HTML(TITLE)
    gr.Markdown(INTRODUCTION_TEXT, elem_classes="markdown-text")
    summary_html = gr.HTML(INITIAL_OUTPUTS[0])

    dataset_value_dfs: list[gr.Dataframe] = []
    dataset_rank_dfs: list[gr.HTML] = []
    domain_status_mds: list[gr.Markdown] = []
    dataset_idx = 0

    with gr.Tabs(elem_classes="tab-buttons"):
        # ── Overall Tab ──────────────────────────────────────────────────────
        with gr.TabItem("Overall"):
            gr.Markdown(
                "**Latest evaluation snapshot**",
                elem_classes="markdown-text",
            )
            gr.Markdown(
                "**Overall summary.** Geometric mean of absolute metric values across the latest dataset snapshot. MSE/CRPS are lower-is-better; RankScore is higher-is-better.",
                elem_classes="markdown-text",
            )
            overall_values_df = _leaderboard_dataframe(
                value=INITIAL_OUTPUTS[1],
                label="Metric values and RankScore",
            )
            gr.Markdown(
                "**Latest snapshot ranks.** Per-metric ranks derived from the snapshot values above; lower rank is better.",
                elem_classes="markdown-text",
            )
            overall_ranks_df = gr.HTML(
                value=INITIAL_OUTPUTS[2],
                elem_classes="rank-html-output",
            )
            gr.Markdown(
                "**Evaluated timestamps by domain** · hover for dataset names",
                elem_classes="markdown-text",
            )
            domain_pie_plot = gr.Plot(
                value=INITIAL_OUTPUTS[3],
                label="Domain distribution",
            )

            gr.Markdown(
                "**Pairwise historical ranking.** Every Ranking table uses all shared releases completed through the current cutoff, balanced equally across datasets. The 24-hour, 7-day, and 30-day windows apply only to their Metrics tables. Official ranks require at least 30 shared releases, 5 shared datasets, 7 days of shared coverage, 3 eligible opponents, and membership in the main comparison component. Absolute metrics remain descriptive and do not determine rank.",
                elem_classes="markdown-text",
            )
            gr.Markdown(
                "**Last 24 hours** · cumulative historical ranking; metrics use releases whose target window ended in the trailing 24 hours.",
                elem_classes="markdown-text",
            )
            gr.Markdown("**Ranking**", elem_classes="markdown-text")
            last_24h_rank_df = gr.HTML(
                value=INITIAL_OUTPUTS[4],
                elem_classes="rank-html-output",
            )
            gr.Markdown("**Metrics**", elem_classes="markdown-text")
            last_24h_metrics_df = gr.HTML(
                value=INITIAL_OUTPUTS[5],
                elem_classes="rank-html-output",
            )
            gr.Markdown(
                "**Last 7 days** · cumulative historical ranking; metrics use a fixed trailing seven-day window.",
                elem_classes="markdown-text",
            )
            gr.Markdown("**Ranking**", elem_classes="markdown-text")
            last_7d_rank_df = gr.HTML(
                value=INITIAL_OUTPUTS[6],
                elem_classes="rank-html-output",
            )
            gr.Markdown("**Metrics**", elem_classes="markdown-text")
            last_7d_metrics_df = gr.HTML(
                value=INITIAL_OUTPUTS[7],
                elem_classes="rank-html-output",
            )
            gr.Markdown(
                "**Last 30 days** · cumulative historical ranking; metrics use a fixed trailing thirty-day window.",
                elem_classes="markdown-text",
            )
            gr.Markdown("**Ranking**", elem_classes="markdown-text")
            last_30d_rank_df = gr.HTML(
                value=INITIAL_OUTPUTS[8],
                elem_classes="rank-html-output",
            )
            gr.Markdown("**Metrics**", elem_classes="markdown-text")
            last_30d_metrics_df = gr.HTML(
                value=INITIAL_OUTPUTS[9],
                elem_classes="rank-html-output",
            )

        # ── GIFT-style grouped result tables ────────────────────────────────
        with gr.TabItem("GIFT-style Aggregates"):
            aggregate_metadata_md = gr.Markdown(
                INITIAL_OUTPUTS[10],
                elem_classes="markdown-text",
            )
            gr.Markdown(
                "MSE and CRPS are normalized per dataset configuration against Seasonal-Naive (1.0 = baseline). Rank is the mean per-configuration CRPS rank. Lower is better for all three metrics.",
                elem_classes="markdown-text",
            )
            with gr.Tabs():
                with gr.TabItem("Prediction Length"):
                    prediction_length_aggregate_df = _leaderboard_dataframe(
                        value=INITIAL_OUTPUTS[11],
                        label="Results on TSFM_Bench aggregated by Prediction Length",
                    )
                with gr.TabItem("Domain"):
                    domain_aggregate_df = _leaderboard_dataframe(
                        value=INITIAL_OUTPUTS[12],
                        label="Results on TSFM_Bench aggregated by Domain",
                    )
                with gr.TabItem("Frequency"):
                    frequency_aggregate_df = _leaderboard_dataframe(
                        value=INITIAL_OUTPUTS[13],
                        label="Results on TSFM_Bench aggregated by Frequency",
                    )

        # ── Domain Tabs ───────────────────────────────────────────────────────
        domain_idx = 0
        for domain, domain_datasets in DOMAIN_GROUPS.items():
            domain_label = f"{domain} ({len(domain_datasets)})"
            with gr.TabItem(domain_label):
                # Per-domain status line (data source, last eval, refresh intervals)
                domain_status_mds.append(
                    gr.Markdown(
                        INITIAL_OUTPUTS[DOMAIN_OUTPUT_OFFSET + domain_idx],
                        elem_classes="markdown-text",
                    )
                )
                domain_idx += 1

                for dataset in domain_datasets:
                    gr.Markdown(
                        dataset_section_title(dataset, RESULTS_PATH),
                        elem_classes="dataset-heading",
                    )
                    dataset_value_dfs.append(
                        _leaderboard_dataframe(
                            value=INITIAL_OUTPUTS[DOMAIN_OUTPUT_OFFSET + N_DOMAINS + dataset_idx],
                            show_label=False,
                        )
                    )
                    gr.Markdown(
                        dataset_section_title(dataset, RESULTS_PATH, ranks=True),
                        elem_classes="dataset-heading",
                    )
                    dataset_rank_dfs.append(
                        gr.HTML(
                            value=INITIAL_OUTPUTS[DOMAIN_OUTPUT_OFFSET + N_DOMAINS + len(DATASETS) + dataset_idx],
                            elem_classes="rank-html-output",
                        )
                    )
                    dataset_key = dataset.split("/")[0]
                    forecast_model_choices = _available_forecast_models(dataset_key)
                    default_forecast_model = forecast_model_choices[0] if forecast_model_choices else None
                    forecast_model_dropdown = gr.Dropdown(
                        choices=forecast_model_choices,
                        value=default_forecast_model,
                        label="Model",
                        interactive=bool(forecast_model_choices),
                    )
                    forecast_plot = gr.Plot(
                        value=_forecast_plot_for(dataset_key, default_forecast_model),
                        label="Forecast snapshot",
                    )
                    forecast_model_dropdown.change(
                        fn=_make_forecast_plotter(dataset_key),
                        inputs=[forecast_model_dropdown],
                        outputs=[forecast_plot],
                    )

                    dataset_idx += 1

    with gr.Accordion("About", open=False):
        gr.Markdown(LLM_BENCHMARKS_TEXT, elem_classes="markdown-text")

    with gr.Accordion("Citation", open=False):
        gr.Textbox(
            value=CITATION_BUTTON_TEXT,
            label=CITATION_BUTTON_LABEL,
            lines=12,
            show_copy_button=True,
        )

    refresh_outputs = [
        summary_html,
        overall_values_df,
        overall_ranks_df,
        domain_pie_plot,
        last_24h_rank_df,
        last_24h_metrics_df,
        last_7d_rank_df,
        last_7d_metrics_df,
        last_30d_rank_df,
        last_30d_metrics_df,
        aggregate_metadata_md,
        prediction_length_aggregate_df,
        domain_aggregate_df,
        frequency_aggregate_df,
        *domain_status_mds,
        *dataset_value_dfs,
        *dataset_rank_dfs,
    ]

    demo.load(refresh_leaderboard, inputs=[], outputs=refresh_outputs)
    gr.Timer(value=REFRESH_SECONDS).tick(refresh_leaderboard, inputs=[], outputs=refresh_outputs)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=20).launch()