"""Print B0/B1 activity results in percent, with one explicit AP protocol."""

import argparse
import json
from pathlib import Path
import re
import statistics


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=Path("results/easycom_local_b1"))
    parser.add_argument("--protocol", choices=("causal_valid_only", "egospeak_full_timeline"), default="causal_valid_only")
    parser.add_argument("--dataset", choices=("easycom", "ego4d"), help="Reject results from a different dataset")
    parser.add_argument("--paired-deltas", action="store_true", help="Compare matching seeds to baseline, plus AV minus audio")
    args = parser.parse_args()
    files = sorted(args.root.glob("*/metrics.json"))
    if not files:
        raise SystemExit(f"No */metrics.json in {args.root}")
    print(f"Protocol: {args.protocol}; all AP/mAP numbers are percentages")
    print(f"{'run':<23} {'0.2s mAP':>11} {'BG AP':>11} {'wearer AP':>11} {'others AP':>11} {'mean mAP':>11}")
    groups, by_seed = {}, {}
    comparison_signature = None
    for path in files:
        metrics = json.loads(path.read_text())
        if args.dataset is not None and metrics.get("dataset") != args.dataset:
            raise ValueError(f"Dataset mismatch in {path}: {metrics.get('dataset')}")
        if args.paired_deltas:
            # Candidates may select different best epochs on the same dev set.
            # Match the maximum training budget, not the selected epoch.
            training_epochs = metrics.get(
                "checkpoint_training_epochs",
                metrics.get("checkpoint_epoch"),
            )
            signature = (metrics.get("dataset"), metrics.get("sessions"), training_epochs,
                         metrics.get("checkpoint_train_sessions"), metrics.get("checkpoint_dev_sessions"),
                         metrics.get("loss"), metrics.get("layer_fusion"))
            if comparison_signature is not None and signature != comparison_signature:
                raise ValueError(f"Paired runs have different splits/budget/loss/fusion: {path}")
            comparison_signature = signature
        score = metrics[args.protocol]
        first = score["horizons"][0]
        ap = first["class_ap"]
        values = [100 * value for value in (first["map"], ap["background"],
                  ap["camera_wearer_speaking"], ap["other_people_speaking"], score["mean_map"])]
        name = path.parent.name
        print(f"{name:<23}" + "".join(f" {v:11.5f}" for v in values))
        group = re.sub(r"_seed\d+$", "", name)
        groups.setdefault(group, []).append(values)
        match = re.fullmatch(r"(.+)_seed(\d+)", name)
        if match:
            by_seed.setdefault(match[1], {})[int(match[2])] = values
    print("\nMean ± sample standard deviation (AP/mAP percentage points)")
    for group, rows in sorted(groups.items()):
        stats = []
        for column in zip(*rows):
            spread = f"{statistics.stdev(column):.5f}" if len(column) > 1 else "n/a"
            stats.append(f"{statistics.mean(column):.5f} ± {spread}")
        print(f"{group:<23}" + " | ".join(stats))
    if args.paired_deltas:
        print("\nPaired deltas (percentage points; same column order; matching seeds only)")
        for treatment, control in (("audio", "baseline"), ("visual", "baseline"),
                                   ("av", "baseline"), ("av", "audio"),
                                   ("track_gru", "pooled"),
                                   ("person_adapter_v2", "pooled"),
                                   ("egospeak_gru", "pooled"),
                                   ("egospeak_gru", "person_adapter_v2"),
                                   ("qwen_audio", "global"),
                                   ("traditional_audio", "global"),
                                   ("highrate", "traditional_audio"),
                                   ("highrate", "qwen_audio"),
                                   ("stage3", "vap_only")):
            left, right = by_seed.get(treatment, {}), by_seed.get(control, {})
            seeds = sorted(left.keys() & right.keys())
            if not seeds:
                print(f"{treatment} - {control}: no matching seeds")
                continue
            differences = [[a - b for a, b in zip(left[seed], right[seed])] for seed in seeds]
            stats = []
            for column in zip(*differences):
                spread = f"{statistics.stdev(column):.5f}" if len(column) > 1 else "n/a"
                stats.append(f"{statistics.mean(column):+.5f} ± {spread}")
            print(f"{treatment} - {control} (seeds={seeds}): " + " | ".join(stats))


if __name__ == "__main__":
    main()