backend/tests/test_simulation.py

"""Simulation engine: movement, capacity, queueing, rerouting, reproducibility."""

from __future__ import annotations

import numpy as np
import pytest

from flowtwin.config import SETTINGS
from flowtwin.simulation.agents import (
    POLICY_ADAPTIVE,
    POLICY_SHORTEST,
    STATUS_ARRIVED,
    STATUS_ON_EDGE,
    STATUS_WAITING,
)
from flowtwin.simulation.engine import RunOverrides, Simulator
from flowtwin.simulation.movement import CapacityBudget, admit, weidmann_speed
from flowtwin.venue import compile_venue, load_scenario


@pytest.fixture(scope="module")
def venue():
    return compile_venue("circuit_alpha")


@pytest.fixture(scope="module")
def scenario():
    return load_scenario("circuit_alpha_post_race")


def make_sim(venue, scenario, seed=42193, crowd=6000, **kw):
    return Simulator(venue, scenario, SETTINGS, seed=seed,
                     overrides=RunOverrides(crowd_size=crowd, **kw))


# ── movement model ───────────────────────────────────────────────────

def test_speed_falls_monotonically_with_density():
    cfg = SETTINGS.movement
    densities = np.array([0.1, 0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 5.4])
    speeds = weidmann_speed(densities, cfg)
    assert np.all(np.diff(speeds) <= 1e-9), "walking speed must not rise with density"
    assert speeds[0] == pytest.approx(cfg.free_speed_mps)
    assert speeds[-1] < 0.3


def test_capacity_budget_delivers_the_nominal_rate():
    budget = CapacityBudget(np.array([90.0]))       # 90 people/minute
    admitted = 0
    for _ in range(60):                              # one minute at dt=1s
        allow = budget.accrue(1.0)
        used = min(int(allow[0]), 5)
        budget.consume(np.array([float(used)]))
        admitted += used
    assert 88 <= admitted <= 92, f"expected ~90 admissions per minute, got {admitted}"


def test_admission_is_first_come_first_served():
    group = np.array([0, 0, 0, 1])
    priority = np.array([30.0, 10.0, 20.0, 5.0])     # join times
    allowance = np.array([2, 1])
    ok = admit(group, priority, allowance)
    assert ok.tolist() == [False, True, True, True]


# ── population ───────────────────────────────────────────────────────

def test_population_matches_requested_size(venue, scenario):
    sim = make_sim(venue, scenario, crowd=5000)
    assert sim.n_agents == 5000
    assert np.all(sim.pop.status == STATUS_WAITING)
    assert np.all(sim.pop.compliance >= 0) and np.all(sim.pop.compliance <= 1)


def test_agents_spawn_move_and_arrive(venue, scenario):
    sim = make_sim(venue, scenario, crowd=4000)
    sim.run_for(120)
    assert np.any(sim.pop.status == STATUS_ON_EDGE), "no agent entered the network"
    sim.run_for(900)
    assert np.any(sim.pop.status == STATUS_ARRIVED), "no agent reached a destination"
    arrived = sim.pop.status == STATUS_ARRIVED
    travel = sim.pop.arrive_t[arrived] - sim.pop.enter_t[arrived]
    assert np.all(travel[~np.isnan(travel)] > 0)


def test_every_agent_eventually_reaches_a_destination(venue, scenario):
    sim = make_sim(venue, scenario, crowd=3000)
    sim.run_until_complete(3600)
    assert sim.remaining == 0, f"{sim.remaining} agents never arrived"


def test_congestion_forms_and_capacity_binds(venue, scenario):
    """The scripted Exit B failure must produce a measurable queue there."""
    sim = make_sim(venue, scenario, crowd=40000)
    sim.run_for(1100)
    exit_b = venue.node_index["EXIT_B"]
    assert sim.node_budget.multiplier[exit_b] == pytest.approx(0.5), \
        "the scripted capacity reduction did not fire"
    assert sim.state.node_queue[exit_b] > 500, "no queue formed at the degraded exit"
    approach = venue.edge_index["X_E_EXITB"]
    assert sim.state.edge_density[approach] > venue.venue.warning_density
    assert sim.state.edge_velocity[approach] < SETTINGS.movement.free_speed_mps


def test_density_never_exceeds_the_jam_limit(venue, scenario):
    sim = make_sim(venue, scenario, crowd=40000)
    sim.run_for(1400)
    jam = SETTINGS.movement.jam_density
    assert sim.state.edge_density.max() <= jam * 1.02
    assert sim.state.edge_peak_local_density.max() <= jam * 1.25


def test_agent_sample_stays_within_the_render_budget(venue, scenario):
    sim = make_sim(venue, scenario, crowd=40000)
    sim.run_for(400)
    sample = sim.agent_sample(1000)
    assert sample["sampled"] <= 1000
    assert len(sample["x"]) == len(sample["y"]) == sample["sampled"]
    assert sample["total"] >= sample["sampled"]


# ── reproducibility ──────────────────────────────────────────────────

def test_same_seed_reproduces_identical_output(venue, scenario):
    a = make_sim(venue, scenario, seed=777)
    b = make_sim(venue, scenario, seed=777)
    a.run_for(600)
    b.run_for(600)
    assert a.metrics() == b.metrics()
    assert np.array_equal(a.pop.pos_m, b.pop.pos_m)
    assert np.allclose(a.state.edge_density, b.state.edge_density)


def test_different_seeds_diverge(venue, scenario):
    a = make_sim(venue, scenario, seed=1)
    b = make_sim(venue, scenario, seed=2)
    a.run_for(600)
    b.run_for(600)
    assert not np.array_equal(a.pop.pos_m, b.pop.pos_m)


def test_snapshot_restore_is_exact(venue, scenario):
    sim = make_sim(venue, scenario)
    sim.run_for(400)
    snap = sim.snapshot()
    sim.run_for(200)
    first = sim.metrics()
    sim.restore(snap)
    sim.run_for(200)
    assert sim.metrics() == first


def test_branch_does_not_disturb_the_parent(venue, scenario):
    sim = make_sim(venue, scenario)
    sim.run_for(400)
    before = sim.pop.pos_m.copy()
    branch = sim.branch()
    branch.run_for(200)
    assert np.array_equal(sim.pop.pos_m, before)


def test_two_branches_of_the_same_state_are_identical(venue, scenario):
    sim = make_sim(venue, scenario)
    sim.run_for(400)
    a, b = sim.branch(), sim.branch()
    a.run_for(180)
    b.run_for(180)
    assert a.metrics() == b.metrics()


# ── interventions ────────────────────────────────────────────────────

def test_diverting_flow_moves_people_onto_another_route(venue, scenario):
    sim = make_sim(venue, scenario, crowd=40000)
    sim.run_for(700)

    edge = venue.edge_index["X_E_EXITB"]
    pair = int(venue.pair_of[edge])
    node = venue.node_index["EXIT_B"]
    slot = sim.dest_indices.index(venue.node_index["TRANSPORT_BUS"])
    before, _ = sim.tables.path_nodes(POLICY_ADAPTIVE, slot, venue.node_index["CON_EAST"])

    baseline = sim.branch()
    diverted = sim.branch()
    accepted = diverted.divert_flow(0.4, {edge, pair}, {node}, penalty=8.0)

    assert accepted > 0, "nobody was diverted"
    assert np.sum(diverted.pop.policy == POLICY_ADAPTIVE) == accepted
    after, _ = diverted.tables.path_nodes(POLICY_ADAPTIVE, slot, venue.node_index["CON_EAST"])
    assert after != before, "the adaptive plan did not change after the penalty"

    baseline.run_for(300)
    diverted.run_for(300)
    assert diverted.state.node_queue[node] < baseline.state.node_queue[node], \
        "diverting flow did not reduce the queue at the degraded exit"
    assert diverted.total_rerouted > 0


def test_diversion_respects_compliance(venue, scenario):
    """Not everyone obeys: accepted must be below the number instructed."""
    sim = make_sim(venue, scenario, crowd=40000)
    sim.run_for(700)
    edge = venue.edge_index["X_E_EXITB"]
    node = venue.node_index["EXIT_B"]
    accepted = sim.divert_flow(1.0, {edge, int(venue.pair_of[edge])}, {node})
    active = int(np.sum(sim.pop.status != STATUS_ARRIVED))
    assert 0 < accepted < active


def test_staggering_release_delays_departures(venue, scenario):
    sim = make_sim(venue, scenario, crowd=20000)
    sim.run_for(120)
    waiting = sim.pop.status == STATUS_WAITING
    before = sim.pop.release_t[waiting].copy()
    moved = sim.stagger_release(["GS_MAIN"], 0.5, 150.0)
    assert moved > 0
    after = sim.pop.release_t[waiting]
    assert after.sum() > before.sum(), "release times did not move later"


def test_opening_an_alternate_exit_raises_its_throughput(venue, scenario):
    sim = make_sim(venue, scenario, crowd=10000)
    idx = venue.node_index["EXIT_C"]
    before = float(sim.node_budget.multiplier[idx])
    assert sim.open_alternate("EXIT_C", 1.35)
    assert sim.node_budget.multiplier[idx] == pytest.approx(before * 1.35)


def test_capacity_override_is_applied_at_construction(venue, scenario):
    sim = make_sim(venue, scenario, crowd=1000,
                   capacity_overrides={"EXIT_A": 0.25})
    idx = venue.node_index["EXIT_A"]
    assert sim.node_budget.multiplier[idx] == pytest.approx(0.25)


def test_whatif_capacity_slider_retunes_the_scripted_failure(venue, scenario):
    """The What-If capacity control must change what actually happens."""
    idx = venue.node_index["EXIT_B"]

    authored = make_sim(venue, scenario, crowd=3000)
    authored.run_for(300)
    assert authored.node_budget.multiplier[idx] == pytest.approx(0.5)

    harsher = make_sim(venue, scenario, crowd=3000,
                       event_factor_overrides={"EXIT_B": 0.25})
    harsher.run_for(300)
    assert harsher.node_budget.multiplier[idx] == pytest.approx(0.25)
    assert "25%" in harsher.event_log[-1]["label"]

    # Dialling it back to 100% must remove the failure entirely.
    healthy = make_sim(venue, scenario, crowd=3000,
                       event_factor_overrides={"EXIT_B": 1.0})
    healthy.run_for(300)
    assert healthy.node_budget.multiplier[idx] == pytest.approx(1.0)


def test_rejects_an_impossible_crowd_size(venue, scenario):
    with pytest.raises(ValueError):
        Simulator(venue, scenario, SETTINGS,
                  overrides=RunOverrides(crowd_size=SETTINGS.simulation.max_agents + 1))