benchmark_report.py

"""
Benchmark report generator -- reads a benchmark results JSON and prints a full report.

Usage:
  python benchmark_report.py gemini/gemini_baseline
  python benchmark_report.py gemini/gemini_baseline --no-answers   # skip LLM responses
  python benchmark_report.py gemini/gemini_baseline --json         # output as JSON summary
"""

import argparse
import json
import os
import sys
from datetime import datetime

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
RESULTS_BASE_DIR = os.path.join(SCRIPT_DIR, "benchmark_results")

C = {
    "RESET": "\033[0m",
    "DIM": "\033[2m",
    "BOLD": "\033[1m",
    "CYAN": "\033[36m",
    "GREEN": "\033[32m",
    "YELLOW": "\033[33m",
    "MAGENTA": "\033[35m",
    "RED": "\033[31m",
    "WHITE": "\033[97m",
    "UNDERLINE": "\033[4m",
}

# Disable colors if not a terminal
if not sys.stdout.isatty():
    C = {k: "" for k in C}


def load_report(path_arg):
    """Resolve path_arg to a JSON file and load it."""
    # Try as-is first, then under benchmark_results/
    candidates = [
        path_arg,
        path_arg + ".json",
        os.path.join(RESULTS_BASE_DIR, path_arg),
        os.path.join(RESULTS_BASE_DIR, path_arg + ".json"),
    ]
    for p in candidates:
        if os.path.isfile(p):
            with open(p) as f:
                return json.load(f), p
    print(f"Not found: {path_arg}")
    print(f"Looked in: {', '.join(candidates)}")
    sys.exit(1)


def fmt_ms(ms):
    if ms >= 1000:
        return f"{ms/1000:.1f}s"
    return f"{ms:.0f}ms"


def fmt_num(n):
    if isinstance(n, float):
        return f"{n:,.2f}"
    return f"{n:,}"


def bar(value, max_val, width=30, fill_char="=", empty_char=" "):
    """Simple ASCII bar."""
    if max_val == 0:
        return "[" + empty_char * width + "]"
    ratio = min(value / max_val, 1.0)
    filled = int(ratio * width)
    return "[" + fill_char * filled + empty_char * (width - filled) + "]"


def print_header(title):
    print(f"\n{C['BOLD']}{C['CYAN']}{'=' * 80}{C['RESET']}")
    print(f"{C['BOLD']}{C['CYAN']}  {title}{C['RESET']}")
    print(f"{C['BOLD']}{C['CYAN']}{'=' * 80}{C['RESET']}")


def print_section(title):
    print(f"\n{C['BOLD']}{C['YELLOW']}--- {title} ---{C['RESET']}")


def print_report(data, file_path, show_answers=True):
    meta = data["run_metadata"]
    results = data["query_results"]
    params = meta.get("params", {})
    bench_params = meta.get("benchmark_params", {})

    # ======================================================================
    # HEADER
    # ======================================================================
    print_header(f"BENCHMARK REPORT: {data['run_name']}")
    print(f"  {C['DIM']}File:{C['RESET']}       {os.path.basename(file_path)}")
    print(f"  {C['DIM']}Timestamp:{C['RESET']}   {data['timestamp']}")
    print(f"  {C['DIM']}Duration:{C['RESET']}    {data['total_duration_s']:.1f}s ({data['total_duration_s']/60:.1f}min)")
    print(f"  {C['DIM']}Queries:{C['RESET']}     {data['num_queries']}")

    # ======================================================================
    # RUN METADATA
    # ======================================================================
    print_section("Run Configuration")
    print(f"  Vectors:       {fmt_num(meta['num_vectors'])}")
    print(f"  Dimensions:    {meta['dims']}")
    print(f"  Embed cost:    ${meta.get('est_cost_usd', 0):.4f}")
    print(f"  Embed time:    {meta.get('embed_seconds', 0):.0f}s ({meta.get('embed_seconds', 0)/60:.1f}min)")
    print(f"  Pickle size:   {meta.get('pkl_size_mb', 0):.1f} MB")
    embed_model = meta.get("embed_model", "unknown")
    print(f"  Embed model:   {embed_model}")
    if params:
        print(f"  Chunk size:    {params.get('chunk_size', '?')} chars, overlap {params.get('overlap', '?')}")
        print(f"  File limit:    {params.get('limit', '?')}")
    print(f"  Max chunks:    {bench_params.get('max_chunks', '?')}")

    # ======================================================================
    # AGGREGATE STATS
    # ======================================================================
    print_section("Aggregate Statistics")

    all_top1 = [r["similarity_stats"]["top_1"] for r in results]
    all_min = [r["similarity_stats"]["min"] for r in results]
    all_spread = [r["similarity_stats"]["spread"] for r in results]
    all_median = [r["similarity_stats"]["median"] for r in results]
    all_chunks = [r["retrieval_stats"]["chunks_used"] for r in results]
    all_sources = [r["retrieval_stats"]["unique_sources_in_context"] for r in results]
    all_boilerplate = [r["retrieval_stats"]["boilerplate_in_context"] for r in results]
    all_latency = [r["latency"]["total_ms"] for r in results]
    all_llm_latency = [r["latency"]["llm_generation_ms"] for r in results]
    all_tokens = [r["llm_prompt"]["est_tokens"] for r in results]

    n = len(results)
    avg = lambda lst: sum(lst) / len(lst) if lst else 0

    print(f"  Similarity (top-1):    avg={avg(all_top1):.4f}  min={min(all_top1):.4f}  max={max(all_top1):.4f}")
    print(f"  Similarity (floor):    avg={avg(all_min):.4f}  min={min(all_min):.4f}  max={max(all_min):.4f}")
    print(f"  Similarity (median):   avg={avg(all_median):.4f}")
    print(f"  Spread (top1-floor):   avg={avg(all_spread):.4f}  min={min(all_spread):.4f}  max={max(all_spread):.4f}")
    print(f"  Chunks used/query:     {all_chunks[0]} (fixed)")
    print(f"  Unique sources/query:  avg={avg(all_sources):.0f}  min={min(all_sources)}  max={max(all_sources)}")
    print(f"  Boilerplate chunks:    total={sum(all_boilerplate)}  avg={avg(all_boilerplate):.1f}/query")
    print(f"  Est. tokens/query:     avg={avg(all_tokens):.0f}  min={min(all_tokens)}  max={max(all_tokens)}  total={sum(all_tokens):,}")
    print(f"  Latency (total):       avg={fmt_ms(avg(all_latency))}  min={fmt_ms(min(all_latency))}  max={fmt_ms(max(all_latency))}")
    print(f"  Latency (LLM only):    avg={fmt_ms(avg(all_llm_latency))}  min={fmt_ms(min(all_llm_latency))}  max={fmt_ms(max(all_llm_latency))}")

    # ======================================================================
    # PER-QUERY TABLE
    # ======================================================================
    print_section("Per-Query Overview")

    hdr = f"  {'ID':<4} {'Category':<18} {'Sources':>7} {'Top1':>7} {'Floor':>7} {'Spread':>7} {'Median':>7} {'Tokens':>7} {'Latency':>8} {'Ans':>5}"
    print(f"{C['BOLD']}{hdr}{C['RESET']}")
    print(f"  {'-' * (len(hdr) - 2)}")

    for r in results:
        qid = r["query_id"]
        cat = r["category"]
        rs = r["retrieval_stats"]
        ss = r["similarity_stats"]
        lat = r["latency"]
        prompt = r["llm_prompt"]

        print(f"  Q{qid:<3} {cat:<18} {rs['unique_sources_in_context']:>7} {ss['top_1']:>7.4f} {ss['min']:>7.4f} {ss['spread']:>7.4f} {ss['median']:>7.4f} {prompt['est_tokens']:>7,} {fmt_ms(lat['total_ms']):>8} {r['llm_response_chars']:>5}")

    # ======================================================================
    # SIMILARITY ANALYSIS
    # ======================================================================
    print_section("Similarity Analysis")
    print(f"  {C['DIM']}Spread = top-1 minus floor (chunk #200). Small spread = flat band, model can't discriminate.{C['RESET']}")
    print()

    max_spread = max(all_spread) if all_spread else 1

    for r in results:
        qid = r["query_id"]
        ss = r["similarity_stats"]

        spread_bar = bar(ss["spread"], max_spread, width=25)
        spread_indicator = ""
        if ss["spread"] < 0.02:
            spread_indicator = f" {C['RED']}FLAT{C['RESET']}"
        elif ss["spread"] < 0.04:
            spread_indicator = f" {C['YELLOW']}narrow{C['RESET']}"

        print(f"  Q{qid:<3} top1={ss['top_1']:.4f} floor={ss['min']:.4f} spread={ss['spread']:.4f}{spread_indicator}  {spread_bar}")

    flat_count = sum(1 for r in results if r["similarity_stats"]["spread"] < 0.02)
    narrow_count = sum(1 for r in results if 0.02 <= r["similarity_stats"]["spread"] < 0.04)
    print(f"\n  {C['DIM']}FLAT (<0.02 spread): {flat_count}/{n} queries | Narrow (0.02-0.04): {narrow_count}/{n}{C['RESET']}")

    # ======================================================================
    # CATEGORY BREAKDOWN
    # ======================================================================
    print_section("Category Breakdown")
    categories = {}
    for r in results:
        cat = r["category"]
        if cat not in categories:
            categories[cat] = []
        categories[cat].append(r)

    cat_header = f"  {'Category':<20} {'Queries':>7} {'Avg Top1':>9} {'Avg Spread':>11} {'Avg Latency':>12} {'Avg Tokens':>11}"
    print(f"{C['BOLD']}{cat_header}{C['RESET']}")
    print(f"  {'-' * (len(cat_header) - 2)}")

    for cat, cat_results in categories.items():
        cn = len(cat_results)
        cat_avg_top1 = sum(r["similarity_stats"]["top_1"] for r in cat_results) / cn
        cat_avg_spread = sum(r["similarity_stats"]["spread"] for r in cat_results) / cn
        cat_avg_lat = sum(r["latency"]["total_ms"] for r in cat_results) / cn
        cat_avg_tok = sum(r["llm_prompt"]["est_tokens"] for r in cat_results) / cn
        print(f"  {cat:<20} {cn:>7} {cat_avg_top1:>9.4f} {cat_avg_spread:>11.4f} {fmt_ms(cat_avg_lat):>12} {cat_avg_tok:>11,.0f}")

    # ======================================================================
    # LATENCY BREAKDOWN
    # ======================================================================
    print_section("Latency Breakdown")
    print(f"  {C['DIM']}Where time is spent per query: embedding the query, cosine search, LLM generation.{C['RESET']}")
    print()

    total_embed = sum(r["latency"]["embed_query_ms"] for r in results)
    total_search = sum(r["latency"]["cosine_search_ms"] for r in results)
    total_llm = sum(r["latency"]["llm_generation_ms"] for r in results)
    total_all = total_embed + total_search + total_llm

    pct = lambda v: (v / total_all * 100) if total_all > 0 else 0

    print(f"  Embed query:    {fmt_ms(total_embed / n):>8} avg ({pct(total_embed):>5.1f}% of total)")
    print(f"  Cosine search:  {fmt_ms(total_search / n):>8} avg ({pct(total_search):>5.1f}% of total)")
    print(f"  LLM generation: {fmt_ms(total_llm / n):>8} avg ({pct(total_llm):>5.1f}% of total)")
    print(f"  Total:          {fmt_ms(total_all / n):>8} avg")

    max_lat = max(r["latency"]["total_ms"] for r in results)
    print()
    for r in results:
        lat = r["latency"]
        total = lat["total_ms"]
        llm_pct = (lat["llm_generation_ms"] / total * 100) if total > 0 else 0
        latency_bar = bar(total, max_lat, width=30, fill_char="#")
        print(f"  Q{r['query_id']:<3} {fmt_ms(total):>7}  {latency_bar}  LLM={llm_pct:.0f}%")

    # ======================================================================
    # PER-QUERY DETAILS + ANSWERS
    # ======================================================================
    if show_answers:
        print_section("Full Query Results")

        for r in results:
            qid = r["query_id"]
            rs = r["retrieval_stats"]
            ss = r["similarity_stats"]
            lat = r["latency"]
            prompt = r["llm_prompt"]

            print(f"\n  {C['BOLD']}{C['WHITE']}Q{qid} [{r['category']}]: {r['query']}{C['RESET']}")
            print(f"  {C['DIM']}Tests: {r['tests']}{C['RESET']}")
            print()

            # Retrieval stats
            print(f"    Retrieval:")
            print(f"      Chunks sent to LLM:   {rs['chunks_used']}")
            print(f"      Unique job sources:    {rs['unique_sources_in_context']}")
            print(f"      Boilerplate chunks:    {rs['boilerplate_in_context']}")
            print(f"      Context size:          {prompt['context_chars']:,} chars (~{prompt['est_tokens']:,} tokens)")

            # Similarity
            print(f"    Similarity:")
            print(f"      Top-1: {ss['top_1']:.4f}  Floor: {ss['min']:.4f}  Spread: {ss['spread']:.4f}  Median: {ss['median']:.4f}")

            # Source diversity
            us = r["unique_sources"]
            print(f"    Source diversity:")
            print(f"      Top-5: {us['top_5']} unique  Top-10: {us['top_10']} unique  Top-20: {us['top_20']} unique")

            # Latency
            print(f"    Latency:")
            print(f"      Embed: {fmt_ms(lat['embed_query_ms'])}  Search: {fmt_ms(lat['cosine_search_ms'])}  LLM: {fmt_ms(lat['llm_generation_ms'])}  Total: {fmt_ms(lat['total_ms'])}")

            # Top-3 sources
            top3 = r["top_20"][:3]
            print(f"    Top-3 chunks:")
            for t in top3:
                bp_flag = f" {C['RED']}[BOILERPLATE]{C['RESET']}" if t["is_boilerplate"] else ""
                print(f"      #{t['rank']} sim={t['similarity']:.4f} src={t['source']} ({t['char_count']} chars){bp_flag}")
                print(f"         {C['DIM']}{t['preview'][:120]}{C['RESET']}")

            # LLM Answer
            print(f"\n    {C['GREEN']}Answer ({r['llm_response_chars']} chars):{C['RESET']}")
            response = r["llm_response"]
            lines = response.split("\n")
            for line in lines:
                while len(line) > 100:
                    split_at = line[:100].rfind(" ")
                    if split_at == -1:
                        split_at = 100
                    print(f"    {line[:split_at]}")
                    line = line[split_at:].lstrip()
                print(f"    {line}")

            print(f"\n  {C['DIM']}{'- ' * 40}{C['RESET']}")

    # ======================================================================
    # FOOTER
    # ======================================================================
    print_header("END OF REPORT")
    print(f"  {data['run_name']} | {n} queries | {data['total_duration_s']:.0f}s | {sum(all_tokens):,} total est. tokens")
    print()


def generate_json_summary(data):
    """Generate a compact JSON summary for programmatic use."""
    results = data["query_results"]
    n = len(results)
    avg = lambda lst: round(sum(lst) / len(lst), 4) if lst else 0

    summary = {
        "run_name": data["run_name"],
        "timestamp": data["timestamp"],
        "total_duration_s": data["total_duration_s"],
        "num_queries": n,
        "embed_model": data["run_metadata"].get("embed_model", "unknown"),
        "num_vectors": data["run_metadata"]["num_vectors"],
        "dims": data["run_metadata"]["dims"],
        "aggregate": {
            "avg_top1_sim": avg([r["similarity_stats"]["top_1"] for r in results]),
            "avg_floor_sim": avg([r["similarity_stats"].get("min", r.get("retrieval_stats", {}).get("min_sim_used", r["similarity_stats"].get("top_20_min", 0))) for r in results]),
            "avg_spread": avg([r["similarity_stats"].get("spread", round(r["similarity_stats"]["top_1"] - r["similarity_stats"].get("top_20_min", r.get("retrieval_stats", {}).get("min_sim_used", 0)), 6)) for r in results]),
            "avg_chunks_used": avg([r["retrieval_stats"]["chunks_used"] for r in results]),
            "total_boilerplate": sum(r["retrieval_stats"]["boilerplate_in_context"] for r in results),
            "avg_est_tokens": avg([r["llm_prompt"]["est_tokens"] for r in results]),
            "avg_latency_ms": avg([r["latency"]["total_ms"] for r in results]),
            "avg_llm_latency_ms": avg([r["latency"]["llm_generation_ms"] for r in results]),
        },
        "queries": [],
    }

    for r in results:
        ss = r["similarity_stats"]
        summary["queries"].append({
            "id": r["query_id"],
            "category": r["category"],
            "query": r["query"],
            "chunks_used": r["retrieval_stats"]["chunks_used"],
            "sources": r["retrieval_stats"]["unique_sources_in_context"],
            "top1_sim": ss["top_1"],
            "floor_sim": ss.get("min", r.get("retrieval_stats", {}).get("min_sim_used", ss.get("top_20_min", 0))),
            "spread": ss.get("spread", round(ss["top_1"] - ss.get("top_20_min", r.get("retrieval_stats", {}).get("min_sim_used", 0)), 6)),
            "est_tokens": r["llm_prompt"]["est_tokens"],
            "latency_ms": r["latency"]["total_ms"],
            "answer_chars": r["llm_response_chars"],
        })

    return summary


def main():
    parser = argparse.ArgumentParser(description="Generate report from benchmark results JSON")
    parser.add_argument("report", type=str, help="Path to results JSON (e.g., gemini/gemini_baseline)")
    parser.add_argument("--no-answers", action="store_true", help="Skip printing full LLM responses")
    parser.add_argument("--json", action="store_true", help="Output compact JSON summary instead of formatted report")
    args = parser.parse_args()

    data, file_path = load_report(args.report)

    if args.json:
        summary = generate_json_summary(data)
        print(json.dumps(summary, indent=2))
    else:
        print_report(data, file_path, show_answers=not args.no_answers)


if __name__ == "__main__":
    main()