core/sector.py

"""
sector.py — Individual sector / department state.

Each Sector holds its baseline plus transient per-round state
(demand, thresholds, allocation, revenue factor, revenue, consumption).
The game orchestrator resets per-round fields each step.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from .config import (
    CRITICAL_RATIO,
    INITIAL_POPULATION,
    INITIAL_PRODUCTIVITY,
    RF_MAX,
    SURPLUS_RATIO,
    WASTAGE_RATIO,
)
from .revenue import compute_thresholds, revenue_factor


class _ThresholdAccessor:
    """Callable numeric threshold for compatibility with stored attributes."""

    def __init__(self, sector: "Sector", threshold_name: str) -> None:
        self._sector = sector
        self._threshold_name = threshold_name
        self._value = 0.0

    def set(self, value: float) -> None:
        self._value = value

    def __call__(self, population: int) -> float:
        critical, demand, surplus, wastage = self._sector.thresholds(population)
        return {
            "critical": critical,
            "demand": demand,
            "surplus": surplus,
            "wastage": wastage,
        }[self._threshold_name]

    def __float__(self) -> float:
        return self._value

    def __round__(self, ndigits: int | None = None) -> float:
        return round(self._value, ndigits) if ndigits is not None else round(self._value)

    def __repr__(self) -> str:
        return repr(self._value)

    def __eq__(self, other: object) -> bool:
        try:
            return self._value == float(other)  # type: ignore[arg-type]
        except (TypeError, ValueError):
            return False

    def __lt__(self, other: Any) -> bool:
        return self._value < float(other)

    def __le__(self, other: Any) -> bool:
        return self._value <= float(other)

    def __gt__(self, other: Any) -> bool:
        return self._value > float(other)

    def __ge__(self, other: Any) -> bool:
        return self._value >= float(other)

    def __add__(self, other: Any) -> float:
        return self._value + float(other)

    def __radd__(self, other: Any) -> float:
        return float(other) + self._value

    def __sub__(self, other: Any) -> float:
        return self._value - float(other)

    def __rsub__(self, other: Any) -> float:
        return float(other) - self._value

    def __mul__(self, other: Any) -> float:
        return self._value * float(other)

    def __rmul__(self, other: Any) -> float:
        return float(other) * self._value

    def __truediv__(self, other: Any) -> float:
        return self._value / float(other)

    def __rtruediv__(self, other: Any) -> float:
        return float(other) / self._value


@dataclass
class Sector:
    """Mutable state for one government department."""

    # ── Identity (set once at episode start) ────────────────────
    name: str
    baseline: float
    full_name: str = ""
    description: str = ""

    # ── Per-round computed state ────────────────────────────────
    event_multiplier: float = 1.0

    demand: _ThresholdAccessor = field(init=False)
    critical: _ThresholdAccessor = field(init=False)
    surplus: _ThresholdAccessor = field(init=False)
    wastage: _ThresholdAccessor = field(init=False)

    allocation: float = 0.0
    revenue_factor_value: float = 0.0
    revenue: float = 0.0
    consumption: float = 0.0

    # ── Config ratios (injected from GameConfig) ────────────────
    _critical_ratio: float = field(default=CRITICAL_RATIO, repr=False)
    _surplus_ratio: float = field(default=SURPLUS_RATIO, repr=False)
    _wastage_ratio: float = field(default=WASTAGE_RATIO, repr=False)
    _rf_max: float = field(default=RF_MAX, repr=False)

    def __post_init__(self) -> None:
        self.demand = _ThresholdAccessor(self, "demand")
        self.critical = _ThresholdAccessor(self, "critical")
        self.surplus = _ThresholdAccessor(self, "surplus")
        self.wastage = _ThresholdAccessor(self, "wastage")
        self.update_thresholds(INITIAL_POPULATION, INITIAL_POPULATION)

    # ── Lifecycle ───────────────────────────────────────────────

    def reset_round(self) -> None:
        """Clear transient state before a new round."""
        self.event_multiplier = 1.0
        self.allocation = 0.0
        self.revenue_factor_value = 0.0
        self.revenue = 0.0
        self.consumption = 0.0

    def update_thresholds(self, population: int, pop_0: int) -> None:
        """Recalculate demand/critical/surplus/wastage for this round."""
        critical, demand, surplus, wastage = compute_thresholds(
            baseline=self.baseline,
            population=population,
            pop_0=pop_0,
            event_multiplier=self.event_multiplier,
            critical_ratio=self._critical_ratio,
            surplus_ratio=self._surplus_ratio,
            wastage_ratio=self._wastage_ratio,
        )
        self.critical.set(critical)
        self.demand.set(demand)
        self.surplus.set(surplus)
        self.wastage.set(wastage)

    def thresholds(self, population: int) -> tuple[float, float, float, float]:
        """Return (critical, demand, surplus, wastage) for a population."""
        return compute_thresholds(
            baseline=self.baseline,
            population=population,
            pop_0=INITIAL_POPULATION,
            event_multiplier=self.event_multiplier,
            critical_ratio=self._critical_ratio,
            surplus_ratio=self._surplus_ratio,
            wastage_ratio=self._wastage_ratio,
        )

    def allocate(self, amount: float) -> None:
        """Store the current round allocation."""
        self.allocation = amount

    def compute_revenue(
        self,
        population: int | None = None,
        productivity: float = INITIAL_PRODUCTIVITY,
        *,
        allocation: float | None = None,
    ) -> float | None:
        """
        Compute revenue using the current or provided allocation.

        Returns
        -------
        float
            Revenue generated by this sector.
        None
            If allocation is below the critical threshold (episode must end).
        """
        if allocation is not None:
            self.allocate(allocation)
        if population is not None:
            self.update_thresholds(population=population, pop_0=INITIAL_POPULATION)

        rf = revenue_factor(
            allocation=self.allocation,
            critical=float(self.critical),
            demand=float(self.demand),
            surplus=float(self.surplus),
            wastage=float(self.wastage),
            rf_max=self._rf_max,
        )

        if rf is None:
            self.revenue_factor_value = 0.0
            self.revenue = 0.0
            return None  # CRITICAL FAILURE

        self.revenue_factor_value = rf
        self.revenue = self.allocation * rf * productivity
        return self.revenue

    def is_critical_failure(self, population: int) -> bool:
        """Return true when allocation is below the critical threshold."""
        return self.allocation < self.critical(population)

    def compute_consumption(self) -> float:
        """
        Consumption = min(allocation, demand).
        Surplus (allocation - consumption) will be returned to treasury.

        Returns the surplus amount.
        """
        self.consumption = min(self.allocation, float(self.demand))
        return self.allocation - self.consumption  # surplus

    # ── Serialisation helpers ───────────────────────────────────

    def to_dict(self) -> dict:
        """Snapshot of sector state for observations / logging."""
        return {
            "name": self.name,
            "full_name": self.full_name,
            "baseline": self.baseline,
            "demand": round(self.demand, 4),
            "critical": round(self.critical, 4),
            "surplus": round(self.surplus, 4),
            "wastage": round(self.wastage, 4),
            "allocation": round(self.allocation, 4),
            "revenue_factor": round(self.revenue_factor_value, 6),
            "revenue": round(self.revenue, 4),
            "consumption": round(self.consumption, 4),
            "event_multiplier": round(self.event_multiplier, 4),
        }

    @classmethod
    def from_config(cls, name: str, info: dict, config) -> "Sector":
        """Build a Sector from sectors.json data + a GameConfig."""
        return cls(
            name=name,
            baseline=info["baseline"] if isinstance(info, dict) else info,
            full_name=(info.get("full_name", name) if isinstance(info, dict) else name),
            description=(info.get("description", "") if isinstance(info, dict) else ""),
            _critical_ratio=config.CRITICAL_RATIO,
            _surplus_ratio=config.SURPLUS_RATIO,
            _wastage_ratio=config.WASTAGE_RATIO,
            _rf_max=config.RF_MAX,
        )