# tests/runtime/test_predator_chase_memory.py
"""Redesigned predator_chase memory: longer pre-kill preroll + wider, more
varied pre-kill roaming.

Baseline (pre-redesign, seed 7):
  - first "eaten" event at turn_idx 71
  - max distractor Manhattan displacement from start during the pre-kill phase
    was only 12 cells (uniform random walk drifts back near the start).

These tests pin measurable improvements:
  - the first kill must land meaningfully later (preroll floor),
  - distractors must travel noticeably farther from their starts before the
    first kill (wider action radius).

Hard invariants (a0 survivor, 3 genuine eaten, determinism) are asserted in
test_director_predator_chase.py and reinforced here as guardrails.
"""
from proteus.game.runtime.multiagent_director import author_predator_chase

_STARTS = [(10, 30), (14, 38), (18, 24), (12, 46)]
_PRED = (54, 31)

# Justified thresholds vs baseline (seed 7):
#   first-eat baseline = 71 -> require >= 100 (the redesign lands it at 104,
#     +33 turns / ~1.46x later; the floor sits a few turns below the achieved
#     value because first-eat is non-monotonic in the grace/roam interaction).
#   max pre-kill distractor displacement baseline = 12 -> require >= 20
#     (the persistent-heading roam achieves 40, vs the random walk's 12).
_FIRST_EAT_FLOOR = 100
_SPREAD_FLOOR = 20


def _run():
    return author_predator_chase(seed=7, agent_starts=_STARTS, predator_start=_PRED)


def _first_eat_turn(ck):
    for t in ck.memory_turns:
        for e in t.events:
            if "eaten" in e:
                return t.turn_idx
    return None


def test_first_kill_happens_later():
    """First genuine 'eaten' event must occur at a clearly later turn."""
    first = _first_eat_turn(_run())
    assert first is not None, "expected at least one 'eaten' event"
    assert first >= _FIRST_EAT_FLOOR, (
        f"first eat at turn {first}, want >= {_FIRST_EAT_FLOOR}"
    )


def test_distractors_roam_wider_before_first_kill():
    """Max distractor displacement from start during the pre-kill phase must
    clearly exceed the uniform-random-walk baseline."""
    ck = _run()
    first = _first_eat_turn(ck)
    assert first is not None
    starts = {f"a{i}": p for i, p in enumerate(_STARTS) if i != 0}
    max_disp = 0
    for t in ck.memory_turns:
        if t.turn_idx >= first:
            break
        for a in t.agents:
            if a.id in starts:
                s = starts[a.id]
                d = abs(a.pos[0] - s[0]) + abs(a.pos[1] - s[1])
                max_disp = max(max_disp, d)
    assert max_disp >= _SPREAD_FLOOR, (
        f"max pre-kill distractor displacement {max_disp}, want >= {_SPREAD_FLOOR}"
    )


def _pred_center(turn):
    for a in turn.agents:
        if a.kind == "predator":
            return (a.pos[0] + a.size // 2, a.pos[1] + a.size // 2)
    return None


def _agent_center(a):
    return (a.pos[0] + a.size // 2, a.pos[1] + a.size // 2)


# Window (in turns) right after the first kill over which the pack must visibly
# move AWAY from the predator. Kept short so edge-cornering of one agent cannot
# make the assertion vacuous or fragile.
_SCATTER_WINDOW = 6
# One-step slack: agent-agent collision (the no-overlap invariant) can block a
# survivor's optimal flee move for a single transition, costing at most ~2 cells
# of Manhattan distance before it recovers. We allow that much and no more.
# Measured separation between the two regimes for seed 7:
#   roaming (pre-change): nearest survivor is steadily approached, min distance
#     drops by 6 over the window  -> exceeds the slack -> FAILS.
#   fleeing (post-change): the pack holds the predator off, min distance drops
#     by at most 2 (one blocked step) then holds -> within slack -> PASSES.
_BLOCK_SLACK = 2


def _min_survivor_dist(turn, ids):
    """Minimum Manhattan distance from any alive distractor in *ids* to the
    predator on *turn*, or None if no predator / none of *ids* alive."""
    pred = _pred_center(turn)
    if pred is None:
        return None
    dists = [
        abs(_agent_center(a)[0] - pred[0]) + abs(_agent_center(a)[1] - pred[1])
        for a in turn.agents
        if a.kind == "agent" and a.id in ids and a.alive
    ]
    return min(dists) if dists else None


def test_pack_scatters_after_first_kill():
    """After the FIRST 'eaten' event the whole pack scatters and FLEES, so the
    predator cannot keep closing the gap: the MINIMUM surviving-distractor
    distance to the predator never drops more than one blocked step below its
    value at the kill turn across the short window right after the kill.

    Pre-change the survivors keep roaming (predator-unaware), so the predator
    steadily corners the nearest one and the minimum distance collapses well
    past the slack -> this assertion FAILS. With the scatter change every
    survivor flees, the pack holds the predator off, and the minimum distance is
    maintained -> it PASSES."""
    ck = _run()
    first = _first_eat_turn(ck)
    assert first is not None
    by_turn = {t.turn_idx: t for t in ck.memory_turns}

    # The kill-turn frame is recorded PRE-move, so the agent being eaten is still
    # alive in it. Restrict to the genuine SURVIVORS (alive non-chosen distractors
    # that are NOT named in any 'eaten' event on the kill turn) so the baseline
    # reflects the pack that must flee, not the agent about to die.
    kill_turn = by_turn[first]
    eaten_now = {e.split()[0] for e in kill_turn.events if "eaten" in e}
    survivor_ids = {
        a.id for a in kill_turn.agents
        if a.kind == "agent" and a.id != "a0" and a.alive and a.id not in eaten_now
    }
    assert survivor_ids, "expected surviving distractors at the first-kill turn"

    kill_min = _min_survivor_dist(kill_turn, survivor_ids)
    assert kill_min is not None
    floor = kill_min - _BLOCK_SLACK

    checked = 0
    for off in range(1, _SCATTER_WINDOW + 1):
        t = by_turn.get(first + off)
        if t is None:
            break
        m = _min_survivor_dist(t, survivor_ids)
        if m is None:  # final settled frame (predator removed) or no survivors
            break
        checked += 1
        assert m >= floor, (
            f"nearest survivor closed on the predator after the kill: "
            f"min distance {m} at turn {t.turn_idx} < floor {floor} "
            f"(kill-turn min {kill_min} - slack {_BLOCK_SLACK}); the pack did not flee"
        )
    assert checked >= 3, f"scatter window too short to be meaningful ({checked} frames)"


def test_guardrails_survivor_and_three_genuine_eaten():
    """a0 is the sole survivor, alive every turn, and exactly 3 distractors are
    eaten by GENUINE predator contact (not the silent end safeguard)."""
    ck = _run()
    for t in ck.memory_turns:
        chosen = next(a for a in t.agents if a.id == "a0")
        assert chosen.alive
    eaten = [e for t in ck.memory_turns for e in t.events if "eaten" in e]
    assert len(eaten) == 3, f"expected 3 genuine eaten events, got {eaten}"
    last = ck.memory_turns[-1]
    alive = [a for a in last.agents if a.kind == "agent" and a.alive]
    assert len(alive) == 1 and alive[0].id == "a0"