#!/usr/bin/env python3
"""
world_gen.py — Generate a simulatable world from text / image / video.

Pipeline (mirrors the robot intent path: language → plan → generate):
  input + a STRUCTURED world spec from System 2 (planner.planWorldSpec, passed as
  --spec) — or, offline, a heuristic parse of the brief —
    → normalise into a typed world (environment template + surface + structures +
      objects with shape/colour/size/placement)
    → emit:
        world_spec.json   — the normalised world
        scene.xml         — a MuJoCo scene built from it, of the PROPER TYPE
                            (warehouse→shelves/pallets, kitchen→counter, …) that
                            actually LOADS, so the world is simulatable
        preview.gif       — procedural preview (PIL, CPU, no GPU)
        dream.mp4         — LTX-Video neural generation when torch+CUDA+weights
                            are present (optional).

Usage:
  python scripts/world_gen.py --mode text --input "<brief>" [--spec spec.json] --out <dir>
"""
from __future__ import annotations

import argparse
import json
import logging
import math
import pathlib
import time
from typing import List, Optional

log = logging.getLogger("world_gen")

# ── vocabulary ────────────────────────────────────────────────────────────────
_COLORS = {
    "red": (0.85, 0.2, 0.2), "blue": (0.3, 0.4, 0.9), "green": (0.2, 0.7, 0.3),
    "yellow": (0.9, 0.85, 0.2), "orange": (0.95, 0.55, 0.15),
    "purple": (0.6, 0.3, 0.8), "white": (0.9, 0.9, 0.9),
    "black": (0.12, 0.12, 0.14), "gray": (0.5, 0.5, 0.52), "grey": (0.5, 0.5, 0.52),
}
_SHAPES = {
    "cube": "box", "box": "box", "block": "box", "crate": "box", "carton": "box",
    "brick": "box", "container": "box", "package": "box", "parcel": "box",
    "bin": "box", "book": "box", "tile": "box",
    "ball": "sphere", "sphere": "sphere", "orb": "sphere", "marble": "sphere",
    "egg": "sphere",
    "cylinder": "cylinder", "can": "cylinder", "bottle": "cylinder",
    "cup": "cylinder", "mug": "cylinder", "pillar": "cylinder",
    "barrel": "cylinder", "drum": "cylinder", "keg": "cylinder", "tin": "cylinder",
    "jar": "cylinder", "vase": "cylinder", "tube": "cylinder", "pipe": "cylinder",
    "capsule": "capsule", "pill": "capsule",
}
_SIZES = {"small": 0.018, "medium": 0.028, "large": 0.045, "tiny": 0.014, "big": 0.05}
_NUMS = {"a": 1, "an": 1, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5,
         "six": 6, "couple": 2, "few": 3, "several": 4, "pair": 2}

# environment → (background rgb, surface type)
_ENVIRONMENTS = {
    "warehouse": ((0.20, 0.22, 0.26), "floor"),
    "factory":   ((0.24, 0.26, 0.29), "floor"),
    "garage":    ((0.28, 0.28, 0.30), "floor"),
    "kitchen":   ((0.86, 0.84, 0.80), "counter"),
    "lab":       ((0.82, 0.86, 0.90), "table"),
    "office":    ((0.78, 0.76, 0.72), "table"),
    "room":      ((0.80, 0.78, 0.74), "floor"),
    "outdoor":   ((0.45, 0.65, 0.85), "ground"),
    "tabletop":  ((0.30, 0.50, 0.70), "table"),
    "table":     ((0.30, 0.50, 0.70), "table"),
}
# surface → top height (z) objects rest on, and the usable area (centre + half-extent)
_SURFACES = {
    "table":   dict(z=0.42, cx=0.36, cy=0.0, hx=0.16, hy=0.32),
    "counter": dict(z=0.45, cx=0.45, cy=0.0, hx=0.30, hy=0.18),
    "shelf":   dict(z=0.50, cx=0.40, cy=0.0, hx=0.18, hy=0.30),
    "floor":   dict(z=0.00, cx=0.60, cy=0.0, hx=0.35, hy=0.45),
    "ground":  dict(z=0.00, cx=0.60, cy=0.0, hx=0.40, hy=0.50),
}
_DEFAULT_STRUCTURES = {
    "warehouse": ["rack", "pallet"], "factory": ["rack", "crate-stack"],
    "garage": ["shelf"], "kitchen": ["counter"], "room": [], "outdoor": [],
    "lab": [], "office": [], "tabletop": [], "table": [],
}


def _size_m(v) -> float:
    if isinstance(v, (int, float)):
        return max(0.01, min(0.06, float(v)))
    return _SIZES.get(str(v).lower(), 0.028)


def _rgba(color: str):
    return list(_COLORS.get(str(color).lower(), (0.6, 0.6, 0.62))) + [1.0]


def _layout(objects: List[dict], surface: dict):
    """Place objects in a centred grid on the surface; set pos + resting z."""
    n = max(1, len(objects))
    cols = int(math.ceil(math.sqrt(n)))
    rows = int(math.ceil(n / cols))
    for i, o in enumerate(objects):
        r, cc = divmod(i, cols)
        fx = (cc - (cols - 1) / 2.0) / max(1, cols) * 2.0   # -1..1
        fy = (r - (rows - 1) / 2.0) / max(1, rows) * 2.0
        s = o["size_m"]
        o["pos"] = [round(surface["cx"] + fx * surface["hx"] * 0.8, 3),
                    round(surface["cy"] + fy * surface["hy"] * 0.8, 3),
                    round(surface["z"] + s + 0.005, 3)]


# ── spec construction (LLM path + heuristic fallback) ─────────────────────────

def normalize_spec(llm: dict, mode: str, src: str) -> dict:
    """Coerce a System-2 JSON world spec into the internal, buildable spec."""
    env = str(llm.get("environment", "tabletop")).lower()
    if env not in _ENVIRONMENTS:
        env = "tabletop"
    bg, default_surface = _ENVIRONMENTS[env]
    surf_type = str(llm.get("surface", default_surface)).lower()
    if surf_type not in _SURFACES:
        surf_type = default_surface
    surface = dict(_SURFACES[surf_type]); surface["type"] = surf_type

    objects: List[dict] = []
    for o in llm.get("objects", []):
        shape = _SHAPES.get(str(o.get("shape", "box")).lower(), "box")
        color = str(o.get("color", "gray")).lower()
        size_m = _size_m(o.get("size", "medium"))
        count = int(o.get("count", 1) or 1)
        for _ in range(max(1, min(6, count))):
            objects.append({"name": o.get("name", f"{color} {shape}"),
                            "shape": shape, "color": color,
                            "rgba": _rgba(color), "size_m": size_m,
                            "material": o.get("material", "plastic")})
    objects = objects[:12] or [{"name": "red cube", "shape": "box", "color": "red",
                                "rgba": _rgba("red"), "size_m": 0.028, "material": "plastic"}]
    _layout(objects, surface)

    structures = [s for s in llm.get("structures", _DEFAULT_STRUCTURES.get(env, []))]
    return _finish_spec(llm.get("title", ""), env, bg, surface, structures, objects,
                        str(llm.get("lighting", "bright")), mode, src)


def parse_world_spec(prompt: str, mode: str, src: str) -> dict:
    """Offline heuristic: ground a brief into the internal spec (LLM fallback)."""
    t = (prompt or "").lower()
    env = next((e for e in _ENVIRONMENTS if e in t), "tabletop")
    bg, default_surface = _ENVIRONMENTS[env]
    surf_type = default_surface
    for s in _SURFACES:
        if s in t:
            surf_type = s; break
    surface = dict(_SURFACES[surf_type]); surface["type"] = surf_type

    objects: List[dict] = []
    words = t.replace(",", " ").split()
    for i, w in enumerate(words):
        cand = [w, w[:-1] if w.endswith("s") else w, w[:-2] if w.endswith("es") else w]
        shape = next((_SHAPES[c] for c in cand if c in _SHAPES), None)
        if not shape:
            continue
        color, count, size = "gray", 1, "medium"
        for j in range(max(0, i - 4), i):
            if words[j] in _COLORS: color = words[j]
            if words[j] in _NUMS:   count = _NUMS[words[j]]
            if words[j] in _SIZES:  size = words[j]
        for _ in range(count):
            objects.append({"name": f"{color} {shape}", "shape": shape, "color": color,
                            "rgba": _rgba(color), "size_m": _size_m(size), "material": "plastic"})
    if not objects:
        objects = [{"name": "red cube", "shape": "box", "color": "red", "rgba": _rgba("red"),
                    "size_m": 0.028, "material": "plastic"},
                   {"name": "blue ball", "shape": "sphere", "color": "blue", "rgba": _rgba("blue"),
                    "size_m": 0.028, "material": "plastic"}]
    objects = objects[:12]
    _layout(objects, surface)
    structures = _DEFAULT_STRUCTURES.get(env, [])
    light = "dim" if env in ("warehouse", "factory", "garage") else "bright"
    return _finish_spec(prompt, env, bg, surface, structures, objects, light, mode, src)


def _finish_spec(title, env, bg, surface, structures, objects, light, mode, src) -> dict:
    return {
        "title": (title or "generated world").strip()[:80],
        "source": {"mode": mode, "input": src},
        "environment": env,
        "background_rgb": list(bg),
        "surface": surface,
        "structures": list(structures)[:4],
        "lighting": light,
        "objects": objects,
        "camera": {"distance": 2.2, "azimuth": 150, "elevation": -20},
        "created": time.strftime("%Y-%m-%dT%H:%M:%S"),
    }


# ── MuJoCo scene of the proper type ───────────────────────────────────────────

def _surface_geom(surface: dict) -> str:
    st = surface["type"]
    cx, cy, z = surface["cx"], surface["cy"], surface["z"]
    if st in ("floor", "ground"):
        return ""   # the world plane is the surface
    hx, hy = surface["hx"] + 0.02, surface["hy"] + 0.02
    top = z + 0.02
    leg = (f'<geom type="box" size="0.02 0.02 {z/2:.3f}" pos="0 0 {-z/2:.3f}" material="wood"/>'
           if z > 0.05 else "")
    return (f'    <body name="surface" pos="{cx} {cy} {z}">\n'
            f'      <geom name="surface_top" type="box" size="{hx:.3f} {hy:.3f} 0.02" material="wood"/>\n'
            f'      {leg}\n    </body>')


def _structures_geom(spec: dict) -> str:
    """Environment furniture appropriate to the world type."""
    out = []
    n = 0
    for s in spec.get("structures", []):
        n += 1
        if s in ("rack", "shelf"):
            # back rack with two shelf levels
            out.append(
                f'    <body name="rack{n}" pos="0.78 0 0">\n'
                f'      <geom type="box" size="0.03 0.45 0.45" pos="0 0 0.45" rgba="0.4 0.4 0.45 1"/>\n'
                f'      <geom type="box" size="0.18 0.45 0.02" pos="-0.18 0 0.35" rgba="0.5 0.5 0.55 1"/>\n'
                f'      <geom type="box" size="0.18 0.45 0.02" pos="-0.18 0 0.70" rgba="0.5 0.5 0.55 1"/>\n'
                f'    </body>')
        elif s == "pallet":
            out.append(
                f'    <body name="pallet{n}" pos="0.55 0.28 0.03">\n'
                f'      <geom type="box" size="0.16 0.12 0.03" rgba="0.55 0.40 0.22 1"/>\n'
                f'    </body>')
        elif s in ("crate-stack", "bin"):
            out.append(
                f'    <body name="crates{n}" pos="0.55 -0.30 0.06">\n'
                f'      <geom type="box" size="0.08 0.08 0.06" rgba="0.6 0.45 0.25 1"/>\n'
                f'      <geom type="box" size="0.07 0.07 0.05" pos="0 0 0.11" rgba="0.62 0.47 0.27 1"/>\n'
                f'    </body>')
        elif s == "counter":
            pass   # the surface already is the counter
    return "\n".join(out)


def _obj_geom(o: dict) -> str:
    r, g, b, a = o["rgba"]; s = o["size_m"]
    if o["shape"] == "sphere":
        return f'<geom type="sphere" size="{s}" rgba="{r} {g} {b} {a}" mass="0.1"/>'
    if o["shape"] == "cylinder":
        return f'<geom type="cylinder" size="{s} {s}" rgba="{r} {g} {b} {a}" mass="0.1"/>'
    if o["shape"] == "capsule":
        return f'<geom type="capsule" size="{s} {s}" rgba="{r} {g} {b} {a}" mass="0.1"/>'
    return f'<geom type="box" size="{s} {s} {s}" rgba="{r} {g} {b} {a}" mass="0.1"/>'


def spec_to_mjcf(spec: dict) -> str:
    bg = spec["background_rgb"]
    light = spec.get("lighting", "bright")
    diff = {"dim": "0.45 0.45 0.5", "warm": "0.95 0.85 0.7", "cool": "0.7 0.8 0.95",
            "natural": "0.9 0.9 0.85", "bright": "1 1 1"}.get(light, "0.9 0.9 0.9")

    bodies = []
    for i, o in enumerate(spec["objects"]):
        px, py, pz = o["pos"]
        bodies.append(
            f'    <body name="obj{i}_{o["color"]}_{o["shape"]}" pos="{px} {py} {pz}">\n'
            f'      <freejoint/>\n      {_obj_geom(o)}\n    </body>')
    objects_xml = "\n".join(bodies)
    surface_xml = _surface_geom(spec["surface"])
    structures_xml = _structures_geom(spec)
    target = "surface" if spec["surface"]["type"] not in ("floor", "ground") else "obj0_" + \
        f'{spec["objects"][0]["color"]}_{spec["objects"][0]["shape"]}'

    return f"""<mujoco model="generated_world">
  <option timestep="0.002" gravity="0 0 -9.81" integrator="implicitfast"/>
  <visual><global offwidth="896" offheight="512"/></visual>
  <asset>
    <texture name="sky" type="skybox" builtin="gradient"
             rgb1="{bg[0]} {bg[1]} {bg[2]}" rgb2="0 0 0" width="512" height="512"/>
    <texture name="grid" type="2d" builtin="checker" rgb1="0.2 0.3 0.4"
             rgb2="0.1 0.15 0.2" width="512" height="512"/>
    <material name="grid" texture="grid" texrepeat="6 6" reflectance="0.1"/>
    <material name="wood" rgba="0.7 0.5 0.3 1"/>
  </asset>
  <worldbody>
    <light name="top" pos="0 0 3" dir="0 0 -1" diffuse="{diff}" specular="0.2 0.2 0.2"/>
    <geom name="floor" type="plane" size="5 5 0.1" material="grid"/>
    <camera name="overview" pos="1.6 -0.7 1.2" mode="targetbody" target="{target}" fovy="55"/>
{surface_xml}
{structures_xml}
{objects_xml}
  </worldbody>
</mujoco>
"""


# ── procedural preview (PIL) ──────────────────────────────────────────────────

def render_preview_gif(spec: dict, out: pathlib.Path, frames: int = 16) -> Optional[pathlib.Path]:
    try:
        from PIL import Image, ImageDraw
    except Exception:
        return None
    W = H = 320
    bg = tuple(int(255 * c) for c in spec["background_rgb"])
    floor = (90, 65, 40) if spec["surface"]["type"] in ("table", "counter", "shelf") else (70, 70, 75)
    objs = spec["objects"]; n = len(objs)
    imgs = []
    for f in range(frames):
        im = Image.new("RGB", (W, H), bg); d = ImageDraw.Draw(im)
        d.rectangle([0, int(H * 0.70), W, H], fill=floor)
        if "rack" in spec["structures"] or "shelf" in spec["structures"]:
            d.rectangle([W - 40, 40, W - 20, int(H * 0.72)], fill=(110, 110, 120))
        for k, o in enumerate(objs):
            col = tuple(int(255 * c) for c in o["rgba"][:3])
            cx = int(W * (0.18 + 0.64 * (k + 0.5) / max(1, n)))
            cy = int(H * 0.58) + int(4 * math.sin(2 * math.pi * f / frames + k))
            r = int(14 + 360 * o["size_m"])
            if o["shape"] == "sphere":
                d.ellipse([cx - r, cy - r, cx + r, cy + r], fill=col, outline=(0, 0, 0))
            else:
                d.rectangle([cx - r, cy - r, cx + r, cy + r], fill=col, outline=(0, 0, 0))
        d.text((8, 8), f'{spec["environment"]}: {spec["title"][:32]}', fill=(235, 235, 235))
        imgs.append(im)
    path = out / "preview.gif"
    imgs[0].save(path, save_all=True, append_images=imgs[1:], duration=80, loop=0)
    return path


# ── optional LTX neural generation ────────────────────────────────────────────

def try_ltx(spec, mode, src, out, frames, steps) -> Optional[pathlib.Path]:
    try:
        import torch
        if not torch.cuda.is_available():
            log.info("LTX: no CUDA — skipping neural video (procedural world still produced).")
            return None
    except Exception:
        log.info("LTX: torch not installed — skipping neural video.")
        return None
    try:
        import numpy as np, imageio.v2 as imageio, sys
        from PIL import Image
        sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))
        from rx1_brain.world_model_ltx import LTXWorldModel
        wm = LTXWorldModel(num_frames=frames, steps=steps)
        prompt = spec["title"]
        if mode == "image":
            ego = np.asarray(Image.open(src).convert("RGB").resize((wm.width, wm.height)))
            video = wm.imagine(ego, prompt)
        elif mode == "video":
            seq = [np.asarray(Image.fromarray(fr).convert("RGB")) for fr in imageio.get_reader(src)]
            _, video = wm.predict_next(seq, prompt)
        else:
            ego = np.full((wm.height, wm.width, 3), 30, np.uint8)
            video = wm.imagine(ego, prompt)
        path = out / "dream.mp4"; imageio.mimsave(path, list(video), fps=16)
        return path
    except Exception as e:
        log.warning("LTX generation failed (%s); procedural world still produced.", str(e)[:160])
        return None


# ── main ───────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(description="Generate a simulatable world")
    ap.add_argument("--mode", choices=["text", "image", "video"], default="text")
    ap.add_argument("--input", required=True)
    ap.add_argument("--prompt", default="")
    ap.add_argument("--spec", default="", help="path to a System-2 JSON world spec")
    ap.add_argument("--out", default="")
    ap.add_argument("--backend", choices=["auto", "stub", "ltx"], default="auto")
    ap.add_argument("--frames", type=int, default=33)
    ap.add_argument("--steps", type=int, default=28)
    args = ap.parse_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(name)s  %(levelname)s  %(message)s",
                        datefmt="%H:%M:%S")

    out = pathlib.Path(args.out) if args.out else (
        pathlib.Path.home() / ".zelpi-worlds" / time.strftime("%Y%m%d-%H%M%S"))
    out.mkdir(parents=True, exist_ok=True)

    # Prefer the System-2 structured spec; fall back to heuristic parsing.
    spec = None
    if args.spec and pathlib.Path(args.spec).exists():
        try:
            spec = normalize_spec(json.loads(pathlib.Path(args.spec).read_text()), args.mode, args.input)
            log.info("WORLD-GEN using System-2 structured spec")
        except Exception as e:
            log.warning("spec parse failed (%s); using heuristic parser", str(e)[:120])
    if spec is None:
        spec = parse_world_spec(args.prompt or args.input, args.mode, args.input)
        log.info("WORLD-GEN heuristic parse")

    log.info("WORLD env=%s surface=%s structures=%s objects=%s",
             spec["environment"], spec["surface"]["type"], spec["structures"],
             [f'{o["color"]} {o["shape"]}' for o in spec["objects"]])

    (out / "world_spec.json").write_text(json.dumps(spec, indent=2))
    mjcf = spec_to_mjcf(spec)
    scene_path = out / "scene.xml"; scene_path.write_text(mjcf)
    try:
        import mujoco
        m = mujoco.MjModel.from_xml_string(mjcf)
        log.info("SCENE ok — loads in MuJoCo: nbody=%d ngeom=%d (simulatable)", m.nbody, m.ngeom)
    except Exception as e:
        log.warning("scene.xml did not validate: %s", str(e)[:200])

    gif = render_preview_gif(spec, out)
    if gif:
        log.info("PREVIEW %s", gif)
    dream = try_ltx(spec, args.mode, args.input, out, args.frames, args.steps) \
        if args.backend in ("auto", "ltx") else None
    if dream:
        log.info("DREAM (LTX) %s", dream)

    result = {"out_dir": str(out), "world_spec": str(out / "world_spec.json"),
              "scene": str(scene_path), "preview": str(gif) if gif else None,
              "dream": str(dream) if dream else None,
              "backend": "ltx" if dream else "procedural",
              "objects": len(spec["objects"]), "environment": spec["environment"],
              "surface": spec["surface"]["type"], "structures": spec["structures"]}
    (out / "result.json").write_text(json.dumps(result, indent=2))
    print("WORLD_GEN_RESULT " + json.dumps(result))


if __name__ == "__main__":
    main()