"""Compare win/loss outcome features within one active setup route."""

from __future__ import annotations

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

from compare_setup_routes import first_right_approach_index, route_for_record


Trace = dict[str, Any]


def read_jsonl(path: Path) -> list[Trace]:
    rows = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def write_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")


def active_right_indices(window: list[Trace]) -> list[int]:
    indices = []
    for index, row in enumerate(window):
        if (
            row.get("dx") is not None
            and row["dx"] > 0
            and row["ball_x"] >= 120
            and row["ball_y"] != 0
            and row["ball_x"] not in (0, 205)
            and row.get("mode") != "inactive"
        ):
            indices.append(index)
    return indices


def row_bucket(row: Trace | None) -> str:
    if row is None:
        return "missing"
    return (
        f"x={row.get('ball_x')};y={row.get('ball_y')};"
        f"dx={row.get('dx')};dy={row.get('dy')};"
        f"target={round(float(row['target_y']), 2)};"
        f"pad={row.get('self_paddle_y')};"
        f"err={round(float(row['self_paddle_y']) - float(row['target_y']), 2)};"
        f"action={row.get('action_name')};mode={row.get('mode')};"
        f"score={row.get('self_score')}-{row.get('opponent_score')}"
    )


def compact_row(row: Trace | None) -> dict[str, Any] | None:
    if row is None:
        return None
    return {
        "ball_x": row.get("ball_x"),
        "ball_y": row.get("ball_y"),
        "dx": row.get("dx"),
        "dy": row.get("dy"),
        "target_y": round(float(row["target_y"]), 2),
        "paddle_y": row.get("self_paddle_y"),
        "paddle_error": round(float(row["self_paddle_y"]) - float(row["target_y"]), 2),
        "action": row.get("action_name"),
        "mode": row.get("mode"),
        "score": f"{row.get('self_score')}-{row.get('opponent_score')}",
    }


def numeric(values: list[float]) -> dict[str, Any]:
    if not values:
        return {"count": 0, "mean": None, "min": None, "max": None}
    return {"count": len(values), "mean": mean(values), "min": min(values), "max": max(values)}


def top_counts(counter: Counter[str], limit: int) -> list[dict[str, Any]]:
    return [{"value": value, "count": count} for value, count in counter.most_common(limit)]


def sample_offsets(
    records: list[Trace],
    anchor: str,
    offsets: list[int],
    top: int,
) -> list[dict[str, Any]]:
    rows = []
    for offset in offsets:
        buckets: Counter[str] = Counter()
        errors: list[float] = []
        examples = []
        for record in records:
            window = record.get("window") or []
            indices = active_right_indices(window)
            if not indices:
                row = None
            else:
                base = indices[0] if anchor == "first_right" else indices[-1]
                index = base + offset
                row = window[index] if 0 <= index < len(window) else None
            buckets[row_bucket(row)] += 1
            if row is not None:
                errors.append(float(row["self_paddle_y"]) - float(row["target_y"]))
                if len(examples) < top:
                    examples.append(
                        {
                            "seed": record["seed"],
                            "point": record["point_index"],
                            "outcome": record["outcome"],
                            "pattern": record["pattern"],
                            "row": compact_row(row),
                        }
                    )
        rows.append(
            {
                "offset": offset,
                "top_buckets": top_counts(buckets, top),
                "paddle_error": numeric(errors),
                "examples": examples,
            }
        )
    return rows


def summarize_group(records: list[Trace], first_offsets: list[int], last_offsets: list[int], top: int) -> dict[str, Any]:
    return {
        "count": len(records),
        "seeds": sorted({record["seed"] for record in records}),
        "outcomes": dict(Counter(record["outcome"] for record in records)),
        "patterns": top_counts(Counter(record["pattern"] for record in records), top),
        "first_right_offsets": sample_offsets(records, "first_right", first_offsets, top),
        "last_right_offsets": sample_offsets(records, "last_right", last_offsets, top),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--jsonl", type=Path, required=True)
    parser.add_argument("--route", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--first-offsets", default="-8,-4,0,4,8,12,16,20,24,28,32,36,40")
    parser.add_argument("--last-offsets", default="-20,-16,-12,-8,-4,0,1,2,3,4")
    parser.add_argument("--top", type=int, default=10)
    args = parser.parse_args()

    first_offsets = [int(part.strip()) for part in args.first_offsets.split(",") if part.strip()]
    last_offsets = [int(part.strip()) for part in args.last_offsets.split(",") if part.strip()]
    records = [record for record in read_jsonl(args.jsonl) if route_for_record(record) == args.route]

    by_group: dict[str, list[Trace]] = defaultdict(list)
    for record in records:
        group = record["pattern"] if record["outcome"] == "loss" else "win"
        by_group[group].append(record)

    payload = {
        "source": str(args.jsonl),
        "route": args.route,
        "count": len(records),
        "groups": {
            group: summarize_group(items, first_offsets, last_offsets, args.top)
            for group, items in sorted(by_group.items())
        },
        "interpretation": (
            "Groups share the same active setup route. first_right offsets show route-entry phase; "
            "last_right offsets show contact and point-resolution phase."
        ),
    }
    write_json(args.output, payload)
    for group, summary in payload["groups"].items():
        print(f"{group}: count={summary['count']} seeds={summary['seeds']}")
    print(f"wrote: {args.output}")


if __name__ == "__main__":
    main()
