sdr-wideband-suite/internal/detector/detector.go

package detector

import (
	"math"
	"sort"
	"time"

	"sdr-visual-suite/internal/classifier"
)

type Event struct {
	ID        int64                      `json:"id"`
	Start     time.Time                  `json:"start"`
	End       time.Time                  `json:"end"`
	CenterHz  float64                    `json:"center_hz"`
	Bandwidth float64                    `json:"bandwidth_hz"`
	PeakDb    float64                    `json:"peak_db"`
	SNRDb     float64                    `json:"snr_db"`
	FirstBin  int                        `json:"first_bin"`
	LastBin   int                        `json:"last_bin"`
	Class     *classifier.Classification `json:"class,omitempty"`
}

type Detector struct {
	ThresholdDb     float64
	MinDuration     time.Duration
	Hold            time.Duration
	EmaAlpha        float64
	HysteresisDb    float64
	MinStableFrames int
	GapTolerance    time.Duration
	CFAREnabled     bool
	CFARGuardCells  int
	CFARTrainCells  int
	CFARRank        int
	CFARScaleDb     float64

	binWidth   float64
	nbins      int
	sampleRate int

	ema    []float64
	active map[int64]*activeEvent
	nextID int64
}

type activeEvent struct {
	id         int64
	start      time.Time
	lastSeen   time.Time
	centerHz   float64
	bwHz       float64
	peakDb     float64
	snrDb      float64
	firstBin   int
	lastBin    int
	class      *classifier.Classification
	stableHits int
}

type Signal struct {
	FirstBin int                        `json:"first_bin"`
	LastBin  int                        `json:"last_bin"`
	CenterHz float64                    `json:"center_hz"`
	BWHz     float64                    `json:"bw_hz"`
	PeakDb   float64                    `json:"peak_db"`
	SNRDb    float64                    `json:"snr_db"`
	Class    *classifier.Classification `json:"class,omitempty"`
}

func New(thresholdDb float64, sampleRate int, fftSize int, minDur, hold time.Duration, emaAlpha, hysteresis float64, minStable int, gapTolerance time.Duration, cfarEnabled bool, cfarGuard, cfarTrain, cfarRank int, cfarScaleDb float64) *Detector {
	if minDur <= 0 {
		minDur = 250 * time.Millisecond
	}
	if hold <= 0 {
		hold = 500 * time.Millisecond
	}
	if emaAlpha <= 0 || emaAlpha > 1 {
		emaAlpha = 0.2
	}
	if hysteresis <= 0 {
		hysteresis = 3
	}
	if minStable <= 0 {
		minStable = 3
	}
	if gapTolerance <= 0 {
		gapTolerance = hold
	}
	if cfarGuard < 0 {
		cfarGuard = 2
	}
	if cfarTrain <= 0 {
		cfarTrain = 16
	}
	if cfarScaleDb <= 0 {
		cfarScaleDb = 6
	}
	if cfarRank <= 0 || cfarRank > 2*cfarTrain {
		cfarRank = int(math.Round(0.75 * float64(2*cfarTrain)))
		if cfarRank <= 0 {
			cfarRank = 1
		}
	}
	return &Detector{
		ThresholdDb:     thresholdDb,
		MinDuration:     minDur,
		Hold:            hold,
		EmaAlpha:        emaAlpha,
		HysteresisDb:    hysteresis,
		MinStableFrames: minStable,
		GapTolerance:    gapTolerance,
		CFAREnabled:     cfarEnabled,
		CFARGuardCells:  cfarGuard,
		CFARTrainCells:  cfarTrain,
		CFARRank:        cfarRank,
		CFARScaleDb:     cfarScaleDb,
		binWidth:        float64(sampleRate) / float64(fftSize),
		nbins:           fftSize,
		sampleRate:      sampleRate,
		ema:             make([]float64, fftSize),
		active:          map[int64]*activeEvent{},
		nextID:          1,
	}
}

func (d *Detector) Process(now time.Time, spectrum []float64, centerHz float64) ([]Event, []Signal) {
	signals := d.detectSignals(spectrum, centerHz)
	finished := d.matchSignals(now, signals)
	return finished, signals
}

// UpdateClasses refreshes active event classes from current signals.
func (d *Detector) UpdateClasses(signals []Signal) {
	for _, s := range signals {
		for _, ev := range d.active {
			if overlapHz(s.CenterHz, s.BWHz, ev.centerHz, ev.bwHz) && math.Abs(s.CenterHz-ev.centerHz) < (s.BWHz+ev.bwHz)/2.0 {
				if s.Class != nil {
					if ev.class == nil || s.Class.Confidence >= ev.class.Confidence {
						ev.class = s.Class
					}
				}
			}
		}
	}
}

func (d *Detector) detectSignals(spectrum []float64, centerHz float64) []Signal {
	n := len(spectrum)
	if n == 0 {
		return nil
	}
	smooth := d.smoothSpectrum(spectrum)
	thresholds := d.cfarThresholds(smooth)
	noiseGlobal := median(smooth)
	var signals []Signal
	in := false
	start := 0
	peak := -1e9
	peakBin := 0
	for i := 0; i < n; i++ {
		v := smooth[i]
		thresholdOn := d.ThresholdDb
		if thresholds != nil && !math.IsNaN(thresholds[i]) {
			thresholdOn = thresholds[i]
		}
		thresholdOff := thresholdOn - d.HysteresisDb
		if v >= thresholdOn {
			if !in {
				in = true
				start = i
				peak = v
				peakBin = i
			} else if v > peak {
				peak = v
				peakBin = i
			}
		} else if in && v < thresholdOff {
			noise := noiseGlobal
			if thresholds != nil && peakBin >= 0 && peakBin < len(thresholds) && !math.IsNaN(thresholds[peakBin]) {
				noise = thresholds[peakBin] - d.CFARScaleDb
			}
			signals = append(signals, d.makeSignal(start, i-1, peak, peakBin, noise, centerHz, smooth))
			in = false
		}
	}
	if in {
		noise := noiseGlobal
		if thresholds != nil && peakBin >= 0 && peakBin < len(thresholds) && !math.IsNaN(thresholds[peakBin]) {
			noise = thresholds[peakBin] - d.CFARScaleDb
		}
		signals = append(signals, d.makeSignal(start, n-1, peak, peakBin, noise, centerHz, smooth))
	}
	return signals
}

func (d *Detector) makeSignal(first, last int, peak float64, peakBin int, noise float64, centerHz float64, spectrum []float64) Signal {
	centerBin := float64(first+last) / 2.0
	centerFreq := centerHz + (centerBin-float64(d.nbins)/2.0)*d.binWidth
	bw := float64(last-first+1) * d.binWidth
	snr := peak - noise
	return Signal{
		FirstBin: first,
		LastBin:  last,
		CenterHz: centerFreq,
		BWHz:     bw,
		PeakDb:   peak,
		SNRDb:    snr,
	}
}

func (d *Detector) cfarThresholds(spectrum []float64) []float64 {
	if !d.CFAREnabled || d.CFARTrainCells <= 0 {
		return nil
	}
	n := len(spectrum)
	train := d.CFARTrainCells
	guard := d.CFARGuardCells
	totalTrain := 2 * train
	if totalTrain <= 0 {
		return nil
	}
	rank := d.CFARRank
	if rank <= 0 || rank > totalTrain {
		rank = int(math.Round(0.75 * float64(totalTrain)))
		if rank <= 0 {
			rank = 1
		}
	}
	rankIdx := rank - 1
	thresholds := make([]float64, n)
	buf := make([]float64, totalTrain)
	for i := 0; i < n; i++ {
		leftStart := i - guard - train
		rightEnd := i + guard + train
		if leftStart < 0 || rightEnd >= n {
			thresholds[i] = math.NaN()
			continue
		}
		copy(buf[:train], spectrum[leftStart:leftStart+train])
		copy(buf[train:], spectrum[i+guard+1:i+guard+1+train])
		sort.Float64s(buf)
		noise := buf[rankIdx]
		thresholds[i] = noise + d.CFARScaleDb
	}
	return thresholds
}

func (d *Detector) smoothSpectrum(spectrum []float64) []float64 {
	if d.ema == nil || len(d.ema) != len(spectrum) {
		d.ema = make([]float64, len(spectrum))
		copy(d.ema, spectrum)
		return d.ema
	}
	alpha := d.EmaAlpha
	for i := range spectrum {
		v := spectrum[i]
		d.ema[i] = alpha*v + (1-alpha)*d.ema[i]
	}
	// IMPORTANT: caller must not modify returned slice
	return d.ema
}

func (d *Detector) matchSignals(now time.Time, signals []Signal) []Event {
	used := make(map[int64]bool, len(d.active))
	for _, s := range signals {
		var best *activeEvent
		var candidates []struct {
			ev   *activeEvent
			dist float64
		}
		for _, ev := range d.active {
			if overlapHz(s.CenterHz, s.BWHz, ev.centerHz, ev.bwHz) && math.Abs(s.CenterHz-ev.centerHz) < (s.BWHz+ev.bwHz)/2.0 {
				candidates = append(candidates, struct {
					ev   *activeEvent
					dist float64
				}{ev: ev, dist: math.Abs(s.CenterHz - ev.centerHz)})
			}
		}
		if len(candidates) > 0 {
			sort.Slice(candidates, func(i, j int) bool { return candidates[i].dist < candidates[j].dist })
			best = candidates[0].ev
		}
		if best == nil {
			id := d.nextID
			d.nextID++
			d.active[id] = &activeEvent{
				id:         id,
				start:      now,
				lastSeen:   now,
				centerHz:   s.CenterHz,
				bwHz:       s.BWHz,
				peakDb:     s.PeakDb,
				snrDb:      s.SNRDb,
				firstBin:   s.FirstBin,
				lastBin:    s.LastBin,
				class:      s.Class,
				stableHits: 1,
			}
			continue
		}
		used[best.id] = true
		best.lastSeen = now
		best.stableHits++
		best.centerHz = (best.centerHz + s.CenterHz) / 2.0
		if s.BWHz > best.bwHz {
			best.bwHz = s.BWHz
		}
		if s.PeakDb > best.peakDb {
			best.peakDb = s.PeakDb
		}
		if s.SNRDb > best.snrDb {
			best.snrDb = s.SNRDb
		}
		if s.FirstBin < best.firstBin {
			best.firstBin = s.FirstBin
		}
		if s.LastBin > best.lastBin {
			best.lastBin = s.LastBin
		}
		if s.Class != nil {
			if best.class == nil || s.Class.Confidence >= best.class.Confidence {
				best.class = s.Class
			}
		}
	}

	var finished []Event
	for id, ev := range d.active {
		if used[id] {
			continue
		}
		if now.Sub(ev.lastSeen) < d.GapTolerance {
			continue
		}
		duration := ev.lastSeen.Sub(ev.start)
		if duration < d.MinDuration || ev.stableHits < d.MinStableFrames {
			delete(d.active, id)
			continue
		}
		finished = append(finished, Event{
			ID:        ev.id,
			Start:     ev.start,
			End:       ev.lastSeen,
			CenterHz:  ev.centerHz,
			Bandwidth: ev.bwHz,
			PeakDb:    ev.peakDb,
			SNRDb:     ev.snrDb,
			FirstBin:  ev.firstBin,
			LastBin:   ev.lastBin,
			Class:     ev.class,
		})
		delete(d.active, id)
	}
	return finished
}

func overlapHz(c1, b1, c2, b2 float64) bool {
	l1 := c1 - b1/2.0
	r1 := c1 + b1/2.0
	l2 := c2 - b2/2.0
	r2 := c2 + b2/2.0
	return l1 <= r2 && l2 <= r1
}

func median(vals []float64) float64 {
	if len(vals) == 0 {
		return 0
	}
	cpy := append([]float64(nil), vals...)
	sort.Float64s(cpy)
	mid := len(cpy) / 2
	if len(cpy)%2 == 0 {
		return (cpy[mid-1] + cpy[mid]) / 2.0
	}
	return cpy[mid]
}