// Package signature identifies a stretch of audio by the landmarks in it, so the same
// advertisement can be recognised in an episode nobody has compared it against.
//
// Why not just correlate. Correlation is what detection already does and it is accurate, but it
// costs an FFT over the whole episode per stored signature. Measured on the live corpus: 500
// signatures against 23 episodes is about an hour that way, and 98 seconds this way. The library
// is only possible because looking something up does not mean scanning for it.
//
// How it works. Reduce audio to a sparse set of spectral peaks, pair nearby peaks, and hash each
// pair together with the time between them. A real repeat produces many hashes that agree on one
// time offset; a coincidence produces scattered hashes that agree on nothing. Looking up is then
// a hash probe, and the answer is the tallest column in a histogram of offsets.
package signature

import (
	"math"

	"github.com/terst/listen/server/internal/detect"
)

// SampleRate is what audio must be decoded at before fingerprinting.
//
// 2kHz, so a Nyquist of 1kHz. Speech carries its distinguishing formants well below that, and
// the rate sets everything downstream -- an hour at 8kHz is 29 million samples to transform and
// at 2kHz it is 7 million.
const SampleRate = 2000

const (
	frameSize = 512 // 256ms at SampleRate
	hopSize   = 256 // 128ms, so just under eight frames a second
	fanOut    = 3   // how many later peaks each anchor is paired with
	minGap    = 1   // frames
	maxGap    = 20  // frames, so pairs span up to ~2.5s
)

// bandEdges splits the usable bins into bands that are narrow low down and wide up top, which is
// where speech actually differs. One peak is kept per band per frame, so a loud low rumble
// cannot crowd out everything above it.
var bandEdges = [...]int{0, 8, 16, 32, 56, 96, 256}

// peakFloorFactor is the multiple of a frame's own mean bin magnitude that a peak must clear.
//
// Relative to the frame, not absolute: a quiet passage must not hash like a loud one, and
// silence must produce nothing at all. Two recordings agree perfectly on silence, so hashing it
// would make every pause collide with every other.
const peakFloorFactor = 2.0

// Fingerprint is the landmarks of one stretch of audio: a hash per peak pair, and the frame the
// pair started at. The two slices are parallel.
type Fingerprint struct {
	Hashes []uint32
	Frames []int32
}

// Take fingerprints audio that has been decoded at [SampleRate].
func Take(samples []float64) Fingerprint {
	return pairs(peaks(samples))
}

type peak struct {
	frame int32
	bin   int32
}

func peaks(samples []float64) []peak {
	if len(samples) < frameSize {
		return nil
	}
	window := hann(frameSize)
	re := make([]float64, frameSize)
	im := make([]float64, frameSize)
	mags := make([]float64, frameSize/2)

	var out []peak
	for start, frame := 0, int32(0); start+frameSize <= len(samples); start, frame = start+hopSize, frame+1 {
		for i := 0; i < frameSize; i++ {
			re[i] = samples[start+i] * window[i]
			im[i] = 0
		}
		detect.Transform(re, im, false)
		total := 0.0
		for i := range mags {
			mags[i] = math.Hypot(re[i], im[i])
			total += mags[i]
		}
		if total <= 0 {
			continue
		}
		floor := peakFloorFactor * total / float64(len(mags))
		for b := 0; b+1 < len(bandEdges); b++ {
			best, bestMag := int32(-1), floor
			for i := bandEdges[b]; i < bandEdges[b+1] && i < len(mags); i++ {
				if mags[i] > bestMag {
					best, bestMag = int32(i), mags[i]
				}
			}
			if best >= 0 {
				out = append(out, peak{frame: frame, bin: best})
			}
		}
	}
	return out
}

func hann(n int) []float64 {
	w := make([]float64, n)
	for i := range w {
		w[i] = 0.5 - 0.5*math.Cos(2*math.Pi*float64(i)/float64(n-1))
	}
	return w
}

// pairs hashes each peak against the few that follow it.
//
// A pair rather than a single peak because one peak is far too common to be worth looking up --
// a hash has to describe a relationship between two moments to be rare enough to mean something,
// and that is also what makes it survive a re-encode.
func pairs(ps []peak) Fingerprint {
	var fp Fingerprint
	for i, anchor := range ps {
		paired := 0
		for j := i + 1; j < len(ps) && paired < fanOut; j++ {
			gap := ps[j].frame - anchor.frame
			if gap < minGap {
				continue
			}
			if gap > maxGap {
				break
			}
			fp.Hashes = append(fp.Hashes, uint32(anchor.bin&0xFF)<<16|uint32(ps[j].bin&0xFF)<<8|uint32(gap&0xFF))
			fp.Frames = append(fp.Frames, anchor.frame)
			paired++
		}
	}
	return fp
}

// FrameMs is how much audio one frame advance covers.
const FrameMs = hopSize * 1000 / SampleRate

type posting struct {
	entry int32
	frame int32
}

// Index maps each hash to everywhere it has been seen.
type Index struct {
	postings map[uint32][]posting
	ids      []string
}

// NewIndex returns an empty index.
func NewIndex() *Index { return &Index{postings: map[uint32][]posting{}} }

// Add puts one fingerprint in, under the given id.
func (ix *Index) Add(id string, fp Fingerprint) {
	entry := int32(len(ix.ids))
	ix.ids = append(ix.ids, id)
	for i, h := range fp.Hashes {
		ix.postings[h] = append(ix.postings[h], posting{entry: entry, frame: fp.Frames[i]})
	}
}

// Len is how many fingerprints the index holds.
func (ix *Index) Len() int { return len(ix.ids) }

// Hit is one indexed entry's best agreement with a query.
type Hit struct {
	ID    string
	Votes int
	// OffsetMs is where the query's audio starts inside the indexed entry.
	OffsetMs int64
}

// Search reports, for each indexed entry that agrees at all, the offset it agrees on most.
//
// The histogram is the whole trick. Every hash that matches casts one vote for the time offset
// it implies. Audio that genuinely repeats puts nearly all of its votes in one bin; coincidental
// collisions scatter theirs. The caller compares the tallest column against the noise and
// decides -- this reports, it does not judge.
func (ix *Index) Search(fp Fingerprint) []Hit {
	type key struct {
		entry  int32
		offset int32
	}
	votes := map[key]int{}
	for i, h := range fp.Hashes {
		for _, p := range ix.postings[h] {
			votes[key{entry: p.entry, offset: p.frame - fp.Frames[i]}]++
		}
	}
	best := map[int32]Hit{}
	for k, n := range votes {
		if cur, ok := best[k.entry]; !ok || n > cur.Votes {
			best[k.entry] = Hit{ID: ix.ids[k.entry], Votes: n, OffsetMs: int64(k.offset) * FrameMs}
		}
	}
	out := make([]Hit, 0, len(best))
	for _, h := range best {
		out = append(out, h)
	}
	return out
}