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

package signal

import (
	"math"

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

// Sample-rate conversion. Decimation and rational resampling run a
// linear-phase FIR anti-alias filter in the time domain, so they suit
// aperiodic streams; ResampleFourier is the exact band-limited
// resample of the Fourier definition and suits whole records.

// kaiserI0 evaluates the modified Bessel function of the first kind
// at a real order zero by its power series, which converges in a few
// dozen terms across the window betas any reasonable design uses.
func kaiserI0(x float64) float64 {
	sum, term := 1.0, 1.0
	half := x / 2
	for k := 1; k < 64; k++ {
		term *= (half * half) / float64(k*k)
		sum += term
		if term < 1e-18*sum {
			break
		}
	}
	return sum
}

// kaiserSinc builds the odd-length FIR low-pass of a windowed sinc:
// gain 1 at DC, cutoff fc normalised to the sampling rate (half the
// output rate after conversion), Kaiser taper with beta 8.6, which
// puts the stopband around 80 dB. The cutoff rides the middle of the
// transition band, so the design slightly attenuates the top of the
// passband by construction; tests pin that below a dB.
func kaiserSinc(taps int, fc float64) []float64 {
	if taps%2 == 0 {
		taps++
	}
	// A 1-tap kernel has no window to apply: the ratio inside r
	// divides by taps−1 = 0 and turns every coefficient NaN. The
	// kernel is the identity, gain 1 at DC, which is exactly what a
	// one-tap low-pass means: keep the sample, filter nothing.
	if taps == 1 {
		return []float64{1}
	}
	const beta = 8.6
	i0b := kaiserI0(beta)
	h := make([]float64, taps)
	centre := (taps - 1) / 2
	sum := 0.0
	for i := range taps {
		r := 2*float64(i)/float64(taps-1) - 1
		w := kaiserI0(beta*math.Sqrt(math.Max(0, 1-r*r))) / i0b
		arg := 2 * fc * float64(i-centre)
		s := 1.0
		if arg != 0 {
			s = math.Sin(math.Pi*arg) / (math.Pi * arg)
		}
		h[i] = 2 * fc * s * w
		sum += h[i]
	}
	for i := range h {
		h[i] /= sum
	}
	return h
}

// Decimate reduces the sample rate by the integer factor: a Kaiser
// tapered FIR anti-alias filter runs first, the filter's group delay
// is compensated, and every factor-th sample of the compensated
// series is kept. The passband ends at nine tenths of the new
// Nyquist; the transition band then reaches its stopband floor of
// about 80 dB before the new Nyquist, so everything that would fold
// into the output band is suppressed. Taps sets the filter length; a
// non-positive value means 32·factor+1, which is the length that
// fits that transition. Output starts once the filter has full
// context, so the result holds about (n − taps)/factor samples of an
// input of n; when the tap count leaves no sample with full context,
// the result is empty rather than read past the end of the signal.
func Decimate(data *core.Array, factor, taps int) (*core.Array, error) {
	const name = "Decimate"
	if data.NDim() != 1 {
		return nil, base.Errf("%s: the series must be a vector, got shape %s", name, base.ShapeText(data.Shape()))
	}
	if data.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex series are not supported", name)
	}
	if factor < 2 {
		return nil, base.Errf("%s: the factor must be at least 2, got %d", name, factor)
	}
	n := data.Len()
	if n == 0 {
		return nil, base.Errf("%s: the series must not be empty", name)
	}
	if taps <= 0 {
		taps = 32*factor + 1
	}
	if taps%2 == 0 {
		taps++
	}
	if taps >= n {
		return nil, base.Errf("%s: %d taps against %d samples leaves nothing after the filter delay", name, taps, n)
	}
	h := kaiserSinc(taps, 0.45/float64(factor))
	filtered, err := FilterApply(h, []float64{1}, data)
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	delay := (taps - 1) / 2
	// Sample the compensated series (compensated index m lives at
	// filtered[m + delay]) from where its filter context is complete:
	// m + delay ≥ taps − 1 means m ≥ delay, so the first kept index
	// is the first multiple of factor at or past delay, and the last
	// is the last multiple whose compensated read still lies inside
	// the signal. When no multiple qualifies, the result is empty:
	// a negative numerator must not be left to truncate toward zero,
	// which used to turn "no fully filtered sample" into one read
	// past the end of the signal.
	skip := (delay + factor - 1) / factor
	outLen := 0
	if span := filtered.Len() - 1 - delay; span >= 0 {
		outLen = max(span/factor-skip+1, 0)
	}
	out := core.New(core.Float, outLen)
	vals := out.RawFloats()
	// FilterApply keeps the input dtype, so a float32 series comes
	// back with a float32 payload: read it through widenFloats (the
	// FloatAt widening the rest of the package reads with) instead of
	// the nil float64 payload.
	src := widenFloats(filtered)
	for i := range outLen {
		vals[i] = src[skip*factor+delay+i*factor]
	}
	return out, nil
}

// Resample converts the sample rate by the rational factor up/down:
// the series is filtered on the up-sampled grid by a Kaiser tapered
// FIR at the tighter of the two Nyquists with the up gain folded in,
// and every down-th sample of the compensated result is kept. Up and
// down must be at least 1 and not both 1; taps sets the kernel length
// on the up-sampled grid, a non-positive value meaning
// 32·max(up, down)+1, the length that fits the anti-alias transition.
// At the series ends the filter reads past the data: the up-sampled
// grid counts as zeros outside the input, so the head and the tail
// after the delay compensation are the sums over those zeros, not
// wrapped or skipped samples.
func Resample(data *core.Array, up, down, taps int) (*core.Array, error) {
	const name = "Resample"
	if data.NDim() != 1 {
		return nil, base.Errf("%s: the series must be a vector, got shape %s", name, base.ShapeText(data.Shape()))
	}
	if data.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex series are not supported", name)
	}
	if up < 1 || down < 1 || (up == 1 && down == 1) {
		return nil, base.Errf("%s: the rate change up/down must not be the identity %d/%d", name, up, down)
	}
	n := data.Len()
	if n == 0 {
		return nil, base.Errf("%s: the series must not be empty", name)
	}
	if taps <= 0 {
		taps = 32*max(up, down) + 1
	}
	if taps%2 == 0 {
		taps++
	}
	if taps > up*n {
		return nil, base.Errf("%s: %d taps against %d up-sampled samples leaves nothing after the filter delay", name, taps, up*n)
	}
	// The source is read through widenFloats: a float32 or int series
	// carries no float64 payload, and widenFloats is the exact FloatAt
	// widening the rest of the package reads with.
	src := widenFloats(data)
	// Cutoff on the up-sampled grid: the smaller Nyquist of input and
	// output, in cycles per up-sampled sample.
	fc := 0.5 * math.Min(1/float64(up), 1/float64(down))
	h := kaiserSinc(taps, fc)
	delay := (taps - 1) / 2
	outLen := (n*up + down - 1) / down
	out := core.New(core.Float, outLen)
	vals := out.RawFloats()
	for m := range outLen {
		// Output sample m is up-sample index m·down; the filter
		// centred there sums the inputs within its support. Only k
		// with centre − k·up inside [0, taps) contributes, so the loop
		// runs over that window alone, in the same ascending order.
		centre := delay + m*down
		kmin := max(0, (centre-taps+1+up-1)/up)
		kmax := min(n-1, centre/up)
		total := 0.0
		for k := kmin; k <= kmax; k++ {
			total += h[centre-k*up] * src[k]
		}
		// j runs over the kernel, so the compensated sample sits at
		// up-index centre − delay; scale by up for the zero-stuffed
		// grid's unit gain.
		vals[m] = float64(up) * total
	}
	return out, nil
}

// ResampleFourier resamples a series to exactly size samples by the
// Fourier (band-limited) definition: the spectrum's bins are kept,
// padded with zeros or truncated at the fold, and scaled so the
// result carries the same tone amplitudes as the input. It is exact
// for series band-limited below the new Nyquist and treats the input
// as one period, the same convention AnalyticSignal uses. Signals
// with energy past the new Nyquist lose it, which is the brick-wall
// anti-alias this resample implies.
func ResampleFourier(data *core.Array, size int) (*core.Array, error) {
	const name = "ResampleFourier"
	if data.NDim() != 1 {
		return nil, base.Errf("%s: the series must be a vector, got shape %s", name, base.ShapeText(data.Shape()))
	}
	n := data.Len()
	if n == 0 {
		return nil, base.Errf("%s: the series must not be empty", name)
	}
	if size < 1 {
		return nil, base.Errf("%s: the target size must be at least 1, got %d", name, size)
	}
	spec, err := RFFT(data)
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	oldHalf := spec.Len()
	newHalf := size/2 + 1
	scaled := make([]complex128, newHalf)
	ratio := float64(size) / float64(n)
	for k := range min(oldHalf, newHalf) {
		scaled[k] = complex(ratio, 0) * spec.RawComplexes()[k]
	}
	target, err := core.FromComplexes(scaled, newHalf)
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	out, err := IRFFT(target, size)
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	return out, nil
}