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

package signal

import (
	"slices"

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

// Daubechies discrete wavelets: db2 through db8 beside the
// Haar transform the package already carries in DWT and IDWT. Each
// family is an orthonormal two-channel filter bank driven over levels
// of halving blocks, packed in the same [A_L, D_L, …, D_1] layout the
// Haar transform uses; each dbN pair has N vanishing moments and 2N
// taps, so higher families resolve smoother signals into sparser
// details but tie the boundary condition to longer blocks.
//
// The tables below are not trusted on authority: they were produced
// by the spectral factorisation of Daubechies' polynomial
// P(y) = Σ_k C(N−1+k, k)·y^k, and the tests certify them against the
// defining conditions, the unit norm, the shift-2 orthogonality, the
// N vanishing moments and the spectral identity
// |H(ω)|² = 2·cos^{2N}(ω/2)·P(sin²(ω/2)), so a mistyped digit cannot
// survive.

// Daubechies names a member of the Daubechies wavelet family: dbN
// carries N vanishing moments and a 2N-tap filter pair. Haar remains
// the db1 case and keeps its own exact transform in DWT and IDWT.
type Daubechies string

const (
	// DB2 is the 4-tap Daubechies wavelet with 2 vanishing moments.
	DB2 Daubechies = "db2"
	// DB3 is the 6-tap Daubechies wavelet with 3 vanishing moments.
	DB3 Daubechies = "db3"
	// DB4 is the 8-tap Daubechies wavelet with 4 vanishing moments.
	DB4 Daubechies = "db4"
	// DB5 is the 10-tap Daubechies wavelet with 5 vanishing moments.
	DB5 Daubechies = "db5"
	// DB6 is the 12-tap Daubechies wavelet with 6 vanishing moments.
	DB6 Daubechies = "db6"
	// DB7 is the 14-tap Daubechies wavelet with 7 vanishing moments.
	DB7 Daubechies = "db7"
	// DB8 is the 16-tap Daubechies wavelet with 8 vanishing moments.
	DB8 Daubechies = "db8"
)

// The scaling filters h, the low-pass analysis side, normalised to
// Σ h = √2. Treat them as read-only.
var (
	db2Coeffs = []float64{
		0.48296291314453421, 0.83651630373780794, 0.22414386804201339, -0.1294095225512604,
	}
	db3Coeffs = []float64{
		0.33267055295008258, 0.80689150931109266, 0.45987750211849149, -0.13501102001025453,
		-0.085441273882026644, 0.035226291885709533,
	}
	db4Coeffs = []float64{
		0.23037781330889651, 0.71484657055291556, 0.63088076792985903, -0.027983769416859948,
		-0.18703481171909306, 0.030841381835560764, 0.032883011666885203, -0.010597401785069032,
	}
	db5Coeffs = []float64{
		0.1601023979741929, 0.60382926979718965, 0.72430852843777283, 0.13842814590132091,
		-0.24229488706638203, -0.032244869584638361, 0.077571493840045691, -0.0062414902127982726,
		-0.012580751999081997, 0.0033357252854737704,
	}
	db6Coeffs = []float64{
		0.11154074335010949, 0.49462389039845328, 0.75113390802109559, 0.31525035170919741,
		-0.22626469396543983, -0.12976686756726177, 0.097501605587322904, 0.027522865530305803,
		-0.031582039317486044, 0.00055384220116149461, 0.0047772575109455116, -0.0010773010853084798,
	}
	db7Coeffs = []float64{
		0.07785205408500917, 0.39653931948191723, 0.72913209084623509, 0.46978228740519357,
		-0.14390600392856487, -0.22403618499387515, 0.071309219266830329, 0.080612609151083051,
		-0.0380299369350144, -0.016574541630666881, 0.012550998556099839, 0.00042957797292136738,
		-0.0018016407040474906, 0.00035371379997452024,
	}
	db8Coeffs = []float64{
		0.054415842243104001, 0.3128715909143, 0.67563073629728954, 0.58535468365420718,
		-0.015829105256349327, -0.28401554296154741, 0.00047248457391386436, 0.12874742662047808,
		-0.017369301001807447, -0.044088253930794685, 0.013981027917398216, 0.0087460940474057974,
		-0.0048703529934515776, -0.00039174037337694672, 0.00067544940645056933, -0.00011747678412476953,
	}
)

// daubechiesTaps resolves a family name to its scaling filter.
func daubechiesTaps(name string, family Daubechies) ([]float64, error) {
	switch family {
	case DB2:
		return db2Coeffs, nil
	case DB3:
		return db3Coeffs, nil
	case DB4:
		return db4Coeffs, nil
	case DB5:
		return db5Coeffs, nil
	case DB6:
		return db6Coeffs, nil
	case DB7:
		return db7Coeffs, nil
	case DB8:
		return db8Coeffs, nil
	default:
		return nil, base.Errf("%s: unknown Daubechies family %q (want db2 through db8)", name, string(family))
	}
}

// DWTMode picks the boundary treatment of the Daubechies transform.
type DWTMode int

const (
	// DWTPeriodic treats the signal as one period of a periodic
	// sequence: the default. Every block keeps its exact energy and
	// the length must offer every level a multiple of the filter
	// length to work on.
	DWTPeriodic DWTMode = iota
	// DWTZeroPad extends the signal with zeros at the tail to the
	// next length the level tree needs, so any length transforms;
	// the coefficients past the signal's own span carry the
	// response of the step to zero at the seam.
	DWTZeroPad
)

// DaubechiesDWT returns the discrete Daubechies wavelet transform of
// the rank-1 real signal x over levels scales, packed as
// [A_levels, D_levels, …, D_1]: the deepest approximation first, each
// detail band after it, the finest last. mode picks the boundary
// treatment, DWTPeriodic by default: it refuses a length that does
// not leave every live block a multiple of the 2N-tap filter (the
// deepest block must hold at least one filter length), while
// DWTZeroPad zero-pads the tail to the next such length first. The
// periodic transform conserves energy exactly; a zero-padded one
// conserves the energy of the padded signal. levels must be at
// least 1.
func DaubechiesDWT(x *core.Array, family Daubechies, levels int, mode DWTMode) (*core.Array, error) {
	const name = "DaubechiesDWT"
	src, h, levels, mode, err := dwtPrepare(name, x, family, levels, mode)
	if err != nil {
		return nil, err
	}
	L := len(h)
	n := x.Len()
	// The periodic contract gates the length; the zero-padded one
	// extends the tail to the smallest length that offers every
	// level a multiple of the filter length instead.
	length := n
	if mode == DWTPeriodic {
		if err := dwtLengthGate(name, n, L, levels); err != nil {
			return nil, err
		}
	} else {
		// The padded block must stay inside the allocator's reach: the
		// shift wraps past the addressable range first, and a span that
		// survives the wrap but tops the makeslice ceiling of 2^45
		// float64 elements panics in the make below instead of refusing
		// here.
		span := L << (levels - 1)
		if span <= 0 || span >= 1<<45 {
			return nil, base.Errf("%s: %d levels need a padded length no machine could hold", name, levels)
		}
		length = ((n-1)/span + 1) * span
	}
	work := make([]float64, length)
	copy(work, src)
	g := daubechiesHighPass(h)
	// Each pass halves the live block: the approximation and the
	// detail land in scratch, then the block rearranges into
	// [A, D] exactly as the packed layout wants.
	scratch := make([]float64, length)
	m := length
	for range levels {
		half := m / 2
		a, d := scratch[:half], scratch[half:m]
		dwtForwardLevel(work[:m], a, d, h, g)
		copy(work[:half], a)
		copy(work[half:m], d)
		m = half
	}
	return core.FromFloats(work, length)
}

// DaubechiesIDWT inverts DaubechiesDWT over the same family, level
// count and mode, undoing the packed layout deepest band first. The
// coefficient vector must satisfy the length contract every forward
// transform produced: a multiple of the filter length at every level
// of the tree, whichever mode the forward ran. A zero-padded forward
// transform therefore inverts to the padded length, whose head is
// the original signal; the mode argument gates exactly this
// contract, the inverse bank itself is the periodic one.
func DaubechiesIDWT(coef *core.Array, family Daubechies, levels int, mode DWTMode) (*core.Array, error) {
	const name = "DaubechiesIDWT"
	src, h, levels, _, err := dwtPrepare(name, coef, family, levels, mode)
	if err != nil {
		return nil, err
	}
	if err := dwtLengthGate(name, coef.Len(), len(h), levels); err != nil {
		return nil, err
	}
	g := daubechiesHighPass(h)
	work := make([]float64, len(src))
	copy(work, src)
	m := len(work) >> levels
	for range levels {
		half := m
		m = 2 * m
		a := slices.Clone(work[:half])
		d := slices.Clone(work[half:m])
		for j := range m {
			var av, dv float64
			for k := range h {
				// The synthesis taps land at even steps only:
				// e/2 is the band index the tap reads.
				e := j - k
				if e < 0 {
					e += m
				}
				if e&1 == 1 {
					continue
				}
				av += h[k] * a[e>>1]
				dv += g[k] * d[e>>1]
			}
			work[j] = av + dv
		}
	}
	return core.FromFloats(work, len(work))
}

// dwtPrepare runs the shared contract of the Daubechies transforms
// ahead of the length question: rank-1 real non-empty input, a known
// family, at least one level, a known mode. It returns the widened
// input, the scaling filter and the validated level count and mode.
func dwtPrepare(name string, x *core.Array, family Daubechies, levels int, mode DWTMode) (src []float64, h []float64, levelsOut int, modeOut DWTMode, err error) {
	if x.NDim() != 1 {
		return nil, nil, 0, 0, base.Errf("%s: needs a rank-1 signal, got shape %s", name, base.ShapeText(x.Shape()))
	}
	if x.Dtype() == core.Complex {
		return nil, nil, 0, 0, base.Errf("%s: complex signals are not supported", name)
	}
	if x.Len() == 0 {
		return nil, nil, 0, 0, base.Errf("%s: an empty signal has no transform", name)
	}
	h, err = daubechiesTaps(name, family)
	if err != nil {
		return nil, nil, 0, 0, err
	}
	if levels < 1 {
		return nil, nil, 0, 0, base.Errf("%s: levels must be at least 1, got %d", name, levels)
	}
	switch mode {
	case DWTPeriodic, DWTZeroPad:
	default:
		return nil, nil, 0, 0, base.Errf("%s: unknown boundary mode %d", name, int(mode))
	}
	return widenFloats(x), h, levels, mode, nil
}

// dwtLengthGate enforces the tree contract on a length: the deepest
// block, n>>levels−1 samples, must be a whole multiple of the filter
// length and at least one filter long, which the halvings before it
// then inherit. Shifts past the word width answer 0 and fail the
// gate, which is exactly the refusal a too deep tree wants.
func dwtLengthGate(name string, n, L, levels int) error {
	blk := n >> (levels - 1)
	if levels-1 >= 64 || (blk<<(levels-1)) != n || blk < L || blk%L != 0 {
		return base.Errf("%s: %d levels need every live block to hold a multiple of the %d-tap filter, which %d samples do not offer",
			name, levels, L, n)
	}
	return nil
}

// daubechiesHighPass derives the wavelet (high-pass) half of the bank
// from the scaling filter: g[k] = (−1)^k·h[L−1−k].
func daubechiesHighPass(h []float64) []float64 {
	L := len(h)
	g := make([]float64, L)
	for k := range L {
		sign := 1.0
		if k%2 == 1 {
			sign = -1
		}
		g[k] = sign * h[L-1-k]
	}
	return g
}

// dwtForwardLevel runs one periodic analysis level over the block
// x of even length m: a[i] = Σ h[k]·x[(2i+k) mod m] and the same with
// g for d. Outputs whose taps never wrap walk the block directly;
// the ones near the block's end correct the few negative indices.
func dwtForwardLevel(x, a, d, h, g []float64) {
	m := len(x)
	L := len(h)
	half := m / 2
	limit := (m - L) / 2
	for i := 0; i <= limit; i++ {
		base := 2 * i
		var av, dv float64
		for k := range L {
			xv := x[base+k]
			av += h[k] * xv
			dv += g[k] * xv
		}
		a[i], d[i] = av, dv
	}
	for i := limit + 1; i < half; i++ {
		base := 2*i - m
		var av, dv float64
		for k := range L {
			idx := base + k
			if idx < 0 {
				idx += m
			}
			xv := x[idx]
			av += h[k] * xv
			dv += g[k] * xv
		}
		a[i], d[i] = av, dv
	}
}