from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, Mapping


@dataclass(frozen=True)
class SiteState:
    site_id: str
    region: str
    on_hand: float
    forecast_demand: float
    equity_weight: float = 1.0
    continuity_weight: float = 1.0
    triage_priority: float = 1.0
    minimum_service: float = 0.0


@dataclass(frozen=True)
class HumanitarianPolicy:
    equity_weight: float = 1.0
    continuity_weight: float = 1.0
    triage_weight: float = 1.0
    regional_floor: float = 0.0

    def site_prior(self, site: SiteState) -> float:
        return (
            site.equity_weight * self.equity_weight
            + site.continuity_weight * self.continuity_weight
            + site.triage_priority * self.triage_weight
        )


@dataclass(frozen=True)
class AllocationResult:
    allocations: dict[str, float]
    unmet_demand: dict[str, float]
    total_allocated: float
    iterations: int


class DistributedAllocator:
    def __init__(self, iterations: int = 8, step_size: float = 0.5) -> None:
        self.iterations = iterations
        self.step_size = step_size

    @staticmethod
    def _project_to_sum(values: list[float], total: float, lowers: list[float], uppers: list[float]) -> list[float]:
        if not values:
            return []
        low = min(v - u for v, u in zip(values, uppers)) - abs(total)
        high = max(v - l for v, l in zip(values, lowers)) + abs(total)
        for _ in range(80):
            shift = (low + high) / 2.0
            projected = [min(upper, max(lower, value - shift)) for value, lower, upper in zip(values, lowers, uppers)]
            current = sum(projected)
            if abs(current - total) < 1e-8:
                return projected
            if current > total:
                low = shift
            else:
                high = shift
        shift = (low + high) / 2.0
        return [min(upper, max(lower, value - shift)) for value, lower, upper in zip(values, lowers, uppers)]

    def allocate(self, sites: Iterable[SiteState], total_available: float, policy: HumanitarianPolicy | None = None) -> AllocationResult:
        policy = policy or HumanitarianPolicy()
        sites = list(sites)
        if not sites:
            return AllocationResult({}, {}, 0.0, 0)

        lowers = [max(0.0, site.minimum_service) for site in sites]
        uppers = [max(lower, min(site.forecast_demand, site.on_hand + total_available)) for site, lower in zip(sites, lowers)]
        weights = [max(0.1, policy.site_prior(site)) for site in sites]
        total = min(total_available, sum(uppers))
        x = [min(upper, max(lower, total * weight / sum(weights))) for weight, lower, upper in zip(weights, lowers, uppers)]
        duals = [0.0 for _ in sites]

        for _ in range(self.iterations):
            proposals = []
            for idx, site in enumerate(sites):
                target = min(site.forecast_demand, site.on_hand + total_available)
                prior = policy.site_prior(site)
                local_update = x[idx] + self.step_size * ((target - x[idx]) / max(1.0, target) + prior - duals[idx])
                proposals.append(local_update)
            x = self._project_to_sum(proposals, total, lowers, uppers)
            duals = [dual + proposal - alloc for dual, proposal, alloc in zip(duals, proposals, x)]

        allocations = {site.site_id: round(value, 4) for site, value in zip(sites, x)}
        unmet = {site.site_id: round(max(0.0, site.forecast_demand - allocations[site.site_id]), 4) for site in sites}
        return AllocationResult(allocations=allocations, unmet_demand=unmet, total_allocated=round(sum(x), 4), iterations=self.iterations)