"""Terrain curricula for the Go1 transfer study.

mjlab's ROUGH_TERRAINS_CFG uses 7 of the 18 shipped presets and caps stair
height at 0.1 m, so a policy trained on it never encounters a gap it has to
place a foot across. HARD adds the unused presets and raises the ceiling.

Difficulty rises along terrain rows; num_rows is the curriculum depth.
"""

from mjlab.terrains import TerrainGeneratorCfg
from mjlab.terrains.config import (
  discrete_obstacles,
  flat,
  hf_pyramid_slope,
  narrow_beams,
  pyramid_stairs,
  pyramid_stairs_inv,
  random_rough,
  stepping_stones,
  wave_terrain,
)

# What mjlab ships, for reference in the transfer matrix.
from mjlab.terrains.config import ROUGH_TERRAINS_CFG  # noqa: F401

FLAT_CFG = TerrainGeneratorCfg(
  size=(8.0, 8.0),
  border_width=20.0,
  num_rows=10,
  num_cols=20,
  sub_terrains={"flat": flat(proportion=1.0)},
  add_lights=True,
)

HARD_CFG = TerrainGeneratorCfg(
  size=(8.0, 8.0),
  border_width=20.0,
  num_rows=10,
  num_cols=20,
  sub_terrains={
    # Discontinuous ground: the robot must choose where to put a foot,
    # which the shipped curriculum never requires.
    "stepping_stones": stepping_stones(
      proportion=0.2, stone_distance_range=(0.1, 0.4), stone_height_variation=0.15
    ),
    "narrow_beams": narrow_beams(proportion=0.15, beam_width_range=(0.15, 0.5)),
    "discrete_obstacles": discrete_obstacles(
      proportion=0.15, obstacle_height_range=(0.05, 0.35), num_obstacles=60
    ),
    # Same families as ROUGH but past its ceiling: 0.2 m steps, not 0.1.
    "pyramid_stairs": pyramid_stairs(proportion=0.15, step_height_range=(0.05, 0.2)),
    "pyramid_stairs_inv": pyramid_stairs_inv(
      proportion=0.15, step_height_range=(0.05, 0.2)
    ),
    "random_rough": random_rough(proportion=0.1, noise_range=(0.05, 0.18)),
    "hf_pyramid_slope": hf_pyramid_slope(proportion=0.05, slope_range=(0.0, 1.2)),
    "wave_terrain": wave_terrain(proportion=0.05),
  },
  add_lights=True,
)

TERRAINS = {"flat": FLAT_CFG, "rough": ROUGH_TERRAINS_CFG, "hard": HARD_CFG}