"""Task registry and dual Actor observation schema checks (no mjlab needed)."""

import math

import torch

from s10_supermimic.core.command_curriculum import stages_for
from s10_supermimic.core.policy import Actor, LEGACY_FRAME_DIM, FRAME_DIM, frame, push_history
from s10_supermimic.core.wheel_intrusion import pair_body_intrusion, wheel_body_intrusion
from s10_supermimic.tasks import get_task, task_names


def main():
  assert task_names() == ("s10-v4", "s10-v4-upward", "s10-v7"), task_names()
  v4, v7, upward = get_task("s10-v4"), get_task("s10-v7"), get_task("s10-v4-upward")
  assert (v4.terrain_revision, v7.terrain_revision) == (4, 7)
  assert upward.terrain_revision == 4
  assert upward.actor.frame_dim == LEGACY_FRAME_DIM
  assert upward.actor.wheel_vel
  assert upward.courses is False and upward.collisions is False
  assert upward.defaults["iterations"] == 30000
  assert upward.defaults["action_noise"] == 0.22
  assert upward.defaults["horizon"] == 64
  assert (v4.actor.frame_dim, v7.actor.frame_dim) == (LEGACY_FRAME_DIM, LEGACY_FRAME_DIM)
  assert (v4.actor.history_dim(5), v7.actor.history_dim(5)) == (265, 265)
  assert v4.actor.wheel_vel and v7.actor.wheel_vel
  assert v7.actor.name == "s10-v7-53"
  assert v4.courses is False and v4.collisions is False
  assert v7.courses and v7.collisions
  assert v7.defaults["imag_batch"] == 2048, "v7 keeps the 16 GB batch-size exception"
  assert v7.defaults["iterations"] == 40000
  assert v7.defaults["command_curriculum"] is True, "v7 uses stair-gated stages capped at 1.5/0.8/1.2"
  assert stages_for(7)[-1][1:] == ((-1.5, 1.5), (-.8, .8), (-1.2, 1.2))
  assert stages_for(4)[-1][1:] == ((-2., 2.), (-1., 1.), (-3., 3.))
  assert v4.defaults["iterations"] == 20000

  state = torch.zeros(2, 38)
  action = torch.zeros(2, 16)
  command = torch.zeros(2, 3)
  bias = torch.zeros(2, 12)
  for history, terms, expected in ((5, v4.actor.terms, 265), (5, v7.actor.terms, 265)):
    actor = Actor(history, actor_wheel_vel=True)
    latest = frame(state, action, command, bias,
                   wheel_vel=actor.actor_wheel_vel)
    assert latest.shape == (2, actor.frame_dim)
    assert push_history(torch.zeros(2, expected), latest, history, terms).shape == (2, expected)
  legacy = Actor(5, actor_wheel_vel=False)
  assert legacy.frame_dim == FRAME_DIM, "legacy 49-value checkpoints remain loadable"

  # v7 safe stance: default hipx 0.05 / hipy 25 deg / knee 1.0 has zero
  # intrusion; a wheel tucked under the body crosses the 1/4 boundary.
  state = torch.zeros(1, 38)
  state[0, 0], state[0, 3] = 0.44, -1.0
  def default_q(hipy, knee=1.0):
    return torch.tensor([0.05, -hipy, knee, -0.05, -hipy, knee,
                         0.05, hipy, -knee, -0.05, hipy, -knee])
  assert wheel_body_intrusion(state, default_q(math.radians(25.0))).item() == 0.0
  # 0.8 rad splay puts the wheel near x=0.134 m: inside the new outer-quarter
  # boundary (0.15 m), outside the old outer-third boundary (0.10 m).
  assert wheel_body_intrusion(state, default_q(.8)).item() > 0.0
  normal = wheel_body_intrusion(state, default_q(.5)).item()
  tucked = wheel_body_intrusion(state, default_q(1.0)).item()
  assert normal == 0.0 < tucked, (normal, tucked)

  # Pair 5/6 restriction: zero during the 0.5 s grace, then active only when
  # both same-side legs are inside. A single leg never triggers it.
  timers = torch.zeros(1, 4)
  for _ in range(24):
    cost, timers = pair_body_intrusion(state, default_q(1.0), timers, .02)
    assert cost.item() == 0.0
  cost, timers = pair_body_intrusion(state, default_q(1.0), timers, .02)
  assert cost.item() > 0.0, cost
  # Front-left alone stays inside; the pair timer stays unarmed.
  one_leg = default_q(math.radians(25.0))
  one_leg[1] = -1.0
  timers = torch.zeros(1, 4)
  for _ in range(40):
    cost, timers = pair_body_intrusion(state, one_leg, timers, .02)
    assert cost.item() == 0.0
  assert timers[0, 1].item() == 0.0, timers
  # Leaving the region resets that wheel's timer.
  timers = torch.zeros(1, 4)
  for _ in range(30):
    _, timers = pair_body_intrusion(state, default_q(1.0), timers, .02)
  cost, timers = pair_body_intrusion(state, default_q(math.radians(25.0)), timers, .02)
  assert cost.item() == 0.0 and timers[0, 0].item() == 0.0, (cost, timers)
  print("Task registry / Actor schema / v7 safety checks passed")


if __name__ == "__main__":
  main()