from __future__ import annotations

from dataclasses import dataclass

from PythonModels.components.amesim_mechanical import (
    AmesimElasticEndstop,
    AmesimMassFrictionEndstops,
    AmesimPistonGeometry,
    circular_area,
    mm_to_m,
)
from PythonModels.reporting.amesim_results import AmesimResults
from PythonModels.reporting.test_mql_variables import (
    TestMqlVariableCatalog,
    build_test_mql_variable_catalog,
)
from PythonModels.systems.test_mql_config import TestMqlConfig, TestMqlResolvedComponent


MM_TO_M = 1.0e-3
N_PER_MM_TO_N_PER_M = 1.0e3
N_PER_MM_PER_S_TO_N_PER_M_PER_S = 1.0e3


@dataclass(frozen=True)
class TestMqlPistonSpec:
    alias: str
    piston_diameter_m: float
    rod_diameter_m: float
    zero_displacement_m: float
    piston_area_m2: float
    rod_area_m2: float
    annulus_area_m2: float
    data_paths: tuple[str, ...]

    def geometry(self) -> AmesimPistonGeometry:
        return AmesimPistonGeometry(
            piston_diameter_m=self.piston_diameter_m,
            rod_diameter_m=self.rod_diameter_m,
            zero_length_m=self.zero_displacement_m,
        )


@dataclass(frozen=True)
class TestMqlMassEndstopSpec:
    alias: str
    mass_kg: float
    xmin_m: float
    xmax_m: float
    min_stiffness_n_per_m: float
    max_stiffness_n_per_m: float
    min_damping_n_per_m_per_s: float
    max_damping_n_per_m_per_s: float
    min_penetration_m: float
    max_penetration_m: float
    stiction_force_n: float
    coulomb_friction_n: float
    viscous_friction_n_per_m_per_s: float
    windage_n_per_m2_per_s2: float
    stick_velocity_threshold_m_s: float
    reset_velocity_threshold_m_s: float
    rest_coeff: float
    stribeck_constant_m_s: float
    use_friction: bool
    stop_type: int
    initial_velocity_m_s: float
    initial_displacement_m: float
    data_paths: tuple[str, ...]

    def endstop(self) -> AmesimMassFrictionEndstops:
        return AmesimMassFrictionEndstops(
            mass_kg=self.mass_kg,
            lower_limit_m=self.xmin_m,
            upper_limit_m=self.xmax_m,
            lower_stiffness_n_per_m=self.min_stiffness_n_per_m,
            upper_stiffness_n_per_m=self.max_stiffness_n_per_m,
            lower_damping_n_per_m_per_s=self.min_damping_n_per_m_per_s,
            upper_damping_n_per_m_per_s=self.max_damping_n_per_m_per_s,
            viscous_friction_n_per_m_per_s=self.viscous_friction_n_per_m_per_s,
            coulomb_friction_n=self.coulomb_friction_n,
            stiction_force_n=self.stiction_force_n,
            windage_n_per_m2_per_s2=self.windage_n_per_m2_per_s2,
        )


@dataclass(frozen=True)
class TestMqlElasticEndstopSpec:
    alias: str
    gap_m: float
    contact_stiffness_n_per_m: float
    contact_damping_n_per_m_per_s: float
    spring_diameter_m: float
    wire_diameter_m: float
    data_paths: tuple[str, ...]

    def endstop(self) -> AmesimElasticEndstop:
        return AmesimElasticEndstop(
            contact_stiffness_n_per_m=self.contact_stiffness_n_per_m,
            contact_damping_n_per_m_per_s=self.contact_damping_n_per_m_per_s,
            gap0_m=self.gap_m,
        )


@dataclass(frozen=True)
class TestMqlMechanicalNodeSpec:
    alias: str
    port_count: int
    sum_mode: int
    data_paths: tuple[str, ...]


@dataclass(frozen=True)
class TestMqlMechanicalAssembly:
    pistons: dict[str, TestMqlPistonSpec]
    masses: dict[str, TestMqlMassEndstopSpec]
    elastic_endstops: dict[str, TestMqlElasticEndstopSpec]
    mechanical_nodes: dict[str, TestMqlMechanicalNodeSpec]
    zero_force_sources: tuple[str, ...]
    force_connectors: tuple[str, ...]

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

    @property
    def aliases(self) -> tuple[str, ...]:
        return tuple(
            [
                *self.pistons,
                *self.masses,
                *self.elastic_endstops,
                *self.mechanical_nodes,
                *self.zero_force_sources,
                *self.force_connectors,
            ]
        )


@dataclass(frozen=True)
class TestMqlMechanicalMassState:
    alias: str
    velocity_m_s: float
    displacement_m: float

    def as_vector(self) -> list[float]:
        return [self.velocity_m_s, self.displacement_m]


@dataclass(frozen=True)
class TestMqlMechanicalNodeKinematics:
    alias: str
    velocities_m_s: dict[int, float]
    displacements_m: dict[int, float]


@dataclass(frozen=True)
class TestMqlMechanicalMassSnapshot:
    states: tuple[TestMqlMechanicalMassState, ...]
    node_kinematics_by_alias: dict[str, TestMqlMechanicalNodeKinematics]

    @property
    def state_count(self) -> int:
        return 2 * len(self.states)


class TestMqlMechanicalMassClosure:
    def __init__(self, assembly: TestMqlMechanicalAssembly) -> None:
        self.assembly = assembly
        self.mass_aliases = tuple(assembly.masses)

    def initial_state_vector(self) -> list[float]:
        state: list[float] = []
        for alias in self.mass_aliases:
            spec = self.assembly.masses[alias]
            state.extend([spec.initial_velocity_m_s, spec.initial_displacement_m])
        return state

    def snapshot(self, state_vector: list[float] | None = None) -> TestMqlMechanicalMassSnapshot:
        values = self.initial_state_vector() if state_vector is None else list(state_vector)
        if len(values) != 2 * len(self.mass_aliases):
            raise ValueError("mechanical mass state vector requires two values per mass")
        states = tuple(
            TestMqlMechanicalMassState(
                alias=alias,
                velocity_m_s=values[2 * index],
                displacement_m=values[2 * index + 1],
            )
            for index, alias in enumerate(self.mass_aliases)
        )
        return TestMqlMechanicalMassSnapshot(
            states=states,
            node_kinematics_by_alias=self._node_kinematics_by_alias(states),
        )

    def _node_kinematics_by_alias(
        self,
        states: tuple[TestMqlMechanicalMassState, ...],
    ) -> dict[str, TestMqlMechanicalNodeKinematics]:
        state_by_alias = {state.alias: state for state in states}
        front = state_by_alias["mass_friction_endstops_18"]
        rear = state_by_alias["mass_friction_endstops_19"]
        return {
            "dynamic_mechanical_node_alternative_2": TestMqlMechanicalNodeKinematics(
                alias="dynamic_mechanical_node_alternative_2",
                velocities_m_s={port: -front.velocity_m_s for port in range(1, 9)},
                displacements_m={port: -front.displacement_m for port in range(1, 9)},
            ),
            "dynamic_mechanical_node_alternative_3": TestMqlMechanicalNodeKinematics(
                alias="dynamic_mechanical_node_alternative_3",
                velocities_m_s={port: rear.velocity_m_s for port in range(1, 9)},
                displacements_m={port: rear.displacement_m for port in range(1, 9)},
            ),
        }

    def rhs(
        self,
        state_vector: list[float],
        *,
        force_by_mass_alias: dict[str, float] | None = None,
    ) -> list[float]:
        snapshot = self.snapshot(state_vector)
        force_by_mass_alias = force_by_mass_alias or {}
        derivatives: list[float] = []
        for state in snapshot.states:
            mass = self.assembly.masses[state.alias].endstop()
            acceleration, velocity = mass.derivatives(
                velocity_m_s=state.velocity_m_s,
                displacement_m=state.displacement_m,
                port_1_force_n=force_by_mass_alias.get(state.alias, 0.0),
            )
            derivatives.extend([acceleration, velocity])
        return derivatives


def build_test_mql_mechanical_assembly(
    config: TestMqlConfig | None = None,
    amesim_results: AmesimResults | None = None,
    variable_catalog: TestMqlVariableCatalog | None = None,
) -> TestMqlMechanicalAssembly:
    config = config or TestMqlConfig.from_amesim_specs()
    if variable_catalog is None and amesim_results is not None:
        variable_catalog = build_test_mql_variable_catalog(amesim_results)

    pistons = {
        component.alias: _build_piston(component, variable_catalog)
        for component in config.components_by_submodel("PNRP17")
    }
    masses = {
        component.alias: _build_mass(component, variable_catalog, amesim_results)
        for component in config.components_by_submodel("MECMAS21")
    }
    elastic_endstops = {
        component.alias: _build_elastic_endstop(component, variable_catalog)
        for component in config.components_by_submodel("LSTP00A")
    }
    mechanical_nodes = {
        component.alias: _build_mechanical_node(component, variable_catalog)
        for component in config.components_by_submodel("LMECHN1")
    }
    zero_force_sources = tuple(component.alias for component in config.components_by_submodel("F000"))
    force_connectors = tuple(component.alias for component in config.components_by_submodel("FORC"))
    return TestMqlMechanicalAssembly(
        pistons=pistons,
        masses=masses,
        elastic_endstops=elastic_endstops,
        mechanical_nodes=mechanical_nodes,
        zero_force_sources=zero_force_sources,
        force_connectors=force_connectors,
    )


def _build_piston(
    component: TestMqlResolvedComponent,
    variable_catalog: TestMqlVariableCatalog | None,
) -> TestMqlPistonSpec:
    geometry = AmesimPistonGeometry(
        piston_diameter_m=mm_to_m(component.parameter_value("dp")),
        rod_diameter_m=mm_to_m(component.parameter_value("dr")),
        zero_length_m=mm_to_m(component.parameter_value("x0")),
    )
    return TestMqlPistonSpec(
        alias=component.alias,
        piston_diameter_m=geometry.piston_diameter_m,
        rod_diameter_m=geometry.rod_diameter_m,
        zero_displacement_m=geometry.zero_length_m,
        piston_area_m2=geometry.piston_area_m2,
        rod_area_m2=geometry.rod_area_m2,
        annulus_area_m2=geometry.annulus_area_m2,
        data_paths=_data_paths(variable_catalog, component.alias),
    )


def _build_mass(
    component: TestMqlResolvedComponent,
    variable_catalog: TestMqlVariableCatalog | None,
    amesim_results: AmesimResults | None,
) -> TestMqlMassEndstopSpec:
    return TestMqlMassEndstopSpec(
        alias=component.alias,
        mass_kg=component.parameter_value("mass"),
        xmin_m=component.parameter_value("xmin"),
        xmax_m=component.parameter_value("xmax"),
        min_stiffness_n_per_m=n_per_mm_to_n_per_m(component.parameter_value("Kbmin")),
        max_stiffness_n_per_m=n_per_mm_to_n_per_m(component.parameter_value("Kbmax")),
        min_damping_n_per_m_per_s=n_per_mm_per_s_to_n_per_m_per_s(component.parameter_value("Dbmin")),
        max_damping_n_per_m_per_s=n_per_mm_per_s_to_n_per_m_per_s(component.parameter_value("Dbmax")),
        min_penetration_m=mm_to_m(component.parameter_value("Pdmin")),
        max_penetration_m=mm_to_m(component.parameter_value("Pdmax")),
        stiction_force_n=component.parameter_value("fstick"),
        coulomb_friction_n=component.parameter_value("fcoul"),
        viscous_friction_n_per_m_per_s=component.parameter_value("rvisc"),
        windage_n_per_m2_per_s2=component.parameter_value("wind"),
        stick_velocity_threshold_m_s=component.parameter_value("dvel"),
        reset_velocity_threshold_m_s=component.parameter_value("restdvel"),
        rest_coeff=component.parameter_value("restcoeff"),
        stribeck_constant_m_s=component.parameter_value("astrib"),
        use_friction=bool(int(component.parameter_value("useFriction"))),
        stop_type=int(component.parameter_value("stoptype")),
        initial_velocity_m_s=_initial_value(amesim_results, f"v1@{component.alias}"),
        initial_displacement_m=_initial_value(amesim_results, f"x1@{component.alias}"),
        data_paths=_data_paths(variable_catalog, component.alias),
    )


def _build_elastic_endstop(
    component: TestMqlResolvedComponent,
    variable_catalog: TestMqlVariableCatalog | None,
) -> TestMqlElasticEndstopSpec:
    return TestMqlElasticEndstopSpec(
        alias=component.alias,
        gap_m=mm_to_m(component.parameter_value("gap0")),
        contact_stiffness_n_per_m=component.parameter_value("kcont"),
        contact_damping_n_per_m_per_s=component.parameter_value("rcont"),
        spring_diameter_m=mm_to_m(component.parameter_value("sdiam")),
        wire_diameter_m=mm_to_m(component.parameter_value("wdiam")),
        data_paths=_data_paths(variable_catalog, component.alias),
    )


def _build_mechanical_node(
    component: TestMqlResolvedComponent,
    variable_catalog: TestMqlVariableCatalog | None,
) -> TestMqlMechanicalNodeSpec:
    return TestMqlMechanicalNodeSpec(
        alias=component.alias,
        port_count=int(component.parameter_value("v1")),
        sum_mode=int(component.parameter_value("sum")),
        data_paths=_data_paths(variable_catalog, component.alias),
    )


def n_per_mm_to_n_per_m(value: float) -> float:
    return value * N_PER_MM_TO_N_PER_M


def n_per_mm_per_s_to_n_per_m_per_s(value: float) -> float:
    return value * N_PER_MM_PER_S_TO_N_PER_M_PER_S


def _initial_value(amesim_results: AmesimResults | None, data_path: str) -> float:
    if amesim_results is None:
        return 0.0
    return float(amesim_results.series(data_path)[0])


def _data_paths(
    variable_catalog: TestMqlVariableCatalog | None,
    alias: str,
) -> tuple[str, ...]:
    if variable_catalog is None:
        return ()
    return variable_catalog.data_paths_for_owner(alias)