path: root/generat3d/scene.py

# generat3D - evaluate Daz DSON scenes and write .dbz files
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import math
import posixpath
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np

from . import Error, warn
from .formula import Graph
from .library import parse_url, read_dson

AXES = "xyz"
REST_CHANNELS = ("center_point", "end_point", "orientation")
POSE_REST = {"rotation": 0.0, "translation": 0.0, "scale": 1.0, "general_scale": 1.0}


def key(node, ident, channel):
    return f"{node}:{ident}?{channel}"


@dataclass
class Node:
    """The figure root or one of its bones."""
    id: str
    name: str
    label: str
    rotation_order: str


@dataclass
class Figure:
    id: str
    name: str
    label: str
    verts: np.ndarray
    material_groups: list
    nodes: list  # nodes[0] is the figure root, then bones in file order
    morphs: dict = field(default_factory=dict)  # value key -> (indices, deltas)


@dataclass
class Scene:
    path: Path
    figure: Figure
    graph: Graph


def _channel(node, name):
    comps = {c["id"]: c.get("current_value", c.get("value", 0.0)) for c in node.get(name, [])}
    return [comps.get(a, 0.0) for a in AXES]


def _lookup(library, doc, url, section):
    """Resolve url to (file path, entry in that file's section). Local refs use doc."""
    _, path, ident, _ = parse_url(url)
    src = library.load(path) if path else doc
    for entry in src.get(section, []):
        if entry.get("id") == ident:
            return path, entry
    raise Error(f"{url}: no {section} entry {ident!r}")


def _number(url, v):
    if isinstance(v, bool) or not isinstance(v, (int, float)):
        raise Error(f"{url}: dial value is not a number")
    return v


class _Loader:
    def __init__(self, path, library):
        self.path = Path(path)
        self.library = library
        self.doc = read_dson(self.path)
        self.graph = Graph()
        self.loaded = set()
        self.hd = 0

    def run(self):
        nodes = self.doc.get("scene", {}).get("nodes", [])
        figs = []
        posed = False
        for n in nodes:
            if "url" not in n:
                continue
            fpath, asset = _lookup(self.library, self.doc, n["url"], "node_library")
            kind = asset.get("type")
            if kind == "bone":
                posed = posed or any(v for ch in ("rotation", "translation")
                                     for v in _channel(n, ch))
            elif kind != "figure":
                warn(f"skipped scene node {n.get('name') or n.get('id')!r} ({kind or 'no type'})")
            elif "conform_target" in n:
                warn(f"skipped follower {n.get('name') or n.get('id')!r} (not handled in v1)")
            else:
                figs.append((n, fpath, asset))
        if not figs:
            raise Error(f"{self.path}: no figure in scene")
        if len(figs) > 1:
            raise Error(f"{self.path}: v1 handles one figure per scene, found {len(figs)}")
        self.figure = self._figure(*figs[0])
        self.ids = {n.id for n in self.figure.nodes} | {n.name for n in self.figure.nodes}
        self._morphs(figs[0][1])
        posed = self._dials() or posed
        if posed:
            warn("pose in scene ignored: the .dbz holds rest data")
        if self.hd:
            warn(f"{self.hd} HD morphs: detail ignored (v1)")
        return Scene(self.path, self.figure, self.graph)

    def _figure(self, inst, fpath, asset):
        fdoc = self.library.load(fpath)
        geoms = inst.get("geometries", [])
        if geoms and "url" in geoms[0]:
            _, geom = _lookup(self.library, fdoc, geoms[0]["url"], "geometry_library")
        elif fdoc.get("geometry_library"):
            geom = fdoc["geometry_library"][0]
        else:
            raise Error(f"{fpath}: figure has no geometry")
        bones = [n for n in fdoc.get("node_library", []) if n.get("type") == "bone"]
        nodes = []
        for n in [asset] + bones:
            nodes.append(Node(n["id"], n.get("name", n["id"]), n.get("label", n["id"]),
                              n.get("rotation_order", "XYZ")))
            for ch in REST_CHANNELS:
                for a, v in zip(AXES, _channel(n, ch)):
                    self.graph.add_prop(key(n["id"], n["id"], f"{ch}/{a}"), v)
            for f in n.get("formulas", []):
                self._formula(f, fpath)
        return Figure(asset["id"], inst.get("name", asset.get("name", asset["id"])),
                      inst.get("label", asset.get("label", asset["id"])),
                      np.array(geom["vertices"]["values"], dtype=float),
                      list(geom.get("polygon_material_groups", {}).get("values", [])),
                      nodes)

    def _key(self, url):
        node, _, ident, channel = parse_url(url)
        if node not in self.ids or node == self.figure.name:
            node = self.figure.id
        return key(node, ident, channel)

    def _formula(self, f, source):
        ops = []
        for op in f.get("operations", []):
            if op.get("op") == "push" and "url" in op:
                channel = parse_url(op["url"])[3].split("/")[0]
                if channel in POSE_REST:
                    ops.append({"op": "push", "val": POSE_REST[channel]})
                else:
                    ops.append({"op": "push", "key": self._key(op["url"])})
            else:
                ops.append(op)
        self.graph.add_formula(self._key(f["output"]), ops, f.get("stage", "sum"), source)

    def _modifiers(self, path, doc):
        nverts = len(self.figure.verts)
        for mod in doc.get("modifier_library", []):
            parent = parse_url(mod.get("parent", ""))[2]
            if parent and parent != self.figure.id:
                continue
            ch = mod.get("channel", {})
            k = key(self.figure.id, mod["id"], "value")
            self.graph.add_prop(k, ch.get("value", 0.0), ch.get("min", -math.inf),
                                ch.get("max", math.inf), ch.get("clamped", False))
            morph = mod.get("morph")
            if morph:
                if "hd_url" in morph:
                    self.hd += 1
                d = morph.get("deltas", {}).get("values", [])
                if morph.get("vertex_count", nverts) != nverts:
                    warn(f"{path}#{mod['id']}: vertex count mismatch, morph skipped")
                elif d:
                    arr = np.array(d, dtype=float)
                    if (arr.ndim != 2 or arr.shape[1] != 4
                            or (arr[:, 0] != np.floor(arr[:, 0])).any()
                            or (arr[:, 0] < 0).any() or (arr[:, 0] >= nverts).any()):
                        raise Error(f"{path}#{mod['id']}: bad morph deltas")
                    self.figure.morphs[k] = (arr[:, 0].astype(int), arr[:, 1:4])
            for f in mod.get("formulas", []):
                self._formula(f, f"{path}#{mod['id']}")

    def _load_file(self, path):
        if path.lower() in self.loaded:
            return
        self.loaded.add(path.lower())
        doc = self.library.load(path)
        try:
            self._modifiers(path, doc)
        except (KeyError, TypeError, ValueError, IndexError, AttributeError) as e:
            raise Error(f"{path}: malformed morph data: {e!r}") from None

    def _morphs(self, fpath):
        morph_dir = posixpath.join(posixpath.dirname(fpath), "Morphs")
        for path in self.library.walk(morph_dir):
            try:
                self.library.load(path)
            except Error as e:
                warn(f"skipped unreadable morph file: {e}")
                continue
            self._load_file(path)

    def _dials(self):
        """Apply .duf dial values. Returns True if the scene animates a pose channel."""
        scene = self.doc.get("scene", {})
        values, missing, posed = {}, [], False
        for m in scene.get("modifiers", []):
            if "url" not in m:
                continue
            k = key(self.figure.id, parse_url(m["url"])[2], "value")
            v = _number(m["url"], m.get("channel", {}).get("current_value", 0.0))
            if k not in self.graph.props:
                path = parse_url(m["url"])[1]
                if path and self.library.find(path):
                    self._load_file(path)
            if k in self.graph.props:
                values[k] = v
            elif v:
                missing.append(m["url"])
        for a in scene.get("animations", []):
            channel = parse_url(a["url"])[3]
            v = _number(a["url"], a.get("keys", [[0, 0.0]])[0][1])
            if channel.split("/")[0] in POSE_REST:
                posed = posed or bool(v)
            elif channel == "value":
                k = self._key(a["url"])
                if k in self.graph.props:
                    values[k] = v
                elif v:
                    missing.append(a["url"])
        if missing:
            raise Error("missing morphs used by the scene:\n  " + "\n  ".join(missing))
        for k, v in values.items():
            self.graph.set_value(k, v)
        return posed


def load(path, library):
    """Load a .duf scene and everything its figure needs."""
    try:
        return _Loader(path, library).run()
    except (KeyError, TypeError, IndexError, ValueError, AttributeError) as e:
        raise Error(f"{path}: malformed scene data: {e!r}") from None