"""S10 asset adapter: twelve position-controlled legs and four velocity wheels.

The MJCF and default pose come from the Deep Robotics S10 deployment SDK
(``S10_description/s10_mjcf``). The copied XML has its torque-motor block
removed, so mjlab can add native position/velocity actuators. v4 keeps
the SDK/IsaacLab leg gains (kp=80, kd=2, effort=50 N·m); v7 uses a softer
kp=20 for deployment safety. Wheels use velocity damping 0.8, effort=14 N·m.
"""

import math
from pathlib import Path

import mujoco

from mjlab.actuator import (
  BuiltinPositionActuatorCfg,
  BuiltinVelocityActuatorCfg,
)
from mjlab.entity import EntityArticulationInfoCfg, EntityCfg
from mjlab.utils.spec_config import CollisionCfg

S10_XML = Path(__file__).parent / "model" / "s10.xml"
LEGS = ("fl", "fr", "hl", "hr")
LEG_JOINT_NAMES = tuple(
  f"{leg}_{joint}_joint" for leg in LEGS for joint in ("hipx", "hipy", "knee")
)
WHEEL_JOINT_NAMES = tuple(f"{leg}_wheel_joint" for leg in LEGS)
WHEEL_BODY_NAMES = tuple(f"{leg}_wheel" for leg in LEGS)
WHEEL_GEOM_NAMES = tuple(f"{leg}_wheel_collision" for leg in LEGS)
# Knee-joint housing, shin link and knee hub; v7 excludes these from the
# non-wheel contact termination so stair shin/knee grazing does not end runs.
KNEE_GEOM_NAMES = tuple(
  f"{leg}_{part}_collision" for leg in LEGS
  for part in ("knee_joint", "shin", "knee_hub"))

# v4 keeps the SDK neutral pose. v7 uses a more crouched, stair-friendly
# pose closer to the gait it needs (hipx 0.05, hipy ±25°, knee 1.0) and a
# softer leg actuator (kp=20) for safer deployment. Knee 1.0 keeps the wheel
# centres outside the body projection even with the ±25° splay.
DEFAULT_HIPY = 0.35
V7_HIPX = 0.05
V7_HIPY = math.radians(25.0)
V7_KNEE = 1.0
V7_LEG_KP = 20.0


def get_spec() -> mujoco.MjSpec:
  spec = mujoco.MjSpec.from_file(str(S10_XML))
  spec.modelname = "s10"
  # Scene solver options are owned by the task, not the attached robot.
  spec.option.cone = mujoco.mjtCone.mjCONE_PYRAMIDAL
  spec.option.impratio = 1.0
  # Name every collision so domain randomisation can select actual contacts.
  for i, geom in enumerate(spec.geoms):
    if geom.contype or geom.conaffinity:
      # Warp MULTICCD cannot handle nonzero margins on these cylinder/box pairs.
      geom.margin = 0.0
      if not geom.name:
        geom.name = f"body_{i}_collision"
  return spec


def get_s10_robot_cfg(hipy: float | None = None, hipx: float | None = None,
                      knee: float | None = None, leg_kp: float = 80.0,
                      leg_kd: float = 2.0) -> EntityCfg:
  """Return a fresh S10 config with mjlab native position/velocity actuators.

  ``hipx``, ``hipy`` and ``knee`` are absolute default joint angles. Omit them
  for the legacy v4 pose; v7 passes its crouched stair pose. ``leg_kp``/``kd``
  let v7 use a softer actuator without changing v4 checkpoints. The joint
  action offset and posture penalty both read ``default_joint_pos``, so these
  values move the action target and the penalised posture together.
  """
  hipx = 0.1 if hipx is None else float(hipx)
  hipy = DEFAULT_HIPY if hipy is None else float(hipy)
  knee = 0.65 if knee is None else float(knee)
  return EntityCfg(
    spec_fn=get_spec,
    init_state=EntityCfg.InitialStateCfg(
      # Default leg pose puts the wheel bottom at ~0.434 m; 0.44 m gives a
      # small margin before the reset z offset is applied.
      pos=(0.0, 0.0, 0.44),
      joint_pos={
        "fl_hipx_joint": hipx, "hl_hipx_joint": hipx,
        "fr_hipx_joint": -hipx, "hr_hipx_joint": -hipx,
        "fl_hipy_joint": -hipy, "fr_hipy_joint": -hipy,
        "hl_hipy_joint": hipy, "hr_hipy_joint": hipy,
        "fl_knee_joint": knee, "fr_knee_joint": knee,
        "hl_knee_joint": -knee, "hr_knee_joint": -knee,
        ".*_wheel_joint": 0.0,
      },
      joint_vel={".*": 0.0},
    ),
    collisions=(CollisionCfg(
      geom_names_expr=(".*_collision",), contype=1, conaffinity=1,
      condim={".*_wheel_collision": 3, "body_.*_collision": 1, ".*": 1},
      priority={".*_wheel_collision": 1, "body_.*_collision": 0, ".*": 0},
    ),),
    articulation=EntityArticulationInfoCfg(
      actuators=(
        # TODO(s10): the SDK/IsaacLab S10 config also randomizes actuator delay
        # (0-5 physics steps) and PD gains. v7 uses a softer kp=20 for safe
        # deployment; v4 keeps the SDK kp=80 recipe.
        BuiltinPositionActuatorCfg(
          target_names_expr=LEG_JOINT_NAMES,
          stiffness=float(leg_kp), damping=float(leg_kd), effort_limit=50.0,
        ),
        BuiltinVelocityActuatorCfg(
          target_names_expr=WHEEL_JOINT_NAMES,
          damping=0.8, effort_limit=14.0, armature=0.0005,
        ),
      ),
      soft_joint_pos_limit_factor=0.9,
    ),
  )