allocation.py

"""Sparse allocation representation for the hybrid optimizer.

An ``Allocation`` is a mapping ``{(source_id, target_id): ships}``. It is the
single object passed between Stages 1/2/3 of the optimizer
(``docs/hybrid_optimization_plan.md``). Allocations are treated as immutable —
every mutation method returns a new instance.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Iterable, Iterator, Mapping

from orbit_model import GameState

ArcKey = tuple[int, int]


@dataclass
class Allocation:
    shipments: dict[ArcKey, int] = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.shipments = {
            (int(src), int(tgt)): int(s)
            for (src, tgt), s in self.shipments.items()
            if int(s) > 0
        }

    def __hash__(self) -> int:  # type: ignore[override]
        return hash(tuple(sorted(self.shipments.items())))

    def __eq__(self, other: object) -> bool:  # type: ignore[override]
        return isinstance(other, Allocation) and self.shipments == other.shipments

    def __len__(self) -> int:
        return len(self.shipments)

    def __iter__(self) -> Iterator[ArcKey]:
        return iter(self.shipments)

    def __contains__(self, key: ArcKey) -> bool:
        return key in self.shipments

    def items(self):
        return self.shipments.items()

    def get(self, key: ArcKey, default: int = 0) -> int:
        return self.shipments.get(key, default)

    def copy(self) -> "Allocation":
        return Allocation(dict(self.shipments))

    def with_shipment(self, source: int, target: int, ships: int) -> "Allocation":
        new = dict(self.shipments)
        if ships <= 0:
            new.pop((source, target), None)
        else:
            new[(source, target)] = int(ships)
        return Allocation(new)

    def with_delta(self, source: int, target: int, delta: int) -> "Allocation":
        cur = self.shipments.get((source, target), 0)
        return self.with_shipment(source, target, cur + delta)

    def without(self, keys: Iterable[ArcKey]) -> "Allocation":
        drop = set(keys)
        return Allocation({k: v for k, v in self.shipments.items() if k not in drop})

    def merge(self, other: "Allocation") -> "Allocation":
        new = dict(self.shipments)
        for k, v in other.shipments.items():
            new[k] = new.get(k, 0) + v
        return Allocation(new)

    def total_from(self, source: int) -> int:
        return sum(s for (src, _), s in self.shipments.items() if src == source)

    def to_target(self, target: int) -> dict[int, int]:
        return {src: s for (src, t), s in self.shipments.items() if t == target}

    def from_source(self, source: int) -> dict[int, int]:
        return {t: s for (src, t), s in self.shipments.items() if src == source}

    def sources(self) -> set[int]:
        return {src for (src, _) in self.shipments}

    def targets(self) -> set[int]:
        return {tgt for (_, tgt) in self.shipments}


def is_feasible(allocation: Allocation, state: GameState) -> bool:
    """Return True iff the allocation respects per-source ship budgets and basic structure.

    Does not check sun-route validity — that's handled when emitting moves.
    """

    totals: dict[int, int] = {}
    for (src, tgt), ships in allocation.shipments.items():
        if ships < 1:
            return False
        if src == tgt:
            return False
        if src not in state.planets or tgt not in state.planets:
            return False
        if state.planets[src].owner != state.player:
            return False
        totals[src] = totals.get(src, 0) + ships

    for src, total in totals.items():
        if total > state.planets[src].ships:
            return False
    return True


def available_ships(state: GameState, source: int, allocation: Allocation | None = None) -> int:
    """Ships at ``source`` minus what allocation already commits."""

    planet = state.planets.get(source)
    if planet is None or planet.owner != state.player:
        return 0
    used = allocation.total_from(source) if allocation is not None else 0
    return max(0, planet.ships - used)


__all__ = ["Allocation", "ArcKey", "is_feasible", "available_ships"]