"""Summarize ordered point-loss chains from traced Pong point records."""

from __future__ import annotations

import argparse
import json
from collections import Counter, defaultdict
from pathlib import Path
from statistics import mean


KNOWN_PROBES = {
    ("loss/low_up/paddle_below_target", 205, 118, -8): [
        "round56_low_up_score00_y118_local_actions_rejected"
    ],
    ("loss/high_down/paddle_above_target", 202, 205, 8): [
        "round60_high_down_score11_y205_local_actions_rejected"
    ],
    ("loss/low_up/near_target", 202, 54, -4): [
        "round58_round59_low_up_score12_y54_local_and_route_actions_rejected"
    ],
    ("loss/low_up/paddle_below_target", 203, 111, -8): [
        "round19_low_up_y111_accepted"
    ],
    ("loss/mid_up/paddle_below_target", 203, 164, -8): [
        "round57_mid_up_y164_local_actions_rejected"
    ],
    ("loss/mid_down/paddle_below_target", 203, 154, 4): [
        "round51_score00_mid_down_y154_final_noop_accepted"
    ],
}


def read_records(path: Path) -> list[dict]:
    records = []
    for line in path.read_text(encoding="utf-8").splitlines():
        if line.strip():
            records.append(json.loads(line))
    return records


def loss_key(record: dict) -> tuple[str, int | None, int | None, int | None]:
    return (
        record.get("pattern"),
        record.get("final_ball_x"),
        record.get("final_ball_y"),
        record.get("final_dy"),
    )


def loss_signature(record: dict) -> str:
    pattern, x, y, dy = loss_key(record)
    return (
        f"{pattern};x={x};y={y};dy={dy};"
        f"score_after={record.get('self_score')}-{record.get('opponent_score')};"
        f"target={record.get('final_target_y')};"
        f"pad={record.get('final_self_paddle_y')};"
        f"err={record.get('final_paddle_error')}"
    )


def compact_signature(record: dict) -> str:
    pattern, x, y, dy = loss_key(record)
    return (
        f"{pattern.split('/', 1)[1] if isinstance(pattern, str) and '/' in pattern else pattern}"
        f"@x{x}y{y}dy{dy}"
        f"->{record.get('self_score')}-{record.get('opponent_score')}"
    )


def first_right_frame(record: dict) -> dict | None:
    for row in record.get("window", []):
        if row.get("dx") is not None and row["dx"] > 0 and row.get("ball_x", 0) >= 120:
            return row
    return None


def summarize_seed(seed: int, records: list[dict]) -> dict:
    losses = [record for record in records if record.get("outcome") == "loss"]
    chain = [loss_signature(record) for record in losses]
    compact_chain = [compact_signature(record) for record in losses]
    uncovered = []
    covered = []
    for record in losses:
        key = loss_key(record)
        probe_labels = KNOWN_PROBES.get(key, [])
        item = {
            "point_index": record.get("point_index"),
            "step": record.get("step"),
            "signature": loss_signature(record),
            "first_right_frame": first_right_frame(record),
            "known_probes": probe_labels,
        }
        if probe_labels:
            covered.append(item)
        else:
            uncovered.append(item)
    final_score = records[-1].get("score_diff") if records else 0
    return {
        "seed": seed,
        "points": len(records),
        "wins": sum(1 for record in records if record.get("outcome") == "win"),
        "losses": len(losses),
        "final_score": final_score,
        "chain_key": "|".join(compact_chain),
        "loss_chain": chain,
        "covered_losses": covered,
        "uncovered_losses": uncovered,
    }


def summarize(path: Path, final_score_filter: int | None) -> dict:
    by_seed: dict[int, list[dict]] = defaultdict(list)
    for record in read_records(path):
        by_seed[int(record["seed"])].append(record)

    seed_summaries = [
        summarize_seed(seed, records) for seed, records in sorted(by_seed.items())
    ]
    if final_score_filter is not None:
        seed_summaries = [
            summary
            for summary in seed_summaries
            if summary["final_score"] == final_score_filter
        ]

    chain_counts = Counter(summary["chain_key"] for summary in seed_summaries)
    uncovered_counts = Counter()
    for summary in seed_summaries:
        for item in summary["uncovered_losses"]:
            uncovered_counts[item["signature"]] += 1

    chain_groups = []
    for chain_key, count in chain_counts.most_common():
        members = [summary for summary in seed_summaries if summary["chain_key"] == chain_key]
        chain_groups.append(
            {
                "count": count,
                "seeds": [summary["seed"] for summary in members],
                "chain_key": chain_key,
                "loss_chain": members[0]["loss_chain"] if members else [],
                "uncovered_signatures": [
                    {"count": item_count, "signature": signature}
                    for signature, item_count in Counter(
                        item["signature"]
                        for summary in members
                        for item in summary["uncovered_losses"]
                    ).most_common()
                ],
                "covered_probe_labels": sorted(
                    {
                        label
                        for summary in members
                        for item in summary["covered_losses"]
                        for label in item["known_probes"]
                    }
                ),
            }
        )

    return {
        "source": str(path),
        "final_score_filter": final_score_filter,
        "seed_count": len(seed_summaries),
        "mean_losses_per_seed": mean(summary["losses"] for summary in seed_summaries)
        if seed_summaries
        else None,
        "chain_groups": chain_groups,
        "uncovered_loss_counts": [
            {"count": count, "signature": signature}
            for signature, count in uncovered_counts.most_common()
        ],
        "seeds": seed_summaries,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("jsonl", type=Path)
    parser.add_argument("--final-score", type=int)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()

    result = summarize(args.jsonl, args.final_score)
    text = json.dumps(result, indent=2, sort_keys=True)
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(text + "\n", encoding="utf-8")
    print(text)


if __name__ == "__main__":
    main()
