"""Six-bar (icosahedral) tensegrity robot model for MuJoCo.

Each of the 24 cables is a spatial tendon driven by a tension-only position
actuator: force = stiffness * (ctrl - length), clipped to <= 0 (cables can only
pull). So ctrl[i] is the cable's rest length and lowering it = contracting it.
"""
from __future__ import annotations

import itertools
import math

import mujoco
import numpy as np

GOLDEN_RATIO = (1 + math.sqrt(5)) / 2


def node_positions(bar_length: float) -> np.ndarray:
    """12 vertices of an icosahedron: (0, ±1, ±φ) and its cyclic permutations."""
    scale = bar_length / (2 * GOLDEN_RATIO)
    nodes = []
    for axis in range(3):  # bars parallel to z, x, y in turn
        for side in (1, -1):
            for end in (1, -1):
                point = [0.0, side * 1.0, end * GOLDEN_RATIO]
                nodes.append(np.roll(point, axis) * scale)
    return np.array(nodes)


def bar_pairs() -> list[tuple[int, int]]:
    """Node indices joined by each of the 6 rigid bars."""
    return [(2 * bar, 2 * bar + 1) for bar in range(6)]


def cable_pairs(nodes: np.ndarray) -> list[tuple[int, int]]:
    """24 cables = icosahedron edges minus the 6 edges between parallel bars."""
    bar_of = {node: bar for bar, pair in enumerate(bar_pairs()) for node in pair}
    distances = {
        (i, j): np.linalg.norm(nodes[i] - nodes[j])
        for i, j in itertools.combinations(range(len(nodes)), 2)
    }
    edge_length = min(distances.values())
    cables = []
    for (i, j), distance in distances.items():
        same_axis_pair = bar_of[i] // 2 == bar_of[j] // 2
        if abs(distance - edge_length) < 1e-6 and not same_axis_pair:
            cables.append((i, j))
    assert len(cables) == 24, len(cables)
    return cables


def build_xml(
    stiffness: float = 300.0,
    slope_deg: float = 0.0,
    bar_length: float = 1.0,
    bar_mass: float = 0.5,
    prestretch: float = 0.9,
    max_cable_force: float = 200.0,
    friction: float = 1.0,
) -> str:
    """Return MJCF text. The slope is modelled by tilting gravity: uphill is +x.

    Bars only collide with the ground (not with each other) to keep the
    simulation stable when strongly contracted cables pull bars close together.
    """
    nodes = node_positions(bar_length)
    cables = cable_pairs(nodes)
    lift = np.array([0, 0, -nodes[:, 2].min() + 0.03])  # drop from just above ground
    gravity_x = -9.81 * math.sin(math.radians(slope_deg))
    gravity_z = -9.81 * math.cos(math.radians(slope_deg))

    bodies = []
    for bar, (a, b) in enumerate(bar_pairs()):
        center = (nodes[a] + nodes[b]) / 2 + lift
        half_a, half_b = nodes[a] - (nodes[a] + nodes[b]) / 2, nodes[b] - (nodes[a] + nodes[b]) / 2
        fmt = lambda v: " ".join(f"{x:.5f}" for x in v)  # noqa: E731
        bodies.append(f"""
    <body name="bar{bar}" pos="{fmt(center)}">
      <freejoint/>
      <geom type="capsule" size="0.02" fromto="{fmt(half_a)} {fmt(half_b)}" mass="{bar_mass}" contype="1" conaffinity="2" rgba="0.85 0.55 0.2 1"/>
      <site name="n{a}" pos="{fmt(half_a)}" size="0.03" rgba="0.2 0.2 0.2 1"/>
      <site name="n{b}" pos="{fmt(half_b)}" size="0.03" rgba="0.2 0.2 0.2 1"/>
    </body>""")

    tendons, actuators = [], []
    for k, (i, j) in enumerate(cables):
        tendons.append(
            f'<spatial name="c{k}" width="0.006" rgba="0.2 0.5 0.9 1" damping="5">'
            f'<site site="n{i}"/><site site="n{j}"/></spatial>'
        )
        actuators.append(
            f'<position name="a{k}" tendon="c{k}" kp="{stiffness}" '
            f'ctrlrange="0.05 {bar_length}" forcerange="-{max_cable_force} 0"/>'
        )

    return f"""<mujoco model="six_bar_tensegrity">
  <option timestep="0.002" integrator="implicitfast" gravity="{gravity_x:.4f} 0 {gravity_z:.4f}"/>
  <default><geom friction="{friction} 0.01 0.001" condim="3"/></default>
  <worldbody>
    <light pos="0 0 4"/>
    <geom name="ground" type="plane" contype="2" conaffinity="1" size="20 20 0.1" rgba="0.8 0.8 0.8 1"/>{''.join(bodies)}
  </worldbody>
  <tendon>{''.join(tendons)}</tendon>
  <actuator>{''.join(actuators)}</actuator>
</mujoco>"""


def build_model(**kwargs) -> mujoco.MjModel:
    return mujoco.MjModel.from_xml_string(build_xml(**kwargs))