src/world3d/physics.py

"""Authoritative 3D rigid-body physics via MuJoCo (V6).

Character pattern: dynamic capsule body, planar velocity servo set from the
action each step, vertical free (gravity integrates). Contacts are resolved
by the solver: walls block, floors support, dynamic bodies get pushed.
Position is NEVER teleported by movement code (only save/restore sets state).

Capabilities verified in tests/test_v6_physics.py: gravity, floor contact +
stable rest, wall blocking, object pushing, raycast queries, save/restore.
"""
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

try:
    import mujoco
    _MUJOCO = True
except Exception:
    mujoco = None  # type: ignore
    _MUJOCO = False

PHYSICS_DT = 0.01          # 100 Hz substeps
CHARACTER_R = 0.3
CHARACTER_H = 0.9          # capsule cylinder section
EYE_HEIGHT = 1.45


def _box_wall(cx, cy, sx, sy, h, name):
    return (f'<geom name="{name}" type="box" pos="{cx} {cy} {h / 2}" '
            f'size="{sx / 2} {sy / 2} {h / 2}" rgba="0.45 0.4 0.5 1"/>')


def build_model_xml(spec: Dict[str, Any],
                    extra_bodies: Optional[List[Dict[str, Any]]] = None) -> str:
    """Compile the world spec (+ characters/foods) to a MuJoCo model."""
    size = float(spec.get("size", 40.0))
    parts = [f'<mujoco><compiler angle="degree"/>',
             f'<option timestep="0.01" gravity="0 0 -9.81"/>',
             f'<worldbody>',
             f'<geom name="ground" type="plane" size="{size / 2} {size / 2} 0.1" '
             f'material="ground"/>']
    home = spec["home"]
    cx, cy, w, d, h, t = (home["cx"], home["cy"], home["w"], home["d"],
                          home["wall_h"], home["wall_t"])
    dw = home["door_w"]
    # north wall (full), south wall split for door gap, east/west full
    parts.append(_box_wall(cx, cy + d / 2, w, t, h, "wall_north"))
    parts.append(_box_wall(cx - (dw / 2 + (w - dw) / 4), cy - d / 2, (w - dw) / 2, t, h,
                           "wall_south_l"))
    parts.append(_box_wall(cx + (dw / 2 + (w - dw) / 4), cy - d / 2, (w - dw) / 2, t, h,
                           "wall_south_r"))
    parts.append(_box_wall(cx - w / 2, cy, t, d, h, "wall_west"))
    parts.append(_box_wall(cx + w / 2, cy, t, d, h, "wall_east"))
    for f in spec.get("furniture", []):
        if f["kind"] == "static_box":
            parts.append(_box_wall(f["x"], f["y"], f["sx"], f["sy"], f["sz"],
                                   "furn_" + str(f["id"])))
    for tr in spec.get("trees", []):
        parts.append(f'<geom name="trunk_{tr["id"]}" type="cylinder" '
                     f'pos="{tr["x"]} {tr["y"]} 1.0" size="0.25 2.0" '
                     f'rgba="0.35 0.25 0.15 1"/>')
    for rk in spec.get("rocks", []):
        parts.append(f'<geom name="rock_{rk["id"]}" type="sphere" '
                     f'pos="{rk["x"]} {rk["y"]} {rk["r"]}" size="{rk["r"]}" '
                     f'rgba="0.5 0.5 0.55 1"/>')
    for st in spec.get("structures", []):
        sid = st.get("id", "struct")
        x, y, z = st.get("x", 0.0), st.get("y", 0.0), st.get("z", 0.0)
        sx, sy, sz = st.get("sx", 1.0), st.get("sy", 1.0), st.get("sz", 0.5)
        rgba = st.get("rgba", "0.55 0.35 0.15 1")
        gtype = st.get("geom_type", "box")
        if not st.get("dynamic", False):
            parts.append(f'<geom name="struct_{sid}" type="{gtype}" '
                         f'pos="{x} {y} {z + sz / 2}" size="{sx / 2} {sy / 2} {sz / 2}" '
                         f'rgba="{rgba}"/>')
        else:
            parts.append(f'<body name="bod_struct_{sid}" pos="{x} {y} {z + sz / 2}">'
                         f'<freejoint/><geom name="dyn_struct_{sid}" type="{gtype}" '
                         f'size="{sx / 2} {sy / 2} {sz / 2}" mass="{st.get("mass", 50.0)}" '
                         f'rgba="{rgba}"/></body>')

    # dynamic bodies: crates, doors (hinged), foods, characters
    for f in spec.get("furniture", []):
        if f["kind"] == "dynamic_box":
            parts.append(
                f'<body name="bod_{f["id"]}" pos="{f["x"]} {f["y"]} {f["sz"] / 2}">'
                f'<freejoint/><geom name="dyn_{f["id"]}" type="box" '
                f'size="{f["sx"] / 2} {f["sy"] / 2} {f["sz"] / 2}" mass="{f.get("mass", 4.0)}" '
                f'rgba="0.6 0.45 0.2 1"/></body>')
    for door in spec.get("doors", []):
        parts.append(
            f'<body name="bod_{door["id"]}" pos="{door["x"] - door["w"] / 2} {door["y"]} 1.0">'
            f'<joint name="hinge_{door["id"]}" type="hinge" axis="0 0 1" damping="2.0"/>'
            f'<geom name="dyn_{door["id"]}" type="box" '
            f'pos="{door["w"] / 2} 0 0" size="{door["w"] / 2} 0.05 {door["h"] / 2}" '
            f'mass="6.0" rgba="0.55 0.35 0.2 1"/></body>')
    for fd in spec.get("foods", []):
        parts.append(
            f'<body name="bod_{fd["id"]}" pos="{fd["x"]} {fd["y"]} 0.25">'
            f'<freejoint/><geom name="food_{fd["id"]}" type="sphere" size="0.18" '
            f'mass="0.3" rgba="0.2 0.8 0.3 1"/></body>')
    for ch in (extra_bodies or []):
        parts.append(
            f'<body name="{ch["name"]}" pos="{ch["x"]} {ch["y"]} 1.2">'
            f'<freejoint/><geom name="char_{ch["name"]}" type="capsule" '
            f'size="{CHARACTER_R} {CHARACTER_H / 2}" mass="60.0" '
            f'rgba="0.2 0.7 0.9 1"/>'
            f'<geom name="head_{ch["name"]}" type="sphere" pos="0 0 0.85" size="0.22" '
            f'mass="5.0" rgba="0.9 0.75 0.6 1"/>'
            f'<site name="eye_{ch["name"]}" pos="0 0.18 0.95"/></body>')
    parts.append('</worldbody>')
    parts.append('<asset><material name="ground" rgba="0.25 0.35 0.22 1"/></asset>')
    parts.append('</mujoco>')
    return "\n".join(parts)


class PhysicsWorld:
    """Owns MuJoCo model+data, steps, queries, save/restore."""

    def __init__(self, spec: Dict[str, Any],
                 characters: Optional[List[Dict[str, Any]]] = None):
        if not _MUJOCO:
            raise RuntimeError("mujoco unavailable: 3D physics cannot run here")
        self.spec = spec
        self.xml = build_model_xml(spec, characters)
        self.model = mujoco.MjModel.from_xml_string(self.xml)
        self.data = mujoco.MjData(self.model)
        mujoco.mj_forward(self.model, self.data)
        self.time = 0.0
        # Character bodies under posture stabilization (documented embodiment:
        # strong angular damping keeps the capsule upright; no teleporting,
        # no kinematic freezing — contacts and gravity still fully simulated).
        self._characters: List[str] = [c["name"] for c in (characters or [])]
        self.posture_damping: float = 0.25
        self._cmd: Dict[str, List[float]] = {c["name"]: [0.0, 0.0]
                                             for c in (characters or [])}

    # ---- characters ----
    def _body_id(self, name: str) -> int:
        bid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, name)
        if bid < 0:
            raise KeyError(f"no body {name!r}")
        return bid

    def _qpos_adr(self, name: str) -> int:
        bid = self._body_id(name)
        jnt = int(self.model.body_jntadr[bid])
        return int(self.model.jnt_qposadr[jnt])

    def _dof_adr(self, name: str) -> int:
        return int(self.model.body_dofadr[self._body_id(name)])

    def char_state(self, name: str) -> Dict[str, Any]:
        qa, da = self._qpos_adr(name), self._dof_adr(name)
        qpos = self.data.qpos[qa:qa + 7].copy()
        qvel = self.data.qvel[da:da + 6].copy()
        pos = [round(float(qpos[0]), 4), round(float(qpos[1]), 4), round(float(qpos[2]), 4)]
        vel = [round(float(qvel[0]), 4), round(float(qvel[1]), 4), round(float(qvel[2]), 4)]
        return {"pos": pos, "vel": vel, "quat": [round(float(x), 4) for x in qpos[3:7]]}

    def drive_character(self, name: str, vx: float, vy: float) -> None:
        """Velocity servo (called BEFORE step): planar velocity command.
        Contacts + gravity still resolved by the solver inside step()."""
        da = self._dof_adr(name)
        self.data.qvel[da + 0] = float(vx)
        self.data.qvel[da + 1] = float(vy)
        self._cmd[name] = [float(vx), float(vy)]

    def push_character(self, name: str, fx: float, fy: float, fz: float = 0.0) -> None:
        da = self._dof_adr(name)
        self.data.qfrc_applied[da + 0] = float(fx)
        self.data.qfrc_applied[da + 1] = float(fy)
        self.data.qfrc_applied[da + 2] = float(fz)

    def teleport(self, name: str, x: float, y: float, z: float) -> None:
        """SAVE/RESTORE ONLY. Movement code must never call this."""
        qa, da = self._qpos_adr(name), self._dof_adr(name)
        self.data.qpos[qa:qa + 3] = [float(x), float(y), float(z)]
        self.data.qvel[da:da + 6] = 0.0
        mujoco.mj_forward(self.model, self.data)

    def _stabilize(self, name: str) -> None:
        """Attitude PD controller (real torques, heading left free).

        Drives body-up toward world-up while preserving yaw. Contacts,
        gravity and pushes still act fully — this only resists toppling,
        like vestibular balance. Gains are part of the embodiment model.
        """
        qa, da = self._qpos_adr(name), self._dof_adr(name)
        qw, qx, qy, qz = (float(v) for v in self.data.qpos[qa + 3:qa + 7])
        n = (qw * qw + qx * qx + qy * qy + qz * qz) ** 0.5 or 1.0
        qw, qx, qy, qz = qw / n, qx / n, qy / n, qz / n
        # current yaw; target = yaw-only orientation (upright, same heading)
        yaw = np.arctan2(2.0 * (qw * qz + qx * qy), 1.0 - 2.0 * (qy * qy + qz * qz))
        tw, tx, ty, tz = np.cos(yaw / 2), 0.0, 0.0, np.sin(yaw / 2)
        # error quat q_err = q_target * conj(q): vector part ~ rotation axis*angle/2
        ex = tw * -qx + tx * qw + ty * -qz + tz * qy
        ey = tw * -qy + tx * qz + ty * qw + tz * -qx
        ez = tw * -qz + tx * -qy + ty * qx + tz * qw
        ew = tw * qw + tx * qx + ty * qy + tz * qz
        s = 1.0 if ew >= 0 else -1.0
        kp, kd = 300.0, 20.0
        # Get-up maneuver: fallen and (nearly) still -> strong righting.
        # Falls are real events (see body damage); recovery is a real maneuver,
        # not a teleport: torques integrate through the solver.
        up_z = 1.0 - 2.0 * (qx * qx + qy * qy)
        cmd = self._cmd.get(name, [0.0, 0.0])
        still = abs(cmd[0]) + abs(cmd[1]) < 0.05
        if up_z < 0.4 and still:
            kp, kd = 1200.0, 40.0
        wx, wy, wz = (float(v) for v in self.data.qvel[da + 3:da + 6])
        # down-weight yaw correction (heading is driven by locomotion)
        tx_, ty_, tz_ = (kp * s * ex - kd * wx, kp * s * ey - kd * wy,
                         0.25 * (kp * s * ez - kd * wz))
        self.data.qfrc_applied[da + 3] += tx_
        self.data.qfrc_applied[da + 4] += ty_
        self.data.qfrc_applied[da + 5] += tz_

    # ---- stepping ----
    def step(self, dt: float = PHYSICS_DT) -> None:
        for name in self._characters:
            try:
                self._stabilize(name)
            except KeyError:
                pass
        mujoco.mj_step(self.model, self.data)
        self.time = float(self.data.time)
        for name in self._characters:
            try:
                # Applied forces are single-step impulses: re-apply each step
                # for continuous push (no hidden persistent forces).
                da = self._dof_adr(name)
                self.data.qfrc_applied[da:da + 6] = 0.0
            except KeyError:
                pass

    def upright(self, name: str) -> float:
        """Up-axis alignment 0..1 (1 = fully upright)."""
        qa = self._qpos_adr(name)
        qw, qx, qy, qz = (float(v) for v in self.data.qpos[qa + 3:qa + 7])
        # rotate world +Z by inverse quat: z-component of body x-axis... use
        # body up vector = R * (0,0,1); take its z via quat formula
        up_z = 1.0 - 2.0 * (qx * qx + qy * qy)
        return round(float(up_z), 4)

    # ---- queries ----
    def contacts(self) -> List[Dict[str, Any]]:
        out = []
        for i in range(self.data.ncon):
            c = self.data.contact[i]
            g1 = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_GEOM, c.geom1) or ""
            g2 = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_GEOM, c.geom2) or ""
            out.append({"geom1": g1, "geom2": g2,
                        "dist": round(float(c.dist), 5)})
        return out

    def char_contacts(self, char_name: str) -> List[str]:
        tag = f"char_{char_name}"
        entities = []
        for c in self.contacts():
            if tag == c["geom1"]:
                entities.append(c["geom2"])
            elif tag == c["geom2"]:
                entities.append(c["geom1"])
        return entities

    def grounded(self, char_name: str) -> bool:
        contacts = self.char_contacts(char_name)
        for name in contacts:
            if "ground" in name or name.startswith("struct_") or "floor" in name:
                return True
        return False


    def raycast(self, origin: List[float], direction: List[float],
                max_dist: float = 20.0, exclude_body: str = "") -> Dict[str, Any]:
        """Single geometric ray (vision primitive). Returns hit entity or miss."""
        geomid = np.zeros(1, dtype=np.int32)
        dist = mujoco.mj_ray(self.model, self.data,
                             np.asarray(origin, dtype=np.float64),
                             np.asarray(direction, dtype=np.float64),
                             None, 1, -1, geomid)
        if dist < 0 or dist > max_dist:
            return {"hit": False, "dist": round(float(max_dist), 3)}
        name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_GEOM, int(geomid[0])) or ""
        if exclude_body and name.endswith(exclude_body):
            return {"hit": False, "dist": round(float(max_dist), 3),
                    "note": "self-hit excluded"}
        return {"hit": True, "dist": round(float(dist), 3), "geom": name,
                "entity": name.split("_", 1)[1] if "_" in name else name}

    # ---- persistence ----
    def snapshot(self) -> Dict[str, Any]:
        return {"time": self.time, "qpos": self.data.qpos.copy().tolist(),
                "qvel": self.data.qvel.copy().tolist(),
                "qacc_warmstart": self.data.qacc_warmstart.copy().tolist(),
                "spec_hash": self.spec.get("spec_hash", "")}

    def restore(self, snap: Dict[str, Any]) -> None:
        if snap.get("spec_hash", "") != self.spec.get("spec_hash", ""):
            raise ValueError("physics snapshot belongs to a different world spec")
        self.data.qpos[:] = np.asarray(snap["qpos"], dtype=np.float64)
        self.data.qvel[:] = np.asarray(snap["qvel"], dtype=np.float64)
        if "qacc_warmstart" in snap:
            self.data.qacc_warmstart[:] = np.asarray(
                snap["qacc_warmstart"], dtype=np.float64)
        mujoco.mj_forward(self.model, self.data)
        self.time = float(snap.get("time", 0.0))