"""S10 terrain revision 4: legacy 53-value Actor frame with measured wheel speed."""

from .base import ActorSpec, TaskSpec, register

ACTOR = ActorSpec(
  name="s10-v4-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=4,
    actor_wheel_vel=ACTOR.wheel_vel)


TASK = register(TaskSpec(
  name="s10-v4",
  description="S10 terrain v4 legacy task: 53-value Actor frame, wheel speed, 20k validated recipe",
  terrain_revision=4,
  actor=ACTOR,
  env_cfg=env_cfg,
  defaults={
    "terrain": True,
    "terrain_revision": 4,
    "num_envs": 4096,
    "iterations": 20000,
    "num_steps_per_env": 1,
    "buffer_len": 512,
    "wm_window": 8,
    "wm_batch": 2048,
    "wm_updates": 4,
    "imag_batch": 2048,
    "horizon": 32,
    "warmup": 50,
    "critic_epochs": 5,
    "critic_minibatches": 4,
    "save_interval": 500,
    "log_interval": 10,
    "command_curriculum": True,
    "command_curriculum_start": 0,
    "prioritized_replay": True,
    "archive_size": 100000,
    "actor_ramp_iterations": 1000,
    "actor_hold_iterations": 0,
    "lr_decay_iterations": 30000,
    "lr_final_frac": 0.03,
    "history": 5,
    "wm_hidden": [512, 512, 512],
    "actor_hidden": [256, 128, 64],
    "critic_hidden": [256, 128, 64],
    "compile_imagination": False,
  },
))