#!/usr/bin/env python3
"""Check routing opportunity in dataset."""

import pandas as pd
import numpy as np

df = pd.read_parquet("data/router_data_qwen3.parquet", columns=["query", "task_id", "llm", "effect"])

# Group by query
query_stats = df.groupby(["query", "task_id"]).agg(
    n_correct=("effect", "sum"),
    n_llms=("effect", "count"),
    min_effect=("effect", "min"),
    max_effect=("effect", "max"),
).reset_index()

n_queries = len(query_stats)

# Categories
all_correct = (query_stats["min_effect"] == 1.0).sum()
all_wrong = (query_stats["max_effect"] == 0.0).sum()
routable = n_queries - all_correct - all_wrong

print("Dataset Composition (Full Dataset):")
print("=" * 50)
print(f"Total queries:              {n_queries:6d}")
print()
print(f"All LLMs correct (trivial): {all_correct:6d} ({100*all_correct/n_queries:5.1f}%)")
print(f"All LLMs wrong (impossible):{all_wrong:6d} ({100*all_wrong/n_queries:5.1f}%)")
print(f"Mixed (routing matters):    {routable:6d} ({100*routable/n_queries:5.1f}%)")
print()

# Distribution of how many LLMs are correct per query
print("Distribution of correct LLMs per query:")
print("-" * 50)
for n in range(8):
    count = (query_stats["n_correct"] == n).sum()
    pct = 100 * count / n_queries
    bar = "█" * int(pct)
    print(f"  {n} LLMs correct: {count:5d} ({pct:5.1f}%) {bar}")
