pp-simulator/pp_simulator/duration.py

from __future__ import annotations

import math
import random
from dataclasses import dataclass
from typing import Any, Mapping


@dataclass(frozen=True)
class OpDurationSpec:
    """Normal-distribution duration config for one operation type."""

    mean: float
    variance: float = 0.0
    min_value: float = 0.0

    def sample(self, rng: random.Random) -> float:
        if self.variance < 0:
            raise ValueError(f"variance must be non-negative, got {self.variance}")
        if self.mean < 0:
            raise ValueError(f"mean must be non-negative, got {self.mean}")

        if self.variance == 0:
            return max(self.min_value, float(self.mean))

        value = rng.gauss(float(self.mean), math.sqrt(float(self.variance)))
        return max(float(self.min_value), value)


def op_type_value(op_type: Any) -> str:
    """Return a stable string value for pp_scheduler.OpType or a plain string."""

    value = getattr(op_type, "value", op_type)
    return str(value)


class DurationSampler:
    def __init__(
        self,
        specs: Mapping[Any, OpDurationSpec | Mapping[str, float]],
        *,
        default_spec: OpDurationSpec | None = None,
        seed: int | None = None,
    ):
        self.rng = random.Random(seed)
        self.default_spec = default_spec
        self.specs = {
            op_type_value(op_type): self._coerce_spec(spec)
            for op_type, spec in specs.items()
        }

    def sample(self, op_type: Any, fallback_duration: float | None = None) -> float:
        op_value = op_type_value(op_type)
        spec = self.specs.get(op_value, self.default_spec)
        if spec is not None:
            return spec.sample(self.rng)
        if fallback_duration is not None:
            return max(0.0, float(fallback_duration))
        raise KeyError(f"No duration spec for op type {op_value!r}")

    @staticmethod
    def _coerce_spec(spec: OpDurationSpec | Mapping[str, float]) -> OpDurationSpec:
        if isinstance(spec, OpDurationSpec):
            return spec
        return OpDurationSpec(
            mean=float(spec["mean"]),
            variance=float(spec.get("variance", 0.0)),
            min_value=float(spec.get("min_value", 0.0)),
        )