"""S10 terrain revision 7: safe low-speed all-terrain locomotion.

V7 keeps the 53-value v4 Actor frame (measured wheel speed is required for
wheel slip/stall feedback), a stair-pass-gated command stage ending at
1.5 m/s, 0.8 m/s and 1.2 rad/s with a 0.05 per-axis deadband, a conservative
initialization (small x/y/yaw only, no random roll/pitch), a hipx 0.05 /
hipy +/-25 deg / knee 1.0 default stance, always-on wheel/body projection
penalties (per-wheel outer-1/4 and paired-leg 5/6 with a 0.5 s grace), reduced
knee/hipy posture stiffness on stairs, and the v5 stair curriculum plus v7
non-wheel collision head (50 N, wheels/knees excluded). Only ``imag_batch`` is
reduced to 2048 to fit a 16 GB GPU.
"""

from .base import ActorSpec, TaskSpec, register

ACTOR = ActorSpec(
  name="s10-v7-53",
  terms=(3, 3, 12, 12, 4, 16, 3),
  frame_terms=(
    {"name": "base_ang_vel", "dim": 3, "scale": 1.0},
    {"name": "projected_gravity", "dim": 3, "scale": 1.0},
    {"name": "joint_pos_rel", "dim": 12, "scale": 1.0},
    {"name": "joint_vel_rel", "dim": 12, "scale": 0.05},
    {"name": "wheel_vel_rel", "dim": 4, "scale": 0.05},
    {"name": "last_action", "dim": 16, "scale": 1.0},
    {"name": "command", "dim": 3, "scale": 1.0},
  ),
  wheel_vel=True,
)


def env_cfg(play: bool, history: int, seed: int):
  from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
  return s10_velocity_terrain_env_cfg(
    play=play, history=history, seed=seed, revision=7,
    actor_wheel_vel=ACTOR.wheel_vel)


TASK = register(TaskSpec(
  name="s10-v7",
  description="S10 terrain v7 safe all-terrain task: 53-value Actor frame, wheel speed, stair-gated low-speed commands, stair curriculum + collision",
  terrain_revision=7,
  actor=ACTOR,
  env_cfg=env_cfg,
  defaults={
    "terrain": True,
    "terrain_revision": 7,
    "num_envs": 4096,
    "iterations": 40000,
    "num_steps_per_env": 1,
    "buffer_len": 512,
    "wm_window": 8,
    "wm_batch": 2048,
    "wm_updates": 4,
    # 8192 is the full-size imagination batch; 2048 keeps the 16 GB footprint.
    "imag_batch": 2048,
    "horizon": 32,
    "warmup": 300,
    "critic_epochs": 5,
    "critic_minibatches": 4,
    "save_interval": 500,
    "log_interval": 50,
    # V7 uses stair-pass-gated stages, capped at its safe envelope.
    "command_curriculum": True,
    "command_curriculum_start": 0,
    "prioritized_replay": True,
    "archive_size": 100000,
    "actor_ramp_iterations": 1000,
    "actor_hold_iterations": 20000,
    "lr_decay_iterations": 20000,
    "lr_final_frac": 0.2,
    "history": 5,
    "wm_hidden": [512, 512, 512],
    "actor_hidden": [256, 128, 64],
    "critic_hidden": [256, 128, 64],
    "compile_imagination": True,
  },
))