"""Always-on wheel/body projection constraints for the safe v7 stance.

The body is 0.60 m long in ``base_link`` (x in [-0.30, 0.30]).

``wheel_body_intrusion`` keeps every individual wheel out of the opposite
three quarters: a front wheel inside the rear three quarters (x < +0.15) or a
rear wheel inside the front three quarters (x > -0.15) shortens the support
polygon and makes the robot tip-prone. Its hinge is exactly zero outside the
outer quarters and grows cubically inside, so a normal extended stance is
untouched while a wheel tucked under the body is penalised quickly.

``pair_body_intrusion`` adds the wider five-sixths restriction requested for
v7: both front wheels may not sit inside the rear five sixths (x < +0.20) at
the same time, and both rear wheels may not sit inside the front five sixths
(x > -0.20). A single leg crossing the line is allowed; a leg only counts as
inside after it has stayed inside continuously for ``PAIR_GRACE_S`` (0.5 s).
The pair penalty is therefore zero during a brief swing-through and applies
immediately once both same-side legs are armed.

Both terms are kinematic and need no terrain scan. Wheels are ordered
fl, fr, hl, hr.
"""
import torch

from .rolling_rewards import wheel_geometry

BODY_HALF_LENGTH = 0.30
# Each wheel must stay in the outer quarter of the body length: x > +0.15 for
# front wheels and x < -0.15 for rear wheels.
WHEEL_STRIP_BOUNDARY = BODY_HALF_LENGTH / 2.0
# Rear/front 5/6 of a 0.60 m body => the boundary sits at x = +-0.20.
PAIR_BOUNDARY = 2.0 * BODY_HALF_LENGTH / 3.0
PAIR_GRACE_S = 0.5


def wheel_body_intrusion(state, default_q, boundary=WHEEL_STRIP_BOUNDARY):
  point, _ = wheel_geometry(state, default_q)
  x = point[..., 0]
  # Front wheels must stay ahead of +L/4, rear wheels behind -L/4.
  violation = torch.cat(((boundary - x[:, :2]).clamp_min(0.),
                         (x[:, 2:] + boundary).clamp_min(0.)), -1)
  normalized = violation / boundary
  return normalized.pow(3).mean(-1)


def pair_region(state, default_q, boundary=PAIR_BOUNDARY):
  """Per-wheel inside mask and normalised depth for the 5/6 restriction."""
  point, _ = wheel_geometry(state, default_q)
  x = point[..., 0]
  inside = torch.cat((x[:, :2] < boundary, x[:, 2:] > -boundary), -1)
  depth = torch.cat(((boundary - x[:, :2]).clamp_min(0.),
                     (x[:, 2:] + boundary).clamp_min(0.)), -1) / boundary
  return inside, depth


def pair_body_intrusion(state, default_q, timers, dt, grace=PAIR_GRACE_S):
  """Penalise both front or both rear wheels staying inside the 5/6 region.

  ``timers`` holds the continuous inside time of each wheel. Returning it lets
  the caller keep the 0.5 s grace across environment steps. A wheel outside
  the boundary resets its own timer, so a single-leg pass-through never
  contributes.
  """
  inside, depth = pair_region(state, default_q)
  timers = torch.where(inside, timers + dt, torch.zeros_like(timers))
  armed = timers >= grace - 1e-6
  front = armed[:, 0] & armed[:, 1]
  rear = armed[:, 2] & armed[:, 3]
  cost = (front.float() * depth[:, :2].mean(-1)
          + rear.float() * depth[:, 2:].mean(-1))
  return cost, timers


def _bridge(env):
  from .terrain import TerrainBridge, TerrainField
  bridge = getattr(env, '_s10_posture_bridge', None)
  if bridge is None:
    bridge = TerrainBridge(env, TerrainField.from_env(env))
    env._s10_posture_bridge = bridge
  return bridge


def env_wheel_body_intrusion(env):
  bridge = _bridge(env)
  return wheel_body_intrusion(
    bridge.state(), bridge.robot.data.default_joint_pos[0, bridge.legs])


def env_pair_body_intrusion(env):
  bridge = _bridge(env)
  timers = getattr(env, '_s10_pair_timer', None)
  if timers is None or timers.shape[0] != env.num_envs:
    timers = torch.zeros(env.num_envs, 4, device=env.device)
  # episode_length_buf is 1 on the first reward step after an auto-reset.
  timers = timers.clone()
  timers[env.episode_length_buf <= 1] = 0.
  cost, timers = pair_body_intrusion(
    bridge.state(), bridge.robot.data.default_joint_pos[0, bridge.legs],
    timers, env.step_dt)
  env._s10_pair_timer = timers
  return cost