// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package signal

import (
	"math"
	"slices"

	"sourcedock.dev/petrbalvin/tensor/internal/base"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// State-space filtering: the Kalman family over the model
// x_{t+1} = f(x_t) + w_t with w ~ N(0, Q), z_t = h(x_t) + v_t with
// v ~ N(0, R). The linear filter runs the standard Riccati recursion
// with a Joseph-form correction, the extended filter linearises f and
// h at the current estimate (analytic Jacobians when supplied, the
// house central-difference helper otherwise), and the unscented filter
// carries the state distribution through the deterministic sigma-point
// set. All three accumulate the exact Gaussian log-likelihood of the
// innovation sequence, and every covariance they publish is mirrored
// into its symmetric average, with positive definiteness enforced
// where the mathematics demands it: a Cholesky factorisation that
// cannot be taken names the step it failed at.

// StateFunc maps one state vector to the image vector the model's
// transition or observation applies: the f of x_{t+1} = f(x_t) + w_t,
// the h of z_t = h(x_t) + v_t. The input array is private to the call
// and may be kept until the call returns; the output must be a real
// rank-1 array of one fixed length.
type StateFunc func(x *core.Array) (*core.Array, error)

// JacobianFunc returns the Jacobian of a StateFunc at x as an
// (m × n) float array, row i holding the partials of output i. An
// analytic Jacobian is always the better instrument; when the option
// is left unset the extended filter builds one by central differences
// at every step.
type JacobianFunc func(x *core.Array) (*core.Array, error)

// KalmanOptions carries the initial condition, the noise levels and
// the filter-specific knobs. An unset (nil) array field takes its
// documented default: a zero initial state, an identity initial
// covariance, zero process noise, an identity measurement noise.
// SigmaAlpha, SigmaBeta and SigmaKappa follow the same rule with the
// defaults 0.001, 2 and 0, the standard scaled unscented choice for a
// Gaussian prior. The extended and unscented filters cannot infer the
// state dimension from their callbacks, so at least one of
// InitialState and InitialCovariance must be set for them.
//
// TransitionJacobian and ObservationJacobian give the extended filter
// the analytic partials it linearises with; a nil one is replaced by
// central differences on the underlying map, two evaluations per
// partial per step.
type KalmanOptions struct {
	InitialState        *core.Array
	InitialCovariance   *core.Array
	ProcessNoise        *core.Array
	MeasurementNoise    *core.Array
	TransitionJacobian  JacobianFunc
	ObservationJacobian JacobianFunc
	SigmaAlpha          float64
	SigmaBeta           float64
	SigmaKappa          float64
}

// KalmanResult holds one filtering pass over the measurement stack.
type KalmanResult struct {
	// States is the (n × d) stack of filtered means x̂_{t|t}, one row
	// per measurement.
	States *core.Array
	// Covariances is the (n × d × d) stack of filtered covariances
	// P_{t|t}, one symmetric positive-definite block per measurement.
	Covariances *core.Array
	// Innovations is the (n × m) stack of one-step prediction errors
	// z_t − h(x̂_{t|t−1}).
	Innovations *core.Array
	// InnovationCovariances is the (n × m × m) stack of the innovation
	// covariances S_t the likelihood reads.
	InnovationCovariances *core.Array
	// LogLikelihood is Σ_t log N(z_t; h(x̂_{t|t−1}), S_t), the exact
	// Gaussian likelihood of the measurement sequence under the model
	// and the quantity noise and parameter estimation maximises.
	LogLikelihood float64
}

// kalmanStepOut carries one step's outputs from a filter's step
// closure into the result stack.
type kalmanStepOut struct {
	mean          []float64
	covariance    []float64
	innovation    []float64
	innovationCov []float64
	logLikelihood float64
}

// kfSymmetryEps is the relative tolerance a covariance's mirror check
// allows, the same reading the multivariate normal takes: matrices
// assembled from products differ from their mirror by an ulp of
// rounding, a genuinely asymmetric pair by far more.
const kfSymmetryEps = 1e-12

// kfFinite refuses the non-finite values a filter would otherwise
// carry silently through every recursion.
func kfFinite(name, what string, vals []float64) error {
	for i, v := range vals {
		if math.IsNaN(v) || math.IsInf(v, 0) {
			return base.Errf("%s: %s holds the non-finite value %g at %d", name, what, v, i)
		}
	}
	return nil
}

// kfVector reads a finite real rank-1 array into a float slice, the
// payload itself for a contiguous float64 array,
// with copies made only where the payload cannot serve, of exactly want entries when want is non-negative.
func kfVector(name, what string, a *core.Array, want int) ([]float64, error) {
	if a.NDim() != 1 {
		return nil, base.Errf("%s: %s must be a vector, got shape %s", name, what, base.ShapeText(a.Shape()))
	}
	if a.Dtype() == core.Complex {
		return nil, base.Errf("%s: %s must be real, got complex", name, what)
	}
	if want >= 0 && a.Len() != want {
		return nil, base.Errf("%s: %s holds %d entries, want %d", name, what, a.Len(), want)
	}
	vals := widenFloats(a)
	if err := kfFinite(name, what, vals); err != nil {
		return nil, err
	}
	return vals, nil
}

// kfMatrix reads a finite real rank-2 array into a fresh row-major
// float slice of exactly rows×cols, either side wildcarded at −1.
func kfMatrix(name, what string, a *core.Array, rows, cols int) ([]float64, error) {
	if a.NDim() != 2 {
		return nil, base.Errf("%s: %s must be rank 2, got shape %s", name, what, base.ShapeText(a.Shape()))
	}
	if a.Dtype() == core.Complex {
		return nil, base.Errf("%s: %s must be real, got complex", name, what)
	}
	shape := a.Shape()
	if (rows >= 0 && shape[0] != rows) || (cols >= 0 && shape[1] != cols) {
		return nil, base.Errf("%s: %s has shape %s, want %d×%d", name, what, base.ShapeText(shape), rows, cols)
	}
	vals := widenFloats(a)
	if err := kfFinite(name, what, vals); err != nil {
		return nil, err
	}
	return vals, nil
}

// kfSymmetric demands every mirror pair agree within a relative
// tolerance, because only one triangle is ever read.
func kfSymmetric(name, what string, a []float64, n int) error {
	for i := range n {
		for j := range i {
			lo, hi := a[i*n+j], a[j*n+i]
			if math.Abs(lo-hi) > kfSymmetryEps*math.Max(math.Abs(lo), math.Abs(hi)) {
				return base.Errf("%s: %s is not symmetric at (%d, %d): %g against %g",
					name, what, i+1, j+1, lo, hi)
			}
		}
	}
	return nil
}

// kfNoise reads one noise covariance: an unset array takes its
// documented default (zeros for the process noise, whose only job is
// to enter sums, an identity for the measurement noise, which is
// factored every step and so must be positive definite), a set one
// must be a finite symmetric matrix of the right size.
func kfNoise(name, what string, a *core.Array, dim int, identityDefault bool) ([]float64, error) {
	if a == nil {
		if identityDefault {
			return kfIdentity(dim), nil
		}
		return make([]float64, dim*dim), nil
	}
	vals, err := kfMatrix(name, what, a, dim, dim)
	if err != nil {
		return nil, err
	}
	if err := kfSymmetric(name, what, vals, dim); err != nil {
		return nil, err
	}
	if identityDefault {
		if _, err := kfCholesky(what, vals, dim); err != nil {
			return nil, base.Errf("%s: %w", name, err)
		}
	}
	return vals, nil
}

// kfInitialCondition reads the initial mean and covariance, zeros and
// identity where the options leave them unset. The initial covariance
// must be symmetric positive definite: the first predict would
// otherwise hand the recursion a structure it cannot factor.
func kfInitialCondition(name string, opts KalmanOptions, d int) (x0, p0 []float64, err error) {
	if opts.InitialState == nil {
		x0 = make([]float64, d)
	} else if x0, err = kfVector(name, "the initial state", opts.InitialState, d); err != nil {
		return nil, nil, err
	}
	if opts.InitialCovariance == nil {
		p0 = kfIdentity(d)
	} else {
		if p0, err = kfMatrix(name, "the initial covariance", opts.InitialCovariance, d, d); err != nil {
			return nil, nil, err
		}
		if err := kfSymmetric(name, "the initial covariance", p0, d); err != nil {
			return nil, nil, err
		}
		if _, err := kfCholesky("the initial covariance", p0, d); err != nil {
			return nil, nil, base.Errf("%s: %w", name, err)
		}
	}
	return x0, p0, nil
}

// kfNonlinearInputs resolves the shared inputs of the extended and
// unscented filters: the state dimension, which the callbacks cannot
// carry, comes from the initial state or the initial covariance, at
// least one of which must be set.
func kfNonlinearInputs(name string, opts KalmanOptions, m int) (d int, x0, p0, q, r []float64, err error) {
	switch {
	case opts.InitialState != nil:
		if x0, err = kfVector(name, "the initial state", opts.InitialState, -1); err != nil {
			return 0, nil, nil, nil, nil, err
		}
		d = len(x0)
	case opts.InitialCovariance != nil:
		if opts.InitialCovariance.NDim() != 2 {
			return 0, nil, nil, nil, nil, base.Errf("%s: the initial covariance must be rank 2, got shape %s",
				name, base.ShapeText(opts.InitialCovariance.Shape()))
		}
		d = opts.InitialCovariance.Shape()[0]
	default:
		return 0, nil, nil, nil, nil, base.Errf("%s: the state dimension must come from an initial state or an initial covariance; neither is set", name)
	}
	if d < 1 {
		return 0, nil, nil, nil, nil, base.Errf("%s: the state dimension must be at least 1, got %d", name, d)
	}
	if x0, p0, err = kfInitialCondition(name, opts, d); err != nil {
		return 0, nil, nil, nil, nil, err
	}
	if q, err = kfNoise(name, "the process noise", opts.ProcessNoise, d, false); err != nil {
		return 0, nil, nil, nil, nil, err
	}
	if r, err = kfNoise(name, "the measurement noise", opts.MeasurementNoise, m, true); err != nil {
		return 0, nil, nil, nil, nil, err
	}
	return d, x0, p0, q, r, nil
}

// kfMeasurements reads the measurement stack: a rank-1 array holds n
// scalar observations, a rank-2 array holds n rows of m channels. At
// least one measurement is needed, every value must be finite.
func kfMeasurements(name string, z *core.Array) (n, m int, rows []float64, err error) {
	if z.NDim() != 1 && z.NDim() != 2 {
		return 0, 0, nil, base.Errf("%s: the measurements must be rank 1 or rank 2, got shape %s",
			name, base.ShapeText(z.Shape()))
	}
	if z.Dtype() == core.Complex {
		return 0, 0, nil, base.Errf("%s: complex measurements are not supported", name)
	}
	if z.NDim() == 1 {
		m = 1
	} else {
		m = z.Shape()[1]
	}
	if m < 1 {
		return 0, 0, nil, base.Errf("%s: the measurement width must be at least 1, got %d", name, m)
	}
	n = z.Len() / m
	if n < 1 {
		return 0, 0, nil, base.Errf("%s: at least one measurement is needed", name)
	}
	rows = widenFloats(z)
	if err := kfFinite(name, "the measurements", rows); err != nil {
		return 0, 0, nil, err
	}
	return n, m, rows, nil
}

// kfIdentity returns the n×n identity, row-major.
func kfIdentity(n int) []float64 {
	a := make([]float64, n*n)
	for i := range n {
		a[i*n+i] = 1
	}
	return a
}

// kfTranspose returns the transpose of the row-major rows×cols
// matrix.
func kfTranspose(a []float64, rows, cols int) []float64 {
	out := make([]float64, rows*cols)
	for i := range rows {
		for j := range cols {
			out[j*rows+i] = a[i*cols+j]
		}
	}
	return out
}

// kfMatVec returns A·x for the row-major rows×cols matrix A.
func kfMatVec(a []float64, rows, cols int, x []float64) []float64 {
	out := make([]float64, rows)
	for i := range rows {
		total := 0.0
		row := a[i*cols : (i+1)*cols]
		for j, v := range row {
			total += v * x[j]
		}
		out[i] = total
	}
	return out
}

// kfMatMul returns A·B for the row-major A of size ra×ca and B of
// size ca×cb.
func kfMatMul(a []float64, ra, ca int, b []float64, cb int) []float64 {
	out := make([]float64, ra*cb)
	for i := range ra {
		row := out[i*cb : (i+1)*cb]
		arow := a[i*ca : (i+1)*ca]
		for k, aik := range arow {
			brow := b[k*cb : (k+1)*cb]
			for j := range cb {
				row[j] += aik * brow[j]
			}
		}
	}
	return out
}

// kfCholesky factors the symmetric positive-definite row-major n×n
// matrix into the lower triangular L with A = L·Lᵀ, only the lower
// mirror read. A non-positive pivot names its row.
func kfCholesky(what string, a []float64, n int) ([]float64, error) {
	l := make([]float64, n*n)
	for i := range n {
		for j := range i + 1 {
			total := a[i*n+j]
			for k := range j {
				total -= l[i*n+k] * l[j*n+k]
			}
			if i == j {
				if !(total > 0) {
					return nil, base.Errf("%s is not positive definite at row %d", what, i+1)
				}
				l[i*n+i] = math.Sqrt(total)
			} else {
				l[i*n+j] = total / l[j*n+j]
			}
		}
	}
	return l, nil
}

// kfCholSolve solves L·Lᵀ·x = b through the forward and backward
// substitutions.
func kfCholSolve(l []float64, n int, b []float64) []float64 {
	x := make([]float64, n)
	for i := range n {
		total := b[i]
		for k := range i {
			total -= l[i*n+k] * x[k]
		}
		x[i] = total / l[i*n+i]
	}
	for i := n - 1; i >= 0; i-- {
		total := x[i]
		for k := i + 1; k < n; k++ {
			total -= l[k*n+i] * x[k]
		}
		x[i] = total / l[i*n+i]
	}
	return x
}

// kfCholSolveMatrix solves L·Lᵀ·X = B for the row-major B of size
// n×cols, row by row.
func kfCholSolveMatrix(l []float64, n int, b []float64, cols int) []float64 {
	x := make([]float64, len(b))
	copy(x, b)
	for i := range n {
		row := x[i*cols : (i+1)*cols]
		for k := range i {
			lk := l[i*n+k]
			krow := x[k*cols : (k+1)*cols]
			for j := range cols {
				row[j] -= lk * krow[j]
			}
		}
		di := l[i*n+i]
		for j := range cols {
			row[j] /= di
		}
	}
	for i := n - 1; i >= 0; i-- {
		row := x[i*cols : (i+1)*cols]
		for k := i + 1; k < n; k++ {
			lk := l[k*n+i]
			krow := x[k*cols : (k+1)*cols]
			for j := range cols {
				row[j] -= lk * krow[j]
			}
		}
		di := l[i*n+i]
		for j := range cols {
			row[j] /= di
		}
	}
	return x
}

// kfLogDet returns the log determinant of the Cholesky factor: twice
// the sum of the log diagonal.
func kfLogDet(l []float64, n int) float64 {
	total := 0.0
	for i := range n {
		total += math.Log(l[i*n+i])
	}
	return 2 * total
}

// kfSymmetrise replaces A by (A + Aᵀ)/2 in place: the rounding drift
// of a recursion that touches a covariance only through symmetric
// expressions cannot survive the mirror average.
func kfSymmetrise(a []float64, n int) {
	for i := range n {
		for j := range i {
			avg := (a[i*n+j] + a[j*n+i]) / 2
			a[i*n+j] = avg
			a[j*n+i] = avg
		}
	}
}

// kalmanRun walks the measurement stack through the step closure,
// which owns one predict-and-correct cycle: it receives the current
// mean and covariance (read-only; it must return fresh slices) and
// the step's measurement row, and returns everything the result stack
// records, the Gaussian log-likelihood contribution included.
func kalmanRun(nMeas, m, d int, meas []float64, x0, p0 []float64,
	step func(t int, zRow, mean, cov []float64) (kalmanStepOut, error),
) (*KalmanResult, error) {
	out := &KalmanResult{
		States:                core.New(core.Float, nMeas, d),
		Covariances:           core.New(core.Float, nMeas, d, d),
		Innovations:           core.New(core.Float, nMeas, m),
		InnovationCovariances: core.New(core.Float, nMeas, m, m),
	}
	mean := slices.Clone(x0)
	cov := slices.Clone(p0)
	total := 0.0
	for t := range nMeas {
		res, err := step(t, meas[t*m:(t+1)*m], mean, cov)
		if err != nil {
			return nil, err
		}
		copy(out.States.RawFloats()[t*d:(t+1)*d], res.mean)
		copy(out.Covariances.RawFloats()[t*d*d:(t+1)*d*d], res.covariance)
		copy(out.Innovations.RawFloats()[t*m:(t+1)*m], res.innovation)
		copy(out.InnovationCovariances.RawFloats()[t*m*m:(t+1)*m*m], res.innovationCov)
		total += res.logLikelihood
		mean, cov = res.mean, res.covariance
	}
	out.LogLikelihood = total
	return out, nil
}

// kfCorrect runs the shared correction of the linear and extended
// filters: the innovation against the predicted observation, its
// covariance S = H·P⁻·Hᵀ + R, the gain K = P⁻·Hᵀ·S⁻¹ through the
// Cholesky solve, the state update and the Joseph-form covariance
// (I−K·H)·P⁻·(I−K·H)ᵀ + K·R·Kᵀ, mirrored into its symmetric average.
// The Joseph form keeps the filtered covariance symmetric positive
// definite by construction, where the plain (I−K·H)·P⁻ recursion
// preserves it only in exact arithmetic.
func kfCorrect(name string, t int, zRow, predicted, xp, pp []float64, d, m int,
	h, ht, r []float64,
) (kalmanStepOut, error) {
	innovation := make([]float64, m)
	for i := range m {
		innovation[i] = zRow[i] - predicted[i]
	}
	hp := kfMatMul(h, m, d, pp, d)
	s := kfMatMul(hp, m, d, ht, m)
	for i := range s {
		s[i] += r[i]
	}
	ls, err := kfCholesky("the innovation covariance", s, m)
	if err != nil {
		return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t, err)
	}
	solved := kfCholSolve(ls, m, innovation)
	quad := 0.0
	for i := range m {
		quad += innovation[i] * solved[i]
	}
	// K = P⁻·Hᵀ·S⁻¹ arrives through the transposed solve: S·X =
	// H·P⁻ gives X = S⁻¹·H·P⁻ = Kᵀ and K = Xᵀ.
	gain := kfTranspose(kfCholSolveMatrix(ls, m, hp, d), m, d)
	mean := make([]float64, d)
	copy(mean, xp)
	for i := range d {
		total := 0.0
		grow := gain[i*m : (i+1)*m]
		for j, y := range innovation {
			total += grow[j] * y
		}
		mean[i] += total
	}
	kh := kfMatMul(gain, d, m, h, d)
	a := make([]float64, d*d)
	for i := range d {
		for j := range d {
			a[i*d+j] = -kh[i*d+j]
		}
		a[i*d+i]++
	}
	at := kfTranspose(a, d, d)
	cov := kfMatMul(kfMatMul(a, d, d, pp, d), d, d, at, d)
	krkt := kfMatMul(kfMatMul(gain, d, m, r, m), d, m, kfTranspose(gain, d, m), d)
	for i := range cov {
		cov[i] += krkt[i]
	}
	kfSymmetrise(cov, d)
	logLik := -0.5 * (float64(m)*math.Log(2*math.Pi) + kfLogDet(ls, m) + quad)
	return kalmanStepOut{mean: mean, covariance: cov, innovation: innovation,
		innovationCov: s, logLikelihood: logLik}, nil
}

// KalmanFilter runs the linear Kalman filter over the measurement
// stack z under the model x_{t+1} = F·x_t + w_t, z_t = H·x_t + v_t:
// transition is the (d × d) state matrix F, observation the (m × d)
// matrix H. Every step predicts with F and corrects with a
// Joseph-form update, and the exact Gaussian log-likelihood of the
// innovation sequence accumulates into the result.
//
// A rank-1 z holds n scalar observations, a rank-2 z holds n rows of
// m channels. Nil option fields take their defaults (see
// KalmanOptions). A non-square transition, a shape mismatch, an
// asymmetric noise covariance, a singular measurement noise, a
// non-finite value anywhere, and an innovation covariance that loses
// positive definiteness (naming the step) are errors.
func KalmanFilter(z *core.Array, transition, observation *core.Array, opts KalmanOptions) (*KalmanResult, error) {
	const name = "KalmanFilter"
	if transition == nil || observation == nil {
		return nil, base.Errf("%s: the transition and observation matrices are required", name)
	}
	if transition.NDim() != 2 {
		return nil, base.Errf("%s: the transition matrix must be rank 2, got shape %s",
			name, base.ShapeText(transition.Shape()))
	}
	shape := transition.Shape()
	if shape[0] != shape[1] {
		return nil, base.Errf("%s: the transition matrix must be square, got %d×%d", name, shape[0], shape[1])
	}
	d := shape[0]
	if d < 1 {
		return nil, base.Errf("%s: the state dimension must be at least 1, got %d", name, d)
	}
	nMeas, m, meas, err := kfMeasurements(name, z)
	if err != nil {
		return nil, err
	}
	f, err := kfMatrix(name, "the transition matrix", transition, d, d)
	if err != nil {
		return nil, err
	}
	h, err := kfMatrix(name, "the observation matrix", observation, m, d)
	if err != nil {
		return nil, err
	}
	q, err := kfNoise(name, "the process noise", opts.ProcessNoise, d, false)
	if err != nil {
		return nil, err
	}
	r, err := kfNoise(name, "the measurement noise", opts.MeasurementNoise, m, true)
	if err != nil {
		return nil, err
	}
	x0, p0, err := kfInitialCondition(name, opts, d)
	if err != nil {
		return nil, err
	}
	ft := kfTranspose(f, d, d)
	ht := kfTranspose(h, m, d)
	return kalmanRun(nMeas, m, d, meas, x0, p0, func(t int, zRow, x, p []float64) (kalmanStepOut, error) {
		// Predict: x⁻ = F·x, P⁻ = F·P·Fᵀ + Q.
		xp := kfMatVec(f, d, d, x)
		pp := kfMatMul(kfMatMul(f, d, d, p, d), d, d, ft, d)
		for i := range pp {
			pp[i] += q[i]
		}
		predicted := kfMatVec(h, m, d, xp)
		return kfCorrect(name, t+1, zRow, predicted, xp, pp, d, m, h, ht, r)
	})
}

// ExtendedKalmanFilter runs the extended Kalman filter over the
// measurement stack z under the nonlinear model x_{t+1} = f(x_t) +
// w_t, z_t = h(x_t) + v_t: every step predicts by propagating the
// mean through f and the covariance through the linearisation F =
// ∂f/∂x at x̂_{t|t}, then corrects through the Jacobian H = ∂h/∂x at
// the predicted mean, with the same Joseph-form update and the same
// exact log-likelihood the linear filter carries. The Jacobians come
// from the options when supplied analytically and from central
// differences otherwise (see JacobianFunc).
//
// The extended filter is the linear one applied to local linear
// models: it inherits the Kalman recursions and with them the first
// order's blindness to the curvature of f and h, so a strongly bent
// observation map wants the unscented filter instead. The state
// dimension must be fixed by an initial state or covariance (see
// KalmanOptions); everything else follows the linear filter's
// contract, with a failing callback or Jacobian reported with the
// step it failed at.
func ExtendedKalmanFilter(z *core.Array, transition, observation StateFunc, opts KalmanOptions) (*KalmanResult, error) {
	const name = "ExtendedKalmanFilter"
	if transition == nil || observation == nil {
		return nil, base.Errf("%s: the transition and observation maps are required", name)
	}
	nMeas, m, meas, err := kfMeasurements(name, z)
	if err != nil {
		return nil, err
	}
	d, x0, p0, q, r, err := kfNonlinearInputs(name, opts, m)
	if err != nil {
		return nil, err
	}
	fjac := opts.TransitionJacobian
	if fjac == nil {
		fjac = func(x *core.Array) (*core.Array, error) {
			return core.Jacobian(transition, x, core.JacobianOptions{})
		}
	}
	hjac := opts.ObservationJacobian
	if hjac == nil {
		hjac = func(x *core.Array) (*core.Array, error) {
			return core.Jacobian(observation, x, core.JacobianOptions{})
		}
	}
	return kalmanRun(nMeas, m, d, meas, x0, p0, func(t int, zRow, x, p []float64) (kalmanStepOut, error) {
		xArr, err := core.FromFloats(x, d)
		if err != nil {
			return kalmanStepOut{}, err
		}
		fOut, err := transition(xArr)
		if err != nil {
			return kalmanStepOut{}, base.Errf("%s: at step %d: the transition: %w", name, t+1, err)
		}
		xp, err := kfVector(name, "the transition output", fOut, d)
		if err != nil {
			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
		}
		fJacArr, err := fjac(xArr)
		if err != nil {
			return kalmanStepOut{}, base.Errf("%s: at step %d: the transition Jacobian: %w", name, t+1, err)
		}
		fMat, err := kfMatrix(name, "the transition Jacobian", fJacArr, d, d)
		if err != nil {
			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
		}
		// Predict through the linearisation at the filtered mean.
		pp := kfMatMul(kfMatMul(fMat, d, d, p, d), d, d, kfTranspose(fMat, d, d), d)
		for i := range pp {
			pp[i] += q[i]
		}
		// Correct through the linearisation at the predicted mean, but
		// innovate against the true observation map.
		xpArr, err := core.FromFloats(xp, d)
		if err != nil {
			return kalmanStepOut{}, err
		}
		hOut, err := observation(xpArr)
		if err != nil {
			return kalmanStepOut{}, base.Errf("%s: at step %d: the observation: %w", name, t+1, err)
		}
		predicted, err := kfVector(name, "the observation output", hOut, m)
		if err != nil {
			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
		}
		hJacArr, err := hjac(xpArr)
		if err != nil {
			return kalmanStepOut{}, base.Errf("%s: at step %d: the observation Jacobian: %w", name, t+1, err)
		}
		hMat, err := kfMatrix(name, "the observation Jacobian", hJacArr, m, d)
		if err != nil {
			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
		}
		return kfCorrect(name, t+1, zRow, predicted, xp, pp, d, m, hMat, kfTranspose(hMat, m, d), r)
	})
}

// UnscentedKalmanFilter runs the unscented Kalman filter over the
// measurement stack z under the nonlinear model x_{t+1} = f(x_t) +
// w_t, z_t = h(x_t) + v_t: the state distribution N(x̂, P) is carried
// through f and h exactly to second order by the deterministic
// sigma-point set x̂ ± sqrt(d+λ)·L[:, i], L the Cholesky factor of P,
// whose weighted moments reconstruct the predicted mean and
// covariance. The update re-draws the set from the predicted
// distribution, so the state-observation cross-covariance carries the
// process noise the innovation covariance does. λ = α²(d+κ) − d, with
// d the state dimension the sigma-point set spans, comes from
// KalmanOptions' SigmaAlpha, SigmaBeta and SigmaKappa; the
// weights are the standard scaled set, with the covariance weight of
// the centre point carrying the 1 − α² + β prior correction.
//
// The correction carries no Joseph form: the gain comes from the
// explicit cross-covariance of state and observation rather than an
// observation matrix, so the covariance leaves the update as
// P⁻ − K·S·Kᵀ mirrored into its symmetric average, and positive
// definiteness is enforced by the next predict's Cholesky
// factorisation, which names the step when it fails. The log-likelihood
// accumulates exactly as in the linear filter. On a linear model the
// sigma transforms are exact and the filter degenerates to the
// Kalman answer to rounding; the state dimension must be fixed by an
// initial state or covariance (see KalmanOptions).
func UnscentedKalmanFilter(z *core.Array, transition, observation StateFunc, opts KalmanOptions) (*KalmanResult, error) {
	const name = "UnscentedKalmanFilter"
	if transition == nil || observation == nil {
		return nil, base.Errf("%s: the transition and observation maps are required", name)
	}
	nMeas, m, meas, err := kfMeasurements(name, z)
	if err != nil {
		return nil, err
	}
	d, x0, p0, q, r, err := kfNonlinearInputs(name, opts, m)
	if err != nil {
		return nil, err
	}
	alpha := opts.SigmaAlpha
	if alpha == 0 {
		alpha = 0.001
	}
	beta := opts.SigmaBeta
	if beta == 0 {
		beta = 2
	}
	kappa := opts.SigmaKappa
	if math.IsNaN(alpha) || alpha < 0 {
		return nil, base.Errf("%s: the sigma alpha must be unset (the default 0.001) or positive, got %g",
			name, opts.SigmaAlpha)
	}
	if math.IsNaN(beta) || beta < 0 {
		return nil, base.Errf("%s: the sigma beta must be unset (the default 2) or positive, got %g",
			name, opts.SigmaBeta)
	}
	if math.IsNaN(kappa) {
		return nil, base.Errf("%s: the sigma kappa must not be NaN, got %g", name, kappa)
	}
	scale := alpha * alpha * (float64(d) + kappa)
	if scale <= 0 {
		return nil, base.Errf("%s: the sigma spread vanishes: alpha %g and kappa %g leave no positive scale for %d states",
			name, alpha, kappa, d)
	}
	points := 2*d + 1
	lambda := scale - float64(d)
	wm := make([]float64, points)
	wc := make([]float64, points)
	wm[0] = lambda / scale
	wc[0] = wm[0] + 1 - alpha*alpha + beta
	for i := 1; i < points; i++ {
		wm[i] = 1 / (2 * scale)
		wc[i] = wm[i]
	}
	spread := math.Sqrt(scale)
	sig := make([]float64, points*d)
	prop := make([]float64, points*d)
	obs := make([]float64, points*m)
	return kalmanRun(nMeas, m, d, meas, x0, p0, func(t int, zRow, x, p []float64) (kalmanStepOut, error) {
		// The sigma set: the mean plus each Cholesky column of the
		// covariance, pushed out by sqrt(d+λ) both ways.
		l, err := kfCholesky("the state covariance", p, d)
		if err != nil {
			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
		}
		copy(sig[:d], x)
		for i := range d {
			plus := sig[(1+2*i)*d : (2+2*i)*d]
			minus := sig[(2+2*i)*d : (3+2*i)*d]
			copy(plus, x)
			copy(minus, x)
			for j := range d {
				v := spread * l[j*d+i]
				plus[j] += v
				minus[j] -= v
			}
		}
		// Predict: every sigma through f, then the weighted moments.
		for s := range points {
			sa, err := core.FromFloats(sig[s*d:(s+1)*d], d)
			if err != nil {
				return kalmanStepOut{}, err
			}
			fa, err := transition(sa)
			if err != nil {
				return kalmanStepOut{}, base.Errf("%s: at step %d: the transition: %w", name, t+1, err)
			}
			fv, err := kfVector(name, "the transition output", fa, d)
			if err != nil {
				return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
			}
			copy(prop[s*d:(s+1)*d], fv)
		}
		xp := make([]float64, d)
		for s := range points {
			for j := range d {
				xp[j] += wm[s] * prop[s*d+j]
			}
		}
		pp := make([]float64, d*d)
		for s := range points {
			w := wc[s]
			prows := prop[s*d : (s+1)*d]
			for a := range d {
				da := prows[a] - xp[a]
				for b := range d {
					pp[a*d+b] += w * da * (prows[b] - xp[b])
				}
			}
		}
		for i := range pp {
			pp[i] += q[i]
		}
		kfSymmetrise(pp, d)
		// The update re-draws the sigma set from the predicted
		// distribution N(x⁻, P⁻), the process noise included: the
		// state-observation cross-covariance must carry the same
		// spread the innovation covariance does, or the gain loses the
		// Q·Hᵀ term and a linear model stops degenerating to the
		// Kalman answer.
		lp2, err := kfCholesky("the predicted state covariance", pp, d)
		if err != nil {
			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
		}
		copy(sig[:d], xp)
		for i := range d {
			plus := sig[(1+2*i)*d : (2+2*i)*d]
			minus := sig[(2+2*i)*d : (3+2*i)*d]
			copy(plus, xp)
			copy(minus, xp)
			for j := range d {
				v := spread * lp2[j*d+i]
				plus[j] += v
				minus[j] -= v
			}
		}
		// Observation: the predicted sigmas through h, the weighted
		// moments, and the state-observation cross-covariance.
		for s := range points {
			pa, err := core.FromFloats(sig[s*d:(s+1)*d], d)
			if err != nil {
				return kalmanStepOut{}, err
			}
			ha, err := observation(pa)
			if err != nil {
				return kalmanStepOut{}, base.Errf("%s: at step %d: the observation: %w", name, t+1, err)
			}
			hv, err := kfVector(name, "the observation output", ha, m)
			if err != nil {
				return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
			}
			copy(obs[s*m:(s+1)*m], hv)
		}
		zp := make([]float64, m)
		for s := range points {
			for j := range m {
				zp[j] += wm[s] * obs[s*m+j]
			}
		}
		s := make([]float64, m*m)
		for k := range points {
			w := wc[k]
			orow := obs[k*m : (k+1)*m]
			for a := range m {
				da := orow[a] - zp[a]
				for b := range m {
					s[a*m+b] += w * da * (orow[b] - zp[b])
				}
			}
		}
		for i := range s {
			s[i] += r[i]
		}
		kfSymmetrise(s, m)
		cross := make([]float64, d*m)
		for k := range points {
			w := wc[k]
			prow := sig[k*d : (k+1)*d]
			orow := obs[k*m : (k+1)*m]
			for a := range d {
				da := prow[a] - xp[a]
				for b := range m {
					cross[a*m+b] += w * da * (orow[b] - zp[b])
				}
			}
		}
		ls, err := kfCholesky("the innovation covariance", s, m)
		if err != nil {
			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
		}
		innovation := make([]float64, m)
		for i := range m {
			innovation[i] = zRow[i] - zp[i]
		}
		solved := kfCholSolve(ls, m, innovation)
		quad := 0.0
		for i := range m {
			quad += innovation[i] * solved[i]
		}
		// K = P_xz·S⁻¹ through the transposed solve: S·X = P_xzᵀ
		// gives X = S⁻¹·P_xzᵀ and K = Xᵀ.
		gain := kfTranspose(kfCholSolveMatrix(ls, m, kfTranspose(cross, d, m), d), m, d)
		mean := make([]float64, d)
		copy(mean, xp)
		for i := range d {
			total := 0.0
			grow := gain[i*m : (i+1)*m]
			for j, y := range innovation {
				total += grow[j] * y
			}
			mean[i] += total
		}
		kskt := kfMatMul(kfMatMul(gain, d, m, s, m), d, m, kfTranspose(gain, d, m), d)
		covariance := make([]float64, d*d)
		for i := range covariance {
			covariance[i] = pp[i] - kskt[i]
		}
		kfSymmetrise(covariance, d)
		logLik := -0.5 * (float64(m)*math.Log(2*math.Pi) + kfLogDet(ls, m) + quad)
		return kalmanStepOut{mean: mean, covariance: covariance, innovation: innovation,
			innovationCov: s, logLikelihood: logLik}, nil
	})
}