"""Probe narrow RAM policy rule candidates without changing the main policy."""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
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 (
    CENTER_Y,
    DOWN,
    LOW_SERVE_PATTERN_STEPS,
    LOW_SERVE_START_Y_MAX,
    LOW_SERVE_TARGET_Y,
    NOOP,
    TARGET_MAX_Y_RAM,
    TARGET_MIN_Y_RAM,
    UP,
    action_from_target,
    is_low_serve_pattern,
    target_y_from_state,
)
from pong_state import PongState, detect_state


@dataclass
class ProbePolicy:
    variant: str
    previous_state: PongState | None = None
    low_serve_steps: int = 0
    last_target_y: float = CENTER_Y
    last_mode: str = "init"
    rule_hits: int = 0
    previous_action: int = NOOP
    previous_paddle_y: int | None = None

    def reset(self) -> None:
        self.previous_state = None
        self.low_serve_steps = 0
        self.last_target_y = CENTER_Y
        self.last_mode = "reset"
        self.rule_hits = 0
        self.previous_action = NOOP
        self.previous_paddle_y = None

    def select_action(self, info: dict) -> int:
        state = detect_state(info)
        target_y = target_y_from_state(state, self.previous_state)
        mode = "intercept"
        dx = None
        dy = None
        if not state.ball_active:
            mode = "inactive"
        elif self.previous_state is None or not self.previous_state.ball_active:
            mode = "initial"
        else:
            dx = state.ball_x - self.previous_state.ball_x
            dy = state.ball_y - self.previous_state.ball_y
            if dx <= 0:
                mode = "center"

        if self.previous_state is not None and dx is not None and dy is not None:
            if is_low_serve_pattern(self.previous_state, dx, dy):
                self.low_serve_steps = LOW_SERVE_PATTERN_STEPS
            if self.low_serve_steps > 0:
                if dx > 0 and dy < 0 and state.ball_y < LOW_SERVE_START_Y_MAX + 2:
                    target_y = LOW_SERVE_TARGET_Y
                    mode = "low_serve"
                self.low_serve_steps -= 1
            target_y, rule_hit = apply_variant(self.variant, state, dx, dy, target_y)
            if rule_hit:
                self.rule_hits += 1
                mode = self.variant

        action = action_from_target(state, target_y)
        if self.previous_state is not None and dx is not None and dy is not None:
            previous_paddle_delta = None
            if self.previous_paddle_y is not None:
                previous_paddle_delta = state.self_paddle_y - self.previous_paddle_y
            action, action_hit = apply_action_variant(
                self.variant,
                state,
                dx,
                dy,
                action,
                self.previous_action,
                previous_paddle_delta,
            )
            if action_hit:
                self.rule_hits += 1
                mode = self.variant
        self.last_target_y = target_y
        self.last_mode = mode
        self.previous_state = state
        self.previous_action = action
        self.previous_paddle_y = state.self_paddle_y
        return action


def apply_variant(
    variant: str,
    state: PongState,
    dx: int,
    dy: int,
    target_y: float,
) -> tuple[float, bool]:
    if variant == "baseline":
        return target_y, False

    if variant == "low_mid_down_plus8":
        if dx > 0 and 184 <= state.ball_x <= 204 and 72 <= state.ball_y <= 145 and 3 <= dy <= 5:
            return target_y + 8, True
        return target_y, False

    if variant == "low_down_y86_plus8":
        if dx > 0 and 196 <= state.ball_x <= 204 and 78 <= state.ball_y <= 96 and 3 <= dy <= 5:
            return target_y + 8, True
        return target_y, False

    if variant == "low_down_y86_plus6":
        if dx > 0 and 196 <= state.ball_x <= 204 and 78 <= state.ball_y <= 96 and 3 <= dy <= 5:
            return target_y + 6, True
        return target_y, False

    if variant == "mid_up_minus8":
        if dx > 0 and 184 <= state.ball_x <= 204 and 145 <= state.ball_y <= 180 and dy <= -6:
            return target_y - 8, True
        return target_y, False

    if variant == "high_down_plus8":
        if dx > 0 and 184 <= state.ball_x <= 204 and state.ball_y >= 165 and dy >= 7:
            return target_y + 8, True
        return target_y, False

    if variant in {
        "low_down_force_down",
        "late_down_no_up",
        "late_down_keep_down",
        "phase_dy4_noop_to_down",
        "phase_dy4_up_to_noop",
        "phase_dyneg8_noop_to_up",
    }:
        return target_y, False

    if variant == "left_center_100":
        if dx <= 0:
            return 100, target_y != 100
        return target_y, False

    if variant == "left_center_140":
        if dx <= 0:
            return 140, target_y != 140
        return target_y, False

    if variant == "left_track_ball_y":
        if dx <= 0:
            tracked = max(TARGET_MIN_Y_RAM, min(TARGET_MAX_Y_RAM, state.ball_y))
            return tracked, target_y != tracked
        return target_y, False

    raise ValueError(f"unknown variant: {variant}")


def apply_action_variant(
    variant: str,
    state: PongState,
    dx: int,
    dy: int,
    action: int,
    previous_action: int,
    previous_paddle_delta: int | None,
) -> tuple[int, bool]:
    if variant == "low_down_force_down":
        if dx > 0 and 190 <= state.ball_x <= 204 and 70 <= state.ball_y <= 100 and 3 <= dy <= 5:
            return DOWN, action != DOWN
        return action, False

    if variant == "late_down_no_up":
        if dx > 0 and 188 <= state.ball_x <= 204 and 70 <= state.ball_y <= 140 and dy > 0 and action == UP:
            return NOOP, True
        return action, False

    if variant == "late_down_keep_down":
        if (
            dx > 0
            and 188 <= state.ball_x <= 204
            and 70 <= state.ball_y <= 140
            and dy > 0
            and action != DOWN
            and previous_action == DOWN
        ):
            return DOWN, True
        return action, False

    if variant == "phase_dy4_noop_to_down":
        if (
            dx > 0
            and 184 <= state.ball_x <= 195
            and dy == 4
            and action == NOOP
            and previous_paddle_delta is not None
            and previous_paddle_delta > 0
        ):
            return DOWN, True
        return action, False

    if variant == "phase_dy4_up_to_noop":
        if (
            dx > 0
            and 184 <= state.ball_x <= 195
            and dy == 4
            and action == UP
            and previous_paddle_delta is not None
            and previous_paddle_delta > 0
            and state.self_paddle_y - state.ball_y > 0
        ):
            return NOOP, True
        return action, False

    if variant == "phase_dyneg8_noop_to_up":
        if (
            dx > 0
            and 184 <= state.ball_x <= 195
            and dy == -8
            and action == NOOP
            and previous_paddle_delta is not None
            and previous_paddle_delta < 0
        ):
            return UP, True
        return action, False

    return action, False


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 run_episode(seed: int, variant: str, max_steps: int) -> dict:
    env = envpool.make("Pong-v5", env_type="gym", num_envs=1, seed=seed)
    try:
        _obs, info = reset_env(env)
        policy = ProbePolicy(variant=variant)
        policy.reset()
        score = 0.0
        for step in range(1, max_steps + 1):
            action = policy.select_action(info)
            _obs, reward, done, info = step_env(env, action)
            score += reward
            if done:
                return {
                    "seed": seed,
                    "variant": variant,
                    "score": score,
                    "steps": step,
                    "done": done,
                    "rule_hits": policy.rule_hits,
                }
        return {
            "seed": seed,
            "variant": variant,
            "score": score,
            "steps": max_steps,
            "done": False,
            "rule_hits": policy.rule_hits,
        }
    finally:
        env.close()


def parse_seeds(raw: str) -> list[int]:
    return [int(part.strip()) for part in raw.split(",") if part.strip()]


def write_jsonl(path: Path, rows: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, separators=(",", ":")))
            handle.write("\n")


def main() -> None:
    parser = argparse.ArgumentParser(description="Probe narrow candidate policy rules.")
    parser.add_argument("--seeds", default="0,1,2,3,4,5,6,7,8,9")
    parser.add_argument("--max-steps", type=int, default=30_000)
    parser.add_argument(
        "--variants",
        default="baseline,low_mid_down_plus8,mid_up_minus8,high_down_plus8",
    )
    parser.add_argument("--jsonl", type=Path, default=Path("diagnostics/candidate_rule_probe.jsonl"))
    args = parser.parse_args()

    seeds = parse_seeds(args.seeds)
    variants = [part.strip() for part in args.variants.split(",") if part.strip()]
    timestamp = datetime.now(timezone.utc).isoformat()
    rows = []
    for variant in variants:
        for seed in seeds:
            row = run_episode(seed=seed, variant=variant, max_steps=args.max_steps)
            row["timestamp"] = timestamp
            row["max_steps"] = args.max_steps
            rows.append(row)

    write_jsonl(args.jsonl, rows)

    print(f"envpool: {envpool.__version__}")
    print(f"seeds: {','.join(str(seed) for seed in seeds)}")
    print(f"max_steps: {args.max_steps}")
    print(f"wrote: {args.jsonl}")
    print("variant,mean,min,max,scores,rule_hits")
    for variant in variants:
        variant_rows = [row for row in rows if row["variant"] == variant]
        scores = [row["score"] for row in variant_rows]
        hits = sum(int(row["rule_hits"]) for row in variant_rows)
        score_text = ";".join(f"{row['seed']}:{row['score']:.0f}" for row in variant_rows)
        print(f"{variant},{mean(scores):.2f},{min(scores):.0f},{max(scores):.0f},{score_text},{hits}")


if __name__ == "__main__":
    main()
