"""Link selected route records to their preceding point outcomes."""

from __future__ import annotations

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


SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

from compare_setup_routes import 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) + "\n", encoding="utf-8")


def score_key(record: Trace) -> str:
    window = record.get("window") or []
    if window:
        first = window[0]
        return f"{first.get('self_score')}-{first.get('opponent_score')}"
    return f"{record.get('self_score')}-{record.get('opponent_score')}"


def phase_key(record: Trace) -> str:
    return f"{score_key(record)}:{record['outcome']}:{record['pattern']}"


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": round(mean(values), 4), "min": min(values), "max": max(values)}


def first_initial_row(record: Trace) -> Trace | None:
    for row in record.get("window") or []:
        if row.get("mode") == "initial" and row.get("ball_x") not in (0, 52, 205):
            return row
    return None


def record_summary(record: Trace | None) -> dict[str, Any] | None:
    if record is None:
        return None
    return {
        "seed": record["seed"],
        "point": record["point_index"],
        "outcome": record["outcome"],
        "pattern": record["pattern"],
        "score_after": f"{record.get('self_score')}-{record.get('opponent_score')}",
        "final_ball_x": record["final_ball_x"],
        "final_ball_y": record["final_ball_y"],
        "final_dx": record["final_dx"],
        "final_dy": record["final_dy"],
        "final_target_y": record["final_target_y"],
        "final_paddle_y": record["final_self_paddle_y"],
        "final_paddle_error": record["final_paddle_error"],
    }


def summarize_links(records: list[tuple[Trace, Trace | None]]) -> dict[str, Any]:
    previous_patterns = Counter(
        "none" if previous is None else f"{previous['outcome']}:{previous['pattern']}"
        for _current, previous in records
    )
    initial_pads = []
    initial_rows = []
    previous_final_errors = []
    previous_final_y = []
    examples = []
    for current, previous in records:
        initial = first_initial_row(current)
        if initial is not None:
            initial_pads.append(float(initial["self_paddle_y"]))
            if len(initial_rows) < 8:
                initial_rows.append(
                    {
                        "seed": current["seed"],
                        "point": current["point_index"],
                        "ball_x": initial["ball_x"],
                        "ball_y": initial["ball_y"],
                        "paddle_y": initial["self_paddle_y"],
                        "target_y": initial["target_y"],
                        "action": initial["action_name"],
                    }
                )
        if previous is not None:
            previous_final_errors.append(float(previous["final_paddle_error"]))
            previous_final_y.append(float(previous["final_ball_y"]))
        if len(examples) < 8:
            examples.append(
                {
                    "current": {
                        "seed": current["seed"],
                        "point": current["point_index"],
                        "phase": phase_key(current),
                        "initial": None if initial is None else {
                            "paddle_y": initial["self_paddle_y"],
                            "target_y": initial["target_y"],
                            "action": initial["action_name"],
                        },
                    },
                    "previous": record_summary(previous),
                }
            )
    return {
        "count": len(records),
        "seeds": sorted({current["seed"] for current, _previous in records}),
        "points": sorted({current["point_index"] for current, _previous in records}),
        "previous_patterns": [
            {"value": value, "count": count}
            for value, count in previous_patterns.most_common()
        ],
        "current_initial_paddle_y": numeric(initial_pads),
        "previous_final_paddle_error": numeric(previous_final_errors),
        "previous_final_ball_y": numeric(previous_final_y),
        "initial_examples": initial_rows,
        "examples": examples,
    }


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)
    args = parser.parse_args()

    rows = read_jsonl(args.jsonl)
    by_seed_point = {(record["seed"], record["point_index"]): record for record in rows}
    links_by_phase: dict[str, list[tuple[Trace, Trace | None]]] = defaultdict(list)
    matched_route_count = 0
    for record in rows:
        if route_for_record(record) != args.route:
            continue
        matched_route_count += 1
        previous = by_seed_point.get((record["seed"], record["point_index"] - 1))
        links_by_phase[phase_key(record)].append((record, previous))

    payload = {
        "source": str(args.jsonl),
        "route": args.route,
        "matched_route_count": matched_route_count,
        "phase_groups": {
            key: summarize_links(records)
            for key, records in sorted(links_by_phase.items())
        },
    }
    write_json(args.output, payload)
    print(f"wrote: {args.output}")
    for key, summary in payload["phase_groups"].items():
        previous = ";".join(f"{item['value']}:{item['count']}" for item in summary["previous_patterns"][:3])
        print(f"{key}: count={summary['count']} initial_pad={summary['current_initial_paddle_y']} previous={previous}")


if __name__ == "__main__":
    main()
