"""Terrain-gated rolling preference, shared by real and imagined rewards."""
import torch
from .world_model import GRAVITY, LEG_POS, LEG_VEL


def rolling_gate(scan, command):
  local = scan.reshape(-1, 21, 13)[:, 5:16, 3:10].flatten(1)
  relief = local.amax(-1) - local.amin(-1)
  # Flat (<2 cm): favor rolling; near obstacles (>=6 cm): allow lifting.
  flat = (1. - (relief - .02) / .04).clamp(0., 1.)
  # Non-steerable wheels may require stepping for lateral/yaw maneuvers.
  straight = torch.exp(-(command[:,1]/.25).square() - (command[:,2]/1.).square())
  return flat * straight


def wheel_geometry(state, default_q):
  """S10 wheel centres in body coordinates and vertical cylinder half-extents.

  Joint order is hipx (abduction), hipy (flexion), knee. Geometry follows
  ``S10_description/s10_mjcf/S10.xml``: hipx at x=±0.2277/y=±0.0574, hipy
  offset y=±0.0779, knee at z=-0.18, wheel at y=±0.045891/z=-0.18,
  wheel radius 0.081 m and half-width 0.0185 m. The analytic result was
  checked against MuJoCo ``mj_forward`` to machine precision.

  The cylinder is symmetric about its axle, so wheel phase is unnecessary.
  """
  q = (state[:,LEG_POS] + default_q).reshape(-1,4,3)
  hipx, hipy, knee = q.unbind(-1)
  side = state.new_tensor([1.,-1.,1.,-1.])
  # Leg sagittal plane: knee link (0.18 m) + wheel link (0.18 m).
  x = .18*knee.sin()
  z = -.18 - .18*knee.cos()
  x3 = hipy.cos()*x - hipy.sin()*z
  z3 = hipy.sin()*x + hipy.cos()*z
  y3 = side*(.0779+.045891)
  y4 = hipx.cos()*y3 + hipx.sin()*z3
  z4 = -hipx.sin()*y3 + hipx.cos()*z3
  point = torch.stack((state.new_tensor([.2277,.2277,-.2277,-.2277]) + x3,
                       state.new_tensor([.0574,-.0574,.0574,-.0574]) + y4,
                       z4),-1)
  up = -torch.nn.functional.normalize(state[:,GRAVITY],dim=-1,eps=1e-6)
  # Wheel axle is local -y, rotated only by hipx; sign is irrelevant after abs.
  axle_z = -up[:,1:2]*hipx.cos()+up[:,2:3]*hipx.sin()
  radius_z = .081*(1.-axle_z.square()).clamp_min(1e-8).sqrt()+.0185*axle_z.abs()
  return point, radius_z


def wheel_bottom_height(state, default_q):
  """Wheel bottoms relative to the ground below the base.

  Only used on near-flat ground by rolling_gate, not as obstacle clearance.
  """
  point, radius_z = wheel_geometry(state, default_q)
  up = -torch.nn.functional.normalize(state[:,GRAVITY],dim=-1,eps=1e-6)
  return state[:,:1] + (point*up[:,None,:]).sum(-1) - radius_z


def rolling_costs(state, command, scan, default_q):
  gate = rolling_gate(scan, command)
  # Tolerate 5 mm of clearance; no positive reward for lifting anywhere.
  lift = ((wheel_bottom_height(state,default_q)-.005).clamp_min(0)/.05).square().sum(-1)
  # Reference mjlab task uses mean squared speed and relaxes during recovery.
  tilt = (state[:,GRAVITY][:,:2].norm(dim=-1)-.08).div(.22).clamp(0.,1.)
  motion = state[:,LEG_VEL].square().mean(-1) * (1.-.8*tilt)
  return {'unnecessary_wheel_lift':gate*lift, 'rolling_leg_motion':gate*motion}


def env_rolling_cost(env, term):
  from .terrain import TerrainField, TerrainBridge
  bridge = getattr(env,'_s10_posture_bridge',None)
  if bridge is None:
    bridge = TerrainBridge(env,TerrainField.from_env(env))
    env._s10_posture_bridge = bridge
  state = bridge.state()
  return rolling_costs(state, env.command_manager.get_command('twist'),
                       bridge.field.scan(state,bridge.pose()),
                       bridge.robot.data.default_joint_pos[0,bridge.legs])[term]