"""Compare windows before the first y86 x166 bounce across score chains."""

from __future__ import annotations

import argparse
import json
import sys
from collections import Counter, defaultdict, deque
from dataclasses import asdict, dataclass
from pathlib import Path
from statistics import mean

import envpool
import numpy as np


REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from pong_policy import ACTION_NAMES, RamPolicy
from pong_state import PongState, detect_state


@dataclass(frozen=True)
class TraceRow:
    step: int
    ball_x: int
    ball_y: int
    self_paddle_y: int
    opponent_paddle_y: int
    dx: int | None
    dy: int | None
    target_y: float
    action_name: str
    mode: str
    self_score: int
    opponent_score: int
    previous_paddle_delta: int | None


@dataclass(frozen=True)
class SeedWindow:
    seed: int
    found: bool
    family: str
    first_bounce_step: int | None
    first_bounce_score: str | None
    rows: list[dict]


def reset_env(env: object) -> tuple[np.ndarray, dict]:
    reset_out = env.reset()
    if isinstance(reset_out, tuple):
        obs, info = reset_out
    else:
        obs, info = reset_out, {}
    return np.asarray(obs), info


def step_env(env: object, action: int) -> tuple[np.ndarray, float, bool, dict]:
    step_out = env.step(np.array([action], dtype=np.int32))
    if len(step_out) == 5:
        obs, reward, terminated, truncated, info = step_out
        done = np.logical_or(terminated, truncated)
    else:
        obs, reward, done, info = step_out
    return (
        np.asarray(obs),
        float(np.asarray(reward).reshape(-1)[0]),
        bool(np.asarray(done).reshape(-1)[0]),
        info,
    )


def is_y86_x166_bounce(state: PongState, previous_state: PongState | None, target_y: float) -> bool:
    if previous_state is None:
        return False
    dx = state.ball_x - previous_state.ball_x
    dy = state.ball_y - previous_state.ball_y
    return (
        dx > 0
        and dy == 5
        and 164 <= state.ball_x <= 168
        and 48 <= state.ball_y <= 52
        and 76 <= state.self_paddle_y <= 80
        and 80 <= target_y <= 85
    )


def family_from_score(score_pair: str | None) -> str:
    if score_pair == "2-2":
        return "accepted_score22_chain"
    if score_pair == "6-0":
        return "pure_score60_floor_chain"
    if score_pair is None:
        return "no_first_bounce"
    return f"other_{score_pair}"


def run_seed(seed: int, max_steps: int, window_size: int) -> SeedWindow:
    env = envpool.make("Pong-v5", env_type="gym", num_envs=1, seed=seed)
    try:
        _obs, info = reset_env(env)
        policy = RamPolicy()
        policy.reset()
        previous_paddle_y: int | None = None
        recent: deque[TraceRow] = deque(maxlen=window_size)

        for step in range(1, max_steps + 1):
            state = detect_state(info)
            previous_state = policy.previous_state
            action = policy.select_action(info)
            dx = dy = None
            if previous_state is not None:
                dx = state.ball_x - previous_state.ball_x
                dy = state.ball_y - previous_state.ball_y
            previous_paddle_delta = None
            if previous_paddle_y is not None:
                previous_paddle_delta = state.self_paddle_y - previous_paddle_y

            row = TraceRow(
                step=step,
                ball_x=state.ball_x,
                ball_y=state.ball_y,
                self_paddle_y=state.self_paddle_y,
                opponent_paddle_y=state.opponent_paddle_y,
                dx=dx,
                dy=dy,
                target_y=policy.last_target_y,
                action_name=ACTION_NAMES.get(action, str(action)),
                mode=policy.last_mode,
                self_score=state.self_score,
                opponent_score=state.opponent_score,
                previous_paddle_delta=previous_paddle_delta,
            )
            recent.append(row)

            if is_y86_x166_bounce(state, previous_state, policy.last_target_y):
                score_pair = f"{state.self_score}-{state.opponent_score}"
                return SeedWindow(
                    seed=seed,
                    found=True,
                    family=family_from_score(score_pair),
                    first_bounce_step=step,
                    first_bounce_score=score_pair,
                    rows=[asdict(item) for item in recent],
                )

            previous_paddle_y = state.self_paddle_y
            _obs, _reward, done, info = step_env(env, action)
            if done:
                break

        return SeedWindow(
            seed=seed,
            found=False,
            family="no_first_bounce",
            first_bounce_step=None,
            first_bounce_score=None,
            rows=[asdict(item) for item in recent],
        )
    finally:
        env.close()


def parse_seeds(raw: str) -> list[int]:
    if ".." in raw:
        start, end = raw.split("..", 1)
        return list(range(int(start), int(end) + 1))
    return [int(part.strip()) for part in raw.split(",") if part.strip()]


def bucket_row(row: dict) -> tuple:
    return (
        row.get("ball_x"),
        row.get("ball_y"),
        row.get("dx"),
        row.get("dy"),
        row.get("target_y"),
        row.get("self_paddle_y"),
        row.get("action_name"),
        row.get("mode"),
        f"{row.get('self_score')}-{row.get('opponent_score')}",
    )


def summarize_offset(rows: list[SeedWindow], offset: int) -> dict:
    by_family: dict[str, list[dict]] = defaultdict(list)
    for item in rows:
        if not item.found or len(item.rows) <= offset:
            continue
        by_family[item.family].append(item.rows[-1 - offset])

    result = {}
    for family, family_rows in sorted(by_family.items()):
        errors = [
            float(row["self_paddle_y"]) - float(row["target_y"])
            for row in family_rows
            if row.get("self_paddle_y") is not None and row.get("target_y") is not None
        ]
        top_bucket, top_count = Counter(bucket_row(row) for row in family_rows).most_common(1)[0]
        result[family] = {
            "count": len(family_rows),
            "top_count": top_count,
            "top_fraction": round(top_count / len(family_rows), 3),
            "mean_paddle_error": round(mean(errors), 3) if errors else None,
            "top": {
                "ball_x": top_bucket[0],
                "ball_y": top_bucket[1],
                "dx": top_bucket[2],
                "dy": top_bucket[3],
                "target_y": top_bucket[4],
                "self_paddle_y": top_bucket[5],
                "action_name": top_bucket[6],
                "mode": top_bucket[7],
                "score": top_bucket[8],
            },
            "action_counts": dict(Counter(row.get("action_name") for row in family_rows)),
            "mode_counts": dict(Counter(row.get("mode") for row in family_rows)),
        }
    return result


def summarize(rows: list[SeedWindow], offsets: list[int]) -> dict:
    return {
        "seed_count": len(rows),
        "family_counts": dict(Counter(row.family for row in rows)),
        "first_bounce_step_counts": dict(Counter(row.first_bounce_step for row in rows if row.found)),
        "offsets": {str(offset): summarize_offset(rows, offset) for offset in offsets},
        "rows": [asdict(row) for row in rows],
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--seeds", default="0..79")
    parser.add_argument("--max-steps", type=int, default=2600)
    parser.add_argument("--window", type=int, default=220)
    parser.add_argument("--offsets", default="0,1,2,4,8,12,20,40,80,120,160,200")
    parser.add_argument("--output", type=Path, default=Path("diagnostics/y86_pre_bounce_windows.json"))
    args = parser.parse_args()

    rows = [run_seed(seed, args.max_steps, args.window) for seed in parse_seeds(args.seeds)]
    offsets = [int(part.strip()) for part in args.offsets.split(",") if part.strip()]
    result = summarize(rows, offsets)
    result["source"] = {
        "script": str(Path(__file__)),
        "seeds": args.seeds,
        "max_steps": args.max_steps,
        "window": args.window,
        "envpool": envpool.__version__,
    }
    text = json.dumps(result, indent=2, sort_keys=True)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(text + "\n", encoding="utf-8")
    print(text)


if __name__ == "__main__":
    main()
