from __future__ import annotations

from dataclasses import dataclass

from PythonModels.reporting.amesim_results import AmesimResults
from PythonModels.reporting.test_mql_chamber_observations import (
    TestMqlChamberObservationCatalog,
    build_test_mql_chamber_observation_catalog,
)
from PythonModels.reporting.test_mql_line_observations import (
    TestMqlLineObservationCatalog,
    build_test_mql_line_observation_catalog,
)
from PythonModels.reporting.test_mql_mechanical_observations import (
    TestMqlMechanicalObservationCatalog,
    build_test_mql_mechanical_observation_catalog,
)
from PythonModels.reporting.test_mql_orifice_observations import (
    TestMqlOrificeObservationCatalog,
    build_test_mql_orifice_observation_catalog,
)
from PythonModels.reporting.test_mql_variables import (
    TestMqlVariableCatalog,
    build_test_mql_variable_catalog,
)


@dataclass(frozen=True)
class TestMqlObservationCatalog:
    variable_catalog: TestMqlVariableCatalog
    chambers: TestMqlChamberObservationCatalog
    orifices: TestMqlOrificeObservationCatalog
    lines: TestMqlLineObservationCatalog
    mechanical: TestMqlMechanicalObservationCatalog

    @property
    def binding_count(self) -> int:
        return (
            len(self.chambers.bindings)
            + len(self.orifices.bindings)
            + self.lines.line_count
            + self.mechanical.binding_count
        )

    def data_paths_by_domain(self) -> dict[str, tuple[str, ...]]:
        return {
            "chambers": _sorted_unique(_chamber_data_paths(self.chambers)),
            "orifices": _sorted_unique(_orifice_data_paths(self.orifices)),
            "lines": _sorted_unique(_line_data_paths(self.lines)),
            "mechanical": _sorted_unique(_mechanical_data_paths(self.mechanical)),
        }

    def data_paths(self) -> tuple[str, ...]:
        paths = []
        for domain_paths in self.data_paths_by_domain().values():
            paths.extend(domain_paths)
        return _sorted_unique(paths)

    def baseline_series_by_data_path(
        self,
        results: AmesimResults,
        data_paths: tuple[str, ...] | list[str] | None = None,
    ) -> dict[str, tuple[float, ...]]:
        selected_paths = tuple(data_paths) if data_paths is not None else self.data_paths()
        _validate_observed_paths(self, selected_paths)
        return {data_path: results.series(data_path) for data_path in selected_paths}


def build_test_mql_observation_catalog(results: AmesimResults) -> TestMqlObservationCatalog:
    variable_catalog = build_test_mql_variable_catalog(results)
    return TestMqlObservationCatalog(
        variable_catalog=variable_catalog,
        chambers=build_test_mql_chamber_observation_catalog(
            results,
            variable_catalog=variable_catalog,
        ),
        orifices=build_test_mql_orifice_observation_catalog(
            results,
            variable_catalog=variable_catalog,
        ),
        lines=build_test_mql_line_observation_catalog(
            results,
            variable_catalog=variable_catalog,
        ),
        mechanical=build_test_mql_mechanical_observation_catalog(
            results,
            variable_catalog=variable_catalog,
        ),
    )


def _chamber_data_paths(catalog: TestMqlChamberObservationCatalog) -> tuple[str, ...]:
    paths = []
    for binding in catalog.bindings:
        paths.extend(
            [
                binding.pressure_path,
                binding.temperature_path,
                binding.gas_mass_path,
                *binding.pressure_duplicate_paths,
                *binding.temperature_duplicate_paths,
            ]
        )
        if binding.volume_path is not None:
            paths.append(binding.volume_path)
    return tuple(paths)


def _orifice_data_paths(catalog: TestMqlOrificeObservationCatalog) -> tuple[str, ...]:
    paths = []
    for binding in catalog.bindings:
        paths.extend(
            [
                binding.primary_mass_flow_path,
                binding.primary_enthalpy_flow_path,
                binding.reversed_mass_flow_path,
                binding.reversed_enthalpy_flow_path,
                binding.mass_flow_parameter_path,
                binding.gas_velocity_path,
            ]
        )
        if binding.opening_path is not None:
            paths.append(binding.opening_path)
    return tuple(paths)


def _line_data_paths(catalog: TestMqlLineObservationCatalog) -> tuple[str, ...]:
    paths = []
    for binding in catalog.bindings:
        paths.extend(binding.mass_flow_paths)
        paths.extend(binding.enthalpy_flow_paths)
        paths.extend(binding.pressure_paths)
        paths.extend(binding.temperature_paths)
        if binding.gas_mass_path is not None:
            paths.append(binding.gas_mass_path)
        paths.extend(
            [
                binding.reynolds_path,
                binding.mass_flow_parameter_path,
                binding.gas_velocity_path,
                binding.friction_factor_path,
            ]
        )
    return tuple(paths)


def _mechanical_data_paths(catalog: TestMqlMechanicalObservationCatalog) -> tuple[str, ...]:
    paths = []
    for binding in catalog.pistons.values():
        paths.extend(
            [
                binding.volume_path,
                binding.volume_rate_path,
                binding.length_path,
                binding.force_port_2_path,
                binding.force_port_3_path,
                binding.displacement_port_2_path,
                binding.velocity_port_2_path,
                binding.displacement_port_3_path,
                binding.velocity_port_3_path,
            ]
        )
    for binding in catalog.masses.values():
        paths.extend(
            [
                binding.displacement_path,
                binding.velocity_path,
                binding.acceleration_path,
                binding.displacement_duplicate_path,
                binding.velocity_duplicate_path,
                binding.acceleration_duplicate_path,
                binding.lower_contact_force_path,
                binding.upper_contact_force_path,
                binding.viscous_friction_force_path,
                binding.dry_friction_force_path,
                binding.stick_flag_path,
            ]
        )
    for binding in catalog.elastic_endstops.values():
        paths.extend(
            [
                binding.force_path,
                binding.duplicate_force_path,
                binding.gap_path,
                binding.stiffness_path,
            ]
        )
    for binding in catalog.zero_force_sources.values():
        paths.append(binding.force_path)
    for binding in catalog.force_connectors.values():
        paths.append(binding.force_path)
    for binding in catalog.mechanical_nodes.values():
        paths.extend(binding.velocity_paths_by_port.values())
        paths.extend(binding.displacement_paths_by_port.values())
        paths.append(binding.total_force_path)
    return tuple(paths)


def _validate_observed_paths(
    catalog: TestMqlObservationCatalog,
    data_paths: tuple[str, ...],
) -> None:
    observed_paths = set(catalog.data_paths())
    missing = [data_path for data_path in data_paths if data_path not in observed_paths]
    if missing:
        raise KeyError(f"Data_Path values are not in the test_mql observation catalog: {missing}")


def _sorted_unique(data_paths: tuple[str, ...] | list[str]) -> tuple[str, ...]:
    return tuple(sorted(set(data_paths)))