fm-rds-tx/internal/offline/watermark_e2e_test.go

package offline

import (
	"math"
	"sort"
	"testing"
	"time"

	"github.com/jan/fm-rds-tx/internal/audio"
	cfgpkg "github.com/jan/fm-rds-tx/internal/config"
	"github.com/jan/fm-rds-tx/internal/dsp"
	"github.com/jan/fm-rds-tx/internal/license"
	"github.com/jan/fm-rds-tx/internal/output"
	"github.com/jan/fm-rds-tx/internal/watermark"
)

// noiseSource provides deterministic broadband stereo material so the
// watermark has energy across many bins and the detector has something real to
// correlate against.
type noiseSource struct {
	stateL uint64
	stateR uint64
	amp    float64
}

func newNoiseSource(seed uint64, amp float64) *noiseSource {
	return &noiseSource{stateL: seed, stateR: seed ^ 0x9e3779b97f4a7c15, amp: amp}
}

func (n *noiseSource) nextNorm(state *uint64) float64 {
	*state = *state*6364136223846793005 + 1442695040888963407
	return float64(int32(*state>>33)) / float64(1<<31)
}

func (n *noiseSource) NextFrame() audio.Frame {
	l := n.amp * n.nextNorm(&n.stateL)
	r := n.amp * n.nextNorm(&n.stateR)
	return audio.NewFrame(audio.Sample(l), audio.Sample(r))
}

func newWatermarkE2EGenerator(t *testing.T, key string) *Generator {
	t.Helper()
	cfg := cfgpkg.Default()
	cfg.RDS.Enabled = false
	cfg.FM.CompositeRateHz = 228000
	cfg.FM.StereoEnabled = true
	cfg.FM.OutputDrive = 0.5
	cfg.FM.LimiterEnabled = true
	cfg.FM.LimiterCeiling = 1.0
	cfg.FM.FMModulationEnabled = false // composite output for inspection
	cfg.Audio.ToneAmplitude = 0        // external source only
	cfg.Audio.Gain = 1.0
	cfg.FM.PreEmphasisTauUS = 50

	gen := NewGenerator(cfg)
	if err := gen.SetExternalSource(newNoiseSource(12345, 0.22)); err != nil {
		t.Fatalf("SetExternalSource: %v", err)
	}
	gen.SetLicense(license.NewState(""))
	gen.ConfigureWatermark(true, key)
	return gen
}

func extractCompositeFrame(frame *output.CompositeFrame) []float64 {
	out := make([]float64, len(frame.Samples))
	for i, s := range frame.Samples {
		out[i] = float64(s.I)
	}
	return out
}

func downsampleForWatermarkDetection(composite []float64, rate float64) ([]float64, float64) {
	decimFactor := int(rate / float64(watermark.WMRate))
	if decimFactor < 1 {
		decimFactor = 1
	}
	actualRate := rate / float64(decimFactor)

	lpf := dsp.NewLPF8(5500, rate)
	filtered := make([]float64, len(composite))
	for i, s := range composite {
		filtered[i] = lpf.Process(s)
	}
	out := make([]float64, len(filtered)/decimFactor)
	for i := range out {
		out[i] = filtered[i*decimFactor]
	}
	return out, actualRate
}

func decodeWatermarkFromComposite(t *testing.T, composite []float64, rate float64) ([watermark.RsDataBytes]byte, [watermark.PayloadBits]float64, bool) {
	t.Helper()
	var zero [watermark.RsDataBytes]byte

	down, actualRate := downsampleForWatermarkDetection(composite, rate)
	if math.Abs(actualRate-float64(watermark.WMRate)) > 1e-9 {
		t.Fatalf("unexpected detector rate %.3f Hz after decimation", actualRate)
	}
	if len(down) <= watermark.FFTSize*2 {
		return zero, [watermark.PayloadBits]float64{}, false
	}

	probe := down[watermark.FFTSize:]
	det := watermark.NewSTFTDetector()
	corrs, bestOffset := det.Detect(probe)

	var sumAbs float64
	var strong int
	for _, c := range corrs {
		ac := math.Abs(c)
		sumAbs += ac
		if ac > 1.0 {
			strong++
		}
	}
	t.Logf("watermark detect: bestOffset=%d avg|c|=%.2f strong=%d/%d", bestOffset, sumAbs/float64(watermark.PayloadBits), strong, watermark.PayloadBits)

	var recv [watermark.RsTotalBytes]byte
	byteConfs := make([]float64, watermark.RsTotalBytes)
	for i := 0; i < watermark.PayloadBits; i++ {
		if corrs[i] < 0 {
			recv[i/8] |= 1 << uint(7-(i%8))
		}
		byteConfs[i/8] += math.Abs(corrs[i])
	}

	type bc struct{ idx int; conf float64 }
	ranked := make([]bc, watermark.RsTotalBytes)
	for i := range ranked {
		ranked[i] = bc{i, byteConfs[i]}
	}
	sort.Slice(ranked, func(a, b int) bool { return ranked[a].conf < ranked[b].conf })

	for ne := 0; ne <= watermark.RsCheckBytes; ne++ {
		erasePos := make([]int, ne)
		for i := 0; i < ne; i++ {
			erasePos[i] = ranked[i].idx
		}
		sort.Ints(erasePos)
		payload, ok := watermark.RSDecode(recv, erasePos)
		if ok {
			return payload, corrs, true
		}
	}
	return zero, corrs, false
}

func TestWatermarkE2E(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping long watermark E2E test in -short mode")
	}
	const key = "test-key-e2e"
	const duration = 45 * time.Second

	gen := newWatermarkE2EGenerator(t, key)
	frame := gen.GenerateFrame(duration)
	if frame == nil {
		t.Fatal("GenerateFrame returned nil")
	}
	composite := extractCompositeFrame(frame)
	payload, _, ok := decodeWatermarkFromComposite(t, composite, frame.SampleRateHz)
	if !ok {
		t.Fatal("RS decode failed — watermark not recoverable from generator composite")
	}
	if !watermark.KeyMatchesPayload(key, payload) {
		t.Fatalf("decoded payload mismatch: key=%q payload=%x", key, payload)
	}
}