"""Summarize point-reset/setup history before the first y86 x166 bounce."""

from __future__ import annotations

import argparse
import json
import sys
from collections import Counter, defaultdict
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 RewardEvent:
    step: int
    reward: int
    score_after: str
    last_ball_x: int
    last_ball_y: int
    last_dx: int | None
    last_dy: int | None
    last_target_y: float
    last_self_paddle_y: int
    last_action_name: str
    last_mode: str


@dataclass(frozen=True)
class SeedSetup:
    seed: int
    found: bool
    family: str
    first_bounce_step: int | None
    first_bounce_score: str | None
    last_reward: dict | None
    reward_history: list[dict]
    steps_since_last_reward: int | None
    first_active_after_reward: dict | None
    first_right_after_reward: dict | None
    first_high_down_after_reward: dict | None
    first_y86_approach_after_reward: dict | None
    bounce_row: dict | None


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 row_dict(row: TraceRow | None) -> dict | None:
    if row is None:
        return None
    return asdict(row)


def reward_event(row: TraceRow, reward: float, state_after: PongState) -> RewardEvent:
    return RewardEvent(
        step=row.step,
        reward=int(reward),
        score_after=f"{state_after.self_score}-{state_after.opponent_score}",
        last_ball_x=row.ball_x,
        last_ball_y=row.ball_y,
        last_dx=row.dx,
        last_dy=row.dy,
        last_target_y=row.target_y,
        last_self_paddle_y=row.self_paddle_y,
        last_action_name=row.action_name,
        last_mode=row.mode,
    )


def update_anchors(
    row: TraceRow,
    first_active: TraceRow | None,
    first_right: TraceRow | None,
    first_high_down: TraceRow | None,
    first_y86_approach: TraceRow | None,
) -> tuple[TraceRow | None, TraceRow | None, TraceRow | None, TraceRow | None]:
    if first_active is None and row.ball_y != 0 and row.ball_x not in (0, 205):
        first_active = row
    if first_right is None and row.dx is not None and row.dx > 0 and row.ball_x >= 120:
        first_right = row
    if (
        first_high_down is None
        and row.dx is not None
        and row.dx > 0
        and row.dy is not None
        and row.dy >= 8
        and 140 <= row.ball_x <= 170
        and row.ball_y >= 45
    ):
        first_high_down = row
    if (
        first_y86_approach is None
        and row.dx is not None
        and row.dx > 0
        and row.dy == -4
        and 80 <= row.ball_x <= 90
        and 115 <= row.ball_y <= 125
    ):
        first_y86_approach = row
    return first_active, first_right, first_high_down, first_y86_approach


def run_seed(seed: int, max_steps: int) -> SeedSetup:
    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
        rewards: list[RewardEvent] = []
        last_reward_step: int | None = None
        first_active: TraceRow | None = None
        first_right: TraceRow | None = None
        first_high_down: TraceRow | None = None
        first_y86_approach: TraceRow | None = None

        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,
            )
            first_active, first_right, first_high_down, first_y86_approach = update_anchors(
                row,
                first_active,
                first_right,
                first_high_down,
                first_y86_approach,
            )

            if is_y86_x166_bounce(state, previous_state, policy.last_target_y):
                score_pair = f"{state.self_score}-{state.opponent_score}"
                return SeedSetup(
                    seed=seed,
                    found=True,
                    family=family_from_score(score_pair),
                    first_bounce_step=step,
                    first_bounce_score=score_pair,
                    last_reward=asdict(rewards[-1]) if rewards else None,
                    reward_history=[asdict(event) for event in rewards[-5:]],
                    steps_since_last_reward=step - last_reward_step if last_reward_step is not None else None,
                    first_active_after_reward=row_dict(first_active),
                    first_right_after_reward=row_dict(first_right),
                    first_high_down_after_reward=row_dict(first_high_down),
                    first_y86_approach_after_reward=row_dict(first_y86_approach),
                    bounce_row=asdict(row),
                )

            previous_paddle_y = state.self_paddle_y
            _obs, reward, done, info = step_env(env, action)
            if reward != 0:
                state_after = detect_state(info)
                rewards.append(reward_event(row, reward, state_after))
                last_reward_step = step
                first_active = None
                first_right = None
                first_high_down = None
                first_y86_approach = None
            if done:
                break

        return SeedSetup(
            seed=seed,
            found=False,
            family="no_first_bounce",
            first_bounce_step=None,
            first_bounce_score=None,
            last_reward=asdict(rewards[-1]) if rewards else None,
            reward_history=[asdict(event) for event in rewards[-5:]],
            steps_since_last_reward=None,
            first_active_after_reward=row_dict(first_active),
            first_right_after_reward=row_dict(first_right),
            first_high_down_after_reward=row_dict(first_high_down),
            first_y86_approach_after_reward=row_dict(first_y86_approach),
            bounce_row=None,
        )
    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 anchor_bucket(row: dict | None) -> str:
    if row is None:
        return "none"
    return (
        f"step_rel_score={row.get('self_score')}-{row.get('opponent_score')};"
        f"x={row.get('ball_x')};y={row.get('ball_y')};"
        f"dx={row.get('dx')};dy={row.get('dy')};"
        f"pad={row.get('self_paddle_y')};target={round(float(row.get('target_y')), 2)};"
        f"action={row.get('action_name')};mode={row.get('mode')}"
    )


def reward_bucket(row: dict | None) -> str:
    if row is None:
        return "none"
    return (
        f"step={row.get('step')};reward={row.get('reward')};score_after={row.get('score_after')};"
        f"x={row.get('last_ball_x')};y={row.get('last_ball_y')};"
        f"dx={row.get('last_dx')};dy={row.get('last_dy')};"
        f"pad={row.get('last_self_paddle_y')};target={round(float(row.get('last_target_y')), 2)};"
        f"action={row.get('last_action_name')};mode={row.get('last_mode')}"
    )


def summarize_anchor(rows: list[SeedSetup], attr: str) -> dict:
    by_family: dict[str, list[str]] = defaultdict(list)
    for row in rows:
        if not row.found:
            continue
        bucket = anchor_bucket(getattr(row, attr))
        by_family[row.family].append(bucket)
    return {
        family: [
            {"count": count, "bucket": bucket}
            for bucket, count in Counter(buckets).most_common(8)
        ]
        for family, buckets in sorted(by_family.items())
    }


def summarize(rows: list[SeedSetup]) -> dict:
    found = [row for row in rows if row.found]
    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 found)),
        "steps_since_last_reward": {
            family: {
                "count": len(values),
                "mean": round(mean(values), 3) if values else None,
                "counts": dict(Counter(values)),
            }
            for family, values in sorted(
                {
                    family: [
                        int(row.steps_since_last_reward)
                        for row in found
                        if row.family == family and row.steps_since_last_reward is not None
                    ]
                    for family in {row.family for row in found}
                }.items()
            )
        },
        "last_reward_buckets": {
            family: [
                {"count": count, "bucket": bucket}
                for bucket, count in Counter(
                    reward_bucket(row.last_reward) for row in found if row.family == family
                ).most_common(8)
            ]
            for family in sorted({row.family for row in found})
        },
        "first_active_after_reward": summarize_anchor(rows, "first_active_after_reward"),
        "first_right_after_reward": summarize_anchor(rows, "first_right_after_reward"),
        "first_high_down_after_reward": summarize_anchor(rows, "first_high_down_after_reward"),
        "first_y86_approach_after_reward": summarize_anchor(rows, "first_y86_approach_after_reward"),
        "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("--output", type=Path, default=Path("diagnostics/y86_chain_setup.json"))
    args = parser.parse_args()

    rows = [run_seed(seed, args.max_steps) for seed in parse_seeds(args.seeds)]
    result = summarize(rows)
    result["source"] = {
        "script": str(Path(__file__)),
        "seeds": args.seeds,
        "max_steps": args.max_steps,
        "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()
