example/ls_eend_mic_gui.py

from __future__ import annotations

import argparse
import queue
import threading
import time
import tkinter as tk
from dataclasses import dataclass
from pathlib import Path
from tkinter import filedialog, messagebox, ttk

import matplotlib

matplotlib.use("TkAgg")

import librosa
import numpy as np
import sounddevice as sd
import soundfile as sf
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.figure import Figure
from scipy.signal import medfilt

from ls_eend_common import (
    ensure_mono,
    save_json,
    write_rttm,
)
from ls_eend_onnx_runtime import ONNXLSEENDInferenceEngine
from ls_eend_streaming_common import StreamingUpdate


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Live microphone GUI for fixed-shape LS-EEND ONNX inference.")
    parser.add_argument("--onnx-model", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, default=Path(__file__).resolve().parent / "artifacts" / "mic_gui")
    parser.add_argument(
        "--onnx-providers",
        nargs="*",
        default=None,
        help="Optional ONNX Runtime providers when --onnx-model is used.",
    )
    parser.add_argument("--num-speakers", type=int, default=None, help="Optional number of displayed speaker tracks.")
    parser.add_argument("--input-device", type=str, default=None, help="Sounddevice input device id or name.")
    parser.add_argument("--input-sample-rate", type=int, default=None, help="Optional microphone sample rate override.")
    parser.add_argument("--block-seconds", type=float, default=0.1, help="Audio callback block size.")
    parser.add_argument("--refresh-seconds", type=float, default=0.1, help="Minimum queued audio delta before dispatching a streaming step.")
    parser.add_argument(
        "--analysis-seconds",
        type=float,
        default=60.0,
        help="Deprecated compatibility flag. The true stateful streaming path ignores it.",
    )
    parser.add_argument("--display-seconds", type=float, default=120.0, help="Timeline window length shown in the GUI.")
    parser.add_argument("--threshold", type=float, default=0.5, help="Binary activity threshold.")
    parser.add_argument("--median", type=int, default=11, help="Median filter width for binary activity display.")
    parser.add_argument("--simulate-audio", type=Path, default=None, help="Optional audio file to stream into the GUI instead of the microphone.")
    parser.add_argument("--simulate-speed", type=float, default=1.0, help="Playback speed for --simulate-audio.")
    parser.add_argument("--auto-start", action="store_true", help="Start capture immediately on launch.")
    parser.add_argument("--list-devices", action="store_true", help="Print input devices and exit.")
    return parser


@dataclass
class InferenceUpdate:
    update: StreamingUpdate | None = None
    finalized: bool = False
    error: str | None = None


class MicrophoneAudioSource:
    def __init__(
        self,
        audio_queue: queue.Queue[np.ndarray],
        status_queue: queue.Queue[str],
        input_device: str | None,
        input_sample_rate: int | None,
        target_sample_rate: int,
        block_seconds: float,
    ) -> None:
        self.audio_queue = audio_queue
        self.status_queue = status_queue
        self.input_device = _resolve_input_device(input_device)
        device_info = sd.query_devices(self.input_device, "input")
        self.device_name = device_info["name"]
        self.sample_rate = int(input_sample_rate or target_sample_rate)
        self.blocksize = max(1, int(round(block_seconds * self.sample_rate)))
        self.stream: sd.InputStream | None = None

    def start(self) -> None:
        self.stream = sd.InputStream(
            samplerate=self.sample_rate,
            blocksize=self.blocksize,
            device=self.input_device,
            channels=1,
            dtype="float32",
            callback=self._callback,
        )
        self.stream.start()
        self.status_queue.put(f"Capturing microphone: {self.device_name} @ {self.sample_rate} Hz")

    def stop(self) -> None:
        if self.stream is None:
            return
        self.stream.stop()
        self.stream.close()
        self.stream = None

    def _callback(self, indata, frames, callback_time, status) -> None:  # type: ignore[no-untyped-def]
        if status:
            self.status_queue.put(f"Audio callback status: {status}")
        self.audio_queue.put(indata[:, 0].copy())


class SimulatedAudioSource:
    def __init__(
        self,
        audio_queue: queue.Queue[np.ndarray],
        status_queue: queue.Queue[str],
        audio_path: Path,
        target_sample_rate: int,
        block_seconds: float,
        speed: float,
    ) -> None:
        self.audio_queue = audio_queue
        self.status_queue = status_queue
        self.audio_path = audio_path
        self.block_seconds = block_seconds
        self.speed = max(speed, 1e-3)
        audio, sample_rate = sf.read(audio_path)
        audio = ensure_mono(audio).astype(np.float32, copy=False)
        if sample_rate != target_sample_rate:
            audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=target_sample_rate).astype(np.float32, copy=False)
            sample_rate = target_sample_rate
        self.audio = audio
        self.sample_rate = int(sample_rate)
        self.device_name = f"Simulated: {audio_path.name}"
        self.blocksize = max(1, int(round(block_seconds * self.sample_rate)))
        self._thread: threading.Thread | None = None
        self._stop_event = threading.Event()

    def start(self) -> None:
        self._stop_event.clear()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()
        self.status_queue.put(f"Streaming simulation: {self.audio_path.name} @ {self.sample_rate} Hz")

    def stop(self) -> None:
        self._stop_event.set()
        if self._thread is not None and self._thread.is_alive():
            self._thread.join(timeout=1.0)
        self._thread = None

    def _run(self) -> None:
        for start in range(0, len(self.audio), self.blocksize):
            if self._stop_event.is_set():
                return
            stop = min(len(self.audio), start + self.blocksize)
            self.audio_queue.put(self.audio[start:stop].copy())
            time.sleep(((stop - start) / self.sample_rate) / self.speed)
        self.status_queue.put("Simulation finished.")


class LSEENDMicGUI:
    def __init__(self, args: argparse.Namespace) -> None:
        self.args = args
        self.engine = ONNXLSEENDInferenceEngine(
            onnx_model_path=args.onnx_model,
            providers=args.onnx_providers,
        )
        backend_label = "ONNX"
        self.output_dir = args.output_dir
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.root = tk.Tk()
        self.root.title(f"LS-EEND Live Microphone GUI ({backend_label})")
        self.root.geometry("1500x920")

        self.audio_queue: queue.Queue[np.ndarray] = queue.Queue()
        self.status_queue: queue.Queue[str] = queue.Queue()
        self.result_queue: queue.Queue[InferenceUpdate] = queue.Queue()
        self.audio_lock = threading.Lock()
        self.pending_audio = np.zeros(0, dtype=np.float32)
        self.total_samples_received = 0
        self.timeline_probabilities = np.zeros((0, 0), dtype=np.float32)
        self.preview_probabilities = np.zeros((0, 0), dtype=np.float32)
        self.preview_start_frame = 0
        self.latest_result = None
        self.session = None
        self.inference_thread: threading.Thread | None = None
        self.inference_in_flight = False
        self.finalize_requested = False
        self.source = None
        self.sample_rate = 0
        self.session_index = 0
        self.display_order: list[int] = []
        self.track_labels: dict[int, tk.StringVar] = {}
        self.swap_left_var = tk.StringVar()
        self.swap_right_var = tk.StringVar()

        self.status_var = tk.StringVar(value="Loading model...")
        self.source_var = tk.StringVar(value="Not started")
        self.buffer_var = tk.StringVar(value="Buffered: 0.0 s")
        self.inference_var = tk.StringVar(value="Inference: idle")
        self.window_seconds_var = tk.DoubleVar(value=float(args.display_seconds))
        self.threshold_var = tk.DoubleVar(value=float(args.threshold))
        self.median_var = tk.IntVar(value=int(args.median))

        self._build_ui()
        self.status_var.set("Ready.")
        self.root.protocol("WM_DELETE_WINDOW", self.on_close)
        self.root.after(50, self._poll_audio)
        self.root.after(75, self._poll_status)
        self.root.after(100, self._poll_results)

        if args.auto_start or args.simulate_audio is not None:
            self.root.after(150, self.start_capture)

    def _build_ui(self) -> None:
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)

        container = ttk.Frame(self.root, padding=8)
        container.grid(row=0, column=0, sticky="nsew")
        container.columnconfigure(0, weight=3)
        container.columnconfigure(1, weight=1)
        container.rowconfigure(0, weight=1)

        plot_frame = ttk.Frame(container)
        plot_frame.grid(row=0, column=0, sticky="nsew", padx=(0, 8))
        plot_frame.columnconfigure(0, weight=1)
        plot_frame.rowconfigure(0, weight=1)

        self.figure = Figure(figsize=(12, 8), dpi=100, constrained_layout=True)
        self.binary_axis = self.figure.add_subplot(211)
        self.prob_axis = self.figure.add_subplot(212, sharex=self.binary_axis)
        self.canvas = FigureCanvasTkAgg(self.figure, master=plot_frame)
        self.canvas.get_tk_widget().grid(row=0, column=0, sticky="nsew")
        self._draw_placeholder()

        side = ttk.Frame(container)
        side.grid(row=0, column=1, sticky="nsew")
        side.columnconfigure(0, weight=1)

        control_frame = ttk.LabelFrame(side, text="Session", padding=8)
        control_frame.grid(row=0, column=0, sticky="ew")
        for column in range(2):
            control_frame.columnconfigure(column, weight=1)

        ttk.Button(control_frame, text="Start", command=self.start_capture).grid(row=0, column=0, sticky="ew", padx=(0, 4))
        ttk.Button(control_frame, text="Stop", command=self.stop_capture).grid(row=0, column=1, sticky="ew", padx=(4, 0))
        ttk.Button(control_frame, text="Reset Timeline", command=self.reset_timeline).grid(row=1, column=0, sticky="ew", pady=(6, 0), padx=(0, 4))
        ttk.Button(control_frame, text="Save RTTM", command=self.save_current_rttm).grid(row=1, column=1, sticky="ew", pady=(6, 0), padx=(4, 0))
        ttk.Button(control_frame, text="Save Heatmap", command=self.save_current_heatmap).grid(row=2, column=0, sticky="ew", pady=(6, 0), padx=(0, 4))
        ttk.Button(control_frame, text="Save Session", command=self.save_session_metadata).grid(row=2, column=1, sticky="ew", pady=(6, 0), padx=(4, 0))

        status_frame = ttk.LabelFrame(side, text="Status", padding=8)
        status_frame.grid(row=1, column=0, sticky="ew", pady=(8, 0))
        ttk.Label(status_frame, textvariable=self.status_var, wraplength=340, justify="left").grid(row=0, column=0, sticky="w")
        ttk.Label(status_frame, textvariable=self.source_var, wraplength=340, justify="left").grid(row=1, column=0, sticky="w", pady=(6, 0))
        ttk.Label(status_frame, textvariable=self.buffer_var, wraplength=340, justify="left").grid(row=2, column=0, sticky="w", pady=(6, 0))
        ttk.Label(status_frame, textvariable=self.inference_var, wraplength=340, justify="left").grid(row=3, column=0, sticky="w", pady=(6, 0))

        display_frame = ttk.LabelFrame(side, text="Display", padding=8)
        display_frame.grid(row=2, column=0, sticky="ew", pady=(8, 0))
        ttk.Label(display_frame, text="Window (s)").grid(row=0, column=0, sticky="w")
        ttk.Spinbox(display_frame, from_=10, to=3600, increment=10, textvariable=self.window_seconds_var, width=10, command=self.refresh_plot).grid(row=0, column=1, sticky="w")
        ttk.Label(display_frame, text="Threshold").grid(row=1, column=0, sticky="w", pady=(6, 0))
        ttk.Spinbox(display_frame, from_=0.05, to=0.95, increment=0.05, textvariable=self.threshold_var, width=10, command=self.refresh_plot).grid(row=1, column=1, sticky="w", pady=(6, 0))
        ttk.Label(display_frame, text="Median").grid(row=2, column=0, sticky="w", pady=(6, 0))
        ttk.Spinbox(display_frame, from_=1, to=51, increment=2, textvariable=self.median_var, width=10, command=self.refresh_plot).grid(row=2, column=1, sticky="w", pady=(6, 0))

        swap_frame = ttk.LabelFrame(side, text="Swap Rows", padding=8)
        swap_frame.grid(row=3, column=0, sticky="ew", pady=(8, 0))
        ttk.Label(swap_frame, text="Row A").grid(row=0, column=0, sticky="w")
        ttk.Label(swap_frame, text="Row B").grid(row=1, column=0, sticky="w", pady=(6, 0))
        self.swap_left = ttk.Combobox(swap_frame, state="readonly", textvariable=self.swap_left_var)
        self.swap_right = ttk.Combobox(swap_frame, state="readonly", textvariable=self.swap_right_var)
        self.swap_left.grid(row=0, column=1, sticky="ew")
        self.swap_right.grid(row=1, column=1, sticky="ew", pady=(6, 0))
        ttk.Button(swap_frame, text="Swap", command=self.swap_selected_rows).grid(row=2, column=0, columnspan=2, sticky="ew", pady=(8, 0))
        ttk.Button(swap_frame, text="Reset Order", command=self.reset_order).grid(row=3, column=0, columnspan=2, sticky="ew", pady=(6, 0))
        swap_frame.columnconfigure(1, weight=1)

        self.speaker_frame = ttk.LabelFrame(side, text="Speakers", padding=8)
        self.speaker_frame.grid(row=4, column=0, sticky="nsew", pady=(8, 0))
        side.rowconfigure(4, weight=1)

    def _draw_placeholder(self) -> None:
        for axis, title in ((self.binary_axis, "Binary Activity"), (self.prob_axis, "Speaker Probability")):
            axis.clear()
            axis.set_title(title)
            axis.text(0.5, 0.5, "No inference yet", ha="center", va="center", transform=axis.transAxes)
            axis.set_yticks([])
        self.prob_axis.set_xlabel("Time (seconds)")
        self.canvas.draw_idle()

    def _create_source(self):
        if self.args.simulate_audio is not None:
            return SimulatedAudioSource(
                audio_queue=self.audio_queue,
                status_queue=self.status_queue,
                audio_path=self.args.simulate_audio,
                target_sample_rate=self.engine.target_sample_rate,
                block_seconds=self.args.block_seconds,
                speed=self.args.simulate_speed,
            )
        return MicrophoneAudioSource(
            audio_queue=self.audio_queue,
            status_queue=self.status_queue,
            input_device=self.args.input_device,
            input_sample_rate=self.args.input_sample_rate,
            target_sample_rate=self.engine.target_sample_rate,
            block_seconds=self.args.block_seconds,
        )

    def start_capture(self) -> None:
        if self.source is not None:
            self.status_var.set("Capture is already running.")
            return
        if self.session is not None and self.session.finalized:
            self.reset_timeline()
        try:
            self.source = self._create_source()
            self.sample_rate = int(self.source.sample_rate)
            self.session = self.engine.create_session(self.sample_rate)
            self.source.start()
        except Exception as exc:
            self.source = None
            self.session = None
            messagebox.showerror("Audio Source Error", str(exc))
            return
        self.finalize_requested = False
        self.source_var.set(f"Source: {self.source.device_name} @ {self.sample_rate} Hz")
        self.status_var.set(
            f"Capture started. Stateful streaming latency is about {self.engine.streaming_latency_seconds:.2f} s."
        )
        self._update_buffer_label()

    def stop_capture(self) -> None:
        if self.source is None and self.session is None:
            return
        if self.source is not None:
            self.source.stop()
            self.source = None
        self.finalize_requested = True
        self.status_var.set("Capture stopped. Flushing the delayed tail...")
        self._maybe_schedule_inference(force=True)

    def reset_timeline(self) -> None:
        if self.source is not None:
            self.source.stop()
            self.source = None
        with self.audio_lock:
            self.pending_audio = np.zeros(0, dtype=np.float32)
            self.total_samples_received = 0
        self.latest_result = None
        self.session = None
        self.timeline_probabilities = np.zeros((0, 0), dtype=np.float32)
        self.preview_probabilities = np.zeros((0, 0), dtype=np.float32)
        self.preview_start_frame = 0
        self.finalize_requested = False
        self.display_order = []
        self._rebuild_speaker_controls()
        self.buffer_var.set("Buffered: 0.0 s")
        self.inference_var.set("Inference: idle")
        self.status_var.set("Timeline reset.")
        self._draw_placeholder()

    def _poll_audio(self) -> None:
        received = False
        while True:
            try:
                chunk = self.audio_queue.get_nowait()
            except queue.Empty:
                break
            received = True
            with self.audio_lock:
                self.pending_audio = np.concatenate([self.pending_audio, chunk], axis=0)
                self.total_samples_received += len(chunk)
            self._update_buffer_label()
        if received:
            self._maybe_schedule_inference(force=False)
        self.root.after(50, self._poll_audio)

    def _poll_status(self) -> None:
        message = None
        while True:
            try:
                message = self.status_queue.get_nowait()
            except queue.Empty:
                break
        if message is not None:
            if message == "Simulation finished.":
                self.stop_capture()
            else:
                self.status_var.set(message)
        self.root.after(100, self._poll_status)

    def _poll_results(self) -> None:
        updated = False
        while True:
            try:
                update = self.result_queue.get_nowait()
            except queue.Empty:
                break
            self.inference_in_flight = False
            if update.error is not None:
                self.status_var.set(update.error)
            else:
                if update.update is not None:
                    self.latest_result = update.update
                    self._merge_result_into_timeline(update.update.start_frame, update.update.probabilities)
                    self._set_preview(update.update.preview_start_frame, update.update.preview_probabilities)
                    self.inference_var.set(
                        f"Inference: {update.update.total_emitted_frames} committed + {update.update.preview_probabilities.shape[0]} preview frames"
                    )
                    self._ensure_track_state(self._combined_probabilities().shape[1])
                    self.refresh_plot()
                    self._update_buffer_label()
                    updated = True
                if update.finalized:
                    self.preview_probabilities = np.zeros((0, self.timeline_probabilities.shape[1] if self.timeline_probabilities.ndim == 2 else 0), dtype=np.float32)
                    self.preview_start_frame = self.timeline_probabilities.shape[0]
                    self._update_buffer_label()
                    self.status_var.set("Streaming tail flushed.")
        if self.session is not None and not self.inference_in_flight:
            with self.audio_lock:
                has_pending_audio = self.pending_audio.size > 0
            if has_pending_audio:
                self._maybe_schedule_inference(force=True)
            elif self.finalize_requested and not self.session.finalized:
                self._start_background_step(finalize=True)
        elif updated:
            self.status_var.set("Inference updated.")
        self.root.after(100, self._poll_results)

    def _maybe_schedule_inference(self, force: bool) -> None:
        if self.session is None or self.sample_rate <= 0:
            return
        minimum_delta = max(1, int(round(self.args.refresh_seconds * self.sample_rate)))
        if self.inference_in_flight:
            return
        with self.audio_lock:
            queued_samples = len(self.pending_audio)
        if queued_samples == 0:
            if force and self.finalize_requested and not self.session.finalized:
                self._start_background_step(finalize=True)
            return
        if not force and queued_samples < minimum_delta:
            return
        self._start_background_step(finalize=False)

    def _start_background_step(self, finalize: bool) -> None:
        self.inference_in_flight = True
        self.inference_var.set("Inference: running...")
        self.inference_thread = threading.Thread(target=self._run_inference, args=(finalize,), daemon=True)
        self.inference_thread.start()

    def _run_inference(self, finalize: bool) -> None:
        try:
            if self.session is None:
                raise RuntimeError("Streaming session is not initialized.")
            if finalize:
                update = self.session.finalize()
                self.result_queue.put(InferenceUpdate(update=update, finalized=True))
                return
            with self.audio_lock:
                chunk = self.pending_audio.copy()
                self.pending_audio = np.zeros(0, dtype=np.float32)
            update = self.session.push_audio(chunk)
            if update is not None and self.args.num_speakers is not None:
                clip = max(0, min(int(self.args.num_speakers), update.probabilities.shape[1]))
                update.logits = update.logits[:, :clip]
                update.probabilities = update.probabilities[:, :clip]
            self.result_queue.put(InferenceUpdate(update=update))
        except Exception as exc:
            self.result_queue.put(InferenceUpdate(error=f"Inference failed: {exc}"))

    def _merge_result_into_timeline(self, start_frame: int, probabilities: np.ndarray) -> None:
        if probabilities.size == 0:
            return
        end_frame = start_frame + probabilities.shape[0]
        current_frames, current_tracks = self.timeline_probabilities.shape if self.timeline_probabilities.size else (0, 0)
        target_tracks = max(current_tracks, probabilities.shape[1])
        if current_frames < end_frame or current_tracks < target_tracks:
            expanded = np.zeros((max(current_frames, end_frame), target_tracks), dtype=np.float32)
            if current_frames > 0 and current_tracks > 0:
                expanded[:current_frames, :current_tracks] = self.timeline_probabilities
            self.timeline_probabilities = expanded
        self.timeline_probabilities[start_frame:end_frame, : probabilities.shape[1]] = probabilities

    def _set_preview(self, start_frame: int, probabilities: np.ndarray) -> None:
        self.preview_start_frame = int(start_frame)
        self.preview_probabilities = probabilities.astype(np.float32, copy=False)

    def _update_buffer_label(self) -> None:
        received_seconds = self.total_samples_received / max(self.sample_rate, 1)
        committed_seconds = self.timeline_probabilities.shape[0] / self.engine.model_frame_hz
        preview_seconds = (
            (self.preview_start_frame + self.preview_probabilities.shape[0]) / self.engine.model_frame_hz
            if self.preview_probabilities.size
            else committed_seconds
        )
        self.buffer_var.set(
            f"Buffered: {received_seconds:.1f} s received, {committed_seconds:.1f} s committed, {preview_seconds:.1f} s incl preview"
        )

    def _selected_probabilities(self) -> np.ndarray:
        probabilities = self._combined_probabilities()
        if probabilities.size == 0:
            return probabilities
        if not self.display_order:
            return probabilities
        return probabilities[:, self.display_order]

    def _selected_preview_range(self) -> tuple[int | None, int | None]:
        if self.preview_probabilities.size == 0:
            return None, None
        return self.preview_start_frame, self.preview_start_frame + self.preview_probabilities.shape[0]

    def _combined_probabilities(self) -> np.ndarray:
        current_frames, current_tracks = self.timeline_probabilities.shape if self.timeline_probabilities.size else (0, 0)
        preview_frames, preview_tracks = self.preview_probabilities.shape if self.preview_probabilities.size else (0, 0)
        total_tracks = max(current_tracks, preview_tracks)
        total_frames = max(current_frames, self.preview_start_frame + preview_frames)
        if total_frames == 0 or total_tracks == 0:
            return np.zeros((0, 0), dtype=np.float32)
        combined = np.zeros((total_frames, total_tracks), dtype=np.float32)
        if current_frames and current_tracks:
            combined[:current_frames, :current_tracks] = self.timeline_probabilities
        if preview_frames and preview_tracks:
            combined[self.preview_start_frame : self.preview_start_frame + preview_frames, :preview_tracks] = self.preview_probabilities
        return combined

    def _speaker_labels_for_display(self) -> list[str]:
        labels = []
        for track_index in self.display_order:
            label_var = self.track_labels.get(track_index)
            label = label_var.get().strip() if label_var is not None else ""
            labels.append(label or f"Speaker {track_index + 1}")
        return labels

    def _ensure_track_state(self, track_count: int) -> None:
        if track_count <= 0:
            return
        if len(self.display_order) != track_count or set(self.display_order) != set(range(track_count)):
            self.display_order = list(range(track_count))
        for track_index in range(track_count):
            if track_index not in self.track_labels:
                self.track_labels[track_index] = tk.StringVar(value=f"Speaker {track_index + 1}")
        self._rebuild_speaker_controls()

    def _rebuild_speaker_controls(self) -> None:
        for child in self.speaker_frame.winfo_children():
            child.destroy()
        if not self.display_order:
            ttk.Label(self.speaker_frame, text="Run inference to populate speaker tracks.").grid(row=0, column=0, sticky="w")
            self.swap_left["values"] = ()
            self.swap_right["values"] = ()
            self.swap_left_var.set("")
            self.swap_right_var.set("")
            return
        row_options = [str(index + 1) for index in range(len(self.display_order))]
        self.swap_left["values"] = row_options
        self.swap_right["values"] = row_options
        if not self.swap_left_var.get():
            self.swap_left_var.set(row_options[0])
        if len(row_options) > 1 and not self.swap_right_var.get():
            self.swap_right_var.set(row_options[1])
        elif row_options:
            self.swap_right_var.set(row_options[0])

        for row_index, track_index in enumerate(self.display_order):
            ttk.Label(self.speaker_frame, text=f"Row {row_index + 1}").grid(row=row_index, column=0, sticky="w")
            ttk.Entry(self.speaker_frame, textvariable=self.track_labels[track_index], width=18).grid(row=row_index, column=1, sticky="ew", padx=(6, 6))
            ttk.Button(self.speaker_frame, text="Up", command=lambda idx=row_index: self.move_row(idx, -1)).grid(row=row_index, column=2, sticky="ew")
            ttk.Button(self.speaker_frame, text="Down", command=lambda idx=row_index: self.move_row(idx, 1)).grid(row=row_index, column=3, sticky="ew", padx=(4, 0))
        self.speaker_frame.columnconfigure(1, weight=1)

    def move_row(self, row_index: int, delta: int) -> None:
        swap_index = row_index + delta
        if swap_index < 0 or swap_index >= len(self.display_order):
            return
        self.display_order[row_index], self.display_order[swap_index] = self.display_order[swap_index], self.display_order[row_index]
        self._rebuild_speaker_controls()
        self.refresh_plot()

    def swap_selected_rows(self) -> None:
        if not self.display_order:
            return
        try:
            left = int(self.swap_left_var.get()) - 1
            right = int(self.swap_right_var.get()) - 1
        except ValueError:
            return
        if left < 0 or right < 0 or left >= len(self.display_order) or right >= len(self.display_order) or left == right:
            return
        self.display_order[left], self.display_order[right] = self.display_order[right], self.display_order[left]
        self._rebuild_speaker_controls()
        self.refresh_plot()

    def reset_order(self) -> None:
        if not self.display_order:
            return
        self.display_order = list(range(len(self.display_order)))
        self._rebuild_speaker_controls()
        self.refresh_plot()

    def _current_binary(self) -> np.ndarray:
        probabilities = self._selected_probabilities()
        if probabilities.size == 0:
            return probabilities
        binary = (probabilities >= float(self.threshold_var.get())).astype(np.float32)
        median_width = max(1, int(self.median_var.get()))
        if median_width > 1:
            if median_width % 2 == 0:
                median_width += 1
            binary = medfilt(binary, kernel_size=(median_width, 1)).astype(np.float32)
        return binary

    def refresh_plot(self) -> None:
        if self.timeline_probabilities.size == 0:
            self._draw_placeholder()
            return
        probabilities = self._selected_probabilities()
        if probabilities.size == 0:
            self._draw_placeholder()
            return
        binary = self._current_binary()
        frame_hz = self.engine.model_frame_hz
        window_frames = max(1, int(round(float(self.window_seconds_var.get()) * frame_hz)))
        start_frame = max(0, probabilities.shape[0] - window_frames)
        shown_probs = probabilities[start_frame:]
        shown_binary = binary[start_frame:]
        start_seconds = start_frame / frame_hz
        end_seconds = probabilities.shape[0] / frame_hz
        speaker_labels = self._speaker_labels_for_display()
        preview_start_frame, preview_end_frame = self._selected_preview_range()
        preview_start_seconds = None if preview_start_frame is None else preview_start_frame / frame_hz
        preview_end_seconds = None if preview_end_frame is None else preview_end_frame / frame_hz

        self.binary_axis.clear()
        self.binary_axis.imshow(
            shown_binary.T,
            aspect="auto",
            origin="lower",
            interpolation="nearest",
            extent=[start_seconds, end_seconds, -0.5, shown_binary.shape[1] - 0.5],
            cmap="Greys",
            vmin=0.0,
            vmax=1.0,
        )
        self.binary_axis.set_title("Binary Activity")
        self.binary_axis.set_yticks(range(len(speaker_labels)))
        self.binary_axis.set_yticklabels(speaker_labels)
        if preview_start_seconds is not None and preview_end_seconds is not None and preview_end_seconds > start_seconds:
            self.binary_axis.axvspan(
                max(preview_start_seconds, start_seconds),
                min(preview_end_seconds, end_seconds),
                color="orange",
                alpha=0.08,
            )
            self.binary_axis.axvline(preview_start_seconds, color="orange", linestyle="--", linewidth=1.0)

        self.prob_axis.clear()
        image = self.prob_axis.imshow(
            shown_probs.T,
            aspect="auto",
            origin="lower",
            interpolation="nearest",
            extent=[start_seconds, end_seconds, -0.5, shown_probs.shape[1] - 0.5],
            cmap="viridis",
            vmin=0.0,
            vmax=1.0,
        )
        self.prob_axis.set_title("Speaker Probability")
        self.prob_axis.set_yticks(range(len(speaker_labels)))
        self.prob_axis.set_yticklabels(speaker_labels)
        self.prob_axis.set_xlabel("Time (seconds)")
        if preview_start_seconds is not None and preview_end_seconds is not None and preview_end_seconds > start_seconds:
            self.prob_axis.axvspan(
                max(preview_start_seconds, start_seconds),
                min(preview_end_seconds, end_seconds),
                color="orange",
                alpha=0.08,
            )
            self.prob_axis.axvline(preview_start_seconds, color="orange", linestyle="--", linewidth=1.0)
        if len(self.figure.axes) < 3:
            self.figure.colorbar(image, ax=self.prob_axis, fraction=0.02, pad=0.01)
        self.canvas.draw_idle()

    def _next_output_stem(self) -> str:
        self.session_index += 1
        return time.strftime(f"mic_session_%Y%m%d_%H%M%S_{self.session_index:02d}")

    def save_current_rttm(self) -> None:
        if self.timeline_probabilities.size == 0:
            messagebox.showinfo("Save RTTM", "No inference available yet.")
            return
        output_path = self.output_dir / f"{self._next_output_stem()}.rttm"
        write_rttm(
            recording_id=output_path.stem,
            binary_prediction=self._current_binary(),
            output_path=output_path,
            frame_rate=self.engine.model_frame_hz,
            speaker_labels=self._speaker_labels_for_display(),
        )
        self.status_var.set(f"Saved RTTM: {output_path.name}")

    def save_current_heatmap(self) -> None:
        if self.timeline_probabilities.size == 0:
            messagebox.showinfo("Save Heatmap", "No inference available yet.")
            return
        path = filedialog.asksaveasfilename(
            title="Save heatmap",
            defaultextension=".png",
            initialdir=str(self.output_dir),
            initialfile=f"{self._next_output_stem()}.png",
            filetypes=[("PNG image", "*.png")],
        )
        if not path:
            return
        self.figure.savefig(path, dpi=200)
        self.status_var.set(f"Saved heatmap: {Path(path).name}")

    def save_session_metadata(self) -> None:
        if self.timeline_probabilities.size == 0:
            messagebox.showinfo("Save Session", "No inference available yet.")
            return
        output_path = self.output_dir / f"{self._next_output_stem()}.json"
        payload = {
            "backend": "onnx",
            "onnx_model": str(self.args.onnx_model),
            "onnx_providers": self.args.onnx_providers,
            "source": self.source_var.get(),
            "duration_seconds": float(self.timeline_probabilities.shape[0] / self.engine.model_frame_hz),
            "preview_duration_seconds": float(self.preview_probabilities.shape[0] / self.engine.model_frame_hz),
            "frame_hz": float(self.engine.model_frame_hz),
            "display_order": self.display_order,
            "speaker_labels": self._speaker_labels_for_display(),
            "threshold": float(self.threshold_var.get()),
            "median": int(self.median_var.get()),
            "streaming_latency_seconds": float(self.engine.streaming_latency_seconds),
        }
        save_json(payload, output_path)
        self.status_var.set(f"Saved session metadata: {output_path.name}")

    def on_close(self) -> None:
        self.stop_capture()
        self.root.destroy()

    def run(self) -> None:
        self.root.mainloop()


def list_devices() -> None:
    devices = sd.query_devices()
    for index, device in enumerate(devices):
        if device["max_input_channels"] <= 0:
            continue
        print(f"{index}: {device['name']} | input_channels={device['max_input_channels']} | default_samplerate={device['default_samplerate']}")


def _resolve_input_device(device: str | None) -> str | int | None:
    if device is None:
        return None
    stripped = device.strip()
    if stripped.isdigit():
        return int(stripped)
    return device


def main() -> None:
    args = build_parser().parse_args()
    if args.list_devices:
        list_devices()
        return
    app = LSEENDMicGUI(args)
    app.run()


if __name__ == "__main__":
    main()