# 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 bisect
import math
import operator

from . import Error

STAGES = ("sum", "mult")
BINARY = {"add": operator.add, "sub": operator.sub,
          "mult": operator.mul, "div": operator.truediv}
SPLINES = ("spline_constant", "spline_linear", "spline_tcb")


class Prop:
    __slots__ = ("value", "lo", "hi", "clamped")

    def __init__(self, value, lo, hi, clamped):
        self.value, self.lo, self.hi, self.clamped = value, lo, hi, clamped


class Graph:
    """Properties keyed by string, plus the formulas that drive them."""

    def __init__(self):
        self.props = {}
        self.formulas = {}  # output key -> [(stage, ops, source)]

    def add_prop(self, key, value=0.0, lo=-math.inf, hi=math.inf, clamped=False):
        self.props[key] = Prop(value, lo, hi, clamped)

    def set_value(self, key, value):
        self.props[key].value = value

    def add_formula(self, output, ops, stage="sum", source="?"):
        if stage not in STAGES:
            raise Error(f"{source}: unknown formula stage {stage!r} for {output}")
        self.formulas.setdefault(output, []).append((stage, ops, source))

    def evaluate(self):
        """Final value of every property."""
        done = {}
        for key in self.props:
            self._value(key, done, [])
        return done

    def _value(self, key, done, path):
        if key in done:
            return done[key]
        if key in path:
            raise Error("formula cycle: " + " -> ".join(path[path.index(key):] + [key]))
        prop = self.props.get(key)
        if prop is None:
            return None
        path.append(key)
        total, factor = prop.value, 1.0
        try:
            for stage, ops, source in self.formulas.get(key, ()):
                r = self._run(key, ops, done, path, source)
                if r is None:
                    continue
                if stage == "sum":
                    total += r
                else:
                    factor *= r
        except RecursionError:
            # ponytail: recursive evaluation; make it iterative if real content hits this
            raise Error("formula dependency chain too deep") from None
        path.pop()
        v = total * factor
        if prop.clamped:
            v = min(max(v, prop.lo), prop.hi)
        done[key] = v
        return v

    def _run(self, key, ops, done, path, source):
        stack = []
        try:
            for op in ops:
                name = op["op"]
                if name == "push":
                    if "key" in op:
                        v = self._value(op["key"], done, path)
                        if v is None:
                            return None
                        stack.append(v)
                    else:
                        stack.append(op["val"])
                elif name in BINARY:
                    b, a = stack.pop(), stack.pop()
                    stack.append(BINARY[name](a, b))
                elif name == "neg":
                    stack.append(-stack.pop())
                elif name in SPLINES:
                    stack.append(_spline_op(name, stack))
                else:
                    raise Error(f"{source}: unknown formula operation {name!r} for {key}")
            if len(stack) != 1:
                raise Error(f"{source}: malformed formula for {key}")
            r = stack[0]
            if not (isinstance(r, (int, float)) and not isinstance(r, bool)
                    and math.isfinite(r)):
                raise Error(f"{source}: malformed formula for {key}")
            return r
        except (IndexError, KeyError, TypeError, ZeroDivisionError, ValueError,
                OverflowError):
            raise Error(f"{source}: malformed formula for {key}") from None


def _spline_op(name, stack):
    n = int(stack.pop())
    if n < 1 or n > len(stack) - 1:
        raise IndexError
    knots = stack[-n:]
    del stack[-n:]
    return spline(name, stack.pop(), knots)


def spline(kind, x, knots):
    """Evaluate a DSON spline at x. Knots are [x, y] or [x, y, t, c, b], sorted by x."""
    xs = [k[0] for k in knots]
    ys = [k[1] for k in knots]
    if x <= xs[0]:
        return ys[0]
    if x >= xs[-1]:
        return ys[-1]
    i = bisect.bisect_right(xs, x) - 1
    if kind == "spline_constant":
        return ys[i]
    s = (x - xs[i]) / (xs[i + 1] - xs[i])
    if kind == "spline_linear":
        return ys[i] + s * (ys[i + 1] - ys[i])
    # ponytail: Kochanek-Bartels with uniform segment parameterization;
    # if the corpus shows Daz scales tangents by knot spacing, add that here.
    m0 = _tcb_tangent(knots, ys, i, outgoing=True)
    m1 = _tcb_tangent(knots, ys, i + 1, outgoing=False)
    s2, s3 = s * s, s * s * s
    return ((2 * s3 - 3 * s2 + 1) * ys[i] + (s3 - 2 * s2 + s) * m0
            + (-2 * s3 + 3 * s2) * ys[i + 1] + (s3 - s2) * m1)


def _tcb_tangent(knots, ys, k, outgoing):
    t, c, b = (list(knots[k][2:5]) + [0, 0, 0])[:3]
    d_in = ys[k] - ys[max(k - 1, 0)]
    d_out = ys[min(k + 1, len(ys) - 1)] - ys[k]
    if outgoing:
        return (1 - t) * ((1 + b) * (1 + c) * d_in + (1 - b) * (1 - c) * d_out) / 2
    return (1 - t) * ((1 + b) * (1 - c) * d_in + (1 - b) * (1 + c) * d_out) / 2