package main

// The three models behind the wake word (V-487 stage two).
//
// openWakeWord's pipeline, run in a row:
//
//	audio      -> melspectrogram.onnx -> 32-bin mel frames, one per 10ms
//	76 frames  -> embedding_model.onnx -> one 96-dim embedding per 80ms
//	16 embeds  -> maven_wakeword.onnx  -> one score
//
// The first two are frozen and pretrained. Only the last was trained here,
// which is why it is 100KB and the other two are megabytes. The shapes are
// not guesses: 2.0s of 16kHz audio measures 197 mel frames, and 76-frame
// windows at stride 8 give exactly the 16 embeddings the head was fitted on.
//
// This file knows ONNX and nothing about the 80ms cadence. wakeword.go knows
// the cadence and nothing about tensors.

import (
	"fmt"

	ort "github.com/yalue/onnxruntime_go"
)

const (
	// melHop — samples per mel frame. 10ms at 16kHz.
	melHop = 160
	// melBins — mel bins per frame, fixed by melspectrogram.onnx.
	melBins = 32
	// embedFrames — mel frames one embedding is computed over, 760ms.
	embedFrames = 76
	// embedStride — mel frames between embeddings, 80ms.
	embedStride = 8
	// embedDim — the embedding width.
	embedDim = 96
	// headWindow — embeddings the head scores at once, 1.28s of audio.
	headWindow = 16

	// melContext — samples of history prepended to each incremental mel
	// call, chosen so the eight frames this call yields continue exactly
	// where the previous call's eight stopped.
	//
	// melspectrogram.onnx returns N/160-3 frames for N samples, and frame i
	// covers [i*160, i*160+400). With 480 samples of history the buffer is
	// 1760 samples, which is 8 frames, and the oldest of them starts one hop
	// after the newest of the previous call. Less history leaves a gap: the
	// first frames of a bare chunk would be computed against silence.
	melContext = 480

	// chunkSamples — audio per embedding step, 80ms.
	chunkSamples = embedStride * melHop
)

// wakeModels holds the three ONNX sessions. It runs on CPU threads beside
// silero and never touches the GPU. That is a rule, not a result: a wake word
// that waits on card admission is not a wake word.
type wakeModels struct {
	mel  *ort.DynamicAdvancedSession
	emb  *ort.DynamicAdvancedSession
	head *ort.DynamicAdvancedSession
}

// newWakeModels loads all three. melPath and embedPath are openWakeWord's
// frozen feature models; headPath is the keyword head trained for "Мэйвен".
func newWakeModels(melPath, embedPath, headPath, libPath string) (*wakeModels, error) {
	if !ort.IsInitialized() {
		if libPath != "" {
			ort.SetSharedLibraryPath(libPath)
		}
		if err := ort.InitializeEnvironment(); err != nil {
			return nil, fmt.Errorf("wake word: onnx runtime: %w", err)
		}
	}
	// One thread per session, not the default of every core. Measured on
	// workpc: the default took mavwaked from 68% of one core to 335% of
	// three, for three graphs that each run in well under 80ms single
	// threaded. An always-on gate that eats a quarter of the workstation is
	// not a gate he will leave running.
	opts, err := ort.NewSessionOptions()
	if err != nil {
		return nil, fmt.Errorf("wake word: session options: %w", err)
	}
	defer opts.Destroy()
	if err := opts.SetIntraOpNumThreads(1); err != nil {
		return nil, fmt.Errorf("wake word: intra-op threads: %w", err)
	}
	if err := opts.SetInterOpNumThreads(1); err != nil {
		return nil, fmt.Errorf("wake word: inter-op threads: %w", err)
	}
	open := func(p string, in, out []string) (*ort.DynamicAdvancedSession, error) {
		s, err := ort.NewDynamicAdvancedSession(p, in, out, opts)
		if err != nil {
			return nil, fmt.Errorf("wake word: load %s: %w", p, err)
		}
		return s, nil
	}
	m := &wakeModels{}
	if m.mel, err = open(melPath, []string{"input"}, []string{"output"}); err != nil {
		return nil, err
	}
	if m.emb, err = open(embedPath, []string{"input_1"}, []string{"conv2d_19"}); err != nil {
		m.Close()
		return nil, err
	}
	if m.head, err = open(headPath, []string{"embeddings"}, []string{"score"}); err != nil {
		m.Close()
		return nil, err
	}
	return m, nil
}

// Close releases the three sessions.
func (m *wakeModels) Close() {
	if m == nil {
		return
	}
	for _, s := range []*ort.DynamicAdvancedSession{m.mel, m.emb, m.head} {
		if s != nil {
			s.Destroy()
		}
	}
	m.mel, m.emb, m.head = nil, nil, nil
}

// melFrames runs one buffer of samples and returns the mel frames it yielded.
func (m *wakeModels) melFrames(buf []float32) ([][melBins]float32, error) {
	in, err := ort.NewTensor(ort.NewShape(1, int64(len(buf))), buf)
	if err != nil {
		return nil, err
	}
	defer in.Destroy()

	n := int64(len(buf)/melHop - 3)
	if n < 1 {
		return nil, fmt.Errorf("wake word: %d samples yield no mel frames", len(buf))
	}
	out, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 1, n, melBins))
	if err != nil {
		return nil, err
	}
	defer out.Destroy()

	if err := m.mel.Run([]ort.Value{in}, []ort.Value{out}); err != nil {
		return nil, err
	}
	data := out.GetData()
	frames := make([][melBins]float32, n)
	for i := range frames {
		for j := 0; j < melBins; j++ {
			// The scaling openWakeWord applies between the two feature
			// models, and the head was fitted on its output.
			frames[i][j] = data[i*melBins+j]/10.0 + 2.0
		}
	}
	return frames, nil
}

// embedding runs embedFrames mel frames through the frozen embedder.
func (m *wakeModels) embedding(mels [][melBins]float32) ([embedDim]float32, error) {
	var e [embedDim]float32
	flat := make([]float32, 0, embedFrames*melBins)
	for _, f := range mels {
		flat = append(flat, f[:]...)
	}
	in, err := ort.NewTensor(ort.NewShape(1, embedFrames, melBins, 1), flat)
	if err != nil {
		return e, err
	}
	defer in.Destroy()
	out, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 1, 1, embedDim))
	if err != nil {
		return e, err
	}
	defer out.Destroy()
	if err := m.emb.Run([]ort.Value{in}, []ort.Value{out}); err != nil {
		return e, err
	}
	copy(e[:], out.GetData())
	return e, nil
}

// score runs the trained head over headWindow embeddings.
func (m *wakeModels) score(embeds [][embedDim]float32) (float64, error) {
	flat := make([]float32, 0, headWindow*embedDim)
	for _, e := range embeds {
		flat = append(flat, e[:]...)
	}
	in, err := ort.NewTensor(ort.NewShape(1, headWindow, embedDim), flat)
	if err != nil {
		return 0, err
	}
	defer in.Destroy()
	out, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 1))
	if err != nil {
		return 0, err
	}
	defer out.Destroy()
	if err := m.head.Run([]ort.Value{in}, []ort.Value{out}); err != nil {
		return 0, err
	}
	return float64(out.GetData()[0]), nil
}