backend/tests.py

"""
tests.py — canon-rule regression tests for the 4-track engine. Pure, no model needed.
Run:  python3 tests.py     (exits non-zero if anything fails)

Covers: map shape, determinism, every loss/win condition, the chaos schedule + caps, no-stacking,
peak-signal-all, crossover-costs-the-turn, and headway collision avoidance.
"""
from __future__ import annotations
import sys
import engine
import rules
import simulation
from engine import Incident
from doctrine import doctrine
from chaos_players import SmartChaos
from agents import enforce_budget

_passed = _failed = 0


def check(name, cond):
    global _passed, _failed
    if cond:
        _passed += 1; print(f"  PASS  {name}")
    else:
        _failed += 1; print(f"  FAIL  {name}")


def run_oracle_game(seed):
    g = engine.new_game(seed); chaos = SmartChaos(); chaos.reset(seed)
    while not g.over:
        plays, _ = chaos.play(g, rules.legal_cards(g))
        for c, loc in plays:
            if engine.CARDS[c] <= g.energy and rules.card_available(g, c) and rules.station_free(g, loc):
                g.energy -= engine.CARDS[c]; rules.apply_chaos(g, c, loc)
        A = rules.legal_actions(g); chosen, _ = doctrine(g, A)
        ann, pol = simulation.apply(g, A, chosen); simulation.advance(g, ann, pol)
    return g


print("MAP")
check("8 stations", len(engine.STATIONS) == 8)
check("Nalasopara dropped", "Nalasopara" not in engine.STATIONS)
check("majors are 4", engine.MAJORS == {"Dadar", "Andheri", "Borivali", "Churchgate"})
check("interchanges from balance", engine.INTERCHANGES == {"Virar", "Churchgate", "Dadar", "Borivali"})
check("20 rounds / peak 16", engine.B["turns_to_win"] == 20 and engine.B["peak_rush_turn"] == 16)
check("crush threshold 170", engine.B["single_overflow_pct"] == 170)

print("DETERMINISM")
g1, g2 = run_oracle_game(3), run_oracle_game(3)
check("same seed => same outcome", (g1.won, g1.turn, g1.reason, round(g1.score)) == (g2.won, g2.turn, g2.reason, round(g2.score)))

print("WIN / LOSS CONDITIONS")
# win: oracle survives to 20 on a calm seed
gw = run_oracle_game(0)
check("oracle can win (survive 20)", gw.won and gw.turn == 20)

# single-station platform crush at >170
gc = engine.new_game(0); gc.stations[engine.IDX["Virar"]].crowd = gc.stations[engine.IDX["Virar"]].cap * 1.8
simulation.advance(gc, False, set())
check("platform crush fires >170%", gc.over and "platform crush" in gc.reason)
# ...but not at 145%
gc2 = engine.new_game(0); gc2.stations[engine.IDX["Virar"]].crowd = gc2.stations[engine.IDX["Virar"]].cap * 1.45
simulation.advance(gc2, False, set())
check("no crush at 145%", not (gc2.over and "platform crush" in (gc2.reason or "")))

# double overflow: two majors >130 (set to 160% so they stay >130% even after a train discharges
# ~360 and inflow lands — a train starts at each of Dadar/Andheri)
gd = engine.new_game(0)
for m in ("Dadar", "Andheri"):
    gd.stations[engine.IDX[m]].crowd = gd.stations[engine.IDX[m]].cap * 1.6
simulation.advance(gd, False, set())
check("double overflow fires", gd.over and gd.reason == "double overflow")

# network lock: 4 trains stuck 3 rounds
gl = engine.new_game(0)
for t in gl.trains[:4]:
    t.held = True; t.stuck_turns = 2
simulation.advance(gl, False, set())
check("network lock fires", gl.over and gl.reason == "network lock")

# safety collapse
gs = engine.new_game(0); gs.safety = 21
gs.incidents.append(Incident("i1", "signal_failure", "Dadar", "slow_up", severity=3, duration=4))
simulation.advance(gs, False, set())
check("safety collapse fires", gs.over and gs.reason == "safety collapse")

# anger collapse
ga = engine.new_game(0); ga.anger = 99
for s in ga.stations:
    s.crowd = s.cap * 0.9
for k in range(3):
    ga.incidents.append(Incident(f"f{k}", "festival_crowd", "Dadar", "-", severity=2, duration=3))
simulation.advance(ga, False, set())
check("anger collapse fires", ga.over and ga.reason == "city anger collapse")

print("CHAOS SCHEDULE / CAPS")
def avail(card, round_no, plays=0, consec=0):
    g = engine.new_game(0); g.turn = round_no - 1
    g.card_plays[card] = plays; g.consec[card] = consec
    return rules.card_available(g, card)
check("flood locked before R8", not avail("monsoon_flood", 7))
check("flood unlocks R8", avail("monsoon_flood", 8))
check("flood max 3", not avail("monsoon_flood", 10, plays=3))
check("festival locked before R13", not avail("festival_crowd", 12))
check("festival unlocks R13", avail("festival_crowd", 13))
check("festival once", not avail("festival_crowd", 16, plays=1))
check("vip once", not avail("vip_special", 5, plays=1))
check("vip available from R1", avail("vip_special", 1))
check("cow/signal cap at 2 consecutive", (not avail("track_cow", 5, consec=2)) and avail("track_cow", 5, consec=1))

print("NO STACKING / PEAK SIGNAL / CROSSOVER / HEADWAY")
gst = engine.new_game(0); rules.apply_chaos(gst, "track_cow", "Dadar")
check("station_free false after a card", (not rules.station_free(gst, "Dadar")) and rules.station_free(gst, "Andheri"))

gp = engine.new_game(0); gp.phase = "peak"; rules.apply_chaos(gp, "signal_failure", "Dadar")
inc = gp.incidents[-1]
t_slow = engine.Train("x", "x", "slow", +1, engine.IDX["Dadar"], "slow")
t_fast = engine.Train("y", "y", "fast", -1, engine.IDX["Dadar"], "fast")
check("peak signal blocks all lanes", inc.track == "all" and rules.incident_blocks(inc, t_slow) and rules.incident_blocks(inc, t_fast))

gx = engine.new_game(0); tr = gx.trains[0]; tr.pos = engine.IDX["Dadar"]; tr.just_switched = True
before = tr.pos; simulation.advance(gx, False, set())
check("crossover consumes the round (no advance)", tr.pos == before)

gh = engine.new_game(0)
# two trains on slow_down: a held leader at Dadar, a follower one station behind (Mumbai Central)
lead = gh.trains[0]; lead.track = "slow"; lead.direction = +1; lead.pos = engine.IDX["Dadar"]; lead.held = True
foll = gh.trains[2]; foll.track = "slow"; foll.direction = +1; foll.pos = engine.IDX["Mumbai Central"]; foll.held = False
# park everyone else off that lane so they don't interfere
for t in gh.trains:
    if t not in (lead, foll):
        t.track = "fast"
c0 = gh.collisions_avoided; simulation.advance(gh, False, set())
check("headway collision avoided (follower holds)", gh.collisions_avoided > c0 and foll.pos == engine.IDX["Mumbai Central"])

print("LEGAL ACTIONS")
gA = engine.new_game(0)
gA.trains[0].pos = engine.IDX["Bandra"]   # not an interchange
gA.trains[1].pos = engine.IDX["Dadar"]    # interchange
A = rules.legal_actions(gA)
sw_bandra = any(a["type"].startswith("switch_to") and a.get("train") == gA.trains[0].id for a in A)
sw_dadar = any(a["type"].startswith("switch_to") and a.get("train") == gA.trains[1].id for a in A)
check("no switch at non-interchange", not sw_bandra)
check("switch offered at interchange", sw_dadar)
budget_ok = enforce_budget([a["action_id"] for a in A], {a["action_id"]: a for a in A})[0]
check("budget enforcement caps at 5", len(budget_ok) <= engine.B["budget_total"])

print(f"\n{_passed} passed, {_failed} failed")
sys.exit(1 if _failed else 0)