GravityLLM-Space-Demo/utils/scene_tools.py

from __future__ import annotations

import json
import math
import re
from pathlib import Path
from typing import Any, Dict, List, Tuple

import matplotlib.pyplot as plt
from jsonschema import Draft7Validator


SCHEMA_PATH = Path(__file__).resolve().parents[1] / "schemas" / "scene.schema.json"
SCHEMA: Dict[str, Any] = json.loads(SCHEMA_PATH.read_text(encoding="utf-8"))
VALIDATOR = Draft7Validator(SCHEMA)

CLASS_DEFAULTS = {
    "lead_vocal": {"az": 0, "el": 10, "dist": 1.6, "width": 0.15, "gain": 0.0, "rev": 0.18, "er": 0.22},
    "dialogue": {"az": 0, "el": 5, "dist": 1.5, "width": 0.10, "gain": 0.0, "rev": 0.08, "er": 0.18},
    "back_vocal": {"az": -18, "el": 8, "dist": 1.9, "width": 0.25, "gain": -1.2, "rev": 0.16, "er": 0.16},
    "kick": {"az": 0, "el": 0, "dist": 2.2, "width": 0.0, "gain": 0.0, "rev": 0.02, "er": 0.05},
    "snare": {"az": 8, "el": 4, "dist": 2.4, "width": 0.1, "gain": -0.3, "rev": 0.05, "er": 0.07},
    "hihat": {"az": 22, "el": 7, "dist": 2.7, "width": 0.18, "gain": -1.0, "rev": 0.06, "er": 0.08},
    "bass": {"az": 0, "el": -5, "dist": 2.6, "width": 0.05, "gain": -0.5, "rev": 0.03, "er": 0.06},
    "drums_bus": {"az": 0, "el": 2, "dist": 2.8, "width": 0.25, "gain": -0.6, "rev": 0.08, "er": 0.10},
    "pad": {"az": -70, "el": 18, "dist": 4.8, "width": 0.82, "gain": -4.0, "rev": 0.28, "er": 0.12},
    "synth_lead": {"az": 24, "el": 15, "dist": 2.0, "width": 0.35, "gain": -1.0, "rev": 0.12, "er": 0.10},
    "guitar": {"az": -28, "el": 10, "dist": 2.8, "width": 0.28, "gain": -1.5, "rev": 0.09, "er": 0.11},
    "piano": {"az": -14, "el": 8, "dist": 2.6, "width": 0.34, "gain": -1.5, "rev": 0.09, "er": 0.10},
    "fx": {"az": 105, "el": 28, "dist": 6.2, "width": 0.66, "gain": -7.0, "rev": 0.42, "er": 0.08},
    "ambience": {"az": -120, "el": 15, "dist": 8.0, "width": 0.90, "gain": -8.0, "rev": 0.55, "er": 0.10},
    "other": {"az": 0, "el": 8, "dist": 3.0, "width": 0.25, "gain": -2.0, "rev": 0.10, "er": 0.10},
}

SPATIAL_CLASSES = list(CLASS_DEFAULTS)


def clip(value: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, value))


def extract_first_json_block(text: str) -> str:
    match = re.search(r"\{.*\}", text, flags=re.DOTALL)
    return match.group(0).strip() if match else text.strip()


def parse_json_text(text: str) -> Dict[str, Any]:
    return json.loads(extract_first_json_block(text))


def validate_scene(scene: Dict[str, Any]) -> Tuple[bool, List[str]]:
    errors = sorted(VALIDATOR.iter_errors(scene), key=lambda e: list(e.path))
    if not errors:
        return True, []
    messages = []
    for err in errors[:50]:
        path = ".".join(str(x) for x in err.path)
        messages.append(f"{path or '<root>'}: {err.message}")
    return False, messages


def make_motion(az: float, el: float, dist: float, sweep: float = 0.0) -> List[Dict[str, float]]:
    if abs(sweep) < 0.001:
        return [
            {"t": 0.0, "az_deg": round(az, 2), "el_deg": round(el, 2), "dist_m": round(dist, 2)},
            {"t": 1.0, "az_deg": round(az, 2), "el_deg": round(el, 2), "dist_m": round(dist, 2)},
        ]
    return [
        {"t": 0.0, "az_deg": round(az - sweep / 2, 2), "el_deg": round(el, 2), "dist_m": round(dist, 2)},
        {"t": 1.0, "az_deg": round(az + sweep / 2, 2), "el_deg": round(el, 2), "dist_m": round(dist, 2)},
    ]


def find_anchor_rules(payload: Dict[str, Any]) -> Dict[str, Dict[str, float]]:
    anchors: Dict[str, Dict[str, float]] = {}
    for rule in payload.get("rules", []):
        if rule.get("type") == "anchor":
            key = rule.get("track_class")
            if key:
                anchors[key] = {
                    "az": float(rule.get("az_deg", 0.0)),
                    "el": float(rule.get("el_deg", 0.0)),
                    "dist": float(rule.get("dist_m", 1.6)),
                }
    return anchors


def width_preference(payload: Dict[str, Any], track_class: str) -> float | None:
    for rule in payload.get("rules", []):
        if rule.get("type") == "width_pref" and rule.get("track_class") == track_class:
            value = rule.get("min_width")
            if value is not None:
                return float(value)
    return None


def target_layout(payload: Dict[str, Any]) -> str:
    fmt = str(payload.get("target_format", "iamf")).lower()
    if fmt in {"binaural", "iamf", "5.1.4", "7.1.4"}:
        return fmt
    return "iamf"


def room_preset(payload: Dict[str, Any]) -> str:
    style = str(payload.get("style", "neutral")).lower()
    mapping = {
        "club": "club_medium",
        "cinematic": "cinema_large",
        "film": "cinema_large",
        "podcast": "studio_dry",
        "live": "stage_wide",
        "intimate": "studio_small",
    }
    return mapping.get(style, "studio_neutral")


def heuristic_scene(payload: Dict[str, Any]) -> Dict[str, Any]:
    anchors = find_anchor_rules(payload)
    objects: List[Dict[str, Any]] = []
    constraints_applied: List[str] = []
    max_objects = int(payload.get("max_objects", 10))
    style = str(payload.get("style", "neutral")).lower()

    for rule in payload.get("rules", []):
        rtype = rule.get("type")
        if rtype == "anchor":
            constraints_applied.append(
                f"anchor:{rule.get('track_class')}@{rule.get('az_deg')}/{rule.get('el_deg')}/{rule.get('dist_m')}"
            )
        elif rtype == "mono_low_end":
            constraints_applied.append(f"mono_low_end<{rule.get('hz_below', 120)}Hz")
        elif rtype == "width_pref":
            constraints_applied.append(f"{rule.get('track_class')}_width>={rule.get('min_width')}")
        elif rtype == "keep_dialogue_clear":
            band = rule.get("band_hz", [1000, 4000])
            constraints_applied.append(f"keep_dialogue_clear_{band[0]}-{band[1]}Hz")
        elif rtype == "avoid_band_masking":
            band = rule.get("band_hz", [1500, 4500])
            constraints_applied.append(f"avoid_masking_{rule.get('mask_target')}_{band[0]}-{band[1]}Hz")

    for idx, stem in enumerate(payload.get("stems", [])[:max_objects]):
        cls = stem.get("class", "other")
        defaults = dict(CLASS_DEFAULTS.get(cls, CLASS_DEFAULTS["other"]))
        if cls in anchors:
            defaults["az"] = anchors[cls]["az"]
            defaults["el"] = anchors[cls]["el"]
            defaults["dist"] = anchors[cls]["dist"]

        width_min = width_preference(payload, cls)
        if width_min is not None:
            defaults["width"] = max(defaults["width"], width_min)

        leadness = float(stem.get("leadness", 0.0))
        transient = float(stem.get("transient", 0.0))
        lufs = float(stem.get("lufs", -20.0))

        if leadness > 0.8 and cls not in {"kick", "bass", "dialogue", "lead_vocal"}:
            defaults["dist"] = clip(defaults["dist"] - 0.35, 0.8, 15.0)
            defaults["gain"] = clip(defaults["gain"] + 0.8, -60.0, 12.0)

        if transient > 0.7 and cls in {"fx", "synth_lead"}:
            defaults["az"] += 12.0

        if lufs < -24 and cls in {"fx", "ambience", "pad"}:
            defaults["dist"] = clip(defaults["dist"] + 0.5, 0.5, 15.0)

        # Add style-specific touch
        sweep = 0.0
        if style == "club" and cls in {"fx", "synth_lead"}:
            sweep = 28.0 if cls == "fx" else 10.0
        elif style in {"cinematic", "film"} and cls in {"fx", "ambience"}:
            sweep = 18.0
        elif style == "podcast" and cls in {"dialogue", "back_vocal"}:
            sweep = 0.0

        obj = {
            "id": str(stem.get("id", f"obj_{idx+1}")),
            "class": cls if cls in SPATIAL_CLASSES else "other",
            "az_deg": round(clip(defaults["az"], -180, 180), 2),
            "el_deg": round(clip(defaults["el"], -45, 90), 2),
            "dist_m": round(clip(defaults["dist"], 0.5, 15.0), 2),
            "width": round(clip(defaults["width"], 0.0, 1.0), 2),
            "gain_db": round(clip(defaults["gain"], -60, 12), 2),
            "reverb_send": round(clip(defaults["rev"], 0.0, 1.0), 2),
            "early_reflections": round(clip(defaults["er"], 0.0, 1.0), 2),
            "motion": make_motion(defaults["az"], defaults["el"], defaults["dist"], sweep=sweep),
        }

        # Low-end mono safety
        if cls in {"kick", "bass"}:
            obj["az_deg"] = 0.0
            obj["width"] = 0.0 if cls == "kick" else min(obj["width"], 0.08)
            obj["motion"] = make_motion(0.0, obj["el_deg"], obj["dist_m"], sweep=0.0)

        objects.append(obj)

    scene = {
        "version": "1.0",
        "bed": {
            "layout": target_layout(payload),
            "loudness_target_lufs": -16.0 if payload.get("style") == "cinematic" else -14.0,
            "room_preset": room_preset(payload),
        },
        "objects": objects or [
            {
                "id": "placeholder",
                "class": "other",
                "az_deg": 0.0,
                "el_deg": 0.0,
                "dist_m": 2.0,
                "width": 0.2,
                "gain_db": 0.0,
                "reverb_send": 0.1,
                "early_reflections": 0.1,
                "motion": make_motion(0.0, 0.0, 2.0),
            }
        ],
        "constraints_applied": constraints_applied,
    }
    return scene


def scene_stats(scene: Dict[str, Any]) -> Dict[str, Any]:
    objects = scene.get("objects", [])
    dominant = sorted(objects, key=lambda o: float(o.get("gain_db", -99.0)), reverse=True)[:3]
    return {
        "layout": scene.get("bed", {}).get("layout", "unknown"),
        "room_preset": scene.get("bed", {}).get("room_preset", "unknown"),
        "object_count": len(objects),
        "dominant": ", ".join(f"{o.get('id')} ({o.get('class')})" for o in dominant) or "n/a",
    }


def scene_table(scene: Dict[str, Any]) -> List[List[Any]]:
    rows = []
    for obj in scene.get("objects", []):
        rows.append([
            obj.get("id"),
            obj.get("class"),
            obj.get("az_deg"),
            obj.get("el_deg"),
            obj.get("dist_m"),
            obj.get("width"),
            obj.get("gain_db"),
        ])
    return rows


def scene_markdown(scene: Dict[str, Any], valid: bool, errors: List[str], backend_used: str) -> str:
    stats = scene_stats(scene)
    badge = "✅ Schema valid" if valid else "⚠️ Validation issues"
    lines = [
        f"### {badge}",
        "",
        f"- **Backend:** {backend_used}",
        f"- **Layout:** `{stats['layout']}`",
        f"- **Room preset:** `{stats['room_preset']}`",
        f"- **Objects:** `{stats['object_count']}`",
        f"- **Top objects:** {stats['dominant']}",
    ]
    if errors:
        lines.append("")
        lines.append("**Validation messages**")
        for err in errors[:8]:
            lines.append(f"- {err}")
    return "\n".join(lines)


def plot_scene(scene: Dict[str, Any]):
    fig = plt.figure(figsize=(7.0, 7.0))
    ax = fig.add_subplot(111)
    ax.set_facecolor("#f8fbff")
    fig.patch.set_facecolor("#f8fbff")

    # rings
    for r in [2, 4, 6, 8, 10]:
        circle = plt.Circle((0, 0), r, color="#d9e5f7", fill=False, linewidth=1)
        ax.add_artist(circle)

    # axes
    ax.axhline(0, color="#d9e5f7", linewidth=1)
    ax.axvline(0, color="#d9e5f7", linewidth=1)

    # labels
    ax.text(0, 10.7, "Front", ha="center", va="bottom", fontsize=11, color="#4a5b77")
    ax.text(10.7, 0, "Right", ha="left", va="center", fontsize=11, color="#4a5b77")
    ax.text(-10.7, 0, "Left", ha="right", va="center", fontsize=11, color="#4a5b77")
    ax.text(0, -10.7, "Rear", ha="center", va="top", fontsize=11, color="#4a5b77")

    palette = ["#1d4ed8", "#0891b2", "#7c3aed", "#ea580c", "#0f766e", "#be123c", "#0369a1", "#4d7c0f"]

    for idx, obj in enumerate(scene.get("objects", [])):
        az = math.radians(float(obj.get("az_deg", 0.0)))
        dist = float(obj.get("dist_m", 1.0))
        x = dist * math.sin(az)
        y = dist * math.cos(az)
        size = 120 + 220 * float(obj.get("width", 0.2))
        color = palette[idx % len(palette)]
        ax.scatter([x], [y], s=size, c=color, alpha=0.85, edgecolors="white", linewidths=1.5, zorder=3)
        ax.text(x, y + 0.42, f"{obj.get('id')}\n{obj.get('class')}", ha="center", va="bottom", fontsize=9, color="#1f2937")

    ax.set_xlim(-11, 11)
    ax.set_ylim(-11, 11)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_title("Spatial Scene Preview", fontsize=15, color="#15233d", pad=12)
    return fig