ANC.py

"""
Narrowband ANC - Expert Build (v3)
===================================
This version replaces "re-estimate frequency from a fresh FFT snapshot every
block" (v1/v2's core weakness) with a bank of continuously self-correcting
Phase-Locked Loops (PLLs), one per tracked tone. This is the actual fix for
the "cancels briefly, then drifts out of phase" symptom you observed - that
symptom is beat-frequency drift caused by a small, uncorrected frequency
error, and a PLL is the standard engineering solution to exactly that
problem (this is how real tonal/narrowband ANC and adaptive line enhancers
work in industry, e.g., transformer hum or engine-order cancellation).

WHAT'S NEW vs v2:
1. PLL-based tracking (per tone): each tone has a local oscillator (NCO).
   The mic signal is demodulated against that NCO (in-phase/quadrature),
   low-pass filtered, and the resulting phase error continuously steers
   the NCO's frequency via a PI (proportional-integral) controller. This
   is a closed feedback loop - it self-corrects drift instead of hoping a
   single snapshot estimate stays accurate.
2. FFT is now only used for periodic ACQUISITION (finding new tones once a
   second), not for continuous tracking. Once a tone is acquired, the PLL
   takes over and tracks it sample-by-sample.
3. Auto latency calibration routine: plays a chirp and cross-correlates it
   against the recorded input to MEASURE actual round-trip audio latency,
   instead of guessing or manually sliding it. (Honest caveat below.)
4. Per-tone lock indicator in the GUI so you can see which tones are
   actually phase-locked vs still acquiring.

HONEST LIMITS (read this - it matters for your report):
- This is still a SINGLE-MIC, NARROWBAND system. It can only ever cancel
  STATIONARY TONAL noise (hums, whines, motor/fan drones) - not broadband
  noise (traffic, speech, transients). That's a mathematical limitation of
  tracking discrete frequency peaks, not a bug.
- Auto latency calibration requires the mic to actually hear the output
  (acoustic loopback). If you're using headphones for the real cancellation
  test, run calibration ONCE using open speakers first, note the measured
  latency, then switch to headphones for the actual cancellation test using
  that measured value (loaded automatically into the slider).
- No system is "finished" - the next real upgrade beyond this is a true
  multi-mic adaptive FIR filter (FxLMS) with a separate reference mic and
  error mic, which is what real ANC headphones use. That's a materially
  bigger hardware + software project, not a code tweak.
"""

import numpy as np
import sounddevice as sd
import tkinter as tk
from tkinter import ttk
from threading import Lock
from matplotlib.figure import Figure
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

# ---------------- Config ----------------
SAMPLE_RATE = 44100
BLOCK_SIZE = 256                 # small block -> fast PLL correction, low latency
NUM_TONES = 3
ACQUIRE_FFT_SIZE = 8192          # long FFT window for fine frequency resolution
ACQUIRE_INTERVAL_SEC = 1.5       # how often we re-scan for new dominant tones
PLL_KP = 6.0                     # proportional gain: phase error -> frequency correction (Hz/rad)
PLL_KI = 1.5                     # integral gain (Hz/rad per second)
IQ_LPF_ALPHA_SEARCH = 0.02       # wide bandwidth (fast) while pulling into lock
IQ_LPF_ALPHA_LOCKED = 0.004      # narrow bandwidth (slow, clean) once locked - kills 2f0 ripple
MIN_FREQ = 60.0                  # ignore rumble below this
MAX_FREQ = 4000.0                # ignore hiss above this (narrowband tones live here)
LOCK_ERROR_THRESH = 0.15         # rad, phase error below this counts as "good" this block
UNLOCK_ERROR_THRESH = 0.4        # rad, must exceed this to drop lock (hysteresis)
LOCK_STREAK_NEEDED = 15          # consecutive good blocks required to declare LOCKED
REACQUIRE_FREQ_TOLERANCE = 15.0  # Hz - only reset an UNLOCKED tracker if peak drifts more than this


class ToneTracker:
    """One phase-locked loop tracking a single tone."""
    def __init__(self):
        self.freq = 0.0
        self.phase = 0.0          # NCO running phase (radians)
        self.amp = 0.0
        self.integrator = 0.0     # PI controller integral term
        self.I_lpf = 0.0
        self.Q_lpf = 0.0
        self.active = False
        self.locked = False
        self.lock_streak = 0

    def acquire(self, freq, phase, amp):
        self.freq = freq
        self.phase = phase
        self.amp = amp
        self.integrator = 0.0
        self.active = True
        self.locked = False
        self.lock_streak = 0


class ANCEngine:
    def __init__(self):
        self.lock = Lock()
        self.trackers = [ToneTracker() for _ in range(NUM_TONES)]
        self.rolling_buf = np.zeros(ACQUIRE_FFT_SIZE, dtype=np.float32)
        self.window = np.hanning(ACQUIRE_FFT_SIZE)
        self.blocks_since_acquire = 0
        self.acquire_every_n_blocks = int(ACQUIRE_INTERVAL_SEC * SAMPLE_RATE / BLOCK_SIZE)

        self.latency_sec = 0.05
        self.gain = 0.9

        self.last_mic_block = np.zeros(BLOCK_SIZE)
        self.last_anti_block = np.zeros(BLOCK_SIZE)

        # calibration state
        self.calibrating = False
        self.calib_chirp = None
        self.calib_record = []
        self.calib_result_ms = None

    # ---------------- Acquisition (periodic FFT scan) ----------------
    def try_acquire_new_tones(self):
        buf = self.rolling_buf * self.window
        spectrum = np.fft.rfft(buf)
        freqs = np.fft.rfftfreq(ACQUIRE_FFT_SIZE, 1 / SAMPLE_RATE)
        mags = np.abs(spectrum)
        mags[(freqs < MIN_FREQ) | (freqs > MAX_FREQ)] = 0

        if not np.any(mags):
            return

        peak_idx = np.argsort(mags)[-NUM_TONES:][::-1]

        for slot, idx in enumerate(peak_idx):
            if mags[idx] <= 0:
                continue
            # parabolic interpolation around the peak bin for sub-bin frequency accuracy
            f_est = self._parabolic_interp(mags, idx, freqs)
            phase_est = np.angle(spectrum[idx])
            amp_est = mags[idx] / (ACQUIRE_FFT_SIZE / 4)

            tr = self.trackers[slot]
            # NEVER reset a tracker that is currently LOCKED - the PLL is
            # doing its job and a noisy FFT snapshot (which may include our
            # own anti-noise bleeding into the mic) is not a reason to throw
            # away a good lock. Only inactive or still-unlocked slots are
            # eligible for a fresh acquisition, and only if the new peak is
            # meaningfully different from what's already being searched for.
            if tr.locked:
                continue
            if (not tr.active) or abs(tr.freq - f_est) > REACQUIRE_FREQ_TOLERANCE:
                tr.acquire(f_est, phase_est, min(amp_est, 1.0 / NUM_TONES))

    @staticmethod
    def _parabolic_interp(mags, idx, freqs):
        if idx <= 0 or idx >= len(mags) - 1:
            return freqs[idx]
        y0, y1, y2 = mags[idx - 1], mags[idx], mags[idx + 1]
        denom = (y0 - 2 * y1 + y2)
        if denom == 0:
            return freqs[idx]
        offset = 0.5 * (y0 - y2) / denom
        bin_width = freqs[1] - freqs[0]
        return freqs[idx] + offset * bin_width

    # ---------------- Real-time PLL tracking + anti-noise synthesis ----------------
    def process_block(self, mic_in, frames):
        out_block = np.zeros(frames, dtype=np.float32)
        n = np.arange(frames)
        dt_block = frames / SAMPLE_RATE

        for tr in self.trackers:
            if not tr.active:
                continue

            # generate this block's NCO phase ramp
            phase_ramp = tr.phase + 2 * np.pi * tr.freq * n / SAMPLE_RATE

            # --- demodulate mic against NCO (I/Q) ---
            I_inst = mic_in * np.sin(phase_ramp)
            Q_inst = mic_in * np.cos(phase_ramp)

            # Bandwidth switching: wide/fast filter while pulling into lock
            # (so acquisition is quick), narrow/slow filter once locked (so
            # the 2*f0 ripple from demodulation is properly suppressed and
            # phase_error stays clean instead of flickering near threshold).
            alpha = IQ_LPF_ALPHA_LOCKED if tr.locked else IQ_LPF_ALPHA_SEARCH
            for k in range(frames):
                tr.I_lpf += alpha * (I_inst[k] - tr.I_lpf)
                tr.Q_lpf += alpha * (Q_inst[k] - tr.Q_lpf)

            phase_error = np.arctan2(tr.I_lpf, tr.Q_lpf)  # small when locked
            amp_est = 2.0 * np.hypot(tr.I_lpf, tr.Q_lpf)

            # PI controller: steer frequency to drive phase_error -> 0
            tr.integrator += PLL_KI * phase_error * dt_block
            freq_correction = PLL_KP * phase_error + tr.integrator
            tr.freq = np.clip(tr.freq + freq_correction * dt_block, MIN_FREQ, MAX_FREQ)

            # Lock detection with hysteresis: requires a sustained streak of
            # good blocks to declare LOCKED, and a much larger error to drop
            # lock again. This prevents the flicker between LOCKED/acquiring
            # caused by single noisy blocks crossing a single threshold.
            if abs(phase_error) < LOCK_ERROR_THRESH:
                tr.lock_streak += 1
            else:
                tr.lock_streak = 0

            if not tr.locked and tr.lock_streak >= LOCK_STREAK_NEEDED:
                tr.locked = True
            elif tr.locked and abs(phase_error) > UNLOCK_ERROR_THRESH:
                tr.locked = False
                tr.lock_streak = 0

            tr.amp = 0.9 * tr.amp + 0.1 * min(amp_est, 1.0 / NUM_TONES)

            # advance NCO phase for next block
            tr.phase = (tr.phase + 2 * np.pi * tr.freq * frames / SAMPLE_RATE) % (2 * np.pi)

            # --- synthesize anti-phase output ---
            latency_shift = 2 * np.pi * tr.freq * self.latency_sec
            total_offset = np.pi + latency_shift  # 180 deg inversion + latency compensation
            wave = tr.amp * self.gain * np.sin(phase_ramp + total_offset)
            out_block += wave

        max_val = np.max(np.abs(out_block))
        if max_val > 0.95:
            out_block = out_block / max_val * 0.95

        return out_block

    # ---------------- Audio callback ----------------
    def audio_callback(self, indata, outdata, frames, time_info, status):
        if status:
            print(status)

        mic_in = indata[:, 0]

        with self.lock:
            if self.calibrating:
                out_block = self._calibration_step(mic_in, frames)
                outdata[:, 0] = out_block
                return

            self.rolling_buf = np.roll(self.rolling_buf, -frames)
            self.rolling_buf[-frames:] = mic_in

            self.blocks_since_acquire += 1
            if self.blocks_since_acquire >= self.acquire_every_n_blocks:
                self.try_acquire_new_tones()
                self.blocks_since_acquire = 0

            out_block = self.process_block(mic_in, frames)

            self.last_mic_block = mic_in.copy()
            self.last_anti_block = out_block.copy()

        outdata[:, 0] = out_block

    # ---------------- Latency calibration (chirp test) ----------------
    def start_calibration(self):
        dur = 0.3
        t = np.linspace(0, dur, int(SAMPLE_RATE * dur), endpoint=False)
        # short linear chirp 500Hz -> 3000Hz, easy to cross-correlate cleanly
        chirp = 0.8 * np.sin(2 * np.pi * (500 * t + (2500 / (2 * dur)) * t**2))
        self.calib_chirp = chirp.astype(np.float32)
        self.calib_record = []
        self.calib_result_ms = None
        self.calibrating = True

    def _calibration_step(self, mic_in, frames):
        self.calib_record.append(mic_in.copy())
        total_recorded = sum(len(b) for b in self.calib_record)

        chirp_len = len(self.calib_chirp)
        pos = total_recorded - frames
        out_block = np.zeros(frames, dtype=np.float32)
        if pos < chirp_len:
            end = min(pos + frames, chirp_len)
            out_block[: end - pos] = self.calib_chirp[pos:end]

        # stop once we've recorded chirp length + generous margin for round-trip delay
        if total_recorded >= chirp_len + SAMPLE_RATE:  # 1 extra second of margin
            self.calibrating = False
            self._finish_calibration()

        return out_block

    def _finish_calibration(self):
        recorded = np.concatenate(self.calib_record)
        correlation = np.correlate(recorded, self.calib_chirp, mode='valid')
        delay_samples = int(np.argmax(np.abs(correlation)))
        self.calib_result_ms = delay_samples / SAMPLE_RATE * 1000
        self.latency_sec = delay_samples / SAMPLE_RATE


class App:
    def __init__(self, master, engine: ANCEngine):
        self.master = master
        self.engine = engine
        master.title("Narrowband ANC - PLL Tracking (Expert Build)")

        ttk.Button(master, text="Auto-Calibrate Latency (uses speakers, not headphones)",
                   command=self.run_calibration).pack(pady=4)
        self.calib_label = ttk.Label(master, text="Latency not yet auto-calibrated")
        self.calib_label.pack()

        ttk.Label(master, text="Latency compensation (manual override, seconds):").pack()
        self.latency_var = tk.DoubleVar(value=engine.latency_sec)
        ttk.Scale(master, from_=0.0, to=0.3, orient=tk.HORIZONTAL,
                  variable=self.latency_var, command=self.update_latency, length=420).pack()
        self.latency_value_label = ttk.Label(master, text=f"{engine.latency_sec*1000:.1f} ms")
        self.latency_value_label.pack()

        ttk.Label(master, text="Output gain:").pack()
        self.gain_var = tk.DoubleVar(value=engine.gain)
        ttk.Scale(master, from_=0.0, to=1.0, orient=tk.HORIZONTAL,
                  variable=self.gain_var, command=self.update_gain, length=420).pack()

        self.status_label = ttk.Label(master, text="Tones: -", font=("Consolas", 10))
        self.status_label.pack(pady=4)

        fig = Figure(figsize=(8, 4))
        self.ax = fig.add_subplot(111)
        self.line_mic, = self.ax.plot([], [], label="Live Mic", alpha=0.7)
        self.line_anti, = self.ax.plot([], [], label="Anti-Phase Output", alpha=0.9)
        self.ax.set_xlim(0, BLOCK_SIZE)
        self.ax.set_ylim(-1.2, 1.2)
        self.ax.legend()
        self.ax.grid(True)
        self.canvas = FigureCanvasTkAgg(fig, master=master)
        self.canvas.get_tk_widget().pack()

        self.refresh()

    def run_calibration(self):
        self.calib_label.config(text="Calibrating... play nothing, stay quiet for ~1.3s")
        self.engine.start_calibration()
        self.master.after(1500, self.check_calibration_done)

    def check_calibration_done(self):
        if self.engine.calibrating:
            self.master.after(200, self.check_calibration_done)
            return
        if self.engine.calib_result_ms is not None:
            self.calib_label.config(
                text=f"Measured latency: {self.engine.calib_result_ms:.1f} ms (loaded into slider)")
            self.latency_var.set(self.engine.latency_sec)
        else:
            self.calib_label.config(text="Calibration failed - try again in a quieter room")

    def update_latency(self, _):
        self.engine.latency_sec = self.latency_var.get()
        self.latency_value_label.config(text=f"{self.engine.latency_sec*1000:.1f} ms")

    def update_gain(self, _):
        self.engine.gain = self.gain_var.get()

    def refresh(self):
        with self.engine.lock:
            mic = self.engine.last_mic_block
            anti = self.engine.last_anti_block
            tones = [(tr.freq, tr.locked, tr.active) for tr in self.engine.trackers]

        self.line_mic.set_data(np.arange(len(mic)), mic)
        self.line_anti.set_data(np.arange(len(anti)), anti)
        self.canvas.draw_idle()

        parts = []
        for f, locked, active in tones:
            if not active:
                continue
            tag = "LOCKED" if locked else "acquiring"
            parts.append(f"{f:.1f}Hz [{tag}]")
        self.status_label.config(text="Tones: " + (", ".join(parts) if parts else "-"))

        self.master.after(100, self.refresh)


def main():
    engine = ANCEngine()

    print("Devices:")
    print(sd.query_devices())
    print("\nStarting full-duplex PLL-tracking ANC.")
    print("1) Click 'Auto-Calibrate Latency' first (uses open speakers, stay quiet).")
    print("2) Then switch to headphones and play a steady tone/hum near the mic.")
    print("3) Watch the Tones readout - wait for '[LOCKED]' before judging cancellation.\n")

    stream = sd.Stream(
        samplerate=SAMPLE_RATE,
        blocksize=BLOCK_SIZE,
        channels=1,
        dtype='float32',
        callback=engine.audio_callback,
        latency='low',
    )

    root = tk.Tk()
    app = App(root, engine)

    def on_close():
        stream.stop()
        stream.close()
        root.destroy()

    root.protocol("WM_DELETE_WINDOW", on_close)

    with stream:
        root.mainloop()


if __name__ == "__main__":
    main()