"""Robot hardware configuration — loaded from robot.yaml next to the URDF.

Separates robot hardware (actuators, joint tuning) from task config
(rewards, episode length) and from the URDF (clean CAD export).

Usage:
    robot = load_robot_config(Path("assets/rotary_cartpole"))
    # robot.urdf_path  → resolved absolute path to the URDF
    # robot.actuators  → list of ActuatorConfig
    # robot.joints     → dict of per-joint overrides
"""

import dataclasses
import math
from pathlib import Path

import structlog
import torch
import yaml

log = structlog.get_logger()


def _as_pair(val) -> tuple[float, float]:
    """Convert scalar or [pos, neg] list to (pos, neg) tuple."""
    if isinstance(val, (list, tuple)) and len(val) == 2:
        return (float(val[0]), float(val[1]))
    return (float(val), float(val))


@dataclasses.dataclass
class ActuatorConfig:
    """Motor/actuator attached to a joint.

    Asymmetric fields use (positive_dir, negative_dir) tuples.
    A scalar in YAML is expanded to a symmetric pair.

    type:
        motor    — direct torque control (ctrl = normalised torque)
        position — PD position servo  (ctrl = target angle, needs kp)
        velocity — P velocity servo   (ctrl = target velocity, needs kp)
    """
    joint: str = ""
    type: str = "motor"
    gear: tuple[float, float] = (1.0, 1.0)            # torque constant (pos, neg)
    ctrl_range: tuple[float, float] = (-1.0, 1.0)     # (lower, upper) control bounds
    deadzone: tuple[float, float] = (0.0, 0.0)        # min |ctrl| for torque (pos, neg)
    damping: tuple[float, float] = (0.0, 0.0)         # viscous damping (pos, neg)
    frictionloss: tuple[float, float] = (0.0, 0.0)    # Coulomb friction (pos, neg)
    kp: float = 0.0            # proportional gain (position / velocity actuators)
    kv: float = 0.0            # derivative gain  (position actuators)
    filter_tau: float = 0.0    # 1st-order filter time constant (s); 0 = no filter
    viscous_quadratic: float = 0.0   # velocity² drag coefficient
    back_emf_gain: float = 0.0      # back-EMF torque reduction
    stribeck_friction_boost: float = 0.0  # extra static friction at low speed (N·m)
    stribeck_vel: float = 2.0             # Stribeck decay velocity (rad/s)
    action_bias: float = 0.0              # additive ctrl bias (driver asymmetry)

    @property
    def gear_avg(self) -> float:
        return (self.gear[0] + self.gear[1]) / 2.0

    @property
    def has_motor_model(self) -> bool:
        """True if this actuator needs the runtime motor model."""
        return (
            self.gear[0] != self.gear[1]
            or self.deadzone != (0.0, 0.0)
            or self.damping != (0.0, 0.0)
            or self.frictionloss != (0.0, 0.0)
            or self.viscous_quadratic > 0
            or self.back_emf_gain > 0
            or self.stribeck_friction_boost > 0
            or self.action_bias != 0.0
        )

    def transform_ctrl(self, ctrl: float) -> float:
        """Clip to ctrl_range, then apply bias, deadzone and gear compensation.

        Must stay in lock-step with the vectorised JAX version in
        ``src/runners/mjx.py`` (step_fn) — sysid fits parameters against
        THIS function, so any drift breaks the identified model.
        """
        # Clip to ctrl_range first (mirrors firmware PWM saturation).
        ctrl = max(self.ctrl_range[0], min(self.ctrl_range[1], ctrl))

        # Additive driver bias (e.g. H-bridge asymmetry).
        ctrl += self.action_bias

        # Deadzone
        dz_pos, dz_neg = self.deadzone
        if ctrl >= 0 and ctrl < dz_pos:
            return 0.0
        if ctrl < 0 and ctrl > -dz_neg:
            return 0.0

        # Gear compensation: rescale so ctrl × gear_avg ≈ action × gear_dir
        gear_avg = self.gear_avg
        if gear_avg > 1e-8:
            gear_dir = self.gear[0] if ctrl >= 0 else self.gear[1]
            ctrl *= gear_dir / gear_avg

        return ctrl

    def compute_motor_force(self, vel: float, ctrl: float,
                            friction_scale: float = 1.0,
                            damping_scale: float = 1.0) -> float:
        """Asymmetric friction (Coulomb + Stribeck), damping, drag, back-EMF.

        ``friction_scale`` / ``damping_scale`` multiply the friction and
        viscous-damping terms for per-env domain randomization
        (1.0 = no randomization, the default used by sysid).
        """
        torque = 0.0

        # Coulomb + Stribeck friction (direction-dependent).  The Stribeck
        # boost adds extra friction at low speed that decays as exp(-(v/vs)²)
        # — crucial for cheap brushed motors near standstill.
        fl_pos, fl_neg = self.frictionloss
        if abs(vel) > 1e-6:
            fl = fl_pos if vel > 0 else fl_neg
            if self.stribeck_friction_boost > 0:
                fl += self.stribeck_friction_boost * math.exp(
                    -((abs(vel) / self.stribeck_vel) ** 2)
                )
            torque -= math.copysign(fl * friction_scale, vel)

        # Viscous damping (direction-dependent)
        damp = (self.damping[0] if vel > 0 else self.damping[1]) * damping_scale
        torque -= damp * vel

        # Quadratic velocity drag
        if self.viscous_quadratic > 0:
            torque -= self.viscous_quadratic * vel * abs(vel)

        # Back-EMF torque reduction
        if self.back_emf_gain > 0 and abs(ctrl) > 1e-6:
            torque -= self.back_emf_gain * vel * math.copysign(1.0, ctrl)

        return max(-10.0, min(10.0, torque))

    def transform_action(self, action):
        """Vectorised clip + bias + deadzone + gear compensation (torch batch).

        Must produce the same result as ``transform_ctrl`` element-wise.
        """
        action = action.clamp(self.ctrl_range[0], self.ctrl_range[1])
        action = action + self.action_bias

        dz_pos, dz_neg = self.deadzone
        if dz_pos > 0 or dz_neg > 0:
            pos_dead = (action >= 0) & (action < dz_pos)
            neg_dead = (action < 0) & (action > -dz_neg)
            action = action.masked_fill(pos_dead | neg_dead, 0.0)

        gear_avg = self.gear_avg
        if gear_avg > 1e-8 and self.gear[0] != self.gear[1]:
            pos = action >= 0
            action = torch.where(
                pos, action * (self.gear[0] / gear_avg),
                action * (self.gear[1] / gear_avg),
            )

        return action


@dataclasses.dataclass
class JointConfig:
    """Per-joint overrides applied on top of the URDF values."""
    damping: float | None = None
    armature: float | None = None      # reflected rotor inertia (kg·m²)
    frictionloss: float | None = None  # Coulomb/dry friction torque (N·m)


@dataclasses.dataclass
class RobotConfig:
    """Complete robot hardware description."""
    urdf_path: Path = dataclasses.field(default_factory=lambda: Path())
    actuators: list[ActuatorConfig] = dataclasses.field(default_factory=list)
    joints: dict[str, JointConfig] = dataclasses.field(default_factory=dict)


def load_robot_config(robot_dir: str | Path) -> RobotConfig:
    """Load robot.yaml from a directory and resolve the URDF path.

    Expected layout:
        robot_dir/
            robot.yaml          ← hardware config
            some_robot.urdf     ← CAD export
            meshes/             ← optional mesh files
    """
    robot_dir = Path(robot_dir).resolve()
    yaml_path = robot_dir / "robot.yaml"

    if not yaml_path.exists():
        raise FileNotFoundError(f"Robot config not found: {yaml_path}")

    raw = yaml.safe_load(yaml_path.read_text())

    # Resolve URDF path relative to robot.yaml directory
    urdf_filename = raw.get("urdf", "")
    if not urdf_filename:
        raise ValueError(f"robot.yaml must specify 'urdf' filename: {yaml_path}")
    urdf_path = robot_dir / urdf_filename
    if not urdf_path.exists():
        raise FileNotFoundError(f"URDF not found: {urdf_path}")

    # Parse actuators — ignore unknown keys (newer sysid exports may add
    # fields before the loader learns about them) instead of crashing.
    known_fields = {f.name for f in dataclasses.fields(ActuatorConfig)}
    actuators = []
    for a in raw.get("actuators", []):
        unknown = set(a) - known_fields
        if unknown:
            log.warning(
                "robot_yaml_unknown_actuator_keys",
                keys=sorted(unknown), file=str(yaml_path),
            )
            a = {k: v for k, v in a.items() if k in known_fields}
        if "ctrl_range" in a:
            a["ctrl_range"] = tuple(a["ctrl_range"])
        for key in ("gear", "deadzone", "damping", "frictionloss"):
            if key in a:
                a[key] = _as_pair(a[key])
        actuators.append(ActuatorConfig(**a))

    # Parse joint overrides
    joints = {}
    for name, jcfg in raw.get("joints", {}).items():
        joints[name] = JointConfig(**jcfg)

    config = RobotConfig(
        urdf_path=urdf_path,
        actuators=actuators,
        joints=joints,
    )

    log.debug("robot_config_loaded", robot_dir=str(robot_dir),
              urdf=urdf_filename, num_actuators=len(actuators),
              joint_overrides=list(joints.keys()))

    return config
