optimization/resource_optimizer.py

"""
Deterministic Resource Optimizer.

This is a real integer linear program, solved with PuLP (CBC backend) —
not a heuristic dressed up to look like one. No language model output
ever reaches this file except through the NeedProfile.priority_score
field, which was itself computed by pure Python in need_engine.py.

FORMULATION
-----------
Decision variables:
    x[loc][resource] = integer number of units of `resource` assigned to `loc`
    0 <= x[loc][resource] <= required[loc][resource]   (never over-allocate; no waste)

Constraints:
    for each resource type r:
        sum over all locations of x[loc][r]  <=  available[r]

Objective (maximize):
    sum over loc, r of  priority_score[loc] * x[loc][r]

This says: every unit of any resource sent to a higher-priority location is
worth more than the same unit sent to a lower-priority location, and the
solver is free to trade off between resource types and locations to
maximize total weighted need covered, subject to the hard resource caps.

Re-solving takes milliseconds for problems this size (5-8 locations x 3
resource types), which is what makes the live Human Override panel
(section 13 of the spec) genuinely interactive rather than a fake animation.
"""

from __future__ import annotations
import pulp
from core.schemas import NeedProfile, ResourcePool, AllocationPlan, LocationAllocation

RESOURCE_TYPES = ["medical_teams", "rescue_teams", "supply_trucks"]

_REQUIRED_FIELD = {
    "medical_teams": "required_medical_teams",
    "rescue_teams": "required_rescue_teams",
    "supply_trucks": "required_supply_trucks",
}


def solve_allocation(profiles: list[NeedProfile], pool: ResourcePool) -> AllocationPlan:
    if not profiles:
        return AllocationPlan(
            allocations=[], overall_coverage_pct=0.0,
            resources_used=ResourcePool(0, 0, 0), resources_available=pool,
            solver_status="NoLocations", objective_value=0.0,
        )

    prob = pulp.LpProblem("disaster_resource_allocation", pulp.LpMaximize)

    x = {}
    for p in profiles:
        for r in RESOURCE_TYPES:
            required = getattr(p, _REQUIRED_FIELD[r])
            x[(p.location_id, r)] = pulp.LpVariable(
                f"x_{p.location_id}_{r}", lowBound=0, upBound=max(required, 0), cat="Integer"
            )

    # Objective: maximize total priority-weighted units allocated
    prob += pulp.lpSum(
        p.priority_score * x[(p.location_id, r)]
        for p in profiles for r in RESOURCE_TYPES
    )

    # Constraints: cannot exceed available pool per resource type
    available = {"medical_teams": pool.medical_teams, "rescue_teams": pool.rescue_teams,
                 "supply_trucks": pool.supply_trucks}
    for r in RESOURCE_TYPES:
        prob += pulp.lpSum(x[(p.location_id, r)] for p in profiles) <= available[r], f"cap_{r}"

    solver = pulp.PULP_CBC_CMD(msg=False)
    prob.solve(solver)
    status = pulp.LpStatus[prob.status]

    allocations = []
    total_required_units = 0
    total_assigned_units = 0
    used = {"medical_teams": 0, "rescue_teams": 0, "supply_trucks": 0}

    for p in profiles:
        assigned = {r: int(round(x[(p.location_id, r)].value() or 0)) for r in RESOURCE_TYPES}
        required = {r: getattr(p, _REQUIRED_FIELD[r]) for r in RESOURCE_TYPES}
        for r in RESOURCE_TYPES:
            used[r] += assigned[r]

        req_total = sum(required.values())
        assigned_total = sum(assigned.values())
        coverage = round((assigned_total / req_total * 100), 1) if req_total > 0 else 100.0

        total_required_units += req_total
        total_assigned_units += assigned_total

        allocations.append(LocationAllocation(
            location_id=p.location_id,
            assigned_medical_teams=assigned["medical_teams"],
            assigned_rescue_teams=assigned["rescue_teams"],
            assigned_supply_trucks=assigned["supply_trucks"],
            required_medical_teams=required["medical_teams"],
            required_rescue_teams=required["rescue_teams"],
            required_supply_trucks=required["supply_trucks"],
            coverage_pct=coverage,
            unmet_medical=max(0, required["medical_teams"] - assigned["medical_teams"]),
            unmet_rescue=max(0, required["rescue_teams"] - assigned["rescue_teams"]),
            unmet_supply=max(0, required["supply_trucks"] - assigned["supply_trucks"]),
        ))

    overall_coverage = round((total_assigned_units / total_required_units * 100), 1) if total_required_units > 0 else 100.0

    return AllocationPlan(
        allocations=allocations,
        overall_coverage_pct=overall_coverage,
        resources_used=ResourcePool(**used),
        resources_available=pool,
        solver_status=status,
        objective_value=pulp.value(prob.objective) or 0.0,
    )