"""Pure-torch checks for the upward recovery reward and per-env curriculum."""

import math

import torch

from s10_supermimic.core.imagine import continuation_value
from s10_supermimic.core.upward import (
  BAD_ORIENTATION_START_DEG,
  HISTORY_LENGTH,
  PROMOTE_SUCCESSES,
  TILT_STAGES_DEG,
  bad_orientation_mask,
  bad_orientation_value,
  init_upward_state,
  reset_root_state_upward,
  upward_curriculum,
  upward_curriculum_stats,
  upward_video_tilt_limit,
)
from s10_supermimic.core.world_model import GRAVITY


def gravity_at(deg: float, direction: float = 0.0) -> torch.Tensor:
  """Projected gravity for a tilt ``deg`` in a fixed horizontal direction."""
  t = math.radians(deg)
  return torch.tensor([[
    math.sin(t) * math.cos(direction),
    math.sin(t) * math.sin(direction),
    -math.cos(t),
  ]])


def test_severity_is_monotonic() -> None:
  values = torch.cat([bad_orientation_value(gravity_at(d)) for d in
                      torch.arange(0.0, 180.5, 0.5).tolist()])
  assert (values[1:] >= values[:-1] - 1e-6).all(), "severity must be monotonic"
  assert values.max() <= 1.0 + 1e-5, values.max()
  assert values.min() >= 0.0, values.min()
  assert bad_orientation_value(gravity_at(BAD_ORIENTATION_START_DEG)).item() < 1e-6
  assert bad_orientation_value(gravity_at(180.0)).item() > 0.99
  v105 = bad_orientation_value(gravity_at(105.0)).item()
  v110 = bad_orientation_value(gravity_at(110.0)).item()
  v120 = bad_orientation_value(gravity_at(120.0)).item()
  assert v105 < v110 < v120, (v105, v110, v120)


def test_severity_gradient_is_finite() -> None:
  for g in (gravity_at(0.0), gravity_at(90.0), gravity_at(180.0)):
    x = g.clone().requires_grad_(True)
    bad_orientation_value(x).sum().backward()
    assert torch.isfinite(x.grad).all(), x.grad


def test_mask_keeps_bad_set() -> None:
  assert bad_orientation_mask(gravity_at(60.0)).item()
  assert not bad_orientation_mask(gravity_at(30.0)).item()
  assert not bad_orientation_mask(gravity_at(50.0, math.pi / 4)).item()
  assert bad_orientation_mask(gravity_at(180.0)).item()


def test_upward_continuation_does_not_use_v4_fall_proxy() -> None:
  state = torch.zeros(3, 38)
  state[0, GRAVITY] = torch.tensor([0.0, 0.0, -1.0])
  state[1, GRAVITY] = torch.tensor([0.0, 0.0, 0.0])
  state[2, GRAVITY] = torch.tensor([0.0, 0.0, 1.0])
  upward = continuation_value(state, {"upright_mode": "upward"})
  v4 = continuation_value(state, {})
  assert torch.allclose(upward, torch.ones(3)), upward
  assert v4[0] > 0.99 and v4[2] < 1e-4, v4


class _Data:
  def __init__(self, n: int):
    self.root_link_pos_w = torch.zeros(n, 3)
    self.projected_gravity_b = torch.tensor([[0.0, 0.0, -1.0]] * n)


class _Robot:
  def __init__(self, n: int):
    self.data = _Data(n)


class _Terrain:
  def __init__(self, n: int):
    self.terrain_levels = torch.zeros(n, dtype=torch.long)
    self.terrain_types = torch.zeros(n, dtype=torch.long)
    self.terrain_origins = torch.zeros(6, 1, 3)
    self.env_origins = torch.zeros(n, 3)


class _Scene:
  def __init__(self, n: int):
    self.robot = _Robot(n)
    self.terrain = _Terrain(n)
    self.env_origins = torch.zeros(n, 3)

  def __getitem__(self, key):
    return self.robot


class _Termination:
  def __init__(self, n: int):
    self.terminated = torch.zeros(n, dtype=torch.bool)


class _Env:
  def __init__(self, n: int):
    self.num_envs = n
    self.device = "cpu"
    self.scene = _Scene(n)
    self.episode_length_buf = torch.ones(n, dtype=torch.long)
    self.termination_manager = _Termination(n)


def _episode(env: _Env, recovered: torch.Tensor) -> None:
  """Advance one finished episode for every env and run the curriculum."""
  env.termination_manager.terminated[:] = ~recovered
  env.scene.robot.data.projected_gravity_b[:] = torch.where(
    recovered[:, None],
    torch.tensor([0.0, 0.0, -1.0]),
    torch.tensor([0.0, 0.0, 1.0]))
  upward_curriculum(env, torch.arange(env.num_envs))


def test_per_env_rolling_window_advances_and_demotes() -> None:
  env = _Env(4)
  outcomes = torch.zeros(4, dtype=torch.bool)
  for episode in range(HISTORY_LENGTH):
    outcomes[0] = True          # 20/20 -> advance
    outcomes[1] = False         # 0/20 -> stays at level 0
    outcomes[2] = episode < 10  # 10/20 -> neutral
    outcomes[3] = episode < PROMOTE_SUCCESSES  # exactly 16/20 -> advance
    _episode(env, outcomes)
  stats = upward_curriculum_stats(env)
  assert env._upward_tilt_levels.tolist() == [1, 0, 0, 1]
  assert env.scene.terrain.terrain_levels.tolist() == [1, 0, 0, 1]
  assert stats["tilt_level_q75"] == 1
  # Advanced envs reset their window and need 20 new episodes; only envs 1
  # (0/20) and 2 (10/20) are ready after the first block.
  assert abs(stats["recovery_window_rate"] - 0.25) < 1e-6
  assert abs(stats["recovery_window_ready"] - 0.5) < 1e-6
  assert math.isclose(upward_video_tilt_limit(env), math.radians(135.0))

  # The history is per-env: env 1 now gets 20 successes and advances alone;
  # env 2 stays neutral with 10 successes.
  for episode in range(HISTORY_LENGTH):
    outcomes[0] = True
    outcomes[1] = True
    outcomes[2] = episode < 10
    outcomes[3] = True
    _episode(env, outcomes)
  assert env._upward_tilt_levels.tolist() == [2, 1, 0, 2]
  assert env.scene.terrain.terrain_levels.tolist() == [2, 1, 0, 2]

  # A full 20-episode failure window demotes env 0 from level 2 to level 1.
  for _ in range(HISTORY_LENGTH):
    outcomes[:] = False
    outcomes[3] = True  # keep env 3 at level 2
    _episode(env, outcomes)
  assert env._upward_tilt_levels.tolist() == [1, 0, 0, 2]
  # Terrain has six levels and is independent of the three tilt levels.
  assert env.scene.terrain.terrain_levels.tolist() == [1, 0, 0, 3]

  # Before 20 episodes no env is eligible to move.
  fresh = _Env(4)
  stats = upward_curriculum_stats(fresh)
  assert stats["recovery_window_ready"] == 0.0
  assert stats["tilt_level_q75"] == 0
  outcomes = torch.ones(4, dtype=torch.bool)
  for _ in range(HISTORY_LENGTH - 1):
    _episode(fresh, outcomes)
  assert fresh._upward_tilt_levels.tolist() == [0, 0, 0, 0]
  _episode(fresh, outcomes)
  assert fresh._upward_tilt_levels.tolist() == [1, 1, 1, 1]


class _Asset:
  def __init__(self, n: int):
    self.data = type("Data", (), {})()
    self.data.default_root_state = torch.zeros(n, 13)
    self.data.default_root_state[:, 3] = 1.0  # identity wxyz quaternion
    self.pose = None
    self.velocity = None

  def write_root_link_pose_to_sim(self, pose, env_ids):
    self.pose = pose.clone()

  def write_root_link_velocity_to_sim(self, velocity, env_ids):
    self.velocity = velocity.clone()


class _ResetScene:
  def __init__(self, n: int):
    self.asset = _Asset(n)
    self.env_origins = torch.zeros(n, 3)

  def __getitem__(self, key):
    return self.asset


class _ResetEnv:
  def __init__(self, n: int):
    self.num_envs = n
    self.device = "cpu"
    self.scene = _ResetScene(n)


def test_reset_supports_per_env_and_play_fallback_limits() -> None:
  env = _ResetEnv(2)
  pose_range = {"x": (-0.1, 0.1), "y": (-0.1, 0.1), "z": (0.0, 0.2),
                "roll": (-math.pi, math.pi), "pitch": (-math.pi, math.pi),
                "yaw": (-math.pi, math.pi)}
  ids = torch.arange(2)
  reset_root_state_upward(env, ids, pose_range)
  assert env.scene.asset.pose.shape == (2, 7)
  assert env.scene.asset.velocity.shape == (2, 6)
  assert torch.isfinite(env.scene.asset.pose).all()
  env._upward_tilt_limits = torch.tensor(
    [math.radians(60.0), math.radians(180.0)])
  reset_root_state_upward(env, ids, pose_range)
  assert torch.isfinite(env.scene.asset.pose).all()


def main() -> None:
  test_severity_is_monotonic()
  test_severity_gradient_is_finite()
  test_mask_keeps_bad_set()
  test_upward_continuation_does_not_use_v4_fall_proxy()
  test_per_env_rolling_window_advances_and_demotes()
  test_reset_supports_per_env_and_play_fallback_limits()
  print("Upward reward / continuation / per-env curriculum checks passed")


if __name__ == "__main__":
  main()