"""Classify point outcomes by previous terminal state and current reset phase."""

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 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 score_before(record: Trace) -> str:
    window = record.get("window") or []
    if window:
        first = window[0]
        return f"{first.get('self_score')}-{first.get('opponent_score')}"
    return "unknown"


def score_after(record: Trace) -> str:
    return f"{record.get('self_score')}-{record.get('opponent_score')}"


def terminal_key(record: Trace, coarse: bool) -> str:
    if coarse:
        return (
            f"y={record['final_ball_y']};dy={record['final_dy']};"
            f"pad={record['final_self_paddle_y']};"
            f"target={round(float(record['final_target_y']), 1)};"
            f"err={round(float(record['final_paddle_error']), 1)};"
            f"outcome={record['outcome']};pattern={record['pattern']}"
        )
    return (
        f"x={record['final_ball_x']};y={record['final_ball_y']};"
        f"dx={record['final_dx']};dy={record['final_dy']};"
        f"pad={record['final_self_paddle_y']};"
        f"target={round(float(record['final_target_y']), 1)};"
        f"err={round(float(record['final_paddle_error']), 1)};"
        f"outcome={record['outcome']};pattern={record['pattern']}"
    )


def initial_key(initial: Trace | None) -> str:
    if initial is None:
        return "initial=missing"
    return (
        f"initial_x={initial.get('ball_x')};"
        f"initial_y={initial.get('ball_y')};"
        f"initial_pad={initial.get('self_paddle_y')};"
        f"initial_target={round(float(initial.get('target_y')), 1)};"
        f"initial_action={initial.get('action_name')}"
    )


def current_label(record: Trace) -> str:
    if record.get("outcome") == "loss":
        return (
            f"loss:{record.get('pattern')};"
            f"final_y={record.get('final_ball_y')};"
            f"score_after={score_after(record)}"
        )
    return "win"


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 summarize_group(items: list[dict[str, Any]], top: int) -> dict[str, Any]:
    labels = Counter(item["current_label"] for item in items)
    routes = Counter(item["current_route"] for item in items)
    scores_before = Counter(item["current_score_before"] for item in items)
    final_scores = Counter(item["seed_final_score"] for item in items)
    loss_count = sum(1 for item in items if item["current_outcome"] == "loss")
    floor18_count = sum(1 for item in items if item["seed_final_score"] == 18)
    examples = items[:top]
    return {
        "count": len(items),
        "loss_count": loss_count,
        "loss_rate": round(loss_count / len(items), 4) if items else 0,
        "floor18_count": floor18_count,
        "floor18_rate": round(floor18_count / len(items), 4) if items else 0,
        "seeds": sorted({item["seed"] for item in items}),
        "current_labels": [
            {"value": value, "count": count}
            for value, count in labels.most_common(top)
        ],
        "current_routes": [
            {"value": value, "count": count}
            for value, count in routes.most_common(top)
        ],
        "current_scores_before": [
            {"value": value, "count": count}
            for value, count in scores_before.most_common(top)
        ],
        "seed_final_scores": [
            {"value": value, "count": count}
            for value, count in final_scores.most_common(top)
        ],
        "current_initial_paddle_y": numeric(
            [float(item["current_initial_paddle_y"]) for item in items if item["current_initial_paddle_y"] is not None]
        ),
        "examples": examples,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--jsonl", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--top", type=int, default=20)
    parser.add_argument("--coarse-terminal", action="store_true")
    args = parser.parse_args()

    records = read_jsonl(args.jsonl)
    by_seed_point = {(record["seed"], record["point_index"]): record for record in records}
    final_score_by_seed = {
        seed: int(seed_records[-1]["score_diff"])
        for seed, seed_records in (
            (seed, sorted([record for record in records if record["seed"] == seed], key=lambda row: row["point_index"]))
            for seed in sorted({record["seed"] for record in records})
        )
        if seed_records
    }

    groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for record in records:
        previous = by_seed_point.get((record["seed"], record["point_index"] - 1))
        if previous is None:
            continue
        initial = first_initial_row(record)
        key = terminal_key(previous, args.coarse_terminal) + ";" + initial_key(initial)
        groups[key].append(
            {
                "seed": record["seed"],
                "point": record["point_index"],
                "seed_final_score": final_score_by_seed.get(record["seed"]),
                "previous_terminal": terminal_key(previous, args.coarse_terminal),
                "current_initial": initial_key(initial),
                "current_initial_paddle_y": None if initial is None else initial.get("self_paddle_y"),
                "current_score_before": score_before(record),
                "current_score_after": score_after(record),
                "current_outcome": record["outcome"],
                "current_pattern": record["pattern"],
                "current_label": current_label(record),
                "current_route": route_for_record(record) or "none",
                "current_final_y": record["final_ball_y"],
                "current_final_error": record["final_paddle_error"],
            }
        )

    summaries = {
        key: summarize_group(items, args.top)
        for key, items in groups.items()
    }
    ranked_by_floor = sorted(
        ({"classifier_key": key, **summary} for key, summary in summaries.items()),
        key=lambda row: (row["floor18_count"], row["loss_count"], row["count"]),
        reverse=True,
    )[: args.top]
    ranked_by_loss = sorted(
        ({"classifier_key": key, **summary} for key, summary in summaries.items()),
        key=lambda row: (row["loss_count"], row["floor18_count"], row["count"]),
        reverse=True,
    )[: args.top]
    y205_groups = [
        {"classifier_key": key, **summary}
        for key, summary in summaries.items()
        if any("final_y=205" in label["value"] for label in summary["current_labels"])
    ]
    y171_groups = [
        {"classifier_key": key, **summary}
        for key, summary in summaries.items()
        if any("final_y=171" in label["value"] for label in summary["current_labels"])
    ]

    payload = {
        "source": str(args.jsonl),
        "coarse_terminal": args.coarse_terminal,
        "group_count": len(summaries),
        "ranked_by_floor18": ranked_by_floor,
        "ranked_by_loss": ranked_by_loss,
        "y205_groups": sorted(y205_groups, key=lambda row: row["count"], reverse=True),
        "y171_groups": sorted(y171_groups, key=lambda row: row["count"], reverse=True),
        "interpretation": (
            "Classifier key is previous terminal geometry plus current initial frame. "
            "This is read-only and should be used to decide whether phase markers predict failures before policy edits."
        ),
    }
    write_json(args.output, payload)
    print(f"wrote: {args.output}")
    print(f"groups: {len(summaries)}")
    print("top floor18 groups:")
    for row in ranked_by_floor[:5]:
        print(
            f"count={row['count']} floor18={row['floor18_count']} losses={row['loss_count']} "
            f"labels={row['current_labels'][:2]} key={row['classifier_key']}"
        )
    print("y205 groups:")
    for row in payload["y205_groups"][:5]:
        print(f"count={row['count']} seeds={row['seeds']} labels={row['current_labels'][:2]}")
    print("y171 groups:")
    for row in payload["y171_groups"][:5]:
        print(f"count={row['count']} seeds={row['seeds']} labels={row['current_labels'][:2]}")


if __name__ == "__main__":
    main()
