"""PlantEnv: gymnasium-style plant environment (spec §8, FR-8.2.14-17). Standalone-first: requires only a virtual node + datastore — no scenario. The environment is the interface for continuous plant simulation; per deterministic tick (runner-owned clock, never wall-clock arithmetic): state = env.observe() # ALL modules: DI/DO/AI/AO channels, # engineering units, names, parameters values = env.model(state, dt) # plant physics (subclass hook) env.apply(values) # new FIELD state -> DI/AI inputs; # DO/AO via env.force() (master-side) `model()` is the override point for plant physics: read what the PLC commanded (DO/AO), compute what the field now shows (DI/AI), return it. `is_done(state)` gives an optional end criterion; the default is False — environments run until manually stopped (SIGINT/SIGTERM). Attachments: degradation (§8.3), invariants (§8.4) and aliases (scenario-continuous mode) hook onto the same tick when the env is constructed on an engine. """ from __future__ import annotations import os from collections.abc import Callable from dataclasses import dataclass, field from typing import TYPE_CHECKING from ..modules import VirtualNode from .engine import ScenarioError if TYPE_CHECKING: from .engine import ScenarioEngine @dataclass class ChannelState: """One channel of the observed plant state (engineering units).""" kind: str # di | do | ai | ao module: int # per-type module number (FR-4.3.11 v1.5) channel: int # 0-based channel on the module io_code: str # e.g. "DI0.0" name: str | None # configured application name, if any value: object # bool for di/do, int/float for ai/ao range: tuple[int, int] | None = None # (min, max) raw counts — analog only @dataclass class PlantState: """Full observation: every channel of every module + attachments.""" sim_time_ms: float channels: list[ChannelState] = field(default_factory=list) aliases: dict[str, object] = field(default_factory=dict) parameters: dict[str, float] = field(default_factory=dict) def by_code(self, io_code: str) -> object: """Value of one channel by I/O code, or KeyError.""" for ch in self.channels: if ch.io_code == io_code: return ch.value raise KeyError(io_code) def by_kind(self, kind: str) -> dict[str, object]: """{io_code: value} for all channels of one kind.""" return {c.io_code: c.value for c in self.channels if c.kind == kind} def channel(self, key: str) -> ChannelState: """Full ChannelState (name, range, ...) by I/O code or name. Matches the I/O code first ("AI0.0"), then the configured application name ("speed_setpoint", FR-10.2.11). KeyError when neither matches — including a `None` name (unconfigured channels match by code only). """ for ch in self.channels: if ch.io_code == key: return ch for ch in self.channels: if ch.name == key: return ch raise KeyError(key) def channel_range(self, key: str) -> tuple[int, int]: """Display range (min, max) of an analog channel, by I/O code or name (FR-10.2.14). Always defined for AI/AO — the range store is dense — so env models can unpack it without None-handling; digital channels carry no range and raise KeyError. """ rng = self.channel(key).range if rng is None: raise KeyError(key) return rng class PlantEnv: """gym-style plant environment (FR-8.2.14-17). Constructed on a VirtualNode (standalone) or on an engine (the scenario-continuous layer). The runner owns the tick clock; `tick()` advances sim-time deterministically by TICK_MS per call so both paced (serve) and unpaced (sim) runners behave identically. """ TICK_MS = 10.0 def __init__(self, node: VirtualNode | None = None, *, engine: ScenarioEngine | None = None) -> None: if node is None: if engine is None: raise ScenarioError("PlantEnv needs a node or an engine") self.node: VirtualNode = engine.node else: self.node = node self.engine = engine self.datastore = self.node.datastore self.sim_time_ms = 0.0 # attachments (optional; the scenario layer wires these) self.degradation = getattr(engine, "degradation", None) self.invariants = getattr(engine, "invariants", None) self._aliases = getattr(engine, "_aliases", None) if engine else None self._alias_value: Callable[[str], object] | None = ( engine.alias_value if engine is not None else None) # -- observation (FR-8.2.14) ----------------------------------------------- def observe(self) -> PlantState: """Snapshot ALL modules (DI/DO/AI/AO) in engineering units.""" node = self.node ds = self.datastore channels: list[ChannelState] = [] for kind, pool in (("di", node.di), ("do", node.do), ("ai", node.ai), ("ao", node.ao)): for i, loc in enumerate(pool, start=1): if kind == "di": value = ds.digital_input_channel(i) elif kind == "do": value = ds.digital_output_channel(i) elif kind == "ai": value = ds.get_input_word(loc.word) else: value = ds.get_output_word(loc.word) code = node.io_code_at_slot( kind, loc.module_index, loc.channel - 1) channels.append(ChannelState( kind=kind, module=node.kind_index(loc.module_index)[kind], channel=loc.channel - 1, io_code=code, name=node.channel_names.get(code), value=value, range=node.channel_ranges[code] if kind in ("ai", "ao") else None, )) aliases = ({name: self._alias_value(name) for name in (self._aliases or {})} if self._alias_value else {}) parameters: dict[str, float] = {} if self.degradation is not None: parameters = {n: m.value for n, m in self.degradation.models.items()} return PlantState( sim_time_ms=self.sim_time_ms, channels=channels, aliases=aliases, parameters=parameters, ) # -- plant model hook (FR-8.2.16) ------------------------------------------- def model(self, state: PlantState, dt_s: float) -> dict[str, object]: """Plant physics: react to PLC commands, compute the new field. Override in a subclass (FR-8.2.17). Receives the full observed state and the tick delta in seconds; returns {I/O code or alias: new value} for the channels it drives — DI/AI normally (the simulated sensors); DO/AO entries are treated as forcing and go through the master-side write paths. Returning {} (the default) makes this a pure observation environment. """ return {} def is_done(self, state: PlantState) -> bool: """Optional end criterion; default runs until manually stopped.""" return False # -- actions (FR-8.2.15) ------------------------------------------------------ def apply(self, values: dict[str, object], *, force: bool = False) -> None: """Write new channel values into the process images. DI/AI keys (the field) write directly — this is the primary direction. DO/AO keys (PLC outputs) require force=True and go through the datastore's internal Modbus dispatch so mirrors, statistics, and watchdog stay consistent (FR-9.2.6) — for headless runs and fault setup only. Keys: I/O code ("AI0.0") or alias name. Unknown keys raise immediately (FR-8.2.16). """ if not isinstance(values, dict): raise ScenarioError("apply() values must be a mapping") for key, value in values.items(): self._write_one(key, value, force=force) def force(self, values: dict[str, object]) -> None: """Master-side forcing of DO/AO (FR-8.2.15 explicit form).""" self.apply(values, force=True) def _resolve(self, key: str): """-> (kind, ChannelLocation, alias_or_None) for a code or alias.""" # alias name first (scenario layer) if self._aliases and key in self._aliases and self.engine is not None: alias = self._aliases[key] return alias.kind, self.engine._alias_location(alias), alias # any channel by I/O code (per-type codes, FR-4.3.11 v1.5) code_up = str(key).strip().upper() node = self.node for kind, pool in (("di", node.di), ("do", node.do), ("ai", node.ai), ("ao", node.ao)): for loc in pool: if node.io_code_at_slot( kind, loc.module_index, loc.channel - 1) == code_up: return kind, loc, None raise ScenarioError(f"apply(): unknown channel {key!r}") def _write_one(self, key: str, value: object, *, force: bool) -> None: kind, loc, alias = self._resolve(key) ds = self.datastore if kind in ("do", "ao") and not force: raise ScenarioError( f"apply(): {key!r} is an output ({kind.upper()}); use " "force=True / env.force() for master-side forcing") scale = getattr(alias, "scale", 1.0) or 1.0 offset = getattr(alias, "offset", 0.0) or 0.0 if kind == "di": ds.set_digital_input_channel(loc.bit + 1, bool(value)) elif kind == "ai": raw = int(float(value) * scale + offset) & 0xFFFF # type: ignore[arg-type] ds.set_input_word(loc.word, raw) elif kind == "do": result = ds.write_bits(loc.bit, [bool(value)]) if isinstance(result, Exception): raise ScenarioError(f"force() write to {key!r}: {result}") else: # ao raw = int(float(value) * scale + offset) & 0xFFFF # type: ignore[arg-type] result = ds.write_words(loc.word, [raw]) if isinstance(result, Exception): raise ScenarioError(f"force() write to {key!r}: {result}") # -- one full tick (FR-8.2.14 cycle) --------------------------------------------- def tick(self, dt_ms: float | None = None) -> PlantState: """One environment tick: model() -> apply(field) -> advance. Deterministic: sim-time advances by TICK_MS (or dt_ms) per call — never from wall-clock. Attachments advance with the tick: degradation (§8.3) then invariants (§8.4), FR-8.2.12 order. """ dt = dt_ms if dt_ms is not None else self.TICK_MS state = self.observe() if self.is_done(state): return state values = self.model(state, dt / 1000.0) if values: self.apply(values) self.sim_time_ms += dt if self.degradation is not None: self.degradation.tick(dt) if self.invariants is not None: self.invariants.tick() return self.observe() # -- episode control (FR-8.2.16) ------------------------------------------------ def reset(self) -> PlantState: """Fresh episode: zero images, reset wear + parameters + clock (deterministic — identical runs replay identically, NFR-8.7.3).""" self.datastore.reset_images() self.sim_time_ms = 0.0 if self.degradation is not None: for model in self.degradation.models.values(): model.reset() model.diagnosed = False self.degradation._last_truth = {} if self.invariants is not None: self.invariants.episodes = [] for inv in self.invariants.invariants: inv.active = None for rl in self.invariants.rate_limits: rl.samples = [] rl.breached = False rl.peak_rate = 0.0 self.invariants.snapshots = [] self.invariants._last_snapshot_ms = 0.0 return self.observe() # -- clock (FR-8.2.16) --------------------------------------------------------- @property def tick_ms(self) -> float: return self.TICK_MS @property def parameters(self) -> dict[str, float]: if self.degradation is None: return {} return {n: m.value for n, m in self.degradation.models.items()} def load_env(ref: str, node: VirtualNode | None = None, *, engine=None, reload: bool = False) -> PlantEnv: """Load an environment from a module reference (FR-8.2.17). Accepts ``pkg.mod.ClassName``, ``pkg.mod.factory_fn`` (returns a PlantEnv), or a bare ``pkg.mod`` whose module defines either a class named ``Env`` or a function ``make_env(...)``. ``reload=True`` (FR-10.2.17) invalidates the cached module first: Python keeps modules in ``sys.modules``, so re-importing an edited file would silently reuse the old class. Used by the web UI's start/restart; the CLI runners start a fresh process anyway and pass the default. """ import importlib import sys # environments resolve from the WORKING DIRECTORY by contract # (``envs/`` is not part of the installed package). A console # script puts its Scripts dir — not the cwd — on sys.path, so # make the cwd importable before trying (idempotent). cwd = os.getcwd() if cwd not in sys.path: sys.path.insert(0, cwd) parts = ref.split(":") attr_path = None mod = None leaf = None if len(parts) == 2: # explicit module:attr form — ref is the WHOLE module path attr_path = parts[1] or None mod_name = parts[0] else: # no colon: try the whole ref as a module first, then as # pkg.mod containing an attr (e.g. ...:Class omitted) try: mod = importlib.import_module(ref) mod_name = ref except ImportError: mod_name, _, leaf = ref.rpartition(".") if not mod_name or ("." not in mod_name and attr_path is None and leaf is None): raise ScenarioError(f"--env {ref!r}: reference must include a module path") if reload and mod_name in sys.modules: # hot reload (FR-10.2.17): drop the cached module so the # import below reads the file from disk again sys.modules.pop(mod_name, None) mod = None if mod is None: try: mod = importlib.import_module(mod_name) except ImportError as exc: raise ScenarioError( f"--env {ref!r}: cannot import {mod_name!r}: {exc}") obj = None if attr_path is not None: # explicit module:attr — must exist; never fall back silently if not hasattr(mod, attr_path): raise ScenarioError( f"--env {ref!r}: no {attr_path!r} in {mod_name!r}") obj = getattr(mod, attr_path) elif leaf and hasattr(mod, leaf): obj = getattr(mod, leaf) elif hasattr(mod, "Env"): obj = getattr(mod, "Env") elif hasattr(mod, "make_env"): obj = getattr(mod, "make_env") else: raise ScenarioError( f"--env {ref!r}: no environment found in {mod_name!r} " "(expected a class, Env, or make_env())") env: PlantEnv if isinstance(obj, type) and issubclass(obj, PlantEnv): env = obj(node=node, engine=engine) elif callable(obj): made = obj(node=node, engine=engine) if node is not None \ else obj(engine=engine) if not isinstance(made, PlantEnv): raise ScenarioError(f"--env {ref!r}: did not return a PlantEnv") env = made else: raise ScenarioError( f"--env {ref!r}: not a PlantEnv subclass or factory") return env