#!/usr/bin/env python3
"""Terrain editor GUI — draw pyramids, mark spawn zones, export PNG + coordinates.

Usage:
    uv run scripts/terrain_editor.py
"""
import tkinter as tk
from tkinter import ttk, messagebox
import numpy as np, os, cv2

# ═══ defaults ═══
HS = 0.05; VS = 0.005
CELL_M = 8.0; BORDER_M = 5.0
PLATFORM_M = 1.0            # platform 1m
SPAWN_RADIUS_M = 0.5        # spawn zone ±0.5m around center
DEFAULT_STEP_H = 0.20; DEFAULT_STEP_D = 0.20; DEFAULT_NUM_STEPS = 10
DEFAULT_REF_PLANE_CM = 200  # all cells start from same reference height
CELL_PX = int(CELL_M / HS); BORDER_PX = int(BORDER_M / HS)
PLATFORM_PX = int(PLATFORM_M / HS); SPAWN_RADIUS_PX = int(SPAWN_RADIUS_M / HS)


class TerrainEditor:
    def __init__(self, root):
        self.root = root
        self.root.title("Terrain Editor")
        self.rows = 2; self.cols = 4
        self.cell_types = {}
        self._init_defaults()
        self.selected = (0, 0)
        self._dragging = False
        self._build_ui()
        self._sync_params()
        self._redraw_all()

    def _init_defaults(self):
        for r in range(self.rows):
            for c in range(self.cols):
                if r == 0:
                    self.cell_types[(r, c)] = {"type": "flat", "spawn": True, "level": 0,
                                                "ref_plane_cm": DEFAULT_REF_PLANE_CM}
                else:
                    self.cell_types[(r, c)] = {
                        "type": "convex" if c % 2 == 0 else "concave",
                        "step_h": DEFAULT_STEP_H, "step_d": DEFAULT_STEP_D,
                        "num_steps": DEFAULT_NUM_STEPS, "spawn": True, "level": 1,
                        "ref_plane_cm": DEFAULT_REF_PLANE_CM,
                    }

    # ═══ UI ═══
    def _build_ui(self):
        paned = ttk.PanedWindow(self.root, orient=tk.HORIZONTAL)
        paned.pack(fill=tk.BOTH, expand=True)
        left = ttk.Frame(paned); paned.add(left, weight=2)
        right = ttk.Frame(paned); paned.add(right, weight=1)
        self._build_preview_ui(left)
        self._build_params_ui(right)

    def _build_preview_ui(self, parent):
        ttk.Label(parent, text="Terrain Preview (click to select cell, right-click toggle spawn)", font=("", 10)).pack(pady=2)
        self.info_label = ttk.Label(parent, text="")
        self.info_label.pack()
        self.preview = tk.Canvas(parent, bg="#333", width=600, height=400)
        self.preview.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        self.preview.bind("<Button-1>", self._on_click)
        self.preview.bind("<B1-Motion>", self._on_drag)
        self.preview.bind("<Button-3>", self._on_right_click)
        ttk.Label(parent, text="Left-click: select | Right-click: toggle spawn | Drag: select").pack()
        ctrl = ttk.Frame(parent)
        ctrl.pack(pady=5)
        ttk.Label(ctrl, text="Rows:").pack(side=tk.LEFT)
        self.rows_var = tk.IntVar(value=self.rows)
        ttk.Spinbox(ctrl, from_=1, to=10, width=4, textvariable=self.rows_var,
                    command=self._on_grid_size).pack(side=tk.LEFT, padx=2)
        ttk.Label(ctrl, text="Cols:").pack(side=tk.LEFT, padx=(10,0))
        self.cols_var = tk.IntVar(value=self.cols)
        ttk.Spinbox(ctrl, from_=1, to=10, width=4, textvariable=self.cols_var,
                    command=self._on_grid_size).pack(side=tk.LEFT, padx=2)
        ttk.Button(parent, text="Export PNG + Coords", command=self._export).pack(pady=5)

    def _build_params_ui(self, parent):
        f = ttk.Frame(parent); f.pack(padx=10, pady=5, fill=tk.X)
        ttk.Label(f, text="Cell Type:").grid(row=0, column=0, sticky=tk.W)
        self.type_var = tk.StringVar(value="flat")
        ttk.Combobox(f, textvariable=self.type_var, values=["flat", "convex", "concave"],
                     state="readonly", width=10).grid(row=0, column=1, padx=5)
        self.type_var.trace("w", lambda *a: self._on_param_change())
        ttk.Label(f, text="Level:").grid(row=0, column=2, sticky=tk.W, padx=(20,0))
        self.level_var = tk.IntVar(value=0)
        ttk.Spinbox(f, from_=0, to=9, width=3, textvariable=self.level_var,
                    command=self._on_param_change).grid(row=0, column=3)
        self.spawn_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(f, text="Spawn", variable=self.spawn_var,
                        command=self._on_param_change).grid(row=0, column=4, padx=10)

        ttk.Separator(parent, orient=tk.HORIZONTAL).pack(fill=tk.X, pady=5, padx=10)
        ttk.Label(parent, text="Pyramid Params").pack()
        g = ttk.Frame(parent); g.pack(padx=10, pady=5, fill=tk.X)
        ttk.Label(g, text="Step height (m):").grid(row=0, column=0, sticky=tk.W)
        self.sh_var = tk.StringVar(value=str(DEFAULT_STEP_H))
        ttk.Entry(g, textvariable=self.sh_var, width=7).grid(row=0, column=1, padx=5)
        self.sh_var.trace("w", lambda *a: self._on_param_change())
        ttk.Label(g, text="Step tread (m):").grid(row=1, column=0, sticky=tk.W)
        self.sd_var = tk.StringVar(value=str(DEFAULT_STEP_D))
        ttk.Entry(g, textvariable=self.sd_var, width=7).grid(row=1, column=1, padx=5)
        self.sd_var.trace("w", lambda *a: self._on_param_change())
        ttk.Label(g, text="Num steps:").grid(row=2, column=0, sticky=tk.W)
        self.ns_var = tk.StringVar(value=str(DEFAULT_NUM_STEPS))
        ttk.Entry(g, textvariable=self.ns_var, width=7).grid(row=2, column=1, padx=5)
        ttk.Label(g, text="Ref plane (cm):").grid(row=3, column=0, sticky=tk.W)
        self.ref_var = tk.StringVar(value=str(DEFAULT_REF_PLANE_CM))
        ttk.Entry(g, textvariable=self.ref_var, width=7).grid(row=3, column=1, padx=5)
        self.ns_var.trace("w", lambda *a: self._on_param_change())

        ttk.Separator(parent, orient=tk.HORIZONTAL).pack(fill=tk.X, pady=5, padx=10)
        ttk.Label(parent, text="Selected Cell").pack()
        self.cell_label = ttk.Label(parent, text="")
        self.cell_label.pack()

    # ═══ events ═══
    def _cell_at(self, ex, ey):
        m = 5; pw = self.preview.winfo_width(); ph = self.preview.winfo_height()
        cw = (pw - 2*m) // max(self.cols,1); ch = (ph - 2*m) // max(self.rows,1)
        col = (ex - m) // cw; row = (ey - m) // ch
        if 0 <= col < self.cols and 0 <= row < self.rows:
            return row, col, m + col*cw, m + row*ch, cw, ch
        return None

    def _on_click(self, event):
        v = self._cell_at(event.x, event.y)
        if v:
            self.selected = (v[0], v[1])
            self._sync_params(); self._redraw_all()

    def _on_drag(self, event):
        v = self._cell_at(event.x, event.y)
        if v:
            self.selected = (v[0], v[1])
            self._sync_params(); self._redraw_all()

    def _on_right_click(self, event):
        v = self._cell_at(event.x, event.y)
        if v:
            r, c = v[0], v[1]
            ct = self.cell_types.get((r, c), {"type": "flat", "spawn": True, "level": r})
            ct = dict(ct)  # copy before modifying
            ct["spawn"] = not ct.get("spawn", True)
            self.cell_types[(r, c)] = ct
            if (r, c) == self.selected:
                self._sync_params()
            self._redraw_all()

    def _on_grid_size(self):
        try: nr = self.rows_var.get()
        except: nr = self.rows
        try: nc = self.cols_var.get()
        except: nc = self.cols
        if nr == self.rows and nc == self.cols: return
        old = self.cell_types
        self.rows, self.cols = nr, nc
        self.cell_types = {}
        for r in range(nr):
            for c in range(nc):
                self.cell_types[(r,c)] = old.get((r,c), {"type": "flat", "spawn": True, "level": r})
        self._redraw_all()

    # ═══ sync ═══
    def _sync_params(self):
        ct = self.cell_types.get(self.selected, {"type": "flat", "spawn": True, "level": self.selected[0]})
        self.type_var.set(ct.get("type", "flat"))
        self.spawn_var.set(ct.get("spawn", True))
        r = self.selected[0]
        self.level_var.set(ct.get("level", r))
        self.ref_var.set(str(ct.get("ref_plane_cm", DEFAULT_REF_PLANE_CM)))
        self.sh_var.set(str(ct.get("step_h", DEFAULT_STEP_H)))
        self.sd_var.set(str(ct.get("step_d", DEFAULT_STEP_D)))
        self.ns_var.set(str(ct.get("num_steps", DEFAULT_NUM_STEPS)))
        r, c = self.selected
        half_x = BORDER_M + self.cols * CELL_M / 2
        half_y = BORDER_M + self.rows * CELL_M / 2
        cx = -half_x + BORDER_M + c * CELL_M + CELL_M / 2
        cy = half_y - BORDER_M - r * CELL_M - CELL_M / 2
        self.cell_label.config(text=f"({r},{c}) center: x={cx:+.1f} y={cy:+.1f}  type={ct['type']}")

    def _on_param_change(self):
        r, c = self.selected
        try: sh = float(self.sh_var.get()); sd = float(self.sd_var.get())
        except: return
        try: ns = int(self.ns_var.get())
        except: return
        try: ref_cm = float(self.ref_var.get())
        except: ref_cm = DEFAULT_REF_PLANE_CM
        cell = {"type": self.type_var.get(), "spawn": self.spawn_var.get(),
                "level": self.level_var.get(), "ref_plane_cm": ref_cm}
        if cell["type"] != "flat":
            cell.update({"step_h": sh, "step_d": sd, "num_steps": ns})
        self.cell_types[(r, c)] = cell
        self._redraw_all()

    # ═══ draw ═══
    def _redraw_all(self):
        w = self.preview.winfo_width(); h = self.preview.winfo_height()
        if w < 10: w = 600
        if h < 10: h = 400
        self._draw_preview(w, h)
        half_x = BORDER_M + self.cols * CELL_M / 2
        half_y = BORDER_M + self.rows * CELL_M / 2
        self.info_label.config(
            text=f"{self.rows}×{self.cols}  "
                 f"{self.cols*CELL_M+2*BORDER_M:.0f}×{self.rows*CELL_M+2*BORDER_M:.0f}m  "
                 f"spawn_cy = {half_y-BORDER_M-CELL_M/2:.0f} - row*{CELL_M:.0f}")

    def _draw_preview(self, pw, ph):
        cv = self.preview; cv.delete("all")
        m = 5; cw = (pw - 2*m) // max(self.cols, 1); ch = (ph - 2*m) // max(self.rows, 1)
        cw = max(cw, 30); ch = max(ch, 30)
        colors = {"flat": "#5b8c5a", "convex": "#c0392b", "concave": "#2471a3"}

        for r in range(self.rows):
            for c in range(self.cols):
                x1, y1 = m + c*cw, m + r*ch
                x2, y2 = x1 + cw, y1 + ch
                ct = self.cell_types.get((r,c), {"type": "flat", "spawn": True, "level": r})
                cv.create_rectangle(x1, y1, x2, y2, fill=colors.get(ct["type"], "#555"),
                                    outline="#888", width=1)
                # Cell center dot
                cx = (x1+x2)//2; cy = (y1+y2)//2
                cv.create_oval(cx-3, cy-3, cx+3, cy+3, fill="white", outline="")

                # Pyramid stairs rings
                if ct["type"] != "flat":
                    sh = ct.get("step_h", DEFAULT_STEP_H)
                    sd = ct.get("step_d", DEFAULT_STEP_D)
                    ns = ct.get("num_steps", DEFAULT_NUM_STEPS)
                    concave = ct["type"] == "concave"
                    p2 = max(2, cw // 16)
                    step_px = max(1, (cw//2 - p2) // max(ns, 1))
                    h_max = int(sh * ns / VS)
                    for i in range(ns + 1):
                        half = p2 + (ns - i) * step_px
                        if concave:
                            frac = (ns - i) / max(ns, 1)
                        else:
                            frac = i / max(ns, 1)
                        g = int(180 - frac * 100)
                        clr = f"#{g:02x}{g:02x}{g:02x}"
                        cv.create_rectangle(cx - half, cy - half, cx + half, cy + half,
                                            fill=clr, outline="")

                # Spawn zone (green rect)
                if ct.get("spawn", True):
                    sz = max(2, int(SPAWN_RADIUS_M / CELL_M * cw))
                    cv.create_rectangle(cx - sz, cy - sz, cx + sz, cy + sz,
                                        outline="#00ff00", width=2)

                # Level + type label
                lvl = ct.get("level", r)
                lbl = f"L{lvl} {ct['type'][:3]}"
                if ct["type"] != "flat":
                    lbl = f"L{lvl} {ct['type'][:3]}-{sh*100:.0f}cm"
                cv.create_text(x1 + 20, y1 + 10, text=lbl, fill="white",
                               font=("", 8), anchor=tk.NW)

        # Highlight selected cell
        r, c = self.selected
        x1, y1 = m + c*cw, m + r*ch
        x2, y2 = x1 + cw, y1 + ch
        cv.create_rectangle(x1, y1, x2, y2, outline="yellow", width=3)

        # Level labels on right
        for r in range(self.rows):
            y = m + r*ch + ch//2
            cv.create_text(pw - 15, y, text=f"L{r}", fill="white", font=("", 12, "bold"))

    # ═══ generate + export ═══
    def _generate_png(self):
        tot_rows = self.rows * CELL_PX + 2 * BORDER_PX
        tot_cols = self.cols * CELL_PX + 2 * BORDER_PX
        canvas = np.zeros((tot_rows, tot_cols), dtype=np.uint16)
        for r in range(self.rows):
            for c in range(self.cols):
                x0 = BORDER_PX + c * CELL_PX; y0 = BORDER_PX + r * CELL_PX
                ct = self.cell_types.get((r,c), {"type": "flat", "level": r})
                if ct["type"] == "flat": continue
                sh = ct.get("step_h", DEFAULT_STEP_H)
                sd = ct.get("step_d", DEFAULT_STEP_D)
                ns = ct.get("num_steps", DEFAULT_NUM_STEPS)
                concave = ct["type"] == "concave"
                ref_cm = ct.get("ref_plane_cm", DEFAULT_REF_PLANE_CM)
                ref_vs = int(ref_cm / 100.0 / VS)
                h_vs = int(sh / VS); d_px = int(sd / HS)
                p2 = PLATFORM_PX // 2
                cx = x0 + CELL_PX // 2; cy = y0 + CELL_PX // 2
                cv2.rectangle(canvas, (x0, y0), (x0+CELL_PX, y0+CELL_PX), int(ref_vs), -1)
                for i in range(ns + 1):
                    half = p2 + (ns - i) * d_px
                    x1, y1 = cx - half, cy - half; x2, y2 = cx + half, cy + half
                    if concave:
                        h = ref_vs - h_vs * i
                    else:
                        h = ref_vs + h_vs * i
                    cv2.rectangle(canvas, (x1, y1), (x2, y2), int(h), -1)
        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)
        return png, z_range, z_min

    def _export(self):
        png, z_range, z_min = self._generate_png()
        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)
        w_m = (self.cols * CELL_PX + 2 * BORDER_PX) * HS
        h_m = (self.rows * CELL_PX + 2 * BORDER_PX) * HS
        half_x = BORDER_M + self.cols * CELL_M / 2
        half_y = BORDER_M + self.rows * CELL_M / 2

        lines = [
            f"# Terrain: {self.rows}×{self.cols}  {w_m:.0f}×{h_m:.0f}m",
            f"XML: size=\"{w_m/2:.1f} {h_m/2:.1f} {z_range:.3f} {max(z_min,0.001):.3f}\"",
            f"dreamwaq.py: terrain_rows={self.rows}  terrain_cols={self.cols}",
            f"",
            f"# === Cell centers ===",
        ]
        for r in range(self.rows):
            for c in range(self.cols):
                cx = -half_x + BORDER_M + c * CELL_M + CELL_M / 2
                cy = half_y - BORDER_M - r * CELL_M - CELL_M / 2
                ct = self.cell_types.get((r,c), {"type": "flat", "level": r})
                lines.append(f"  ({r},{c}): x={cx:+.1f} y={cy:+.1f}  {ct['type']}")

        lines.append("")
        lines.append("# === Level boundaries (robot out of bounds → reset) ===")
        level_bounds = {}
        for r in range(self.rows):
            for c in range(self.cols):
                lv = self.cell_types.get((r,c), {"level": r})["level"]
                if lv not in level_bounds:
                    level_bounds[lv] = {"rmin": r, "rmax": r, "cmin": c, "cmax": c}
                else:
                    b = level_bounds[lv]
                    b["rmin"] = min(b["rmin"], r)
                    b["rmax"] = max(b["rmax"], r)
                    b["cmin"] = min(b["cmin"], c)
                    b["cmax"] = max(b["cmax"], c)
        for lv in sorted(level_bounds):
            b = level_bounds[lv]
            x_min = -half_x + BORDER_M + b["cmin"] * CELL_M
            x_max = -half_x + BORDER_M + (b["cmax"] + 1) * CELL_M
            y_min = half_y - BORDER_M - (b["rmax"] + 1) * CELL_M
            y_max = half_y - BORDER_M - b["rmin"] * CELL_M
            lines.append(f"  level {lv}: x=[{x_min:+.1f}, {x_max:+.1f}]  "
                         f"y=[{y_min:+.1f}, {y_max:+.1f}]  "
                         f"({b['rmax']-b['rmin']+1}×{b['cmax']-b['cmin']+1} cells)")

        lines.append("")
        lines.append("# === Spawn positions ===")
        for lv in sorted(level_bounds):
            b = level_bounds[lv]
            spawn_cells = [(r,c) for r in range(b["rmin"], b["rmax"]+1)
                           for c in range(b["cmin"], b["cmax"]+1)
                           if self.cell_types.get((r,c), {}).get("spawn", True)]
            if spawn_cells:
                lines.append(f"  level {lv}: {len(spawn_cells)} spawn cells")
                for (rr, cc) in spawn_cells:
                    cx = -half_x + BORDER_M + cc * CELL_M + CELL_M / 2
                    cy = half_y - BORDER_M - rr * CELL_M - CELL_M / 2
                    ct = self.cell_types.get((rr,cc), {})
                    if ct.get("type") == "flat":
                        z_plat = 0
                    else:
                        ref_cm = ct.get("ref_plane_cm", DEFAULT_REF_PLANE_CM)
                        sh = ct.get("step_h", 0)
                        ns = ct.get("num_steps", 0)
                        concave = ct["type"] == "concave"
                        z_plat = (ref_cm - ns * sh * 100) / 100.0 if concave else (ref_cm + ns * sh * 100) / 100.0
                    lines.append(f"    ({rr},{cc}) x={cx:+.1f} y={cy:+.1f}  {ct['type']} "
                                 f"z_plat={z_plat*100:.0f}cm")

        lines.append("")
        lines.append("# === Pyramid tread details ===")
        for r in range(self.rows):
            for c in range(self.cols):
                ct = self.cell_types.get((r,c), {"type": "flat"})
                if ct["type"] == "flat":
                    continue
                cx = -half_x + BORDER_M + c * CELL_M + CELL_M / 2
                cy = half_y - BORDER_M - r * CELL_M - CELL_M / 2
                sh = ct.get("step_h", DEFAULT_STEP_H)
                sd = ct.get("step_d", DEFAULT_STEP_D)
                ns = ct.get("num_steps", DEFAULT_NUM_STEPS)
                concave = ct["type"] == "concave"
                ref_cm = ct.get("ref_plane_cm", DEFAULT_REF_PLANE_CM)
                ref_z = ref_cm / 100.0
                h_vs = int(sh / VS)
                p2 = PLATFORM_PX // 2
                d_px = int(sd / HS)
                plat_z = (ref_cm - ns * sh * 100) / 100.0 if concave else (ref_cm + ns * sh * 100) / 100.0
                lines.append(f"  ({r},{c}) {ct['type']}  center=({cx:+.1f}, {cy:+.1f})  "
                             f"ref_plane={ref_z*100:.0f}cm  platform={plat_z*100:.0f}cm  "
                             f"step_h={sh*100:.0f}cm  tread={sd*100:.0f}cm  steps={ns}")
                for i in range(ns + 1):
                    half_m = (p2 + (ns - i) * d_px) * HS
                    if concave:
                        z = ref_z - (h_vs * i) * VS
                    else:
                        z = ref_z + (h_vs * i) * VS
                    ring_type = "platform" if i == ns else "ring"
                    lines.append(f"    {ring_type} {i}: z={z*100:5.0f}cm  "
                                 f"half={half_m:.2f}m  "
                                 f"x=[{cx-half_m:+.1f},{cx+half_m:+.1f}] "
                                 f"y=[{cy-half_m:+.1f},{cy+half_m:+.1f}]")

        lines.append("")
        lines.append(f"# Training: DREAMWAQ_TERRAIN=flat_stairs "
                     f"uv run scripts/train_dreamwaq_rsl.py --level N")

        info = "\n".join(lines)
        print(info)
        messagebox.showinfo("Exported", f"{out_path}\n\n{info}")


if __name__ == "__main__":
    root = tk.Tk()
    root.geometry("900x550")
    TerrainEditor(root)
    root.mainloop()