#!/usr/bin/env python3
"""
Generate a model-comparison report from experiments/<run-id>/ directories.

Each ``experiments/<run-id>/`` directory is expected to contain:

* ``manifest.json`` with at least ``run_id``, ``model_tag``, ``t0_wall_ms``,
  ``duration_seconds``, ``device_label`` (written by ``tools/run_experiment.sh``).
* One or more ``*.json`` files matching the experiment-log schema produced by
  the sample app's Experiment Mode (``device_model``, ``model_name``, ``run_id``,
  ``started_wall_ms``, ``stopped_wall_ms``, ``events``).

For every run we compute a scorecard (detection rate, mean confidence, class
diversity, per-class jitter), a per-bucket comparison table aligned to each
run's ``t0_wall_ms`` so model A and model B share the same x-axis when run
against the same video, and one matplotlib chart per class plotting confidence
over time. matplotlib is optional — if it's not installed the report still
renders, just without the timeline images.

Usage:
    tools/compare_logs.py experiments/                # report.html + summary.csv in cwd
    tools/compare_logs.py experiments/ -o ./reports/  # write into ./reports/
    tools/compare_logs.py experiments/ --bucket 200   # use 200ms buckets
"""

from __future__ import annotations

import argparse
import base64
import csv
import io
import json
import math
import statistics
import sys
from collections import defaultdict
from html import escape
from pathlib import Path

# matplotlib is optional. If unavailable, the report omits timeline charts.
try:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    HAS_MPL = True
except ImportError:
    HAS_MPL = False


# ----------------------------------------------------------------------------- run loading

def load_runs(experiments_dir: Path):
    """Yield (manifest, events_sorted_by_wall_ms) for each run subdir."""
    for run_dir in sorted(experiments_dir.iterdir()):
        if not run_dir.is_dir():
            continue
        manifest_path = run_dir / "manifest.json"
        if not manifest_path.exists():
            print(f"  skipping {run_dir.name}: no manifest.json", file=sys.stderr)
            continue
        try:
            with open(manifest_path) as f:
                manifest = json.load(f)
        except json.JSONDecodeError as e:
            print(f"  skipping {run_dir.name}: bad manifest.json ({e})", file=sys.stderr)
            continue
        log_files = [f for f in run_dir.glob("*.json") if f.name != "manifest.json"]
        if not log_files:
            print(f"  skipping {run_dir.name}: no log files", file=sys.stderr)
            continue
        events: list[dict] = []
        for log in log_files:
            try:
                with open(log) as f:
                    data = json.load(f)
            except json.JSONDecodeError as e:
                print(f"  warning: bad log {log.name} ({e})", file=sys.stderr)
                continue
            events.extend(data.get("events", []))
        events.sort(key=lambda e: e.get("wall_ms", 0))
        yield manifest, events


# ----------------------------------------------------------------------------- per-run metrics

def compute_scorecard(manifest: dict, events: list[dict]) -> dict:
    """Return summary metrics for one run."""
    t0 = manifest["t0_wall_ms"]
    duration_s = float(manifest.get("duration_seconds") or 0)

    total_events = len(events)
    events_with_objects = sum(1 for e in events if e.get("objects"))

    all_confs: list[float] = []
    classes: set[str] = set()
    bbox_areas: list[float] = []
    for e in events:
        for obj in e.get("objects", []):
            all_confs.append(float(obj.get("confidence", 0)))
            classes.add(obj.get("label", ""))
            bbox = obj.get("bbox", [0, 0, 0, 0])
            if len(bbox) >= 4:
                w = max(0.0, float(bbox[2]) - float(bbox[0]))
                h = max(0.0, float(bbox[3]) - float(bbox[1]))
                bbox_areas.append(w * h)

    # Jitter per class: stddev of consecutive bbox-center deltas. High jitter
    # means the box is bouncing around frame-to-frame; low jitter means a
    # stable lock. We sum stdev_dx + stdev_dy as a coarse single-number score.
    jitter_per_class: dict[str, float] = {}
    for cls in classes:
        centers: list[tuple[float, float]] = []
        for e in events:
            for obj in e.get("objects", []):
                if obj.get("label") == cls:
                    bbox = obj.get("bbox", [0, 0, 0, 0])
                    if len(bbox) >= 4:
                        cx = (float(bbox[0]) + float(bbox[2])) / 2
                        cy = (float(bbox[1]) + float(bbox[3])) / 2
                        centers.append((cx, cy))
                    break  # at most one center per frame for each class
        if len(centers) >= 3:
            dxs = [centers[i + 1][0] - centers[i][0] for i in range(len(centers) - 1)]
            dys = [centers[i + 1][1] - centers[i][1] for i in range(len(centers) - 1)]
            jitter_per_class[cls] = statistics.stdev(dxs) + statistics.stdev(dys)

    return {
        "run_id": manifest["run_id"],
        "model_tag": manifest.get("model_tag", "unknown"),
        "device_label": manifest.get("device_label", "unknown"),
        "video_path": manifest.get("video_path", ""),
        "duration_s": duration_s,
        "t0_wall_ms": t0,
        "total_events": total_events,
        "events_with_objects": events_with_objects,
        "detection_rate": events_with_objects / total_events if total_events else 0.0,
        "events_per_second": total_events / duration_s if duration_s else 0.0,
        "mean_confidence": statistics.mean(all_confs) if all_confs else 0.0,
        "class_diversity": len(classes),
        "classes": sorted(c for c in classes if c),
        "mean_bbox_area": statistics.mean(bbox_areas) if bbox_areas else 0.0,
        "jitter_per_class": jitter_per_class,
    }


# ----------------------------------------------------------------------------- bucketing

def bucket_events(events: list[dict], t0: int, bucket_ms: int) -> dict[int, list[dict]]:
    """Group events into buckets keyed by floor(relative_ms / bucket_ms)."""
    buckets: dict[int, list[dict]] = defaultdict(list)
    for e in events:
        rel = int(e.get("wall_ms", 0)) - t0
        if rel < 0:
            continue
        buckets[rel // bucket_ms].append(e)
    return buckets


def bucket_summary(events_in_bucket: list[dict]) -> dict:
    """Per-bucket / per-model aggregation: count, top-class, top-confidence."""
    confs: list[float] = []
    label_counts: dict[str, int] = defaultdict(int)
    for e in events_in_bucket:
        for obj in e.get("objects", []):
            confs.append(float(obj.get("confidence", 0)))
            label_counts[obj.get("label", "")] += 1
    top_label = max(label_counts.items(), key=lambda kv: kv[1])[0] if label_counts else ""
    return {
        "frames": len(events_in_bucket),
        "detections": sum(label_counts.values()),
        "top_label": top_label,
        "max_conf": max(confs) if confs else 0.0,
        "mean_conf": statistics.mean(confs) if confs else 0.0,
    }


# ----------------------------------------------------------------------------- charts

def render_timeline_png(per_run_series: dict[str, list[tuple[float, float]]], title: str) -> str | None:
    """Build a single matplotlib chart and return it as a base64 data URI."""
    if not HAS_MPL:
        return None
    fig, ax = plt.subplots(figsize=(9, 3))
    has_data = False
    for model_tag, series in per_run_series.items():
        if not series:
            continue
        xs = [s[0] / 1000.0 for s in series]
        ys = [s[1] for s in series]
        ax.plot(xs, ys, marker=".", linestyle="-", linewidth=1, label=model_tag)
        has_data = True
    if not has_data:
        plt.close(fig)
        return None
    ax.set_xlabel("seconds since t0")
    ax.set_ylabel("confidence")
    ax.set_ylim(0, 1)
    ax.set_title(title)
    ax.legend(loc="upper right", fontsize=8)
    ax.grid(True, alpha=0.3)
    buf = io.BytesIO()
    fig.tight_layout()
    fig.savefig(buf, format="png", dpi=110)
    plt.close(fig)
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("ascii")


def class_timeline_data(manifest: dict, events: list[dict], cls: str) -> list[tuple[float, float]]:
    """Return [(relative_ms, top_confidence_for_class)] for one run + class."""
    t0 = manifest["t0_wall_ms"]
    out: list[tuple[float, float]] = []
    for e in events:
        rel = int(e.get("wall_ms", 0)) - t0
        if rel < 0:
            continue
        best = 0.0
        found = False
        for obj in e.get("objects", []):
            if obj.get("label") == cls:
                c = float(obj.get("confidence", 0))
                if c > best:
                    best = c
                found = True
        if found:
            out.append((rel, best))
    return out


# ----------------------------------------------------------------------------- HTML rendering

def fmt_pct(x: float) -> str:
    return f"{x * 100:.1f}%"

def fmt_num(x: float, places: int = 2) -> str:
    return f"{x:.{places}f}"


def render_html(scorecards: list[dict], runs: list[tuple[dict, list[dict]]], bucket_ms: int) -> str:
    style = """
    body { font-family: -apple-system, sans-serif; margin: 24px; color: #222; }
    h1, h2 { font-weight: 600; }
    table { border-collapse: collapse; margin: 8px 0 24px 0; font-size: 13px; }
    th, td { border: 1px solid #ddd; padding: 4px 8px; text-align: right; }
    th { background: #f5f5f5; text-align: left; }
    td.left, th.left { text-align: left; }
    .scorecard td:first-child { font-weight: 600; }
    .small { color: #888; font-size: 11px; }
    img { max-width: 100%; border: 1px solid #ddd; margin: 8px 0; }
    .warn { background: #fff7e0; border: 1px solid #f0c000; padding: 8px 12px; border-radius: 4px; }
    """

    parts: list[str] = []
    parts.append("<!doctype html><html><head><meta charset='utf-8'>")
    parts.append("<title>Model comparison report</title>")
    parts.append(f"<style>{style}</style></head><body>")
    parts.append("<h1>Model comparison report</h1>")
    parts.append(f"<p class='small'>{len(scorecards)} run(s), bucket size {bucket_ms} ms</p>")

    if not HAS_MPL:
        parts.append(
            "<p class='warn'>matplotlib is not installed — timeline charts are omitted. "
            "Install with <code>uv pip install matplotlib</code> (or pip) and re-run.</p>"
        )

    # ----- per-run scorecards
    parts.append("<h2>Per-run scorecard</h2>")
    parts.append("<table class='scorecard'>")
    headers = ["metric"] + [s["model_tag"] for s in scorecards]
    parts.append("<tr>" + "".join(f"<th class='left'>{escape(h)}</th>" for h in headers) + "</tr>")

    def row(label: str, values: list[str]) -> str:
        return "<tr><td class='left'>" + escape(label) + "</td>" + "".join(
            f"<td>{escape(v)}</td>" for v in values
        ) + "</tr>"

    parts.append(row("run_id", [s["run_id"] for s in scorecards]))
    parts.append(row("device", [s["device_label"] for s in scorecards]))
    parts.append(row("duration (s)", [fmt_num(s["duration_s"], 1) for s in scorecards]))
    parts.append(row("total events", [str(s["total_events"]) for s in scorecards]))
    parts.append(row("events / second", [fmt_num(s["events_per_second"], 2) for s in scorecards]))
    parts.append(row("frames with detection", [fmt_pct(s["detection_rate"]) for s in scorecards]))
    parts.append(row("mean confidence", [fmt_num(s["mean_confidence"], 3) for s in scorecards]))
    parts.append(row("classes seen", [str(s["class_diversity"]) for s in scorecards]))
    parts.append(row("mean bbox area (px²)", [fmt_num(s["mean_bbox_area"], 0) for s in scorecards]))
    parts.append(row("class list", [", ".join(s["classes"]) for s in scorecards]))
    parts.append("</table>")

    # ----- per-class jitter table
    all_classes = sorted({c for s in scorecards for c in s["classes"]})
    if all_classes:
        parts.append("<h2>Jitter per class (lower = more stable)</h2>")
        parts.append("<table>")
        parts.append("<tr><th class='left'>class</th>" + "".join(f"<th>{escape(s['model_tag'])}</th>" for s in scorecards) + "</tr>")
        for cls in all_classes:
            row_vals = []
            for s in scorecards:
                v = s["jitter_per_class"].get(cls)
                row_vals.append(fmt_num(v, 1) if v is not None else "—")
            parts.append("<tr><td class='left'>" + escape(cls) + "</td>" + "".join(f"<td>{v}</td>" for v in row_vals) + "</tr>")
        parts.append("</table>")

    # ----- per-bucket comparison table
    parts.append(f"<h2>Per-bucket comparison ({bucket_ms} ms bins)</h2>")
    bucketed_per_run = []
    for s, (manifest, events) in zip(scorecards, runs):
        bucketed_per_run.append(bucket_events(events, manifest["t0_wall_ms"], bucket_ms))
    all_buckets = sorted({b for d in bucketed_per_run for b in d.keys()})
    parts.append("<table>")
    header_cells = ["<th class='left'>bucket (ms)</th>"]
    for s in scorecards:
        tag = escape(s["model_tag"])
        header_cells.append(f"<th>{tag}<br><span class='small'>frames</span></th>")
        header_cells.append(f"<th>{tag}<br><span class='small'>top class</span></th>")
        header_cells.append(f"<th>{tag}<br><span class='small'>max conf</span></th>")
    parts.append("<tr>" + "".join(header_cells) + "</tr>")
    for b in all_buckets:
        cells = [f"<td class='left'>{b * bucket_ms}–{(b + 1) * bucket_ms}</td>"]
        for run_buckets in bucketed_per_run:
            evs = run_buckets.get(b, [])
            if evs:
                summ = bucket_summary(evs)
                cells.append(f"<td>{summ['frames']}</td>")
                cells.append(f"<td>{escape(summ['top_label'])}</td>")
                cells.append(f"<td>{fmt_num(summ['max_conf'], 2)}</td>")
            else:
                cells.append("<td>—</td><td>—</td><td>—</td>")
        parts.append("<tr>" + "".join(cells) + "</tr>")
    parts.append("</table>")

    # ----- per-class timeline charts
    if HAS_MPL:
        parts.append("<h2>Per-class confidence timelines</h2>")
        for cls in all_classes:
            per_run_series: dict[str, list[tuple[float, float]]] = {}
            for s, (manifest, events) in zip(scorecards, runs):
                per_run_series[s["model_tag"]] = class_timeline_data(manifest, events, cls)
            data_uri = render_timeline_png(per_run_series, f"class: {cls}")
            if data_uri:
                parts.append(f"<img alt='{escape(cls)} timeline' src='{data_uri}'/>")

    parts.append("</body></html>")
    return "".join(parts)


# ----------------------------------------------------------------------------- summary csv

def write_summary_csv(scorecards: list[dict], path: Path) -> None:
    fieldnames = [
        "run_id",
        "model_tag",
        "device_label",
        "duration_s",
        "total_events",
        "events_per_second",
        "detection_rate",
        "mean_confidence",
        "class_diversity",
        "classes",
        "mean_bbox_area",
    ]
    with open(path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()
        for s in scorecards:
            w.writerow(
                {
                    "run_id": s["run_id"],
                    "model_tag": s["model_tag"],
                    "device_label": s["device_label"],
                    "duration_s": f"{s['duration_s']:.2f}",
                    "total_events": s["total_events"],
                    "events_per_second": f"{s['events_per_second']:.3f}",
                    "detection_rate": f"{s['detection_rate']:.4f}",
                    "mean_confidence": f"{s['mean_confidence']:.4f}",
                    "class_diversity": s["class_diversity"],
                    "classes": "|".join(s["classes"]),
                    "mean_bbox_area": f"{s['mean_bbox_area']:.1f}",
                }
            )


# ----------------------------------------------------------------------------- main

def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(
        description="Build a model comparison report from experiments/<run-id>/ directories.",
    )
    parser.add_argument("experiments_dir", type=Path, help="path to the experiments/ directory")
    parser.add_argument("-o", "--output-dir", type=Path, default=Path.cwd(),
                        help="where report.html and summary.csv should be written (default: cwd)")
    parser.add_argument("--bucket", type=int, default=100, help="bucket size in milliseconds (default: 100)")
    args = parser.parse_args(argv)

    if not args.experiments_dir.is_dir():
        print(f"error: not a directory: {args.experiments_dir}", file=sys.stderr)
        return 2

    runs = list(load_runs(args.experiments_dir))
    if not runs:
        print("error: no usable runs found.", file=sys.stderr)
        return 2

    scorecards = [compute_scorecard(m, e) for m, e in runs]

    args.output_dir.mkdir(parents=True, exist_ok=True)
    html_path = args.output_dir / "report.html"
    csv_path = args.output_dir / "summary.csv"

    html_path.write_text(render_html(scorecards, runs, args.bucket))
    write_summary_csv(scorecards, csv_path)

    print(f"wrote {html_path}")
    print(f"wrote {csv_path}")
    print(f"runs: {len(runs)}")
    for s in scorecards:
        print(
            f"  {s['model_tag']}: {s['total_events']} events, "
            f"{s['detection_rate'] * 100:.1f}% with detection, "
            f"mean conf {s['mean_confidence']:.3f}"
        )
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
