tests/test_optimizer.py

import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from core.schemas import NeedProfile, ResourcePool
from optimization.resource_optimizer import solve_allocation


def make_profile(loc, priority, req_med=0, req_rescue=0, req_supply=0):
    return NeedProfile(
        location_id=loc, display_name=loc, severity=5, people_affected=10,
        urgency=5, medical_need=5, rescue_need=5, supply_need=5,
        priority_score=priority,
        required_medical_teams=req_med, required_rescue_teams=req_rescue,
        required_supply_trucks=req_supply, damage_type="structural_collapse",
        evidence_caption="test", lat=0.0, lon=0.0,
    )


def test_high_priority_gets_served_first_under_scarcity():
    profiles = [
        make_profile("ZONE_A", priority=90, req_med=5),
        make_profile("ZONE_B", priority=30, req_med=5),
    ]
    pool = ResourcePool(medical_teams=5, rescue_teams=0, supply_trucks=0)
    plan = solve_allocation(profiles, pool)
    a = next(a for a in plan.allocations if a.location_id == "ZONE_A")
    b = next(a for a in plan.allocations if a.location_id == "ZONE_B")
    assert a.assigned_medical_teams == 5, "highest priority zone should be fully served first"
    assert b.assigned_medical_teams == 0, "scarce resource should not spill to lower priority zone"
    print("PASS: high priority zone served first under scarcity")


def test_allocation_changes_when_resources_decrease():
    profiles = [
        make_profile("ZONE_A", priority=90, req_med=5),
        make_profile("ZONE_B", priority=60, req_med=5),
    ]
    plan_before = solve_allocation(profiles, ResourcePool(medical_teams=10, rescue_teams=0, supply_trucks=0))
    plan_after = solve_allocation(profiles, ResourcePool(medical_teams=4, rescue_teams=0, supply_trucks=0))
    assert plan_before.overall_coverage_pct > plan_after.overall_coverage_pct
    b_before = next(a for a in plan_before.allocations if a.location_id == "ZONE_B").assigned_medical_teams
    b_after = next(a for a in plan_after.allocations if a.location_id == "ZONE_B").assigned_medical_teams
    assert b_after < b_before, "reducing resources must reduce allocation to the lower-priority zone"
    print(f"PASS: coverage dropped from {plan_before.overall_coverage_pct}% to {plan_after.overall_coverage_pct}% "
          f"when medical teams cut from 10 to 4")


def test_allocation_changes_when_priority_changes():
    """Simulates severity/urgency change flowing through to a different allocation."""
    low_severity = [make_profile("ZONE_A", priority=20, req_med=5), make_profile("ZONE_B", priority=25, req_med=5)]
    high_severity = [make_profile("ZONE_A", priority=95, req_med=5), make_profile("ZONE_B", priority=25, req_med=5)]
    pool = ResourcePool(medical_teams=5, rescue_teams=0, supply_trucks=0)

    plan_low = solve_allocation(low_severity, pool)
    plan_high = solve_allocation(high_severity, pool)

    a_low = next(a for a in plan_low.allocations if a.location_id == "ZONE_A").assigned_medical_teams
    a_high = next(a for a in plan_high.allocations if a.location_id == "ZONE_A").assigned_medical_teams
    assert a_high > a_low, "raising ZONE_A's priority score must increase its allocation"
    print(f"PASS: ZONE_A allocation rose from {a_low} to {a_high} units when its priority score increased")


def test_human_override_reruns_optimizer_with_new_constraints():
    """Simulates the Human Override panel: user manually shrinks the pool for one resource."""
    profiles = [
        make_profile("ZONE_A", priority=80, req_rescue=4),
        make_profile("ZONE_B", priority=70, req_rescue=4),
        make_profile("ZONE_C", priority=50, req_rescue=4),
    ]
    before = solve_allocation(profiles, ResourcePool(medical_teams=0, rescue_teams=8, supply_trucks=0))
    override = solve_allocation(profiles, ResourcePool(medical_teams=0, rescue_teams=3, supply_trucks=0))
    assert before.overall_coverage_pct != override.overall_coverage_pct
    assert override.resources_used.rescue_teams <= 3
    print(f"PASS: override from 8->3 rescue teams changed coverage {before.overall_coverage_pct}% -> {override.overall_coverage_pct}%")


def test_never_exceeds_available_pool():
    profiles = [make_profile(f"ZONE_{i}", priority=50 + i, req_med=10, req_rescue=10, req_supply=10) for i in range(6)]
    pool = ResourcePool(medical_teams=7, rescue_teams=3, supply_trucks=5)
    plan = solve_allocation(profiles, pool)
    assert plan.resources_used.medical_teams <= 7
    assert plan.resources_used.rescue_teams <= 3
    assert plan.resources_used.supply_trucks <= 5
    print("PASS: solver never exceeds available resource pool across 6 competing locations")


def test_solver_status_optimal():
    profiles = [make_profile("ZONE_A", priority=50, req_med=3)]
    plan = solve_allocation(profiles, ResourcePool(medical_teams=3, rescue_teams=0, supply_trucks=0))
    assert plan.solver_status == "Optimal"
    print("PASS: solver reports Optimal status on a feasible problem")


def test_empty_locations_does_not_crash():
    plan = solve_allocation([], ResourcePool(medical_teams=5, rescue_teams=5, supply_trucks=5))
    assert plan.overall_coverage_pct == 0.0
    print("PASS: empty location list handled without crashing")


if __name__ == "__main__":
    tests = [v for k, v in list(globals().items()) if k.startswith("test_")]
    for t in tests:
        t()
    print(f"\n{len(tests)}/{len(tests)} optimizer tests passed.")