SystemSimulationApp/PythonModels/reporting/test_mql_orifice_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,
)


G_PER_S_TO_KG_PER_S = 1.0e-3


@dataclass(frozen=True)
class TestMqlOrificeObservation:
    time: float
    mass_flow_kg_s: float
    enthalpy_flow_w: float
    mass_flow_parameter: float
    gas_velocity_m_s: float
    opening: float
    effective_area_m2: float


@dataclass(frozen=True)
class TestMqlOrificeBinding:
    alias: str
    submodel: str
    nominal_area_m2: float
    flow_coefficient: float
    primary_mass_flow_path: str
    primary_enthalpy_flow_path: str
    reversed_mass_flow_path: str
    reversed_enthalpy_flow_path: str
    mass_flow_parameter_path: str
    gas_velocity_path: str
    opening_path: str | None

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

    def opening_series(self, results: AmesimResults) -> tuple[float, ...]:
        if self.opening_path is None:
            return tuple(1.0 for _ in results.times)
        return tuple(results.series(self.opening_path))

    def mass_flow_kg_s_series(self, results: AmesimResults) -> tuple[float, ...]:
        return tuple(value * G_PER_S_TO_KG_PER_S for value in results.series(self.primary_mass_flow_path))

    def reversed_mass_flow_kg_s_series(self, results: AmesimResults) -> tuple[float, ...]:
        return tuple(value * G_PER_S_TO_KG_PER_S for value in results.series(self.reversed_mass_flow_path))

    def effective_area_series(self, results: AmesimResults) -> tuple[float, ...]:
        return tuple(self.nominal_area_m2 * max(opening, 0.0) for opening in self.opening_series(results))

    def observation_at(self, results: AmesimResults, index: int) -> TestMqlOrificeObservation:
        opening = self.opening_series(results)[index]
        return TestMqlOrificeObservation(
            time=results.times[index],
            mass_flow_kg_s=results.series(self.primary_mass_flow_path)[index] * G_PER_S_TO_KG_PER_S,
            enthalpy_flow_w=results.series(self.primary_enthalpy_flow_path)[index],
            mass_flow_parameter=results.series(self.mass_flow_parameter_path)[index],
            gas_velocity_m_s=results.series(self.gas_velocity_path)[index],
            opening=opening,
            effective_area_m2=self.nominal_area_m2 * max(opening, 0.0),
        )


@dataclass(frozen=True)
class TestMqlOrificeObservationCatalog:
    bindings: tuple[TestMqlOrificeBinding, ...]

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

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

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


def build_test_mql_orifice_observation_catalog(
    results: AmesimResults,
    *,
    variable_catalog: TestMqlVariableCatalog | None = None,
    assembly: TestMqlPneumaticAssembly | None = None,
) -> TestMqlOrificeObservationCatalog:
    variable_catalog = variable_catalog or build_test_mql_variable_catalog(results)
    assembly = assembly or build_test_mql_pneumatic_assembly()
    bindings = []
    for alias, orifice in {
        **assembly.fixed_orifices,
        **assembly.variable_orifices,
    }.items():
        owner_variables = tuple(
            variable
            for variable in variable_catalog.variables
            if variable.owner_alias == alias
        )
        primary_mass_flow = _find_primary(owner_variables, signal_prefix="dm")
        primary_enthalpy_flow = _find_primary(owner_variables, signal_prefix="dh")
        reversed_mass_flow = _find_reversed(owner_variables, signal_prefix="dm")
        reversed_enthalpy_flow = _find_reversed(owner_variables, signal_prefix="dh")
        mass_flow_parameter = _find_by_signal(owner_variables, "cm")
        gas_velocity = _find_by_signal(owner_variables, "gasvel")
        opening = _find_optional_by_signal(owner_variables, "xv")
        bindings.append(
            TestMqlOrificeBinding(
                alias=alias,
                submodel=primary_mass_flow.submodel,
                nominal_area_m2=orifice.area,
                flow_coefficient=orifice.flow_coefficient,
                primary_mass_flow_path=primary_mass_flow.data_path,
                primary_enthalpy_flow_path=primary_enthalpy_flow.data_path,
                reversed_mass_flow_path=reversed_mass_flow.data_path,
                reversed_enthalpy_flow_path=reversed_enthalpy_flow.data_path,
                mass_flow_parameter_path=mass_flow_parameter.data_path,
                gas_velocity_path=gas_velocity.data_path,
                opening_path=opening.data_path if opening is not None else None,
            )
        )
    return TestMqlOrificeObservationCatalog(
        bindings=tuple(sorted(bindings, key=lambda binding: binding.alias))
    )


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


def _find_reversed(
    variables: tuple[TestMqlVariableBinding, ...],
    *,
    signal_prefix: str,
) -> TestMqlVariableBinding:
    matches = [
        variable
        for variable in variables
        if variable.signal_name.startswith(signal_prefix)
        and "sign reversed duplicate" in variable.label
    ]
    return _single(matches, f"reversed {signal_prefix}")


def _find_by_signal(
    variables: tuple[TestMqlVariableBinding, ...],
    signal_name: str,
) -> TestMqlVariableBinding:
    return _single(
        [variable for variable in variables if variable.signal_name == signal_name],
        signal_name,
    )


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


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