uopi/internal/dsp/nodes.go

// Package dsp provides signal processing node primitives used by the synthetic
// data source.
package dsp

import (
	"errors"
	"fmt"
	"math"
	"strconv"
	"strings"
	"time"
	"unicode"

	lua "github.com/yuin/gopher-lua"
)

// Node is a single processing stage in a synthetic signal pipeline.
type Node interface {
	// Process receives the latest input values (one per upstream signal, in order)
	// and returns the computed output. state is node-local persistent state.
	Process(inputs []float64, state map[string]any) (float64, error)
	// Type returns the node type name used in JSON definitions.
	Type() string
}

// ── GainNode ──────────────────────────────────────────────────────────────────

// GainNode multiplies input[0] by a fixed gain factor.
type GainNode struct {
	Gain float64
}

func (n *GainNode) Type() string { return "gain" }
func (n *GainNode) Process(inputs []float64, _ map[string]any) (float64, error) {
	if len(inputs) == 0 {
		return 0, errors.New("gain: no inputs")
	}
	return inputs[0] * n.Gain, nil
}

// ── OffsetNode ────────────────────────────────────────────────────────────────

// OffsetNode adds a fixed offset to input[0].
type OffsetNode struct {
	Offset float64
}

func (n *OffsetNode) Type() string { return "offset" }
func (n *OffsetNode) Process(inputs []float64, _ map[string]any) (float64, error) {
	if len(inputs) == 0 {
		return 0, errors.New("offset: no inputs")
	}
	return inputs[0] + n.Offset, nil
}

// ── AddNode ───────────────────────────────────────────────────────────────────

// AddNode returns the sum of all inputs.
type AddNode struct{}

func (n *AddNode) Type() string { return "add" }
func (n *AddNode) Process(inputs []float64, _ map[string]any) (float64, error) {
	var sum float64
	for _, v := range inputs {
		sum += v
	}
	return sum, nil
}

// ── SubtractNode ──────────────────────────────────────────────────────────────

// SubtractNode returns input[0] - input[1].
type SubtractNode struct{}

func (n *SubtractNode) Type() string { return "subtract" }
func (n *SubtractNode) Process(inputs []float64, _ map[string]any) (float64, error) {
	if len(inputs) < 2 {
		return 0, errors.New("subtract: need at least 2 inputs")
	}
	return inputs[0] - inputs[1], nil
}

// ── MultiplyNode ──────────────────────────────────────────────────────────────

// MultiplyNode returns the product of all inputs.
type MultiplyNode struct{}

func (n *MultiplyNode) Type() string { return "multiply" }
func (n *MultiplyNode) Process(inputs []float64, _ map[string]any) (float64, error) {
	if len(inputs) == 0 {
		return 0, errors.New("multiply: no inputs")
	}
	product := 1.0
	for _, v := range inputs {
		product *= v
	}
	return product, nil
}

// ── DivideNode ────────────────────────────────────────────────────────────────

// DivideNode returns input[0] / input[1]; returns 0 if denominator is 0.
type DivideNode struct{}

func (n *DivideNode) Type() string { return "divide" }
func (n *DivideNode) Process(inputs []float64, _ map[string]any) (float64, error) {
	if len(inputs) < 2 {
		return 0, errors.New("divide: need at least 2 inputs")
	}
	if inputs[1] == 0 {
		return 0, nil
	}
	return inputs[0] / inputs[1], nil
}

// ── MovingAverageNode ─────────────────────────────────────────────────────────

// MovingAverageNode computes a running mean over the last Window samples.
type MovingAverageNode struct {
	Window int
}

func (n *MovingAverageNode) Type() string { return "moving_average" }
func (n *MovingAverageNode) Process(inputs []float64, state map[string]any) (float64, error) {
	if len(inputs) == 0 {
		return 0, errors.New("moving_average: no inputs")
	}
	w := n.Window
	if w < 1 {
		w = 1
	}

	buf, _ := state["buf"].([]float64)
	buf = append(buf, inputs[0])
	if len(buf) > w {
		buf = buf[len(buf)-w:]
	}
	state["buf"] = buf

	var sum float64
	for _, v := range buf {
		sum += v
	}
	return sum / float64(len(buf)), nil
}

// ── RMSNode ───────────────────────────────────────────────────────────────────

// RMSNode computes the root-mean-square over the last Window samples.
type RMSNode struct {
	Window int
}

func (n *RMSNode) Type() string { return "rms" }
func (n *RMSNode) Process(inputs []float64, state map[string]any) (float64, error) {
	if len(inputs) == 0 {
		return 0, errors.New("rms: no inputs")
	}
	w := n.Window
	if w < 1 {
		w = 1
	}

	buf, _ := state["buf"].([]float64)
	buf = append(buf, inputs[0])
	if len(buf) > w {
		buf = buf[len(buf)-w:]
	}
	state["buf"] = buf

	var sumSq float64
	for _, v := range buf {
		sumSq += v * v
	}
	return math.Sqrt(sumSq / float64(len(buf))), nil
}

// ── DerivativeNode ────────────────────────────────────────────────────────────

// DerivativeNode computes the finite difference (current - previous) / dt where
// dt is in seconds. On the first call it returns 0.
type DerivativeNode struct{}

func (n *DerivativeNode) Type() string { return "derivative" }
func (n *DerivativeNode) Process(inputs []float64, state map[string]any) (float64, error) {
	if len(inputs) == 0 {
		return 0, errors.New("derivative: no inputs")
	}
	now := time.Now()
	if prevVal, ok := state["prev_val"].(float64); ok {
		if prevTime, ok := state["prev_time"].(time.Time); ok {
			dt := now.Sub(prevTime).Seconds()
			if dt <= 0 {
				dt = 1e-9 // avoid division by zero
			}
			result := (inputs[0] - prevVal) / dt
			state["prev_val"] = inputs[0]
			state["prev_time"] = now
			return result, nil
		}
	}
	state["prev_val"] = inputs[0]
	state["prev_time"] = now
	return 0, nil
}

// ── ClampNode ─────────────────────────────────────────────────────────────────

// ClampNode clamps the output to [Min, Max].
type ClampNode struct {
	Min float64
	Max float64
}

func (n *ClampNode) Type() string { return "clamp" }
func (n *ClampNode) Process(inputs []float64, _ map[string]any) (float64, error) {
	if len(inputs) == 0 {
		return 0, errors.New("clamp: no inputs")
	}
	v := inputs[0]
	if v < n.Min {
		return n.Min, nil
	}
	if v > n.Max {
		return n.Max, nil
	}
	return v, nil
}

// ── ThresholdNode ─────────────────────────────────────────────────────────────

// ThresholdNode outputs High when input[0] >= Threshold, Low otherwise.
type ThresholdNode struct {
	Threshold float64
	High      float64
	Low       float64
}

func (n *ThresholdNode) Type() string { return "threshold" }
func (n *ThresholdNode) Process(inputs []float64, _ map[string]any) (float64, error) {
	if len(inputs) == 0 {
		return 0, errors.New("threshold: no inputs")
	}
	if inputs[0] >= n.Threshold {
		return n.High, nil
	}
	return n.Low, nil
}

// ── ExprNode ──────────────────────────────────────────────────────────────────

// ExprNode evaluates a simple arithmetic expression with variables a, b, c, d
// bound to inputs[0..3]. It uses a hand-written recursive descent parser.
type ExprNode struct {
	Expr string
}

func (n *ExprNode) Type() string { return "expr" }
func (n *ExprNode) Process(inputs []float64, _ map[string]any) (float64, error) {
	vars := map[string]float64{}
	names := []string{"a", "b", "c", "d"}
	for i, name := range names {
		if i < len(inputs) {
			vars[name] = inputs[i]
		} else {
			vars[name] = 0
		}
	}
	p := &exprParser{src: strings.TrimSpace(n.Expr), vars: vars}
	result, err := p.parseExpr()
	if err != nil {
		return 0, fmt.Errorf("expr: %w", err)
	}
	if p.pos < len(p.src) {
		return 0, fmt.Errorf("expr: unexpected character %q at position %d", p.src[p.pos], p.pos)
	}
	return result, nil
}

// exprParser is a recursive-descent parser for arithmetic expressions.
// Grammar:
//
//	expr   = term   (('+' | '-') term)*
//	term   = power  (('*' | '/') power)*
//	power  = unary  ('^' power)*            — right-associative exponentiation
//	unary  = '-' unary | call
//	call   = ident '(' expr ')' | factor
//	factor = '(' expr ')' | number | variable
//
// Supported functions: exp, log, log2, log10, sqrt, abs, sin, cos, tan,
//
//	asin, acos, atan, floor, ceil, round, pow
//
// Variables: a, b, c, d (bound to inputs[0..3]).
type exprParser struct {
	src  string
	pos  int
	vars map[string]float64
}

func (p *exprParser) skipSpaces() {
	for p.pos < len(p.src) && unicode.IsSpace(rune(p.src[p.pos])) {
		p.pos++
	}
}

func (p *exprParser) peek() (byte, bool) {
	p.skipSpaces()
	if p.pos >= len(p.src) {
		return 0, false
	}
	return p.src[p.pos], true
}

func (p *exprParser) parseExpr() (float64, error) {
	left, err := p.parseTerm()
	if err != nil {
		return 0, err
	}
	for {
		ch, ok := p.peek()
		if !ok || (ch != '+' && ch != '-') {
			break
		}
		p.pos++
		right, err := p.parseTerm()
		if err != nil {
			return 0, err
		}
		if ch == '+' {
			left += right
		} else {
			left -= right
		}
	}
	return left, nil
}

func (p *exprParser) parseTerm() (float64, error) {
	left, err := p.parsePower()
	if err != nil {
		return 0, err
	}
	for {
		ch, ok := p.peek()
		if !ok || (ch != '*' && ch != '/') {
			break
		}
		p.pos++
		right, err := p.parsePower()
		if err != nil {
			return 0, err
		}
		if ch == '*' {
			left *= right
		} else {
			if right == 0 {
				return 0, nil
			}
			left /= right
		}
	}
	return left, nil
}

// parsePower handles right-associative exponentiation: a^b^c = a^(b^c).
// Unary minus is consumed here so that -a^2 = -(a^2), not (-a)^2 — matching
// standard mathematical and Python/Julia precedence rules.
func (p *exprParser) parsePower() (float64, error) {
	// Collect leading unary minuses; an even count cancels out.
	sign := 1.0
	p.skipSpaces()
	for p.pos < len(p.src) && p.src[p.pos] == '-' {
		p.pos++
		sign = -sign
		p.skipSpaces()
	}

	base, err := p.parseCall()
	if err != nil {
		return 0, err
	}

	p.skipSpaces()
	if p.pos < len(p.src) && p.src[p.pos] == '^' {
		p.pos++
		exp, err := p.parsePower() // right-associative recursion
		if err != nil {
			return 0, err
		}
		return sign * math.Pow(base, exp), nil
	}
	return sign * base, nil
}

// parseCall handles function calls like exp(a) or pow(a, 2).
func (p *exprParser) parseCall() (float64, error) {
	p.skipSpaces()
	if p.pos >= len(p.src) {
		return 0, errors.New("unexpected end of expression")
	}

	// Try to read an identifier (function name or variable).
	if p.src[p.pos] >= 'a' && p.src[p.pos] <= 'z' || p.src[p.pos] >= 'A' && p.src[p.pos] <= 'Z' {
		start := p.pos
		for p.pos < len(p.src) && (p.src[p.pos] >= 'a' && p.src[p.pos] <= 'z' ||
			p.src[p.pos] >= 'A' && p.src[p.pos] <= 'Z' ||
			p.src[p.pos] >= '0' && p.src[p.pos] <= '9' ||
			p.src[p.pos] == '_') {
			p.pos++
		}
		name := p.src[start:p.pos]

		// Check for function call syntax: name(...)
		p.skipSpaces()
		if p.pos < len(p.src) && p.src[p.pos] == '(' {
			p.pos++ // consume '('
			arg1, err := p.parseExpr()
			if err != nil {
				return 0, err
			}
			// Optional second argument for two-argument functions.
			var arg2 float64
			p.skipSpaces()
			if p.pos < len(p.src) && p.src[p.pos] == ',' {
				p.pos++
				arg2, err = p.parseExpr()
				if err != nil {
					return 0, err
				}
				p.skipSpaces()
			}
			if p.pos >= len(p.src) || p.src[p.pos] != ')' {
				return 0, fmt.Errorf("missing ')' after %s(...)", name)
			}
			p.pos++ // consume ')'

			switch name {
			case "exp":
				return math.Exp(arg1), nil
			case "log", "ln":
				return math.Log(arg1), nil
			case "log2":
				return math.Log2(arg1), nil
			case "log10":
				return math.Log10(arg1), nil
			case "sqrt":
				return math.Sqrt(arg1), nil
			case "abs":
				return math.Abs(arg1), nil
			case "sin":
				return math.Sin(arg1), nil
			case "cos":
				return math.Cos(arg1), nil
			case "tan":
				return math.Tan(arg1), nil
			case "asin":
				return math.Asin(arg1), nil
			case "acos":
				return math.Acos(arg1), nil
			case "atan":
				return math.Atan(arg1), nil
			case "atan2":
				return math.Atan2(arg1, arg2), nil
			case "pow":
				return math.Pow(arg1, arg2), nil
			case "floor":
				return math.Floor(arg1), nil
			case "ceil":
				return math.Ceil(arg1), nil
			case "round":
				return math.Round(arg1), nil
			case "min":
				return math.Min(arg1, arg2), nil
			case "max":
				return math.Max(arg1, arg2), nil
			default:
				return 0, fmt.Errorf("unknown function %q", name)
			}
		}

		// Not a function call — must be a single-letter variable.
		if len(name) != 1 {
			return 0, fmt.Errorf("unknown identifier %q (use a–d for variables, or a known function name)", name)
		}
		val, ok := p.vars[name]
		if !ok {
			return 0, fmt.Errorf("unknown variable %q (allowed: a, b, c, d)", name)
		}
		return val, nil
	}

	return p.parseFactor()
}

func (p *exprParser) parseFactor() (float64, error) {
	p.skipSpaces()
	if p.pos >= len(p.src) {
		return 0, errors.New("unexpected end of expression")
	}

	ch := p.src[p.pos]

	// Parenthesised subexpression.
	if ch == '(' {
		p.pos++
		val, err := p.parseExpr()
		if err != nil {
			return 0, err
		}
		p.skipSpaces()
		if p.pos >= len(p.src) || p.src[p.pos] != ')' {
			return 0, errors.New("missing closing parenthesis")
		}
		p.pos++
		return val, nil
	}

	// Number (including optional decimal point and exponent notation).
	if ch >= '0' && ch <= '9' || ch == '.' {
		start := p.pos
		for p.pos < len(p.src) && (p.src[p.pos] >= '0' && p.src[p.pos] <= '9' ||
			p.src[p.pos] == '.' || p.src[p.pos] == 'e' || p.src[p.pos] == 'E' ||
			((p.src[p.pos] == '+' || p.src[p.pos] == '-') && p.pos > start &&
				(p.src[p.pos-1] == 'e' || p.src[p.pos-1] == 'E'))) {
			p.pos++
		}
		f, err := strconv.ParseFloat(p.src[start:p.pos], 64)
		if err != nil {
			return 0, fmt.Errorf("invalid number %q", p.src[start:p.pos])
		}
		return f, nil
	}

	return 0, fmt.Errorf("unexpected character %q at position %d", ch, p.pos)
}

// ── LuaNode ───────────────────────────────────────────────────────────────────

// LuaNode runs a Lua script in a sandboxed gopher-lua VM.
// Inputs are bound to globals a, b, c, d. The script's return value is the output.
// The os, io, package, and debug libraries are disabled.
type LuaNode struct {
	Script string
}

func (n *LuaNode) Type() string { return "lua" }

func (n *LuaNode) Process(inputs []float64, state map[string]any) (result float64, retErr error) {
	// Retrieve or create the Lua VM.
	var L *lua.LState
	if existing, ok := state["L"].(*lua.LState); ok && existing != nil {
		L = existing
	} else {
		L = lua.NewState(lua.Options{SkipOpenLibs: true})
		// Open only safe libs.
		for _, pair := range []struct {
			name string
			fn   lua.LGFunction
		}{
			{lua.LoadLibName, lua.OpenPackage}, // required by other libs
			{lua.BaseLibName, lua.OpenBase},
			{lua.MathLibName, lua.OpenMath},
			{lua.StringLibName, lua.OpenString},
			{lua.TabLibName, lua.OpenTable},
		} {
			if err := L.CallByParam(lua.P{
				Fn:      L.NewFunction(pair.fn),
				NRet:    0,
				Protect: true,
			}, lua.LString(pair.name)); err != nil {
				L.Close()
				return 0, fmt.Errorf("lua init: %w", err)
			}
		}
		// Remove potentially dangerous globals that sneak in via base.
		L.SetGlobal("load", lua.LNil)
		L.SetGlobal("loadfile", lua.LNil)
		L.SetGlobal("dofile", lua.LNil)
		L.SetGlobal("require", lua.LNil)
		state["L"] = L
	}

	// Catch panics from Lua execution.
	defer func() {
		if r := recover(); r != nil {
			retErr = fmt.Errorf("lua panic: %v", r)
		}
	}()

	// Clear the stack.
	L.SetTop(0)

	// Bind inputs.
	names := []string{"a", "b", "c", "d"}
	for i, name := range names {
		if i < len(inputs) {
			L.SetGlobal(name, lua.LNumber(inputs[i]))
		} else {
			L.SetGlobal(name, lua.LNumber(0))
		}
	}

	// Execute the script.
	if err := L.DoString(n.Script); err != nil {
		return 0, fmt.Errorf("lua: %w", err)
	}

	// Read the return value (top of the stack after DoString leaves the chunk's
	// results there).
	top := L.GetTop()
	if top < 1 {
		return 0, errors.New("lua: script did not return a value")
	}
	lv := L.Get(top)
	num, ok := lv.(lua.LNumber)
	if !ok {
		return 0, fmt.Errorf("lua: script returned non-number %T", lv)
	}
	return float64(num), nil
}