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

package signal

import (
	"math"
	"math/big"
	"testing"
)

// The fused radix-4 stages replace the radix-2 walk's arithmetic, so
// the accuracy question is answered against references this package
// does not implement: the quadratic definition in float64 (the one the
// functional tests read) and the same definition carried out in 256-bit
// floating point. The two paths under comparison are called directly,
// the radix-2 walk and the package dispatcher, so the error figures
// attribute cleanly to the stage kernels.

// precisionFixture returns the fixed test signal of length n: small
// integer-like values whose transform is far above the rounding floor
// of either path but far below the tolerances below.
func precisionFixture(n int) []complex128 {
	vals := make([]complex128, n)
	for i := range vals {
		vals[i] = complex(float64(i%17)-8, float64(i%5)-2)
	}
	return vals
}

// TestFFTPathAccuracyAgainstNaive pins the fused radix-4 stages to the
// quadratic float64 reference at lengths the functional tests never
// reach (every length they hold sits below the radix-4 floor). The
// tolerance leaves the quadratic reference's own accumulated rounding
// an order of magnitude of headroom.
func TestFFTPathAccuracyAgainstNaive(t *testing.T) {
	for _, n := range []int{64, 128, 256, 1024} {
		vals := precisionFixture(n)
		want := naiveDFT(vals, -1)
		got := append([]complex128(nil), vals...)
		fftPow2(got, -1)
		for i := range want {
			if math.Abs(real(got[i])-real(want[i])) > 1e-6 {
				t.Fatalf("FFT(%d)[%d]: real error %g against the quadratic reference", n, i, math.Abs(real(got[i])-real(want[i])))
			}
			if math.Abs(imag(got[i])-imag(want[i])) > 1e-6 {
				t.Fatalf("FFT(%d)[%d]: imaginary error %g against the quadratic reference", n, i, math.Abs(imag(got[i])-imag(want[i])))
			}
		}
	}
}

func newBig(prec uint) *big.Float { return new(big.Float).SetPrec(prec) }

// machinPi computes π at the given precision by Machin's formula
// π = 16·atan(1/5) − 4·atan(1/239), the arctangents by their Taylor
// series.
func machinPi(prec uint) *big.Float {
	atan := func(x *big.Float) *big.Float {
		x2 := newBig(prec).Mul(x, x)
		term := newBig(prec).Set(x)
		sum := newBig(prec).Set(x)
		tiny := newBig(prec).SetMantExp(big.NewFloat(1), -int(prec)-12)
		for k := 1; k < 200; k++ {
			term.Mul(term, x2)
			// The denominator moves from 2k−1 to 2k+1 across the step.
			term.Mul(term, big.NewFloat(float64(2*k-1)))
			term.Quo(term, big.NewFloat(float64(2*k+1)))
			if k%2 == 1 {
				sum.Sub(sum, term)
			} else {
				sum.Add(sum, term)
			}
			if abs := newBig(prec).Abs(term); abs.Cmp(tiny) < 0 {
				break
			}
		}
		return sum
	}
	fifth := newBig(prec).Quo(big.NewFloat(1), big.NewFloat(5))
	t239th := newBig(prec).Quo(big.NewFloat(1), big.NewFloat(239))
	pi := atan(fifth)
	pi.Mul(pi, big.NewFloat(16))
	t := atan(t239th)
	t.Mul(t, big.NewFloat(4))
	return pi.Sub(pi, t)
}

// sincosBig evaluates cosine and sine of a finite angle at 256-bit
// precision: the angle is range-reduced into [−π, π] and both series
// are summed to well below the precision floor. The pair comes back
// real part first.
func sincosBig(theta *big.Float, prec uint) [2]*big.Float {
	pi := machinPi(prec)
	twoPi := newBig(prec).Mul(pi, big.NewFloat(2))
	piNeg := newBig(prec).Neg(pi)
	for theta.Cmp(pi) > 0 {
		theta.Sub(theta, twoPi)
	}
	for theta.Cmp(piNeg) < 0 {
		theta.Add(theta, twoPi)
	}
	return [2]*big.Float{sincosSeries(theta, false, prec), sincosSeries(theta, true, prec)}
}

// sincosSeries sums the Taylor series of sine (odd powers) or cosine
// (even powers) of a reduced angle.
func sincosSeries(theta *big.Float, sine bool, prec uint) *big.Float {
	sum := newBig(prec)
	term := newBig(prec)
	if sine {
		term.Set(theta)
	} else {
		term.SetInt64(1)
	}
	theta2 := newBig(prec).Mul(theta, theta)
	tiny := newBig(prec).SetMantExp(big.NewFloat(1), -int(prec)-12)
	for k := range 120 {
		if k%2 == 1 {
			sum.Sub(sum, term)
		} else {
			sum.Add(sum, term)
		}
		d1, d2 := 2*k+1, 2*k+2
		if sine {
			d1, d2 = 2*k+2, 2*k+3
		}
		term.Mul(term, theta2)
		term.Quo(term, big.NewFloat(float64(d1)*float64(d2)))
		if abs := newBig(prec).Abs(term); abs.Cmp(tiny) < 0 {
			break
		}
	}
	return sum
}

// bigFloatDFT evaluates the DFT by its quadratic definition at 256-bit
// precision: the twiddle of every residue is an exact-angle series
// evaluation, the accumulation runs in big.Float complex arithmetic,
// and the result is rounded once into complex128.
func bigFloatDFT(vals []complex128, sign float64) []complex128 {
	const prec = 256
	n := len(vals)
	pi := machinPi(prec)
	twoPi := newBig(prec).Mul(pi, big.NewFloat(2))
	tw := make([][2]*big.Float, n)
	for r := range n {
		angle := newBig(prec).SetInt64(int64(r))
		angle.Mul(angle, twoPi)
		angle.Quo(angle, big.NewFloat(float64(n)))
		if sign < 0 {
			angle.Neg(angle)
		}
		tw[r] = sincosBig(angle, prec)
	}
	out := make([][2]*big.Float, n)
	for k := range out {
		out[k] = [2]*big.Float{newBig(prec), newBig(prec)}
	}
	vr, vi := newBig(prec), newBig(prec)
	pr, pii := newBig(prec), newBig(prec)
	tr, ti := newBig(prec), newBig(prec)
	for k := range n {
		re, im := out[k][0], out[k][1]
		for j := range n {
			t := tw[j*k%n]
			vr.SetFloat64(real(vals[j]))
			vi.SetFloat64(imag(vals[j]))
			tr.Set(t[0])
			ti.Set(t[1])
			pr.Mul(vr, tr)
			pii.Mul(vi, ti)
			pr.Sub(pr, pii)
			pii.Mul(vr, ti)
			ti.Mul(vi, tr)
			pii.Add(pii, ti)
			re.Add(re, pr)
			im.Add(im, pii)
		}
	}
	res := make([]complex128, n)
	for k := range n {
		f64re, _ := out[k][0].Float64()
		f64im, _ := out[k][1].Float64()
		res[k] = complex(f64re, f64im)
	}
	return res
}

// TestFFTPrecisionReport compares the radix-2 walk and the fused
// radix-4 dispatch against the 256-bit quadratic reference and reports
// both maximum errors. The assertion only demands that the fused path
// hold the radix-2 path's accuracy within a small factor; the logged
// figures are the evidence the report quotes.
func TestFFTPrecisionReport(t *testing.T) {
	for _, n := range []int{64, 256} {
		vals := precisionFixture(n)
		ref := bigFloatDFT(vals, -1)
		oldPath := append([]complex128(nil), vals...)
		fftRadix2(oldPath, -1)
		newPath := append([]complex128(nil), vals...)
		fftPow2(newPath, -1)
		errOld, errNew := 0.0, 0.0
		for i := range ref {
			eo := math.Max(math.Abs(real(oldPath[i])-real(ref[i])), math.Abs(imag(oldPath[i])-imag(ref[i])))
			en := math.Max(math.Abs(real(newPath[i])-real(ref[i])), math.Abs(imag(newPath[i])-imag(ref[i])))
			errOld = math.Max(errOld, eo)
			errNew = math.Max(errNew, en)
		}
		t.Logf("n=%d: radix-2 max error %.3e, fused radix-4 max error %.3e, ratio %.3f", n, errOld, errNew, errNew/errOld)
		if errNew > 4*errOld+1e-13 {
			t.Fatalf("n=%d: fused radix-4 error %.3e is worse than 4x the radix-2 error %.3e", n, errNew, errOld)
		}
	}
}

// TestBigFloatReferenceAgreesWithQuadratic guards the precision
// harness itself: at lengths this small the 256-bit reference and the
// float64 quadratic definition must agree to the float64 path's own
// rounding, so a broken series evaluation cannot pass unnoticed into
// the error figures above.
func TestBigFloatReferenceAgreesWithQuadratic(t *testing.T) {
	for _, n := range []int{4, 8} {
		vals := precisionFixture(n)
		ref := bigFloatDFT(vals, -1)
		naive := naiveDFT(vals, -1)
		for i := range ref {
			if math.Abs(real(ref[i])-real(naive[i])) > 1e-12 || math.Abs(imag(ref[i])-imag(naive[i])) > 1e-12 {
				t.Fatalf("n=%d [%d]: 256-bit reference %v, quadratic %v", n, i, ref[i], naive[i])
			}
		}
	}
}