from __future__ import annotations

from dataclasses import dataclass

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


AMESIM_REFERENCE_PRESSURE_PA = 101_300.0
BAR_TO_PA = 1.0e5
DEFAULT_TEST_MQL_TEMPERATURE_K = 293.15
DEFAULT_VARIABLE_CHAMBER_PRESSURE_BAR = 1.0


@dataclass(frozen=True)
class TestMqlStepSignalSpec:
    alias: str
    initial_output: float
    final_output: float
    step_time_s: float
    transition_duration_s: float
    transition_type: int

    def output_at(self, time_s: float) -> float:
        if self.transition_type != 1:
            raise ValueError(
                f"unsupported STEP0 transition type {self.transition_type} on {self.alias}"
            )
        return self.initial_output if time_s < self.step_time_s else self.final_output


@dataclass(frozen=True)
class TestMqlVariableOrificeControl:
    orifice_alias: str
    step: TestMqlStepSignalSpec

    def opening_at(self, time_s: float) -> float:
        return self.step.output_at(time_s)


@dataclass(frozen=True)
class TestMqlPneumaticAssembly:
    fixed_chambers: dict[str, AmesimPneumaticVolume]
    variable_chambers: dict[str, AmesimVariablePneumaticVolume]
    fixed_orifices: dict[str, AmesimPneumaticOrifice]
    variable_orifices: dict[str, AmesimPneumaticOrifice]
    variable_orifice_controls: dict[str, TestMqlVariableOrificeControl]
    fixed_initial_absolute_pressure_pa: float
    variable_initial_absolute_pressure_pa: float

    @property
    def initial_pressure_pa(self) -> float:
        return pressure_to_amesim_gauge_pa(self.fixed_initial_absolute_pressure_pa)

    @property
    def fixed_initial_gauge_pressure_pa(self) -> float:
        return pressure_to_amesim_gauge_pa(self.fixed_initial_absolute_pressure_pa)

    @property
    def variable_initial_gauge_pressure_pa(self) -> float:
        return pressure_to_amesim_gauge_pa(self.variable_initial_absolute_pressure_pa)

    @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 variable_orifice_control_count(self) -> int:
        return len(self.variable_orifice_controls)

    def set_variable_orifice_openings(self, time_s: float) -> None:
        for alias, control in self.variable_orifice_controls.items():
            self.variable_orifices[alias].opening = control.opening_at(time_s)

    @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()
    fixed_initial_absolute_pressure_pa = absolute_pressure_from_amesim_bar_parameter(
        config.global_parameters["P0"]
    )
    variable_initial_absolute_pressure_pa = absolute_pressure_from_amesim_bar_parameter(
        DEFAULT_VARIABLE_CHAMBER_PRESSURE_BAR
    )
    fixed_chambers = {
        component.alias: _build_chamber(
            component,
            volume_parameter="cvol",
            gas=gas,
            initial_pressure_pa=fixed_initial_absolute_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=variable_initial_absolute_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_orifice_controls = _build_variable_orifice_controls(config)
    variable_orifices = {
        component.alias: _build_orifice(
            component,
            area_parameter="area0",
            gas=gas,
            opening=variable_orifice_controls[component.alias].opening_at(0.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,
        variable_orifice_controls=variable_orifice_controls,
        fixed_initial_absolute_pressure_pa=fixed_initial_absolute_pressure_pa,
        variable_initial_absolute_pressure_pa=variable_initial_absolute_pressure_pa,
    )


def _build_variable_orifice_controls(
    config: TestMqlConfig,
) -> dict[str, TestMqlVariableOrificeControl]:
    from PythonModels.systems.test_mql import CONNECTION_SPECS

    components_by_alias = {component.alias: component for component in config.components}
    variable_orifice_aliases = {
        component.alias for component in config.components_by_submodel("PNVO001")
    }
    controls: dict[str, TestMqlVariableOrificeControl] = {}
    for connection in CONNECTION_SPECS:
        if connection["submodel"] != "DIRECT":
            continue
        source_alias = str(connection["source_component"])
        target_alias = str(connection["target_component"])
        if target_alias in variable_orifice_aliases:
            orifice_alias = target_alias
            step_alias = source_alias
        elif source_alias in variable_orifice_aliases:
            orifice_alias = source_alias
            step_alias = target_alias
        else:
            continue
        step_component = components_by_alias.get(step_alias)
        if step_component is None or step_component.submodel != "STEP0":
            continue
        controls[orifice_alias] = TestMqlVariableOrificeControl(
            orifice_alias=orifice_alias,
            step=TestMqlStepSignalSpec(
                alias=step_alias,
                initial_output=step_component.parameter_value("out0"),
                final_output=step_component.parameter_value("out1"),
                step_time_s=step_component.parameter_value("t0"),
                transition_duration_s=step_component.parameter_value("td"),
                transition_type=int(step_component.parameter_value("transitionType")),
            ),
        )
    missing = variable_orifice_aliases - controls.keys()
    if missing:
        raise ValueError(
            "missing STEP0 controls for PNVO001 components: "
            + ", ".join(sorted(missing))
        )
    return controls


def absolute_pressure_from_amesim_bar_parameter(pressure_bar: float) -> float:
    return pressure_bar * BAR_TO_PA


def pressure_to_amesim_gauge_pa(absolute_pressure_pa: float) -> float:
    return absolute_pressure_pa - AMESIM_REFERENCE_PRESSURE_PA


def pressure_from_amesim_bar_parameter(pressure_bar: float) -> float:
    return pressure_to_amesim_gauge_pa(absolute_pressure_from_amesim_bar_parameter(pressure_bar))


def _build_chamber(
    component: TestMqlResolvedComponent,
    *,
    volume_parameter: str,
    gas: AmesimPneumaticGas,
    initial_pressure_pa: float,
) -> AmesimPneumaticVolume:
    if volume_parameter == "cvol0":
        return AmesimVariablePneumaticVolume.from_liters(
            name=component.alias,
            dead_volume_liters=component.parameter_value(volume_parameter),
            gas=gas,
            p0=initial_pressure_pa,
            T0=_component_temperature(component),
            heat_transfer_coefficient=component.parameter_value("kth"),
            heat_transfer_area=component.parameter_value("sth"),
            external_temperature_k=_component_temperature(component),
        )
    return AmesimPneumaticVolume.from_liters(
        name=component.alias,
        volume_liters=component.parameter_value(volume_parameter),
        gas=gas,
        p0=initial_pressure_pa,
        T0=_component_temperature(component),
        heat_transfer_coefficient=component.parameter_value("kth"),
        heat_transfer_area=component.parameter_value("sth"),
        external_temperature_k=_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