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_mechanical import (
    TestMqlMechanicalAssembly,
    build_test_mql_mechanical_assembly,
)


@dataclass(frozen=True)
class TestMqlPistonObservation:
    time: float
    chamber_volume_cm3: float
    chamber_volume_rate_l_min: float
    chamber_length_mm: float
    force_port_2_n: float
    force_port_3_n: float
    displacement_port_2_m: float
    velocity_port_2_m_s: float
    displacement_port_3_m: float
    velocity_port_3_m_s: float


@dataclass(frozen=True)
class TestMqlMassEndstopObservation:
    time: float
    displacement_m: float
    velocity_m_s: float
    acceleration_m_s2: float
    lower_contact_force_n: float
    upper_contact_force_n: float
    viscous_friction_force_n: float
    dry_friction_force_n: float
    stick_flag: float


@dataclass(frozen=True)
class TestMqlElasticEndstopObservation:
    time: float
    force_n: float
    duplicate_force_n: float
    gap_mm: float
    stiffness_n_m: float


@dataclass(frozen=True)
class TestMqlForceSourceObservation:
    time: float
    force_n: float


@dataclass(frozen=True)
class TestMqlForceConnectorObservation:
    time: float
    force_n: float


@dataclass(frozen=True)
class TestMqlMechanicalNodeObservation:
    time: float
    velocities_m_s: dict[int, float]
    displacements_m: dict[int, float]
    total_force_n: float


@dataclass(frozen=True)
class TestMqlPistonObservationBinding:
    alias: str
    volume_path: str
    volume_rate_path: str
    length_path: str
    force_port_2_path: str
    force_port_3_path: str
    displacement_port_2_path: str
    velocity_port_2_path: str
    displacement_port_3_path: str
    velocity_port_3_path: str

    def observation_at(self, results: AmesimResults, index: int) -> TestMqlPistonObservation:
        return TestMqlPistonObservation(
            time=results.times[index],
            chamber_volume_cm3=results.series(self.volume_path)[index],
            chamber_volume_rate_l_min=results.series(self.volume_rate_path)[index],
            chamber_length_mm=results.series(self.length_path)[index],
            force_port_2_n=results.series(self.force_port_2_path)[index],
            force_port_3_n=results.series(self.force_port_3_path)[index],
            displacement_port_2_m=results.series(self.displacement_port_2_path)[index],
            velocity_port_2_m_s=results.series(self.velocity_port_2_path)[index],
            displacement_port_3_m=results.series(self.displacement_port_3_path)[index],
            velocity_port_3_m_s=results.series(self.velocity_port_3_path)[index],
        )


@dataclass(frozen=True)
class TestMqlMassEndstopObservationBinding:
    alias: str
    displacement_path: str
    velocity_path: str
    acceleration_path: str
    displacement_duplicate_path: str
    velocity_duplicate_path: str
    acceleration_duplicate_path: str
    lower_contact_force_path: str
    upper_contact_force_path: str
    viscous_friction_force_path: str
    dry_friction_force_path: str
    stick_flag_path: str

    def observation_at(self, results: AmesimResults, index: int) -> TestMqlMassEndstopObservation:
        return TestMqlMassEndstopObservation(
            time=results.times[index],
            displacement_m=results.series(self.displacement_path)[index],
            velocity_m_s=results.series(self.velocity_path)[index],
            acceleration_m_s2=results.series(self.acceleration_path)[index],
            lower_contact_force_n=results.series(self.lower_contact_force_path)[index],
            upper_contact_force_n=results.series(self.upper_contact_force_path)[index],
            viscous_friction_force_n=results.series(self.viscous_friction_force_path)[index],
            dry_friction_force_n=results.series(self.dry_friction_force_path)[index],
            stick_flag=results.series(self.stick_flag_path)[index],
        )


@dataclass(frozen=True)
class TestMqlElasticEndstopObservationBinding:
    alias: str
    force_path: str
    duplicate_force_path: str
    gap_path: str
    stiffness_path: str

    def observation_at(self, results: AmesimResults, index: int) -> TestMqlElasticEndstopObservation:
        return TestMqlElasticEndstopObservation(
            time=results.times[index],
            force_n=results.series(self.force_path)[index],
            duplicate_force_n=results.series(self.duplicate_force_path)[index],
            gap_mm=results.series(self.gap_path)[index],
            stiffness_n_m=results.series(self.stiffness_path)[index],
        )


@dataclass(frozen=True)
class TestMqlForceSourceObservationBinding:
    alias: str
    force_path: str

    def observation_at(self, results: AmesimResults, index: int) -> TestMqlForceSourceObservation:
        return TestMqlForceSourceObservation(
            time=results.times[index],
            force_n=results.series(self.force_path)[index],
        )


@dataclass(frozen=True)
class TestMqlForceConnectorObservationBinding:
    alias: str
    force_path: str

    def observation_at(self, results: AmesimResults, index: int) -> TestMqlForceConnectorObservation:
        return TestMqlForceConnectorObservation(
            time=results.times[index],
            force_n=results.series(self.force_path)[index],
        )


@dataclass(frozen=True)
class TestMqlMechanicalNodeObservationBinding:
    alias: str
    velocity_paths_by_port: dict[int, str]
    displacement_paths_by_port: dict[int, str]
    total_force_path: str

    def observation_at(self, results: AmesimResults, index: int) -> TestMqlMechanicalNodeObservation:
        return TestMqlMechanicalNodeObservation(
            time=results.times[index],
            velocities_m_s={
                port: results.series(path)[index]
                for port, path in self.velocity_paths_by_port.items()
            },
            displacements_m={
                port: results.series(path)[index]
                for port, path in self.displacement_paths_by_port.items()
            },
            total_force_n=results.series(self.total_force_path)[index],
        )


@dataclass(frozen=True)
class TestMqlMechanicalObservationCatalog:
    pistons: dict[str, TestMqlPistonObservationBinding]
    masses: dict[str, TestMqlMassEndstopObservationBinding]
    elastic_endstops: dict[str, TestMqlElasticEndstopObservationBinding]
    zero_force_sources: dict[str, TestMqlForceSourceObservationBinding]
    force_connectors: dict[str, TestMqlForceConnectorObservationBinding]
    mechanical_nodes: dict[str, TestMqlMechanicalNodeObservationBinding]

    @property
    def binding_count(self) -> int:
        return (
            len(self.pistons)
            + len(self.masses)
            + len(self.elastic_endstops)
            + len(self.zero_force_sources)
            + len(self.force_connectors)
            + len(self.mechanical_nodes)
        )


def build_test_mql_mechanical_observation_catalog(
    results: AmesimResults,
    *,
    variable_catalog: TestMqlVariableCatalog | None = None,
    mechanical_assembly: TestMqlMechanicalAssembly | None = None,
) -> TestMqlMechanicalObservationCatalog:
    variable_catalog = variable_catalog or build_test_mql_variable_catalog(results)
    mechanical_assembly = mechanical_assembly or build_test_mql_mechanical_assembly(
        amesim_results=results,
        variable_catalog=variable_catalog,
    )
    return TestMqlMechanicalObservationCatalog(
        pistons={
            alias: _build_piston_binding(alias, variable_catalog)
            for alias in mechanical_assembly.pistons
        },
        masses={
            alias: _build_mass_binding(alias, variable_catalog)
            for alias in mechanical_assembly.masses
        },
        elastic_endstops={
            alias: _build_elastic_endstop_binding(alias, variable_catalog)
            for alias in mechanical_assembly.elastic_endstops
        },
        zero_force_sources={
            alias: _build_zero_force_source_binding(alias, variable_catalog)
            for alias in mechanical_assembly.zero_force_sources
        },
        force_connectors={
            alias: _build_force_connector_binding(alias, variable_catalog)
            for alias in mechanical_assembly.force_connectors
        },
        mechanical_nodes={
            alias: _build_mechanical_node_binding(alias, variable_catalog)
            for alias in mechanical_assembly.mechanical_nodes
        },
    )


def _build_piston_binding(
    alias: str,
    variable_catalog: TestMqlVariableCatalog,
) -> TestMqlPistonObservationBinding:
    variables = _owner_variables(variable_catalog, alias)
    return TestMqlPistonObservationBinding(
        alias=alias,
        volume_path=_required_path(variables, "vol1", expected_units="cm**3"),
        volume_rate_path=_required_path(variables, "vvol1", expected_units="L/min"),
        length_path=_required_path(variables, "length", expected_units="mm"),
        force_port_2_path=_required_path(variables, "f2", expected_units="N"),
        force_port_3_path=_required_path(variables, "f3", expected_units="N"),
        displacement_port_2_path=_required_path(variables, "x5", expected_units="m"),
        velocity_port_2_path=_required_path(variables, "v5", expected_units="m/s"),
        displacement_port_3_path=_required_path(variables, "x4", expected_units="m"),
        velocity_port_3_path=_required_path(variables, "v4", expected_units="m/s"),
    )


def _build_mass_binding(
    alias: str,
    variable_catalog: TestMqlVariableCatalog,
) -> TestMqlMassEndstopObservationBinding:
    variables = _owner_variables(variable_catalog, alias)
    return TestMqlMassEndstopObservationBinding(
        alias=alias,
        displacement_path=_required_path(variables, "x1", expected_units="m"),
        velocity_path=_required_path(variables, "v1", expected_units="m/s"),
        acceleration_path=_required_path(variables, "acc1", expected_units="m/s/s"),
        displacement_duplicate_path=_required_path(variables, "x1dup", expected_units="m"),
        velocity_duplicate_path=_required_path(variables, "v1dup", expected_units="m/s"),
        acceleration_duplicate_path=_required_path(variables, "acc1dup", expected_units="m/s/s"),
        lower_contact_force_path=_required_path(variables, "Fmin", expected_units="N"),
        upper_contact_force_path=_required_path(variables, "Fmax", expected_units="N"),
        viscous_friction_force_path=_required_path(variables, "Fvisc", expected_units="N"),
        dry_friction_force_path=_required_path(variables, "Ffric", expected_units="N"),
        stick_flag_path=_required_path(variables, "stick", expected_units=None),
    )


def _build_elastic_endstop_binding(
    alias: str,
    variable_catalog: TestMqlVariableCatalog,
) -> TestMqlElasticEndstopObservationBinding:
    variables = _owner_variables(variable_catalog, alias)
    return TestMqlElasticEndstopObservationBinding(
        alias=alias,
        force_path=_required_path(variables, "f1", expected_units="N"),
        duplicate_force_path=_required_path(variables, "f2", expected_units="N"),
        gap_path=_required_path(variables, "gap", expected_units="mm"),
        stiffness_path=_required_path(variables, "kval", expected_units="N/m"),
    )


def _build_zero_force_source_binding(
    alias: str,
    variable_catalog: TestMqlVariableCatalog,
) -> TestMqlForceSourceObservationBinding:
    variables = _owner_variables(variable_catalog, alias)
    return TestMqlForceSourceObservationBinding(
        alias=alias,
        force_path=_required_path(variables, "fzero", expected_units="N"),
    )


def _build_force_connector_binding(
    alias: str,
    variable_catalog: TestMqlVariableCatalog,
) -> TestMqlForceConnectorObservationBinding:
    variables = _owner_variables(variable_catalog, alias)
    return TestMqlForceConnectorObservationBinding(
        alias=alias,
        force_path=_required_path(variables, "force", expected_units="N"),
    )


def _build_mechanical_node_binding(
    alias: str,
    variable_catalog: TestMqlVariableCatalog,
) -> TestMqlMechanicalNodeObservationBinding:
    variables = _owner_variables(variable_catalog, alias)
    velocity_paths_by_port = {}
    displacement_paths_by_port = {}
    for port in range(1, 9):
        velocity_paths_by_port[port] = _required_path(
            variables,
            f"p{port}__vt",
            expected_units="m/s",
        )
        displacement_paths_by_port[port] = _required_path(
            variables,
            f"p{port}__xt",
            expected_units="m",
        )
    return TestMqlMechanicalNodeObservationBinding(
        alias=alias,
        velocity_paths_by_port=velocity_paths_by_port,
        displacement_paths_by_port=displacement_paths_by_port,
        total_force_path=_required_path(variables, "tforce", expected_units="N"),
    )


def _owner_variables(
    variable_catalog: TestMqlVariableCatalog,
    alias: str,
) -> tuple[TestMqlVariableBinding, ...]:
    return tuple(
        variable
        for variable in variable_catalog.variables
        if variable.owner_alias == alias
    )


def _required_path(
    variables: tuple[TestMqlVariableBinding, ...],
    signal_name: str,
    *,
    expected_units: str | None,
) -> str:
    variable = _single(
        tuple(variable for variable in variables if variable.signal_name == signal_name),
        signal_name,
    )
    _assert_units(variable, expected_units)
    return variable.data_path


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}."
        )