from __future__ import annotations

from dataclasses import dataclass
from math import pi


MM_TO_M = 1.0e-3
M_TO_MM = 1.0e3
M3_TO_CM3 = 1.0e6
M3_PER_S_TO_L_PER_MIN = 60_000.0


def circular_area(diameter_m: float) -> float:
    if diameter_m < 0.0:
        raise ValueError("diameter_m must be non-negative.")
    return pi * diameter_m * diameter_m / 4.0


def mm_to_m(value: float) -> float:
    return value * MM_TO_M


def m_to_mm(value: float) -> float:
    return value * M_TO_MM


@dataclass(frozen=True)
class AmesimPistonGeometry:
    """Geometry relations used by AMESim PNRP17 pneumatic piston variables."""

    piston_diameter_m: float
    rod_diameter_m: float = 0.0
    zero_length_m: float = 0.0

    @property
    def piston_area_m2(self) -> float:
        return circular_area(self.piston_diameter_m)

    @property
    def rod_area_m2(self) -> float:
        return circular_area(self.rod_diameter_m)

    @property
    def annulus_area_m2(self) -> float:
        return self.piston_area_m2 - self.rod_area_m2

    def chamber_length_m(self, port4_displacement_m: float, port5_displacement_m: float) -> float:
        return self.zero_length_m + port5_displacement_m - port4_displacement_m

    def chamber_length_mm(self, port4_displacement_m: float, port5_displacement_m: float) -> float:
        return m_to_mm(self.chamber_length_m(port4_displacement_m, port5_displacement_m))

    @property
    def chamber_area_m2(self) -> float:
        return self.annulus_area_m2

    def chamber_volume_m3(self, port4_displacement_m: float, port5_displacement_m: float) -> float:
        return self.chamber_area_m2 * self.chamber_length_m(
            port4_displacement_m,
            port5_displacement_m,
        )

    def chamber_volume_cm3(self, port4_displacement_m: float, port5_displacement_m: float) -> float:
        return self.chamber_volume_m3(port4_displacement_m, port5_displacement_m) * M3_TO_CM3

    def chamber_volume_rate_m3_s(self, port4_velocity_m_s: float, port5_velocity_m_s: float) -> float:
        return self.chamber_area_m2 * (port5_velocity_m_s - port4_velocity_m_s)

    def chamber_volume_rate_l_min(self, port4_velocity_m_s: float, port5_velocity_m_s: float) -> float:
        return self.chamber_volume_rate_m3_s(
            port4_velocity_m_s,
            port5_velocity_m_s,
        ) * M3_PER_S_TO_L_PER_MIN


@dataclass(frozen=True)
class AmesimElasticEndstop:
    """Contact force part of AMESim LSTP00A elastic endstop."""

    contact_stiffness_n_per_m: float
    contact_damping_n_per_m_per_s: float = 0.0
    gap0_m: float = 0.0

    def penetration_m_from_gap_mm(self, gap_mm: float) -> float:
        return max(-(mm_to_m(gap_mm) - self.gap0_m), 0.0)

    def static_contact_force(self, gap_mm: float) -> float:
        return self.contact_stiffness_n_per_m * self.penetration_m_from_gap_mm(gap_mm)

    def contact_force(self, gap_mm: float, penetration_velocity_m_s: float = 0.0) -> float:
        if self.penetration_m_from_gap_mm(gap_mm) <= 0.0:
            return 0.0
        damping_force = self.contact_damping_n_per_m_per_s * penetration_velocity_m_s
        return max(self.static_contact_force(gap_mm) + damping_force, 0.0)


@dataclass(frozen=True)
class AmesimMassFrictionEndstops:
    """Parameter and observable helpers for AMESim MECMAS21 translation masses."""

    mass_kg: float
    lower_limit_m: float
    upper_limit_m: float
    lower_stiffness_n_per_m: float
    upper_stiffness_n_per_m: float
    lower_damping_n_per_m_per_s: float = 0.0
    upper_damping_n_per_m_per_s: float = 0.0
    viscous_friction_n_per_m_per_s: float = 0.0
    coulomb_friction_n: float = 0.0
    stiction_force_n: float = 0.0
    windage_n_per_m2_per_s2: float = 0.0

    def lower_penetration_m(self, displacement_m: float) -> float:
        return max(self.lower_limit_m - displacement_m, 0.0)

    def upper_penetration_m(self, displacement_m: float) -> float:
        return max(displacement_m - self.upper_limit_m, 0.0)

    def lower_static_force_magnitude(self, displacement_m: float) -> float:
        return self.lower_stiffness_n_per_m * self.lower_penetration_m(displacement_m)

    def upper_static_force_magnitude(self, displacement_m: float) -> float:
        return self.upper_stiffness_n_per_m * self.upper_penetration_m(displacement_m)

    def viscous_friction_force(self, velocity_m_s: float) -> float:
        return -self.viscous_friction_n_per_m_per_s * velocity_m_s

    def windage_force(self, velocity_m_s: float) -> float:
        return -self.windage_n_per_m2_per_s2 * velocity_m_s * abs(velocity_m_s)

    def dry_friction_force(self, velocity_m_s: float) -> float:
        if velocity_m_s > 0.0:
            return -self.coulomb_friction_n
        if velocity_m_s < 0.0:
            return self.coulomb_friction_n
        return 0.0

    def limit_contact_force(self, displacement_m: float, velocity_m_s: float) -> float:
        lower_force = self.lower_static_force_magnitude(displacement_m)
        if lower_force > 0.0:
            lower_force += max(-self.lower_damping_n_per_m_per_s * velocity_m_s, 0.0)

        upper_force = self.upper_static_force_magnitude(displacement_m)
        if upper_force > 0.0:
            upper_force += max(self.upper_damping_n_per_m_per_s * velocity_m_s, 0.0)

        return lower_force - upper_force

    def derivatives(
        self,
        *,
        velocity_m_s: float,
        displacement_m: float,
        port_1_force_n: float = 0.0,
        port_2_force_n: float = 0.0,
        external_force_n: float = 0.0,
    ) -> tuple[float, float]:
        total_force = (
            port_1_force_n
            + port_2_force_n
            + external_force_n
            + self.viscous_friction_force(velocity_m_s)
            + self.windage_force(velocity_m_s)
            + self.dry_friction_force(velocity_m_s)
            + self.limit_contact_force(displacement_m, velocity_m_s)
        )
        return total_force / self.mass_kg, velocity_m_s