from __future__ import annotations

from dataclasses import dataclass

from PythonModels.components.amesim_pneumatic import (
    HELIUM_PNEUMATIC_GAS,
    AmesimPneumaticGas,
    AmesimPneumaticOrifice,
    AmesimPneumaticVolume,
)
from PythonModels.systems.test_mql_config import TestMqlConfig, TestMqlResolvedComponent


ATMOSPHERIC_PRESSURE_BAR = 1.013
BAR_TO_PA = 1.0e5
DEFAULT_TEST_MQL_TEMPERATURE_K = 293.15


@dataclass(frozen=True)
class TestMqlPneumaticAssembly:
    fixed_chambers: dict[str, AmesimPneumaticVolume]
    variable_chambers: dict[str, AmesimPneumaticVolume]
    fixed_orifices: dict[str, AmesimPneumaticOrifice]
    variable_orifices: dict[str, AmesimPneumaticOrifice]
    initial_pressure_pa: float

    @property
    def chamber_count(self) -> int:
        return len(self.fixed_chambers) + len(self.variable_chambers)

    @property
    def orifice_count(self) -> int:
        return len(self.fixed_orifices) + len(self.variable_orifices)

    @property
    def component_count(self) -> int:
        return self.chamber_count + self.orifice_count

    @property
    def aliases(self) -> tuple[str, ...]:
        return tuple(
            [
                *self.fixed_chambers,
                *self.variable_chambers,
                *self.fixed_orifices,
                *self.variable_orifices,
            ]
        )


def build_test_mql_pneumatic_assembly(
    config: TestMqlConfig | None = None,
    gas: AmesimPneumaticGas = HELIUM_PNEUMATIC_GAS,
) -> TestMqlPneumaticAssembly:
    config = config or TestMqlConfig.from_amesim_specs()
    initial_pressure_pa = pressure_from_amesim_bar_parameter(config.global_parameters["P0"])
    fixed_chambers = {
        component.alias: _build_chamber(
            component,
            volume_parameter="cvol",
            gas=gas,
            initial_pressure_pa=initial_pressure_pa,
        )
        for component in config.components_by_submodel("PNCH023")
    }
    variable_chambers = {
        component.alias: _build_chamber(
            component,
            volume_parameter="cvol0",
            gas=gas,
            initial_pressure_pa=initial_pressure_pa,
        )
        for component in config.components_by_submodel("PNCH012")
    }
    fixed_orifices = {
        component.alias: _build_orifice(
            component,
            area_parameter="area",
            gas=gas,
            opening=1.0,
        )
        for component in config.components_by_submodel("PNOR001")
    }
    variable_orifices = {
        component.alias: _build_orifice(
            component,
            area_parameter="area0",
            gas=gas,
            opening=1.0,
        )
        for component in config.components_by_submodel("PNVO001")
    }
    return TestMqlPneumaticAssembly(
        fixed_chambers=fixed_chambers,
        variable_chambers=variable_chambers,
        fixed_orifices=fixed_orifices,
        variable_orifices=variable_orifices,
        initial_pressure_pa=initial_pressure_pa,
    )


def pressure_from_amesim_bar_parameter(pressure_bar: float) -> float:
    return (pressure_bar - ATMOSPHERIC_PRESSURE_BAR) * BAR_TO_PA


def _build_chamber(
    component: TestMqlResolvedComponent,
    *,
    volume_parameter: str,
    gas: AmesimPneumaticGas,
    initial_pressure_pa: float,
) -> AmesimPneumaticVolume:
    return AmesimPneumaticVolume.from_liters(
        name=component.alias,
        volume_liters=component.parameter_value(volume_parameter),
        gas=gas,
        p0=initial_pressure_pa,
        T0=_component_temperature(component),
    )


def _build_orifice(
    component: TestMqlResolvedComponent,
    *,
    area_parameter: str,
    gas: AmesimPneumaticGas,
    opening: float,
) -> AmesimPneumaticOrifice:
    return AmesimPneumaticOrifice.from_mm2(
        name=component.alias,
        area_mm2=component.parameter_value(area_parameter),
        flow_coefficient=component.parameter_value("cq"),
        gas=gas,
        opening=opening,
    )


def _component_temperature(component: TestMqlResolvedComponent) -> float:
    parameter = component.parameters.get("extemp")
    if parameter is None or parameter.value is None:
        return DEFAULT_TEST_MQL_TEMPERATURE_K
    return parameter.value