ebur128.go

// Package ebur128 implements EBU R-128 loudness measurement
// (integrated loudness, true peak, loudness range).
//
// Only 48 kHz is supported.
package ebur128

import "fmt"

type ChannelLayout int

const (
	LayoutStereo ChannelLayout = iota // 2ch: L, R
	Layout51                          // 6ch: L, R, C, LFE, Ls, Rs
)

// Result of a loudness measurement. Fields are set to loudnessUnmeasurable
// (-99) when the input is too short or fully gated out.
type Result struct {
	IntegratedLoudness float64 // LUFS (gated)
	TruePeak           float64 // dBTP, max across channels
	LoudnessRange      float64 // LU; 0 if unmeasurable

	// Number of 400ms blocks that passed both gates (0 = unmeasurable).
	GatedBlockCount int
	// Number of 3s windows that passed both gates for LRA (<2 = unmeasurable).
	LRAGatedWindowCount int
}

type Meter struct {
	numCh        int
	subBlockSize int // samples per 100ms sub-block (sampleRate / 10)
	chanWeights  []float64
	filters      []kweightFilter
	tpDet        []truePeakDetector

	chanSqSum   []float64
	sampleCount int
	subBlocks   []float64 // weighted energy per 100ms sub-block
}

func New(layout ChannelLayout, sampleRate int) (*Meter, error) {
	if sampleRate != sampleRate48k {
		return nil, fmt.Errorf("ebur128: sample rate %d not supported (must be 48k)", sampleRate)
	}

	var numCh int
	var weights []float64

	switch layout {
	case LayoutStereo:
		numCh = 2
		weights = []float64{weightFront, weightFront}
	case Layout51:
		numCh = 6
		weights = []float64{
			weightFront, weightFront, weightFront,
			weightLFE,
			weightSurround, weightSurround,
		}
	default:
		return nil, fmt.Errorf("ebur128: unknown layout %d", layout)
	}

	return &Meter{
		numCh:        numCh,
		subBlockSize: sampleRate * subBlockMs / 1000,
		chanWeights:  weights,
		filters:      make([]kweightFilter, numCh),
		tpDet:        make([]truePeakDetector, numCh),
		chanSqSum:    make([]float64, numCh),
	}, nil
}

func (m *Meter) Reserve(seconds float64) {
	n := int(seconds*1000) / subBlockMs
	if n > cap(m.subBlocks) {
		sb := make([]float64, len(m.subBlocks), n)
		copy(sb, m.subBlocks)
		m.subBlocks = sb
	}
}

func (m *Meter) Write(buf []float64) {
	numCh := m.numCh
	for i := 0; i+numCh-1 < len(buf); i += numCh {
		for ch := range numCh {
			s := buf[i+ch]
			m.tpDet[ch].process(s)
			filtered := m.filters[ch].process(s)
			m.chanSqSum[ch] += filtered * filtered
		}

		m.sampleCount++
		if m.sampleCount == m.subBlockSize {
			m.finishSubBlock()
		}
	}
}

func (m *Meter) WriteFloat32(input []float32) {
	const chunkSize = 4096
	var buf [chunkSize]float64
	for len(input) > 0 {
		n := min(len(input), chunkSize)
		for i := range n {
			buf[i] = float64(input[i])
		}
		m.Write(buf[:n])
		input = input[n:]
	}
}

func (m *Meter) Reset() {
	for i := range m.filters {
		m.filters[i] = kweightFilter{}
	}
	for i := range m.tpDet {
		m.tpDet[i] = truePeakDetector{}
	}
	for i := range m.chanSqSum {
		m.chanSqSum[i] = 0
	}
	m.sampleCount = 0
	m.subBlocks = m.subBlocks[:0]
}

// Finalize flushes any incomplete sub-block (zero-padded). Idempotent.
func (m *Meter) Finalize() {
	if m.sampleCount == 0 {
		return
	}
	m.finishSubBlock()
}

func (m *Meter) finishSubBlock() {
	var energy float64
	for ch := 0; ch < m.numCh; ch++ {
		meanSq := m.chanSqSum[ch] / float64(m.subBlockSize)
		energy += m.chanWeights[ch] * meanSq
	}
	m.subBlocks = append(m.subBlocks, energy)

	for ch := range m.chanSqSum {
		m.chanSqSum[ch] = 0
	}
	m.sampleCount = 0
}

// Loudness returns the measurement. Call Finalize first to include trailing samples.
func (m *Meter) Loudness() Result {
	var maxPeak float64
	for ch := 0; ch < m.numCh; ch++ {
		if m.tpDet[ch].peak > maxPeak {
			maxPeak = m.tpDet[ch].peak
		}
	}

	il, gbc := integratedLoudness(m.subBlocks)
	lra, lraWin := loudnessRange(m.subBlocks)

	return Result{
		IntegratedLoudness:  il,
		TruePeak:            truePeakDBTP(maxPeak),
		LoudnessRange:       lra,
		GatedBlockCount:     gbc,
		LRAGatedWindowCount: lraWin,
	}
}