SystemSimulationApp/PythonModels/reporting/test_mql_chamber_observations.py

from __future__ import annotations

from dataclasses import dataclass

from PythonModels.reporting.amesim_results import AmesimResults
from PythonModels.reporting.test_mql_variables import (
    TestMqlVariableBinding,
    TestMqlVariableCatalog,
    build_test_mql_variable_catalog,
)
from PythonModels.systems.test_mql_pneumatic import (
    TestMqlPneumaticAssembly,
    build_test_mql_pneumatic_assembly,
)


@dataclass(frozen=True)
class TestMqlChamberObservation:
    time: float
    pressure_pa: float
    temperature_k: float
    gas_mass_g: float
    volume_cm3: float | None


@dataclass(frozen=True)
class TestMqlChamberBinding:
    alias: str
    submodel: str
    pressure_path: str
    temperature_path: str
    gas_mass_path: str
    pressure_duplicate_paths: tuple[str, ...]
    temperature_duplicate_paths: tuple[str, ...]
    volume_path: str | None

    @property
    def is_variable(self) -> bool:
        return self.volume_path is not None

    def observation_at(self, results: AmesimResults, index: int) -> TestMqlChamberObservation:
        return TestMqlChamberObservation(
            time=results.times[index],
            pressure_pa=results.series(self.pressure_path)[index],
            temperature_k=results.series(self.temperature_path)[index],
            gas_mass_g=results.series(self.gas_mass_path)[index],
            volume_cm3=(
                results.series(self.volume_path)[index]
                if self.volume_path is not None
                else None
            ),
        )


@dataclass(frozen=True)
class TestMqlChamberObservationCatalog:
    bindings: tuple[TestMqlChamberBinding, ...]

    @property
    def fixed_count(self) -> int:
        return sum(1 for binding in self.bindings if binding.submodel == "PNCH023")

    @property
    def variable_count(self) -> int:
        return sum(1 for binding in self.bindings if binding.submodel == "PNCH012")

    def by_alias(self, alias: str) -> TestMqlChamberBinding:
        for binding in self.bindings:
            if binding.alias == alias:
                return binding
        raise KeyError(alias)


def build_test_mql_chamber_observation_catalog(
    results: AmesimResults,
    *,
    variable_catalog: TestMqlVariableCatalog | None = None,
    assembly: TestMqlPneumaticAssembly | None = None,
) -> TestMqlChamberObservationCatalog:
    variable_catalog = variable_catalog or build_test_mql_variable_catalog(results)
    assembly = assembly or build_test_mql_pneumatic_assembly()
    chamber_aliases = {
        **{alias: "PNCH023" for alias in assembly.fixed_chambers},
        **{alias: "PNCH012" for alias in assembly.variable_chambers},
    }
    bindings = []
    for alias, submodel in chamber_aliases.items():
        variables = tuple(
            variable
            for variable in variable_catalog.variables
            if variable.owner_alias == alias
        )
        pressure = _primary_observable(variables, "press", expected_units="Pa")
        temperature = _primary_observable(variables, "temp", expected_units="K")
        gas_mass = _required_path(
            variables,
            "mgas1" if submodel == "PNCH012" else "mgas",
            expected_units="g",
        )
        volume = _optional_path(variables, "vol", expected_units="cm**3")
        bindings.append(
            TestMqlChamberBinding(
                alias=alias,
                submodel=submodel,
                pressure_path=pressure.data_path,
                temperature_path=temperature.data_path,
                gas_mass_path=gas_mass,
                pressure_duplicate_paths=_duplicate_paths(variables, "press", expected_units="Pa"),
                temperature_duplicate_paths=_duplicate_paths(variables, "temp", expected_units="K"),
                volume_path=volume,
            )
        )
    return TestMqlChamberObservationCatalog(
        bindings=tuple(sorted(bindings, key=lambda binding: binding.alias))
    )


def _primary_observable(
    variables: tuple[TestMqlVariableBinding, ...],
    signal_prefix: str,
    *,
    expected_units: str,
) -> TestMqlVariableBinding:
    matches = tuple(
        variable
        for variable in variables
        if variable.signal_name == signal_prefix
        and "duplicate" not in variable.label
    )
    variable = _single(matches, f"primary {signal_prefix}")
    _assert_units(variable, expected_units)
    return variable


def _duplicate_paths(
    variables: tuple[TestMqlVariableBinding, ...],
    signal_prefix: str,
    *,
    expected_units: str,
) -> tuple[str, ...]:
    matches = tuple(
        variable
        for variable in variables
        if variable.signal_name.startswith(signal_prefix)
        and variable.signal_name != signal_prefix
        and "duplicate" in variable.label
    )
    for variable in matches:
        _assert_units(variable, expected_units)
    return tuple(variable.data_path for variable in matches)


def _required_path(
    variables: tuple[TestMqlVariableBinding, ...],
    signal_name: str,
    *,
    expected_units: str,
) -> str:
    variable = _single(
        tuple(variable for variable in variables if variable.signal_name == signal_name),
        signal_name,
    )
    _assert_units(variable, expected_units)
    return variable.data_path


def _optional_path(
    variables: tuple[TestMqlVariableBinding, ...],
    signal_name: str,
    *,
    expected_units: str,
) -> str | None:
    matches = tuple(variable for variable in variables if variable.signal_name == signal_name)
    if not matches:
        return None
    variable = _single(matches, signal_name)
    _assert_units(variable, expected_units)
    return variable.data_path


def _single(
    matches: tuple[TestMqlVariableBinding, ...],
    description: str,
) -> TestMqlVariableBinding:
    if len(matches) != 1:
        raise ValueError(f"Expected one {description} variable, found {len(matches)}.")
    return matches[0]


def _assert_units(variable: TestMqlVariableBinding, expected_units: str) -> None:
    if variable.units != expected_units:
        raise ValueError(
            f"Unexpected units for {variable.data_path}: "
            f"{variable.units!r}, expected {expected_units!r}."
        )