TIR_PROJ/tools/gen_mecanum_wheels.py

"""Generate ShepherdDogMecanum.proto wheel blocks with physical rollers.

Each wheel becomes:
  HingeJoint (motor, axis 0 1 0 = body lateral)
    -> Solid (wheel hub, rotation 0 -1 0 π/2)
       children:
         - WHEEL_VIS (visual, kept as-is for appearance)
         - 8x HingeJoint (passive roller, axis tilted ±45° from wheel rotation
                          axis, tangent to the wheel circumference at the mount
                          point)
             -> Solid (capsule)
       boundingObject: a small Cylinder for the hub (smaller radius than the
                       roller circle so the hub doesn't touch the ground)

X-pattern roller tilt assignment:
  FR, RL -> -45° (wheel-axis-relative)
  FL, RR -> +45°

All math is done in the WHEEL SOLID's local frame. The wheel solid's
rotation `0 -1 0 π/2` takes wheel-local x -> body +z (up),
wheel-local y -> body +y (lateral, = wheel rotation axis),
wheel-local z -> body -x (rearward). Conversely, a body-frame offset
(dx, dy, dz) becomes (dz, dy, -dx) in wheel-local coords.

For a wheel rotating about body y at angle θ (θ=0 = body +x = forward,
θ=π/2 = body +z = top), the roller mount in body frame is
(R*cos(θ), 0, R*sin(θ)) relative to wheel centre. Tangent (radial-perp,
in the wheel-spin plane) is (-sin(θ), 0, cos(θ)); the wheel rotation
axis is (0, 1, 0). Roller axis tilted +45° from tangent toward wheel
axis:
    axis_body(+45°) = (1/√2) * (-sin(θ), +1, cos(θ))
    axis_body(-45°) = (1/√2) * (-sin(θ), -1, cos(θ))

Transformed to wheel-local: (dz, dy, -dx) on each component gives
    mount_local = (R*sin(θ), 0, -R*cos(θ))
    axis_local(+45) = (cos(θ)/√2,  +1/√2, sin(θ)/√2)
    axis_local(-45) = (cos(θ)/√2,  -1/√2, sin(θ)/√2)

The Solid's `rotation` field needs to align the Capsule's default
axis (+y) with that local axis. The minimal axis-angle that does this:
    rotation_axis = (sin(θ), 0, -cos(θ))   (unit)
    rotation_angle = π/4 for +45° tilt, 3π/4 for -45° tilt
"""
import math

WHEEL_NAMES = {
    # Tilt sign refers to roller-axis tilt direction relative to the wheel
    # rotation axis (body +y). X-pattern requires rollers on each wheel to
    # tilt INWARD toward the body centre. For a wheel at +y body coord, that
    # means tilting toward -y; for a wheel at -y, tilting toward +y.
    "fr": ("front right",  +0.14, -0.14, +1),  # +1 = +45° tilt (toward +y, inward)
    "fl": ("front left",   +0.14, +0.14, -1),  # -1 = -45° tilt (toward -y, inward)
    "rr": ("rear right",   -0.14, -0.14, -1),  # -1 (toward -y, "outward"...
    "rl": ("rear left",    -0.14, +0.14, +1),  # +1 (toward +y, "outward"...
    # ...for the rear pair the X-pattern flips so diagonal pairs FL+RR have
    # SAME tilt direction in body frame, FR+RL the other. The signs above
    # encode that: FR/RL both +1, FL/RR both -1.
}

R_ROLLER_OFFSET = 0.031   # roller-centre distance from wheel hub centre
R_ROLLER_RADIUS = 0.007
R_ROLLER_HEIGHT = 0.020
ROLLER_MASS = 0.003
HUB_RADIUS = 0.020        # < R_ROLLER_OFFSET - R_ROLLER_RADIUS so hub doesn't touch
HUB_HEIGHT = 0.022
HUB_MASS = 0.045
N_ROLLERS = 8


def wheel_block(key):
    name, ax, ay, tilt_sign = WHEEL_NAMES[key]
    contact_mat = "MecanumWheelA" if tilt_sign > 0 else "MecanumWheelB"
    safe = name.replace(" ", "_").upper()

    rollers = []
    for k in range(N_ROLLERS):
        theta = 2.0 * math.pi * k / N_ROLLERS
        s, c = math.sin(theta), math.cos(theta)
        # Mount position in wheel-local frame.
        mx = R_ROLLER_OFFSET * s
        my = 0.0
        mz = -R_ROLLER_OFFSET * c
        # Hinge axis in wheel-local frame.
        ax_l = c / math.sqrt(2.0)
        ay_l = tilt_sign / math.sqrt(2.0)
        az_l = s / math.sqrt(2.0)
        # Rotation that maps Capsule default axis (0,1,0) to (ax_l, ay_l, az_l).
        rot_axis = (s, 0.0, -c)
        rot_angle = math.pi / 4.0 if tilt_sign > 0 else 3.0 * math.pi / 4.0
        rollers.append(f"""\
        # Mecanum roller {k+1} (θ={math.degrees(theta):.0f}°)
        HingeJoint {{
          jointParameters HingeJointParameters {{
            axis {ax_l:.6f} {ay_l:.6f} {az_l:.6f}
            anchor {mx:.6f} {my:.6f} {mz:.6f}
          }}
          endPoint Solid {{
            translation {mx:.6f} {my:.6f} {mz:.6f}
            rotation {rot_axis[0]:.6f} {rot_axis[1]:.6f} {rot_axis[2]:.6f} {rot_angle:.6f}
            children [
              Shape {{
                appearance PBRAppearance {{
                  baseColor 0.12 0.12 0.12
                  roughness 0.7
                  metalness 0.1
                }}
                geometry Capsule {{
                  height {R_ROLLER_HEIGHT}
                  radius {R_ROLLER_RADIUS}
                  subdivision 8
                }}
              }}
            ]
            name "{name} roller {k+1}"
            contactMaterial "{contact_mat}"
            boundingObject Capsule {{
              height {R_ROLLER_HEIGHT}
              radius {R_ROLLER_RADIUS}
              subdivision 8
            }}
            physics Physics {{
              density -1
              mass {ROLLER_MASS}
              centerOfMass [
                0 0 0
              ]
            }}
          }}
        }}""")
    rollers_str = "\n".join(rollers)

    return f"""\
      # ========== {name.upper()} WHEEL ==========
      DEF {safe}_WHEEL_JOINT HingeJoint {{
        jointParameters HingeJointParameters {{
          axis 0 1 0
          anchor {ax} {ay} 0.038
        }}
        device [
          RotationalMotor {{
            name "{name} wheel motor"
            maxVelocity 70.0
            maxTorque 20.0
          }}
          PositionSensor {{
            name "{name} wheel sensor"
            resolution 0.00628
          }}
        ]
        endPoint Solid {{
          translation {ax} {ay} 0.038
          rotation 0 -1 0 1.570796
          children [
            # Visual hub only — the rollers below provide ground contact.
            Pose {{
              rotation 1 0 0 -1.5708
              children [
                Shape {{
                  appearance PBRAppearance {{
                    baseColor 0.5 0.5 0.5
                    roughness 0.3
                    metalness 0.7
                  }}
                  geometry Cylinder {{
                    height 0.018
                    radius {HUB_RADIUS - 0.002}
                    subdivision 16
                  }}
                }}
                Shape {{
                  appearance PBRAppearance {{
                    baseColor 0.6 0.6 0.6
                    roughness 0.2
                    metalness 0.8
                  }}
                  geometry Cylinder {{
                    height 0.022
                    radius 0.008
                    subdivision 8
                  }}
                }}
              ]
            }}
{rollers_str}
          ]
          name "{name} wheel"
          boundingObject Pose {{
            rotation 1 0 0 -1.5708
            children [
              Cylinder {{
                height {HUB_HEIGHT}
                radius {HUB_RADIUS}
              }}
            ]
          }}
          physics Physics {{
            density -1
            mass {HUB_MASS}
            centerOfMass [
              0 0 0
            ]
          }}
        }}
      }}"""


if __name__ == "__main__":
    for k in ("fr", "fl", "rr", "rl"):
        print(wheel_block(k))
        print()