"""Compare RAM trajectories after counterfactual action interventions."""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass
from pathlib import Path
from statistics import mean
from typing import Any

import envpool


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

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

from pong_policy import ACTION_NAMES, RamPolicy
from pong_state import detect_state
from probe_counterfactual_actions import (
    StepContext,
    Variant,
    parse_seeds,
    parse_variants,
    reset_env,
    sign,
    step_env,
    trigger_y86_bounce_score22_x166_down,
    trigger_y86_bounce_x166_down,
)


@dataclass(frozen=True)
class TriggerInfo:
    step: int
    previous_paddle_delta: int | None


def context_for_step(
    step: int,
    info: dict,
    policy: RamPolicy,
    previous_paddle_y: int | None,
) -> tuple[StepContext, int]:
    state = detect_state(info)
    previous_state = policy.previous_state
    action = policy.select_action(info)
    dx = None
    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
    return (
        StepContext(
            step=step,
            state=state,
            dx=dx,
            dy=dy,
            action=action,
            target_y=policy.last_target_y,
            mode=policy.last_mode,
            previous_paddle_delta=previous_paddle_delta,
        ),
        action,
    )


def trigger_matches(variant: Variant, ctx: StepContext, y86_score22_seen: bool) -> bool:
    if variant.trigger is None:
        return False
    if variant.name == "repeat_y86_bounce_after_score22_x166_down_to_up_h1":
        return y86_score22_seen and trigger_y86_bounce_x166_down(ctx)
    return variant.trigger(ctx)


def find_trigger(seed: int, variant: Variant, max_steps: int) -> TriggerInfo | None:
    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
        y86_score22_seen = False
        for step in range(1, max_steps + 1):
            ctx, action = context_for_step(step, info, policy, previous_paddle_y)
            if trigger_y86_bounce_score22_x166_down(ctx):
                y86_score22_seen = True
            if trigger_matches(variant, ctx, y86_score22_seen):
                return TriggerInfo(step=step, previous_paddle_delta=ctx.previous_paddle_delta)
            previous_paddle_y = ctx.state.self_paddle_y
            _obs, _reward, done, info = step_env(env, action)
            if done:
                return None
        return None
    finally:
        env.close()


def trace_run(
    seed: int,
    variant: Variant,
    max_steps: int,
    trace_start: int,
    trace_len: int,
) -> dict[str, Any]:
    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
        override_steps = 0
        intervention_count = 0
        interventions: list[dict[str, Any]] = []
        rows: list[dict[str, Any]] = []
        y86_score22_seen = False
        score = 0.0
        done = False

        for step in range(1, max_steps + 1):
            ctx, action = context_for_step(step, info, policy, previous_paddle_y)
            if trigger_y86_bounce_score22_x166_down(ctx):
                y86_score22_seen = True
            if (
                intervention_count < variant.max_interventions
                and variant.override_action is not None
                and trigger_matches(variant, ctx, y86_score22_seen)
            ):
                intervention_count += 1
                override_steps = variant.horizon
                interventions.append(
                    {
                        "step": step,
                        "ball_x": ctx.state.ball_x,
                        "ball_y": ctx.state.ball_y,
                        "self_paddle_y": ctx.state.self_paddle_y,
                        "self_score": ctx.state.self_score,
                        "opponent_score": ctx.state.opponent_score,
                        "baseline_action": ACTION_NAMES.get(action, str(action)),
                        "override_action": ACTION_NAMES.get(variant.override_action, str(variant.override_action)),
                        "horizon": variant.horizon,
                        "target_y": ctx.target_y,
                        "mode": ctx.mode,
                        "previous_paddle_delta": ctx.previous_paddle_delta,
                        "previous_paddle_delta_sign": sign(ctx.previous_paddle_delta),
                    }
                )

            executed_action = action
            if override_steps > 0 and variant.override_action is not None:
                executed_action = variant.override_action
                override_steps -= 1

            should_trace = trace_start <= step < trace_start + trace_len
            if should_trace:
                rows.append(
                    {
                        "step": step,
                        "ball_x": ctx.state.ball_x,
                        "ball_y": ctx.state.ball_y,
                        "self_paddle_y": ctx.state.self_paddle_y,
                        "self_score": ctx.state.self_score,
                        "opponent_score": ctx.state.opponent_score,
                        "dx": ctx.dx,
                        "dy": ctx.dy,
                        "target_y": ctx.target_y,
                        "mode": ctx.mode,
                        "baseline_action": ACTION_NAMES.get(action, str(action)),
                        "executed_action": ACTION_NAMES.get(executed_action, str(executed_action)),
                        "previous_paddle_delta": ctx.previous_paddle_delta,
                    }
                )

            previous_paddle_y = ctx.state.self_paddle_y
            _obs, reward, done, info = step_env(env, executed_action)
            score += reward
            if should_trace:
                rows[-1]["reward"] = reward
                rows[-1]["done_after"] = done
            if done and step >= trace_start + trace_len:
                break

        return {
            "seed": seed,
            "variant": variant.name,
            "score": score,
            "done": done,
            "interventions": interventions,
            "trace_start": trace_start,
            "trace_len": trace_len,
            "rows": rows,
        }
    finally:
        env.close()


def first_divergence(
    baseline_rows: list[dict[str, Any]],
    variant_rows: list[dict[str, Any]],
    keys: tuple[str, ...],
) -> dict[str, Any] | None:
    for index, (base, variant) in enumerate(zip(baseline_rows, variant_rows)):
        changed = {key: [base.get(key), variant.get(key)] for key in keys if base.get(key) != variant.get(key)}
        if changed:
            return {"offset": index, "step": base.get("step"), "changed": changed}
    return None


def first_reconvergence_after_divergence(
    baseline_rows: list[dict[str, Any]],
    variant_rows: list[dict[str, Any]],
    keys: tuple[str, ...],
) -> dict[str, Any] | None:
    divergence = first_divergence(baseline_rows, variant_rows, keys)
    if divergence is None:
        return {"offset": 0, "step": baseline_rows[0].get("step") if baseline_rows else None}
    for index, (base, variant) in enumerate(
        zip(baseline_rows[divergence["offset"] + 1 :], variant_rows[divergence["offset"] + 1 :]),
        start=divergence["offset"] + 1,
    ):
        if all(base.get(key) == variant.get(key) for key in keys):
            return {"offset": index, "step": base.get("step")}
    return None


def divergent_offsets(
    baseline_rows: list[dict[str, Any]],
    variant_rows: list[dict[str, Any]],
    keys: tuple[str, ...],
) -> list[int]:
    offsets = []
    for index, (base, variant) in enumerate(zip(baseline_rows, variant_rows)):
        if any(base.get(key) != variant.get(key) for key in keys):
            offsets.append(index)
    return offsets


def compare(seed: int, variants: list[Variant], max_steps: int, pre_steps: int, trace_len: int) -> dict[str, Any]:
    active_variants = [variant for variant in variants if variant.name != "baseline"]
    trigger_infos = {
        variant.name: find_trigger(seed, variant, max_steps)
        for variant in active_variants
    }
    trigger_steps = [info.step for info in trigger_infos.values() if info is not None]
    if not trigger_steps:
        return {"seed": seed, "error": "no trigger found", "trigger_infos": trigger_infos}
    trace_start = max(1, min(trigger_steps) - pre_steps)
    baseline = trace_run(seed, variants[0], max_steps, trace_start, trace_len)
    variant_results = []
    physical_keys = (
        "ball_x",
        "ball_y",
        "self_paddle_y",
        "self_score",
        "opponent_score",
        "dx",
        "dy",
        "target_y",
        "mode",
        "reward",
        "done_after",
    )
    control_keys = physical_keys + ("baseline_action", "executed_action", "previous_paddle_delta")
    for variant in active_variants:
        result = trace_run(seed, variant, max_steps, trace_start, trace_len)
        physical_divergent_offsets = divergent_offsets(baseline["rows"], result["rows"], physical_keys)
        control_divergent_offsets = divergent_offsets(baseline["rows"], result["rows"], control_keys)
        variant_results.append(
            {
                "variant": variant.name,
                "score": result["score"],
                "interventions": result["interventions"],
                "first_physical_divergence": first_divergence(baseline["rows"], result["rows"], physical_keys),
                "first_control_divergence": first_divergence(baseline["rows"], result["rows"], control_keys),
                "first_physical_reconvergence": first_reconvergence_after_divergence(
                    baseline["rows"], result["rows"], physical_keys
                ),
                "first_control_reconvergence": first_reconvergence_after_divergence(
                    baseline["rows"], result["rows"], control_keys
                ),
                "physical_divergent_offset_count": len(physical_divergent_offsets),
                "physical_last_divergent_offset": physical_divergent_offsets[-1]
                if physical_divergent_offsets
                else None,
                "control_divergent_offset_count": len(control_divergent_offsets),
                "control_last_divergent_offset": control_divergent_offsets[-1]
                if control_divergent_offsets
                else None,
                "last_rows": result["rows"][-5:],
            }
        )
    return {
        "seed": seed,
        "trace_start": trace_start,
        "trace_len": trace_len,
        "trigger_infos": {
            name: None if info is None else {"step": info.step, "previous_paddle_delta": info.previous_paddle_delta}
            for name, info in trigger_infos.items()
        },
        "baseline": {
            "score": baseline["score"],
            "last_rows": baseline["rows"][-5:],
        },
        "variants": variant_results,
    }


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 main() -> None:
    parser = argparse.ArgumentParser(description="Compare trajectories after counterfactual interventions.")
    parser.add_argument("--seeds", required=True)
    parser.add_argument("--variants", required=True)
    parser.add_argument("--max-steps", type=int, default=30_000)
    parser.add_argument("--pre-steps", type=int, default=4)
    parser.add_argument("--trace-len", type=int, default=120)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    seeds = parse_seeds(args.seeds)
    variants = parse_variants(args.variants)
    if not variants or variants[0].name != "baseline":
        raise ValueError("first variant must be baseline")

    seed_results = [
        compare(seed, variants, args.max_steps, args.pre_steps, args.trace_len)
        for seed in seeds
    ]
    scores = [result["baseline"]["score"] for result in seed_results if "baseline" in result]
    payload = {
        "source_variants": [variant.name for variant in variants],
        "max_steps": args.max_steps,
        "pre_steps": args.pre_steps,
        "trace_len": args.trace_len,
        "baseline_mean": round(mean(scores), 4) if scores else None,
        "seeds": seeds,
        "seed_results": seed_results,
    }
    write_json(args.output, payload)
    print(f"wrote: {args.output}")
    for result in seed_results:
        print(f"seed {result.get('seed')}: trace_start={result.get('trace_start')}")
        for variant in result.get("variants", []):
            print(
                f"  {variant['variant']}: physical={variant['first_physical_divergence']} "
                f"control={variant['first_control_divergence']}"
            )


if __name__ == "__main__":
    main()
