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

package signal

import (
	"math"
	"sync"

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

// Savitzky-Golay smoothing: every sample is replaced by the value at
// its position of a polynomial of the given order fitted by least
// squares over a symmetric window of samples. Unlike a moving average
// this preserves the moments of the signal up to the fit order, so
// peaks keep their heights and widths instead of being flattened,
// which is why spectroscopy and chromatography standardised on it.
//
// The interior uses one set of precomputed weights, the polynomial
// fit's hat row for the window centre. Samples within half a window
// of an edge use a window truncated at the boundary, fitted at the
// sample's own position, so the output is complete without padding
// and the ends keep the full polynomial treatment at reduced support.

// SavitzkyGolay smooths a rank-1 signal with a Savitzky-Golay
// polynomial filter: window is the (odd) number of samples per fit,
// order the polynomial degree. The standard constraint order ≤
// window/2 applies: a higher degree leaves the truncated edge windows
// with fewer samples than coefficients, so the fit is underdetermined
// and an order above window/2 is an error. A polynomial of degree ≤
// order passes through unchanged.
func SavitzkyGolay(data *core.Array, window, order int) (*core.Array, error) {
	if data.Dtype() == core.Complex {
		return nil, base.Errf("SavitzkyGolay: complex signals are not supported")
	}
	if data.NDim() != 1 {
		return nil, base.Errf("SavitzkyGolay: needs a rank-1 signal, got shape %s", base.ShapeText(data.Shape()))
	}
	n := data.Len()
	if n == 0 {
		return nil, base.Errf("SavitzkyGolay: the signal must not be empty")
	}
	if window < 3 || window%2 == 0 {
		return nil, base.Errf("SavitzkyGolay: window must be an odd number ≥ 3, got %d", window)
	}
	if window > n {
		return nil, base.Errf("SavitzkyGolay: window %d exceeds the signal length %d", window, n)
	}
	if order < 0 {
		return nil, base.Errf("SavitzkyGolay: order must not be negative, got %d", order)
	}
	if order > window/2 {
		return nil, base.Errf("SavitzkyGolay: order %d exceeds window/2 (%d), the edge fits would be underdetermined",
			order, window/2)
	}
	h := window / 2

	// Weights for a window with offsets offs, evaluating the fit at
	// offset 0 (the centre sample).
	weights := func(offs []int) ([]float64, error) {
		k := order + 1
		m := len(offs)
		// Normal equations of the Vandermonde system.
		ata := make([][]float64, k)
		for i := range k {
			ata[i] = make([]float64, k)
			for j := range k {
				s := 0.0
				for _, o := range offs {
					s += math.Pow(float64(o), float64(i+j))
				}
				ata[i][j] = s
			}
		}
		// Huge windows at high order overflow the power sums: a
		// non-finite entry would flow through the solver into NaN
		// weights published as a smooth signal, so it is refused here.
		for i := range k {
			for j := range k {
				if math.IsInf(ata[i][j], 0) || math.IsNaN(ata[i][j]) {
					return nil, base.Errf("SavitzkyGolay: the normal equations overflow for window %d at order %d", 2*h+1, order)
				}
			}
		}
		// Right-hand side: the coefficient vector of evaluating at 0
		// is e_0, so b_i = 0^i, which is 1 at i = 0 and 0 after.
		b := make([]float64, k)
		b[0] = 1
		sol, err := base.SolveSystem("SavitzkyGolay", ata, [][]float64{b})
		if err != nil {
			return nil, base.Errf("SavitzkyGolay: %w", err)
		}
		coef := sol[0]
		w := make([]float64, m)
		for i, o := range offs {
			s := 0.0
			for c := range k {
				s += coef[c] * math.Pow(float64(o), float64(c))
			}
			if math.IsInf(s, 0) || math.IsNaN(s) {
				return nil, base.Errf("SavitzkyGolay: the weights overflow for window %d at order %d", 2*h+1, order)
			}
			w[i] = s
		}
		return w, nil
	}

	// The interior weights and every edge point's truncated window are
	// a pure function of the window and the order, so the whole fit set
	// is served from a cache keyed by the two: the power sums, the
	// solver and the weight evaluations run once per window-and-order
	// instead of once per call.
	fits, ferr := savitzkyGolayFits(window, order, h, weights)
	if ferr != nil {
		return nil, ferr
	}
	edge := h
	out, oerr := core.Zeros(core.Float, []int{n}...)
	if oerr != nil {
		return nil, oerr
	}
	outRow := out.RawFloats()
	src := widenFloats(data)
	// Every output sample owns its own slot and, within half a window
	// of an edge, its own precomputed weight fit, so the points split
	// across workers once a worker's chunk clears the floor below; the
	// per-point summation over the window keeps the serial order. The
	// regions split so the interior walk carries no per-point branch:
	// each region sees its weights and its slice bounds settled once.
	// Every loop is clamped to the worker's own [ps, pe): an edge wider
	// than one chunk would otherwise let the edge loops cross a chunk
	// boundary, and two workers would write the same output slot at
	// once, a data race whose bits happen to agree but whose access the
	// memory model and the race detector both refuse.
	engine.ParallelMin(n, sgMinPointsPerWorker, func(ps, pe int) {
		sweep := func(i int) {
			var w []float64
			start := max(i-h, 0)
			switch {
			case i < edge:
				w = fits.left[i]
			case i >= n-edge:
				w = fits.right[n-1-i]
			default:
				w = fits.full
			}
			sum := 0.0
			row := src[start : start+len(w)]
			for k, weight := range w {
				sum += weight * row[k]
			}
			outRow[i] = sum
		}
		midLo := max(ps, edge)
		midHi := min(pe, n-edge)
		for i := ps; i < min(midLo, pe); i++ {
			sweep(i)
		}
		if midHi > midLo {
			// Interior: the full weights always apply and the window
			// never truncates, so the loop runs without the fit choice.
			for i := midLo; i < midHi; i++ {
				start := i - h
				sum := 0.0
				row := src[start : start+len(fits.full)]
				for k, weight := range fits.full {
					sum += weight * row[k]
				}
				outRow[i] = sum
			}
		}
		for i := max(midHi, ps); i < pe; i++ {
			sweep(i)
		}
	})
	return out, nil
}

// sgFitKey names one cached fit set: the window and order that
// determine every weight row in it.
type sgFitKey struct {
	window, order int
}

// sgFits holds the interior hat row and the truncated edge fits: for
// left edge point j the window runs over the offsets [-j, h], for the
// mirrored right edge point over [-h, j]. The rows are read-only once
// published.
type sgFits struct {
	full        []float64
	left, right [][]float64
}

// sgFitCacheMax bounds how many window-and-order fit sets stay cached;
// each set is a few hundred floats, so the map stays negligible
// whatever the callers do.
const sgFitCacheMax = 256

var (
	sgFitMu   sync.RWMutex
	sgFitKeys = map[sgFitKey]*sgFits{}
)

// savitzkyGolayFits returns the fit set for a window and order,
// building it through the caller's weights evaluation on first use.
// Every edge point owns a truncated window, fitted at the sample's own
// position, so the parallel sweep cannot fail mid-kernel.
func savitzkyGolayFits(window, order, h int, weights func(offs []int) ([]float64, error)) (*sgFits, error) {
	key := sgFitKey{window: window, order: order}
	sgFitMu.RLock()
	fits, ok := sgFitKeys[key]
	sgFitMu.RUnlock()
	if ok {
		return fits, nil
	}
	full, err := weights(rangeInts(-h, h))
	if err != nil {
		return nil, err
	}
	edge := h
	fits = &sgFits{full: full, left: make([][]float64, edge), right: make([][]float64, edge)}
	for j := range edge {
		lw, lerr := weights(rangeInts(-j, h))
		if lerr != nil {
			return nil, lerr
		}
		fits.left[j] = lw
		rw, rerr := weights(rangeInts(-h, j))
		if rerr != nil {
			return nil, rerr
		}
		fits.right[j] = rw
	}
	sgFitMu.Lock()
	if len(sgFitKeys) < sgFitCacheMax {
		sgFitKeys[key] = fits
	}
	sgFitMu.Unlock()
	return fits, nil
}

// sgMinPointsPerWorker is the smallest per-worker chunk of output
// points the Savitzky-Golay sweep splits for: below it a chunk's
// window arithmetic no longer pays the worker spawn cost.
const sgMinPointsPerWorker = 1 << 10

// widenFloats returns the array's elements widened to float64 exactly
// as FloatAt widens them: a contiguous float64 payload aliases the
// array, every other dtype (and any strided array, whose payload order
// is not its element order) converts into a fresh slice. Treat an
// aliased result as read-only.
func widenFloats(a *core.Array) []float64 {
	if a.Dtype() == core.Float && !a.Strided() {
		return a.RawFloats()
	}
	out := make([]float64, a.Len())
	for i := range out {
		out[i] = a.FloatAt(i)
	}
	return out
}

// rangeInts returns the integers from a to b inclusive.
func rangeInts(a, b int) []int {
	out := make([]int, b-a+1)
	for i := range out {
		out[i] = a + i
	}
	return out
}

// Digital filter design. The Butterworth family is the flat
// magnitude response: |H|² is monotone in frequency with no ripple on
// either side, which is what a data-cleaning low-pass wants. The
// design runs the classical route: prewarp the cutoff with the
// bilinear transform's tangent, place the analog Butterworth poles on
// the left half-plane circle, map them inside the unit circle, and
// read the coefficients off the factored transfer function. Conjugate
// pole pairs keep every coefficient real; odd orders carry one real
// pole.

// ButterworthLowPass designs an order-th low-pass filter for sample
// rate fs, returning the direct-form coefficients b (numerator) and a
// (denominator, a[0] = 1). The order must be at least 1, fs positive,
// and the cutoff strictly inside (0, fs/2). The response is −3 dB at
// the cutoff and unity at DC.
func ButterworthLowPass(order int, fs, cutoff float64) (b, a []float64, err error) {
	return butterworth(order, fs, cutoff, false)
}

// ButterworthHighPass designs the high-pass mirror: the same flat
// response turned around, −3 dB at the same cutoff, unity gain at
// Nyquist. It shares the low-pass poles and carries its zeros at
// z = 1, the bilinear image of the analog prototype's origin zeros.
func ButterworthHighPass(order int, fs, cutoff float64) (b, a []float64, err error) {
	return butterworth(order, fs, cutoff, true)
}

// butterworth carries the shared design: the digital low-pass poles,
// zeros and gain as coefficient polynomials, sign-flipped for the
// high-pass.
func butterworth(order int, fs, cutoff float64, highpass bool) (b, a []float64, err error) {
	const name = "Butterworth"
	if order < 1 {
		return nil, nil, base.Errf("%s: the order must be at least 1, got %d", name, order)
	}
	if !(fs > 0) || math.IsInf(fs, 0) {
		return nil, nil, base.Errf("%s: fs must be positive and finite, got %g", name, fs)
	}
	if !(cutoff > 0) || cutoff >= fs/2 {
		return nil, nil, base.Errf("%s: the cutoff must lie in (0, fs/2), got %g for fs %g", name, cutoff, fs)
	}
	// Prewarped analog cutoff: the bilinear transform maps the digital
	// cutoff onto it exactly, so the −3 dB point lands where asked.
	w := math.Tan(math.Pi * cutoff / fs)
	// Analog poles s_k = w·e^{iθ_k} on the left half-plane circle;
	// odd orders place one pole at θ = π, which is real.
	a = []float64{1}
	for k := range order {
		num := 2*k + order + 1
		var s complex128
		if num == 2*order {
			// The odd order's real pole at θ = π: built exactly, not
			// through sin(π), whose float noise would detour it into
			// the pair branch and fabricate a quadratic factor.
			s = complex(-w, 0)
		} else {
			theta := math.Pi * float64(num) / float64(2*order)
			// One Sincos serves both factors, the same pair the
			// separate calls produced, so the pole keeps its bits.
			sinT, cosT := math.Sincos(theta)
			s = complex(w*cosT, w*sinT)
		}
		z := (1 + s) / (1 - s)
		if imag(s) == 0 {
			a = mulPolyReal(a, []float64{1, -real(z)})
			continue
		}
		if imag(s) < 0 {
			continue // the conjugate partner was consumed on its pass
		}
		zc := (1 + conj(s)) / (1 - conj(s))
		// (1 − z·u)(1 − z̄·u) = 1 − 2Re(z)·u + |z|²·u², real by
		// construction; the rounding noise in the imaginary part drops.
		a = mulPolyReal(a, []float64{1, -2 * real(z), real(z * zc)})
	}
	// Zeros: the low-pass puts them all at z = −1 (numerator
	// (1+u)^order, DC-normalised); the high-pass shares the very same
	// poles and moves the zeros to z = +1 (numerator (1−u)^order,
	// Nyquist-normalised). Reflecting the poles by negating z instead
	// would answer a different question at the wrong frequencies.
	if highpass {
		b = binomialCoeffs(order)
		for i := range b {
			if i%2 == 1 {
				b[i] = -b[i]
			}
		}
		k := polyEvalAtMinusOne(a) / math.Pow(2, float64(order))
		for i := range b {
			b[i] *= k
		}
		return b, a, nil
	}
	b = binomialCoeffs(order)
	k := polyEvalAtOne(a) / math.Pow(2, float64(order))
	for i := range b {
		b[i] *= k
	}
	return b, a, nil
}

// mulPolyReal multiplies two real polynomials in u = z^{-1} (index m
// is the coefficient of u^m).
func mulPolyReal(p, q []float64) []float64 {
	out := make([]float64, len(p)+len(q)-1)
	for i, pv := range p {
		for j, qv := range q {
			out[i+j] += pv * qv
		}
	}
	return out
}

// binomialCoeffs returns (1+u)^order as a coefficient list, built by
// exact rational steps so the symmetry C(n,m) = C(n,n−m) survives
// rounding.
func binomialCoeffs(order int) []float64 {
	out := make([]float64, order+1)
	out[0] = 1
	for k := range order {
		out[k+1] = out[k] * float64(order-k) / float64(k+1)
	}
	return out
}

// polyEvalAtMinusOne evaluates a u-polynomial at z = −1: the
// alternating sum of its coefficients.
func polyEvalAtMinusOne(p []float64) float64 {
	s := 0.0
	for m, v := range p {
		if m%2 == 1 {
			s -= v
		} else {
			s += v
		}
	}
	return s
}

// polyEvalAtOne evaluates a u-polynomial at z = 1: the sum of its
// coefficients.
func polyEvalAtOne(p []float64) float64 {
	s := 0.0
	for _, v := range p {
		s += v
	}
	return s
}

// FilterApply runs a direct-form-II-transposed filter over a rank-1
// real signal: y[n] = Σ b_m·x[n−m] − Σ a_m·y[n−m]. The coefficient
// lists may differ in length (the shorter is zero-padded), a[0] must
// be non-zero and is normalised away, and the output has the same
// length as the input. Float32 and float64 inputs keep their dtype;
// int input widens to float64, because a filtered signal is not
// integral and truncating it would quantise the answer to zero
// wherever the response is small. This is the evaluator the
// Butterworth designs hand their coefficients to, and the sweep
// Filtfilt runs in both directions.
func FilterApply(b, a []float64, x *core.Array) (*core.Array, error) {
	const name = "FilterApply"
	bc, ac, out, err := filterPrepare(name, b, a, x)
	if err != nil {
		return nil, err
	}
	src := widenFloats(x)
	dst := make([]float64, x.Len())
	filterSweep(bc, ac, src, dst)
	outF := out.RawFloats()
	if outF == nil || out.Strided() {
		for i := range dst {
			out.SetFloatAt(i, dst[i])
		}
	} else {
		copy(outF, dst)
	}
	return out, nil
}

// filterPrepare runs the shared contract of the filter evaluators:
// rank-1 real input, non-empty coefficient lists with a non-zero
// leading denominator, coefficients normalised by a[0] and padded to
// one common length, and an output array keeping the input's float
// dtype. The normalised coefficients and the output come back
// together.
func filterPrepare(name string, b, a []float64, x *core.Array) (bc, ac []float64, out *core.Array, err error) {
	if x.NDim() != 1 {
		return nil, nil, nil, base.Errf("%s: needs a rank-1 signal, got shape %s", name, base.ShapeText(x.Shape()))
	}
	if x.Dtype() == core.Complex {
		return nil, nil, nil, base.Errf("%s: complex signals are not supported", name)
	}
	if len(b) == 0 || len(a) == 0 {
		return nil, nil, nil, base.Errf("%s: the coefficient lists must not be empty", name)
	}
	if a[0] == 0 {
		return nil, nil, nil, base.Errf("%s: a[0] must be non-zero", name)
	}
	n := max(len(b), len(a))
	bc = make([]float64, n)
	ac = make([]float64, n)
	for i := range b {
		bc[i] = b[i] / a[0]
	}
	ac[0] = 1
	for i := 1; i < len(a); i++ {
		ac[i] = a[i] / a[0]
	}
	outDT := x.Dtype()
	if outDT != core.Float && outDT != core.Float32 {
		outDT = core.Float
	}
	out, oerr := core.Zeros(outDT, x.Len())
	if oerr != nil {
		return nil, nil, nil, oerr
	}
	return bc, ac, out, nil
}

// filterSweep is the direct-form-II-transposed kernel both filter
// evaluators run: one pass of the normalised coefficients over src,
// written to dst, which must not overlap src. Transposed delays:
// coefficient indices run 0..n−1, so the recurrence carries exactly
// n−1 registers; the last one takes no +state term, which is where a
// phantom extra delay would sneak in.
func filterSweep(bc, ac, src, dst []float64) {
	n := len(bc)
	state := make([]float64, max(n-1, 1))
	for i := range src {
		xv := src[i]
		yv := bc[0] * xv
		if n > 1 {
			yv += state[0]
			for j := 1; j < n-1; j++ {
				state[j-1] = bc[j]*xv - ac[j]*yv + state[j]
			}
			state[n-2] = bc[n-1]*xv - ac[n-1]*yv
		}
		dst[i] = yv
	}
}