#!/usr/bin/env python3
"""Generate standard linear stairs for MuJoCo sim2sim testing.

Each cell: flat 2m approach -> N linear steps -> flat 2m platform -> N steps down -> flat edge.
Treads are horizontal, rises are vertical (1px = 0.1m wide, acceptable for hfield).

Usage:
    uv run scripts/gen_stairs_test.py --step-height 0.07 --step-depth 0.31
"""
import numpy as np, os, argparse
from PIL import Image

HS = 0.1          # horizontal scale [m/px]
VS = 0.005        # vertical scale  [m/unit]
CELL_M = 8.0
PLATFORM_M = 4.0  # bigger flat platform -> fewer steps
CELL_PX = int(CELL_M / HS)  # 80
PLATFORM_PX = int(PLATFORM_M / HS)  # 40
BORDER_M = 2.0
BORDER_PX = int(BORDER_M / HS)  # 20
NUM_CELLS = 2
TOT_PX = NUM_CELLS * CELL_PX + 2 * BORDER_PX
TOTAL_M = TOT_PX * HS

np.random.seed(42)


def make_linear_stairs(step_height_m, step_depth_m=0.31):
    """Linear stairs: flat approach -> N steps up -> flat platform -> edge.
    Each step has a flat horizontal tread and (essentially) vertical rise."""
    t = np.zeros((CELL_PX, CELL_PX), dtype=np.int16)
    sd = int(step_depth_m / HS)   # tread depth in px
    sh = int(step_height_m / VS)  # rise height in pixel units

    # How many steps fit on each side of the platform?
    avail = (CELL_PX - PLATFORM_PX) // 2
    n_steps = avail // max(sd, 1)
    if n_steps < 1:
        n_steps = 1

    edge = (CELL_PX - PLATFORM_PX - n_steps * sd) // 2  # remaining flat on each side

    # Draw steps going UP from left (in +x direction)
    # Each step: flat tread at current height, then rise to next height
    x = edge
    h = 0
    for i in range(n_steps):
        x_next = x + sd
        t[:, x:x_next] = h  # tread at current height
        x = x_next
        h += sh

    # Platform (flat at max height)
    plat_start = x
    plat_end = plat_start + PLATFORM_PX
    t[:, plat_start:plat_end] = h

    # Continue stairs going DOWN on the right (optional: mirror)
    x = plat_end
    for i in range(n_steps):
        h -= sh
        x_next = x + sd
        t[:, x:x_next] = h
        x = x_next

    return t


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--step-height", type=float, default=0.15,
                   help="step rise height in metres (default 0.15)")
    p.add_argument("--step-depth", type=float, default=0.50,
                   help="step tread depth in metres (default 0.50)")
    args = p.parse_args()

    print(f"Linear stairs: step_h={args.step_height:.2f}m tread={args.step_depth:.2f}m platform={PLATFORM_M:.0f}m")

    hf_raw = np.zeros((TOT_PX, TOT_PX), dtype=np.int16)
    for i in range(NUM_CELLS):
        for j in range(NUM_CELLS):
            cell = make_linear_stairs(args.step_height, args.step_depth)
            y0 = BORDER_PX + i * CELL_PX
            x0 = BORDER_PX + j * CELL_PX
            hf_raw[y0:y0 + CELL_PX, x0:x0 + CELL_PX] = cell

    hf_m = hf_raw.astype(np.float32) * VS
    z_min = float(hf_m.min())
    z_max = float(hf_m.max())
    z_range = max(z_max - z_min, 0.001)

    print(f"  height: [{z_min:.3f}, {z_max:.3f}]m  z_scale={z_range:.3f}  max={z_max*100:.0f}cm")

    png = ((hf_m - z_min) / z_range * 65535.0).astype(np.uint16)

    out_dir = os.path.join(os.path.dirname(__file__), "..",
                           "motrix_envs", "src", "motrix_envs", "locomotion",
                           "go1", "xmls", "assets")
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, "stairs_test.png")
    Image.fromarray(png).save(out_path)
    print(f"  saved: {out_path}")
    sbase = max(z_min, 0.001)
    print(f"  XML: size=\"{TOTAL_M/2:.1f} {TOTAL_M/2:.1f} {z_range:.3f} {sbase:.3f}\"")


if __name__ == "__main__":
    main()