openvisengine/src/lib/ava/beat.ts

// AVA — beat detection, tempo tracking & prediction
//
// Strategy (mirrors the planned Rust implementation):
//  1. DETECT   — a beat fires when the smoothed bass energy exceeds
//                `threshold ×` the rolling average (last ~60 frames),
//                gated by `minInterval` and an absolute floor.
//  2. LEARN    — consecutive beat intervals build a BPM estimate; when the
//                intervals are consistent (low spread), the period is locked
//                with an exponential moving average and confidence rises.
//  3. PREDICT  — a phase accumulator advances by dt/period every frame. When
//                the bass drops out, the predictor keeps firing beats on the
//                learned grid so the visuals never lose the beat.
//  4. TAP      — the "Set Beat" button: taps lock the period + phase from the
//                tap intervals (median), teaching the tracker the song's beat.

import type { BeatConfig, BeatFrame, BeatSource } from './types'
import { DEFAULT_BEAT_CONFIG } from './types'

const HISTORY_SIZE = 60
const MAX_INTERVALS = 10
const MAX_TAPS = 8
const TAP_RESET_MS = 2500

function median(values: number[]): number {
  const sorted = [...values].sort((a, b) => a - b)
  const mid = sorted.length >> 1
  return sorted.length % 2 ? sorted[mid] : (sorted[mid - 1] + sorted[mid]) / 2
}

function foldPeriod(ms: number): number {
  // keep the period inside 60..200 BPM equivalent
  let p = ms
  while (p > 1000) p /= 2
  while (p < 300) p *= 2
  return p
}

export interface TapResult {
  bpm: number
  phase: number
  confidence: number
  tapCount: number
  /** true when enough taps have accumulated to lock a tempo */
  locked: boolean
}

export class BeatTracker {
  cfg: BeatConfig
  private history: number[] = []
  private intervals: number[] = []
  private taps: number[] = []
  private lastTapAt = -1e9
  private periodMs = 500
  private phase = 0
  private bpm = 0
  private confidence = 0
  private lastBeatAt = -1e9
  private lastRealBeatAt = -1e9
  private prevRealBeatAt = -1e9
  private beatCount = 0

  constructor(cfg: BeatConfig = { ...DEFAULT_BEAT_CONFIG }) {
    this.cfg = cfg
  }

  reset(): void {
    this.history = []
    this.intervals = []
    this.taps = []
    this.periodMs = 500
    this.phase = 0
    this.bpm = 0
    this.confidence = 0
    this.lastBeatAt = -1e9
    this.lastRealBeatAt = -1e9
    this.prevRealBeatAt = -1e9
    this.beatCount = 0
  }

  getBeatCount(): number {
    return this.beatCount
  }

  /** Constant frame returned while no audio source is playing/loaded. */
  private static readonly STANDBY_FRAME: BeatFrame = {
    beat: false,
    source: 'none',
    bpm: 0,
    phase: 0,
    confidence: 0,
  }

  /**
   * Called by the engine every frame while NO audio source is playing or
   * loaded (v6.4). Beat detection, the prediction grid and any tapped tempo
   * all stop: the indicator goes dark, the BPM readout resets and everything
   * learned is dropped — so a stale tempo can never leak ghost beats while
   * idle (or into the next track once audio starts again).
   */
  standby(): BeatFrame {
    this.reset()
    return BeatTracker.STANDBY_FRAME
  }

  getBpm(): number {
    return this.bpm
  }

  /**
   * Advance the tracker one frame.
   * @param bassAtt smoothed bass level (the detector input)
   * @param nowMs   monotonic clock in milliseconds (performance.now())
   * @param dtSec   frame delta in seconds
   */
  update(bassAtt: number, nowMs: number, dtSec: number, cfg?: BeatConfig): BeatFrame {
    if (cfg) this.cfg = cfg

    // --- rolling energy history -------------------------------------------
    this.history.push(bassAtt)
    if (this.history.length > HISTORY_SIZE) this.history.shift()
    let sum = 0
    for (let k = 0; k < this.history.length; k++) sum += this.history[k]
    const avg = sum / this.history.length

    let beat = false
    let source: BeatSource = 'none'

    // --- real beat detection ----------------------------------------------
    if (
      bassAtt > this.cfg.floorLevel &&
      bassAtt > avg * this.cfg.threshold &&
      nowMs - this.lastBeatAt >= this.cfg.minInterval
    ) {
      beat = true
      source = 'audio'
      this.lastBeatAt = nowMs
      this.lastRealBeatAt = nowMs
      this.beatCount++

      // interval learning
      if (this.prevRealBeatAt > 0) {
        const interval = nowMs - this.prevRealBeatAt
        if (interval >= 200 && interval <= 2500) {
          this.intervals.push(interval)
          if (this.intervals.length > MAX_INTERVALS) this.intervals.shift()
          if (this.intervals.length >= 3) {
            const med = median(this.intervals)
            const mean = this.intervals.reduce((a, b) => a + b, 0) / this.intervals.length
            const variance = this.intervals.reduce((a, b) => a + (b - mean) * (b - mean), 0) / this.intervals.length
            const spread = Math.sqrt(variance) / Math.max(mean, 1)
            if (spread < 0.3) {
              // consistent tempo: lock the period (EMA) and gain confidence
              this.periodMs = this.periodMs * 0.6 + foldPeriod(med) * 0.4
              this.bpm = Math.round((60000 / this.periodMs) * 10) / 10
              this.confidence = Math.min(1, this.confidence + 0.2)
            } else {
              // irregular — trust it less
              this.confidence = Math.max(0, this.confidence - 0.05)
            }
          }
        }
      }
      this.prevRealBeatAt = nowMs
      this.phase = 0
    }

    // --- predictor (phase-locked grid) -------------------------------------
    if (this.periodMs > 50) {
      this.phase += (dtSec * 1000) / this.periodMs
    }
    if (this.phase >= 1) {
      this.phase -= Math.floor(this.phase)
      const bassGone = nowMs - this.lastRealBeatAt > this.periodMs * 1.4
      if (
        !beat &&
        this.cfg.prediction &&
        this.confidence > 0.35 &&
        bassGone &&
        nowMs - this.lastBeatAt >= this.cfg.minInterval
      ) {
        beat = true
        source = 'predicted'
        this.lastBeatAt = nowMs
        this.beatCount++
        // predicted beats slowly spend confidence so a stale lock fades out
        this.confidence = Math.max(0, this.confidence - 0.02)
      }
    }

    // --- confidence decay when audio beats disappear -----------------------
    if (nowMs - this.lastRealBeatAt > 2000 && this.taps.length === 0) {
      this.confidence = Math.max(0, this.confidence - dtSec * 0.06)
    }

    return { beat, source, bpm: this.bpm, phase: this.phase, confidence: this.confidence }
  }

  /**
   * "Set Beat" — user taps to the music. Two or more taps lock the tempo
   * (median of tap intervals) and re-phase the predictor to the taps.
   */
  tap(nowMs: number): TapResult {
    if (nowMs - this.lastTapAt > TAP_RESET_MS) this.taps = []
    this.lastTapAt = nowMs
    this.taps.push(nowMs)
    if (this.taps.length > MAX_TAPS) this.taps.shift()

    // immediate visual feedback + re-phase
    this.beatCount++
    this.lastBeatAt = nowMs
    this.phase = 0

    if (this.taps.length >= 2) {
      const gaps: number[] = []
      for (let k = 1; k < this.taps.length; k++) gaps.push(this.taps[k] - this.taps[k - 1])
      const med = median(gaps)
      if (med >= 250 && med <= 1500) {
        this.periodMs = foldPeriod(med)
        this.bpm = Math.round((60000 / this.periodMs) * 10) / 10
        this.confidence = Math.min(1, 0.45 + 0.11 * this.taps.length)
        // keep the interval learner in sync with the tapped grid
        this.intervals = [...gaps.slice(-MAX_INTERVALS)]
        this.lastRealBeatAt = nowMs
      }
    } else {
      this.confidence = Math.max(this.confidence, 0.25)
    }

    return {
      bpm: this.bpm,
      phase: this.phase,
      confidence: this.confidence,
      tapCount: this.taps.length,
      locked: this.taps.length >= 2 && this.bpm > 0,
    }
  }

  /** Current phase (0..1) — for UI indicators. */
  getPhase(): number {
    return this.phase
  }
}