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_lines import (
    TestMqlLineAssembly,
    build_test_mql_line_assembly,
)


G_PER_S_TO_KG_PER_S = 1.0e-3


@dataclass(frozen=True)
class TestMqlLineObservation:
    time: float
    mass_flows_kg_s: dict[str, float]
    enthalpy_flows_w: dict[str, float]
    pressures_pa: dict[str, float]
    temperatures_k: dict[str, float]
    gas_mass_g: float | None
    reynolds_number: float
    mass_flow_parameter: float
    gas_velocity_m_s: float
    friction_factor: float


@dataclass(frozen=True)
class TestMqlLineObservationBinding:
    alias: str
    submodel: str
    pattern: str
    mass_flow_paths: tuple[str, ...]
    enthalpy_flow_paths: tuple[str, ...]
    pressure_paths: tuple[str, ...]
    temperature_paths: tuple[str, ...]
    gas_mass_path: str | None
    reynolds_path: str
    mass_flow_parameter_path: str
    gas_velocity_path: str
    friction_factor_path: str

    def mass_flow_kg_s_series(
        self,
        results: AmesimResults,
        data_path: str | None = None,
    ) -> tuple[float, ...]:
        path = data_path or self.mass_flow_paths[0]
        if path not in self.mass_flow_paths:
            raise KeyError(path)
        return tuple(value * G_PER_S_TO_KG_PER_S for value in results.series(path))

    def observation_at(self, results: AmesimResults, index: int) -> TestMqlLineObservation:
        return TestMqlLineObservation(
            time=results.times[index],
            mass_flows_kg_s={
                path: results.series(path)[index] * G_PER_S_TO_KG_PER_S
                for path in self.mass_flow_paths
            },
            enthalpy_flows_w={
                path: results.series(path)[index]
                for path in self.enthalpy_flow_paths
            },
            pressures_pa={
                path: results.series(path)[index]
                for path in self.pressure_paths
            },
            temperatures_k={
                path: results.series(path)[index]
                for path in self.temperature_paths
            },
            gas_mass_g=(
                results.series(self.gas_mass_path)[index]
                if self.gas_mass_path is not None
                else None
            ),
            reynolds_number=results.series(self.reynolds_path)[index],
            mass_flow_parameter=results.series(self.mass_flow_parameter_path)[index],
            gas_velocity_m_s=results.series(self.gas_velocity_path)[index],
            friction_factor=results.series(self.friction_factor_path)[index],
        )


@dataclass(frozen=True)
class TestMqlLineObservationCatalog:
    bindings: tuple[TestMqlLineObservationBinding, ...]

    @property
    def line_count(self) -> int:
        return len(self.bindings)

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

    def by_submodel(self, submodel: str) -> tuple[TestMqlLineObservationBinding, ...]:
        return tuple(binding for binding in self.bindings if binding.submodel == submodel)


def build_test_mql_line_observation_catalog(
    results: AmesimResults,
    *,
    variable_catalog: TestMqlVariableCatalog | None = None,
    line_assembly: TestMqlLineAssembly | None = None,
) -> TestMqlLineObservationCatalog:
    variable_catalog = variable_catalog or build_test_mql_variable_catalog(results)
    line_assembly = line_assembly or build_test_mql_line_assembly(results, variable_catalog)
    bindings = []
    for line in line_assembly.lines:
        variables = tuple(
            variable
            for variable in variable_catalog.variables
            if variable.owner_alias == line.alias
        )
        bindings.append(
            TestMqlLineObservationBinding(
                alias=line.alias,
                submodel=line.submodel,
                pattern=line.pattern,
                mass_flow_paths=_paths_with_prefix(variables, "dm", expected_units="g/s"),
                enthalpy_flow_paths=_paths_with_prefix(variables, "dh", expected_units="J/s"),
                pressure_paths=_paths_with_prefix(variables, "p", expected_units="Pa"),
                temperature_paths=_paths_with_prefix(variables, "t", expected_units="K"),
                gas_mass_path=_optional_path(variables, "mgas", expected_units="g"),
                reynolds_path=_required_path(variables, "re", expected_units=None),
                mass_flow_parameter_path=_required_path(
                    variables,
                    "cm",
                    expected_units="(kg*K/J)**(1/2)",
                ),
                gas_velocity_path=_required_path(variables, "v", expected_units="m/s"),
                friction_factor_path=_required_path(variables, "ff", expected_units=None),
            )
        )
    return TestMqlLineObservationCatalog(bindings=tuple(bindings))


def _paths_with_prefix(
    variables: tuple[TestMqlVariableBinding, ...],
    prefix: str,
    *,
    expected_units: str | None,
) -> tuple[str, ...]:
    matches = tuple(
        variable
        for variable in variables
        if variable.signal_name.startswith(prefix)
    )
    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 | None,
) -> str:
    variable = _single_signal(variables, signal_name)
    _assert_units(variable, expected_units)
    return variable.data_path


def _optional_path(
    variables: tuple[TestMqlVariableBinding, ...],
    signal_name: str,
    *,
    expected_units: str | None,
) -> 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_signal(
    variables: tuple[TestMqlVariableBinding, ...],
    signal_name: str,
) -> TestMqlVariableBinding:
    return _single(
        tuple(variable for variable in variables if variable.signal_name == signal_name),
        signal_name,
    )


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) -> None:
    if variable.units != expected_units:
        raise ValueError(
            f"Unexpected units for {variable.data_path}: "
            f"{variable.units!r}, expected {expected_units!r}."
        )