mm.py

"""Hardcoded 6-step walkthrough on outbreak_easy (in-process env).

Run: ``uv run python mm.py``

Prints the action, env-accepted flag, reward, per-region telemetry, and
resource inventory after each step. See ``mm.md`` for the design intent
behind each step.
"""

from __future__ import annotations

from CrisisWorldCortex.models import (
    CrisisworldcortexAction,
    DeployResource,
    Escalate,
    NoOp,
    OuterActionPayload,
    RestrictMovement,
)
from CrisisWorldCortex.server.CrisisWorldCortex_environment import (
    CrisisworldcortexEnvironment,
)


def _summarize(obs) -> str:
    regions = " | ".join(
        f"{r.region}: cases={r.reported_cases_d_ago:>3} "
        f"hosp={r.hospital_load:.2f} comp={r.compliance_proxy:.2f}"
        for r in obs.regions
    )
    res = obs.resources
    restr = (
        ", ".join(f"{r.region}={r.severity}({r.ticks_remaining})" for r in obs.active_restrictions)
        or "(none)"
    )
    return (
        f"  regions:    {regions}\n"
        f"  resources:  kits={res.test_kits} beds={res.hospital_beds_free} "
        f"mu={res.mobile_units} vax={res.vaccine_doses}\n"
        f"  restricts:  {restr}\n"
        f"  tick={obs.tick}/{obs.tick + obs.ticks_remaining}  done={obs.done}"
    )


def main() -> None:
    env = CrisisworldcortexEnvironment()
    obs = env.reset(task_name="outbreak_easy", seed=0)
    print("=== Initial state (outbreak_easy, seed=0) ===")
    print(_summarize(obs))
    print()

    steps: list[tuple[str, OuterActionPayload]] = [
        ("NoOp baseline", NoOp()),
        (
            "Deploy 200 test_kits to R1",
            DeployResource(region="R1", resource_type="test_kits", quantity=200),
        ),
        ("Moderate restriction on R1", RestrictMovement(region="R1", severity="moderate")),
        (
            "Deploy 500 vaccine_doses to R1",
            DeployResource(region="R1", resource_type="vaccine_doses", quantity=500),
        ),
        ("Escalate to national", Escalate(to_authority="national")),
        (
            "Strict restriction on R1 (now legal)",
            RestrictMovement(region="R1", severity="strict"),
        ),
    ]

    for i, (label, payload) in enumerate(steps, start=1):
        obs = env.step(CrisisworldcortexAction(action=payload))
        last_log = obs.recent_action_log[-1]
        reward = obs.reward if obs.reward is not None else 0.0
        print(f"=== Step {i}: {label} ===")
        print(f"  action:     {payload.kind}  accepted={last_log.accepted}")
        print(f"  reward:     {reward:.3f}")
        print(_summarize(obs))
        print()


if __name__ == "__main__":
    main()