"""Small certified temporal-budget runtime for a learned controller."""
from dataclasses import dataclass
from typing import Callable, Optional


@dataclass
class Contract:
    phi0: float
    gamma: float
    delta: float
    t0: float

    @property
    def budget(self) -> float:
        if self.gamma <= 0:
            return float("inf") if self.phi0 >= 0 else -float("inf")
        return self.phi0 / self.gamma - self.delta

    def remaining(self, t: float) -> float:
        """Conservative usable latency budget at time t."""
        return self.budget - (t - self.t0)

    def permits(self, t: float, latency: float) -> bool:
        return self.phi0 >= 0 and latency <= self.remaining(t)


class CertifiedScheduler:
    """Certificate gate between sensing and a learned policy."""
    def __init__(self, gamma: float, delta: float, fallback: Callable):
        if gamma < 0 or delta < 0:
            raise ValueError("gamma and delta must be nonnegative")
        self.gamma, self.delta, self.fallback = gamma, delta, fallback
        self.contract: Optional[Contract] = None
        self.evaluations = 0

    def refresh(self, phi: float, t: float) -> Contract:
        self.contract = Contract(phi, self.gamma, self.delta, t)
        return self.contract

    def choose(self, state, t: float, measured_latency: float,
               policy: Callable):
        if self.contract is None:
            return self.fallback(state), False
        if self.contract.permits(t, measured_latency):
            self.evaluations += 1
            return policy(state), True
        return self.fallback(state), False