#!/usr/bin/env python3
"""Generate 2-level terrain: level 0=flat, level 1=pyramid stairs.

1cm = 1px. OpenCV draws concentric filled rectangles (outside→in, 10 steps).
Higher values overwrite lower = convex pyramid (stairs up toward center).
Lower values overwrite higher = concave pyramid (stairs down from center).

Usage:
    uv run scripts/gen_flat_stairs.py                        # convex (stairs UP)
    uv run scripts/gen_flat_stairs.py --concave              # concave (stairs DOWN)
    uv run scripts/gen_flat_stairs.py --step-h 0.10 --num-steps 5
"""
import cv2, numpy as np, os, argparse

# ═══ fixed params ═══
HS = 0.01             # 1cm/px
VS = 0.005            # height unit = 0.5cm
CELL_M = 8.0          # 8m cell
BORDER_M = 5.0        # 5m border
NUM_ROWS = 2          # flat + stairs
NUM_COLS = 4          # columns

CELL_PX = int(CELL_M / HS)       # 800
BORDER_PX = int(BORDER_M / HS)   # 500
PLATFORM_PX = int(1.0 / HS)      # 1m platform = 100px

TOT_ROWS = NUM_ROWS * CELL_PX + 2 * BORDER_PX
TOT_COLS = NUM_COLS * CELL_PX + 2 * BORDER_PX

# ═══ stairs params (override via CLI) ═══
NUM_STEPS = 10        # 10 steps
STEP_H_CM = 20        # 20cm rise per step
STEP_D_CM = 20        # 20cm tread per step

STEP_H_VS = int(STEP_H_CM / 100.0 / VS)  # 0.20 / 0.005 = 40
STEP_D_PX = int(STEP_D_CM / 100.0 / HS)  # 0.20 / 0.01 = 20


def draw_pyramid(canvas, x0, y0, num_steps, step_d_px, step_h_vs, concave=False):
    """Draw concentric rectangles from outside→in.

    Convex: edge=0 → platform=max (stairs up toward center)
    Concave: raise whole cell to max, then draw pit: edge=max → platform=0
    """
    cx, cy = x0 + CELL_PX // 2, y0 + CELL_PX // 2
    p2 = PLATFORM_PX // 2
    h_max = step_h_vs * num_steps

    # Fill cell to cell boundary with reference-plane height
    half_max = CELL_PX // 2  # extend to cell edge
    cv2.rectangle(canvas, (cx - half_max, cy - half_max),
                  (cx + half_max, cy + half_max), int(h_max), -1)

    if concave:
        # Pit: rings going DOWN from reference plane
        for i in range(num_steps + 1):
            half = p2 + (num_steps - i) * step_d_px
            x1, y1 = cx - half, cy - half
            x2, y2 = cx + half, cy + half
            h = h_max - step_h_vs * i
            cv2.rectangle(canvas, (x1, y1), (x2, y2), int(h), -1)
    else:
        # Mound: rings going UP from reference plane
        for i in range(num_steps + 1):
            half = p2 + (num_steps - i) * step_d_px
            x1, y1 = cx - half, cy - half
            x2, y2 = cx + half, cy + half
            h = h_max + step_h_vs * i
            cv2.rectangle(canvas, (x1, y1), (x2, y2), int(h), -1)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--concave", action="store_true", help="concave pyramid (stairs down from center)")
    p.add_argument("--step-h", type=float, default=0.20, help="step rise (m)")
    p.add_argument("--step-d", type=float, default=0.20, help="step tread (m)")
    p.add_argument("--num-steps", type=int, default=10, help="number of steps")
    args = p.parse_args()

    step_h_vs = int(args.step_h / VS)
    step_d_px = int(args.step_d / HS)
    total_h = step_h_vs * args.num_steps * VS

    print(f"Building {TOT_COLS}×{TOT_ROWS}px ({TOT_COLS*HS:.0f}×{TOT_ROWS*HS:.0f}m)")
    print(f"  type={'concave' if args.concave else 'convex'}  "
          f"steps={args.num_steps}  rise={step_h_vs*VS*100:.0f}cm  "
          f"tread={step_d_px*HS*100:.0f}cm  total_h={total_h*100:.0f}cm")

    canvas = np.zeros((TOT_ROWS, TOT_COLS), dtype=np.uint16)

    for row in range(NUM_ROWS):
        for col in range(NUM_COLS):
            x0 = BORDER_PX + col * CELL_PX
            y0 = BORDER_PX + row * CELL_PX
            if row == 1:
                # alternate convex/concave across cols
                concave_cell = (col % 2 == 1)
                draw_pyramid(canvas, x0, y0, args.num_steps,
                             step_d_px, step_h_vs, concave_cell)

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

    png = ((hf_m - z_min) / z_range * 65535).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, "flat_stairs.png")
    cv2.imwrite(out_path, png)
    print(f"  saved: {out_path}")
    print(f"  XML: size=\"{TOT_COLS*HS/2:.1f} {TOT_ROWS*HS/2:.1f} "
          f"{z_range:.3f} {max(z_min,0.001):.3f}\"")


if __name__ == "__main__":
    main()