main.py

"""Application entry point for the Windows sound visualizer."""

from __future__ import annotations

import sys
import time
import traceback
from copy import deepcopy
from pathlib import Path

import numpy as np
from PyQt6.QtCore import QProcess, QTimer
from PyQt6.QtGui import QAction, QIcon
from PyQt6.QtWidgets import QApplication, QMenu, QMessageBox, QSystemTrayIcon

from audio_capture import AudioCapture
from config import load_config, save_config
from overlay import OverlayWidget
from signal_processing import SpectrumProcessor
from settings import SettingsWindow
from startup import set_startup


class VisualizerApp:
    def __init__(self, config: dict) -> None:
        self.config = config
        if config["input"]["source"] == "__none__":
            config["input"]["source"] = None
            save_config(config)
        self.capture = AudioCapture(
            config["output"], config["input"],
        )
        self.output_processor = self._create_processor(config, "output")
        self.input_processor = self._create_processor(config, "input")
        self.overlay = OverlayWidget(config)
        self.settings: SettingsWindow | None = None
        self.tray = self._create_tray()
        self.timer = QTimer()
        self.timer.timeout.connect(self._update)
        self._onset_times: list[float] = []
        self._last_onset = 0.0
        self._onset_times_by_stream = {"output": [], "input": []}
        self._last_onset_by_stream = {"output": 0.0, "input": 0.0}
        self._bpm_display = {"output": None, "input": None}
        self._last_data_update = 0.0
        self._last_frequency = {"output": None, "input": None}
        self._silence_since = {"output": None, "input": None}
        self._calibration_until = 0.0
        self._calibration_activity: list[np.ndarray] = []

    def _create_processor(self, config: dict, prefix: str) -> SpectrumProcessor:
        settings = config[prefix]
        return SpectrumProcessor(
            sample_rate=settings["sample_rate"],
            fft_size=settings["fft_size"],
            frequency_min=config["general"]["frequency_min"],
            frequency_max=config["general"]["frequency_max"],
            output_count=config["visual"]["bar_count"],
            smoothing=config["general"]["smoothing"],
            decay=config["general"]["decay"],
            magnitude_scale=settings["magnitude_scale"],
            magnitude_mode=settings["magnitude_mode"],
            binning_method=config["general"]["binning_method"],
            interpolate_to_zero=config["general"]["interpolate_to_zero"],
            active_bins=config["general"]["non_empty_bins"],
        )

    def start(self) -> None:
        self.overlay.show()
        self.capture.start()
        self.timer.start(1000 // self.config["visual"]["frame_rate"])

    def stop(self) -> None:
        self.timer.stop()
        self.capture.stop()
        self.overlay.save_state()
        self.overlay.close()

    def _update(self) -> None:
        blocks = self.capture.read_latest()
        output_size = self.config["output"]["block_size"]
        input_size = self.config["input"]["block_size"]
        output = blocks["output"]
        input_block = blocks["input"]
        output_spectrum = None if self.config["output"]["source"] == "__none__" else self.output_processor.process(output if output is not None else np.zeros((output_size, 2)))
        input_spectrum = None if self.config["input"]["source"] == "__none__" else self.input_processor.process(input_block if input_block is not None else np.zeros((input_size, 2)))
        self.overlay.set_spectrum(output_spectrum, input_spectrum)
        self._collect_calibration(output, input_block)
        output_frequency = self._dominant_frequency("output", output_spectrum, output)
        input_frequency = self._dominant_frequency("input", input_spectrum, input_block)
        output_bpm = self._estimate_bpm("output", output)
        input_bpm = self._estimate_bpm("input", input_block)
        output_volume = self._volume(output)
        input_volume = self._volume(input_block)
        now = time.monotonic()
        if now - self._last_data_update >= 1.0 / self.config["data"]["update_rate"]:
            self.overlay.set_data(output_frequency, output_bpm, input_frequency, input_bpm, output_volume, input_volume)
            self._last_data_update = now

    def _dominant_frequency(self, stream: str, spectrum, block: np.ndarray | None) -> float | None:
        if spectrum is None:
            return None
        if block is None:
            return self._last_frequency[stream]
        energy = float(np.sqrt(np.mean(np.asarray(block, dtype=np.float32) ** 2)))
        now = time.monotonic()
        if energy < 0.0005:
            if self._silence_since[stream] is None:
                self._silence_since[stream] = now
            if now - self._silence_since[stream] < 1.0:
                return self._last_frequency[stream]
            self._last_frequency[stream] = 0.0
            return 0.0
        self._silence_since[stream] = None
        if not spectrum.left_magnitudes.size:
            return self._last_frequency[stream]
        magnitudes = np.maximum(spectrum.left_magnitudes, spectrum.right_magnitudes)
        frequency = float(spectrum.frequencies[int(np.argmax(magnitudes))])
        self._last_frequency[stream] = frequency
        return frequency

    def _estimate_bpm(self, stream: str, block: np.ndarray | None) -> float | None:
        if not self.config["data"][f"show_{stream}_bpm"]:
            self._bpm_display[stream] = None
            return None
        if block is None:
            return self._bpm_display[stream]
        energy = float(np.sqrt(np.mean(np.asarray(block, dtype=np.float32) ** 2)))
        now = time.monotonic()
        if energy < 0.0005:
            if self._silence_since[stream] is None:
                self._silence_since[stream] = now
            if now - self._silence_since[stream] < 1.0:
                return self._bpm_display[stream]
            self._onset_times_by_stream[stream].clear()
            self._last_onset_by_stream[stream] = 0.0
            self._bpm_display[stream] = 0.0
            return 0.0
        self._silence_since[stream] = None
        onset_times = self._onset_times_by_stream[stream]
        history = getattr(self, "_energy_history", {"output": [], "input": []})
        self._energy_history = history
        baseline = float(np.median(history[stream][-32:])) if history[stream] else energy
        history[stream].append(energy)
        del history[stream][:-64]
        if energy > max(0.0005, baseline * 1.35) and now - self._last_onset_by_stream[stream] > 0.28:
            onset_times.append(now)
            self._last_onset_by_stream[stream] = now
            del onset_times[:-12]
        if len(onset_times) < 3:
            return self._bpm_display[stream]
        intervals = np.diff(onset_times)
        bpm = 60.0 / float(np.median(intervals))
        while bpm < 60:
            bpm *= 2
        while bpm > 180:
            bpm /= 2
        previous = self._bpm_display[stream]
        self._bpm_display[stream] = bpm if previous is None else previous + (bpm - previous) * 0.12
        return self._bpm_display[stream]

    @staticmethod
    def _volume(block: np.ndarray | None) -> float | None:
        if block is None:
            return None
        return max(0.0, min(1.0, float(np.sqrt(np.mean(np.asarray(block, dtype=np.float32) ** 2)))))

    def _create_tray(self) -> QSystemTrayIcon:
        tray = QSystemTrayIcon(self.overlay)
        icon_path = Path(__file__).with_name("icon.png")
        tray.setIcon(QIcon(str(icon_path)) if icon_path.exists() else self.overlay.windowIcon())
        menu = QMenu()
        toggle = QAction("Show / hide", menu)
        toggle.triggered.connect(self._toggle_overlay)
        menu.addAction(toggle)
        reload_action = QAction("Reload config", menu)
        reload_action.triggered.connect(self._reload_config)
        menu.addAction(reload_action)
        settings_action = QAction("Settings", menu)
        settings_action.triggered.connect(self._show_settings)
        menu.addAction(settings_action)
        menu.addSeparator()
        quit_action = QAction("Quit", menu)
        quit_action.triggered.connect(self.exit)
        menu.addAction(quit_action)
        tray.setContextMenu(menu)
        tray.activated.connect(self._tray_activated)
        tray.show()
        return tray

    def calibrate_non_empty(self) -> None:
        if self.settings is not None:
            self.settings.set_calibrating(True)
        self.config["general"]["binning_method"] = "default"
        self.output_processor = self._create_processor(self.config, "output")
        self.input_processor = self._create_processor(self.config, "input")
        self._calibration_until = time.monotonic() + 5.0
        self._calibration_activity = []

    def _collect_calibration(self, output, input_block) -> None:
        if self._calibration_until <= 0:
            return
        for processor, block in ((self.output_processor, output), (self.input_processor, input_block)):
            if block is not None:
                self._calibration_activity.append(processor.detect_non_empty(block))
        if time.monotonic() < self._calibration_until:
            return
        self._calibration_until = 0.0
        if not self._calibration_activity:
            if self.settings is not None:
                self.settings.set_calibrating(False)
            return
        activity = np.max(np.vstack(self._calibration_activity), axis=0)
        mask = activity.astype(bool).tolist()
        self.config["general"]["non_empty_bins"] = mask
        self.config["general"]["binning_method"] = "non_empty"
        save_config(self.config)
        self.output_processor = self._create_processor(self.config, "output")
        self.input_processor = self._create_processor(self.config, "input")
        self.overlay.apply_config(self.config)
        if self.settings is not None:
            self.settings.config = deepcopy(self.config)
            self.settings.set_calibrating(False)

    def _tray_activated(self, reason: QSystemTrayIcon.ActivationReason) -> None:
        if reason == QSystemTrayIcon.ActivationReason.Trigger:
            if self.settings is not None and self.settings.isVisible():
                self.settings.close()
            else:
                self._show_settings()

    def _toggle_overlay(self) -> None:
        self.overlay.setVisible(not self.overlay.isVisible())

    def _show_settings(self) -> None:
        if self.settings is None:
            self.settings = SettingsWindow(self.config, self._apply_config, self.overlay, self.overlay.recenter, self.overlay.set_data_edit_mode, self._reset_data_layout, self.calibrate_non_empty, self.restart, self.exit, self._apply_settings_state)
        else:
            self.settings.config = deepcopy(self.config)
            self.settings._restore_state()
        self.config["general"]["settings_window_open"] = True
        self.settings.config["general"]["settings_window_open"] = True
        save_config(self.config)
        save_config(self.settings.config)
        self.settings.show()
        self.settings.raise_()
        self.settings.activateWindow()

    def _apply_config(self, config: dict) -> None:
        old_config = self.config
        config["data"]["text_layout"] = deepcopy(self.overlay.config["data"]["text_layout"])
        self.config = config
        if self.settings is not None:
            self.settings.config = deepcopy(config)
        if old_config["general"]["startup_enabled"] != config["general"]["startup_enabled"]:
            set_startup(config["general"]["startup_enabled"])
        self.overlay.apply_config(config)
        self.overlay.update()
        analysis_changed = any(old_config["general"][key] != config["general"][key] for key in ("frequency_min", "frequency_max", "smoothing", "decay", "binning_method", "interpolate_to_zero", "non_empty_bins"))
        if old_config["output"] != config["output"] or old_config["input"] != config["input"] or old_config["visual"]["bar_count"] != config["visual"]["bar_count"] or analysis_changed:
            self.output_processor = self._create_processor(config, "output")
            self.input_processor = self._create_processor(config, "input")
        if old_config["output"] != config["output"] or old_config["input"] != config["input"]:
            self.capture.stop()
            self.capture.output = config["output"]
            self.capture.input = config["input"]
            self.capture.start()
        self.timer.start(1000 // config["visual"]["frame_rate"])

    def _apply_settings_state(self, config: dict) -> None:
        self.config["general"].update({
            key: deepcopy(config["general"][key])
            for key in ("settings_window_monitor", "settings_window_position", "settings_window_size", "settings_window_tab", "settings_window_open")
        })
        self.overlay.config = self.config
        if self.settings is not None:
            self.settings.config = deepcopy(self.config)
        save_config(self.config)

    def _reset_data_layout(self) -> None:
        self.overlay.reset_data_layout()
        self.config = self.overlay.config

    def restart(self) -> None:
        if self.settings is not None:
            state = deepcopy(self.settings.config)
            state["general"]["settings_window_open"] = self.settings.isVisible()
            self._apply_settings_state(state)
        self.overlay.save_state()
        QProcess.startDetached(sys.executable, sys.argv)
        QApplication.quit()

    def exit(self) -> None:
        if self.settings is not None:
            state = deepcopy(self.settings.config)
            state["general"]["settings_window_open"] = self.settings.isVisible()
            self._apply_settings_state(state)
        self.overlay.save_state()
        QApplication.quit()

    def _reload_config(self) -> None:
        try:
            self._apply_config(load_config())
        except ValueError as error:
            QMessageBox.warning(self.overlay, "Sound Visualizer", str(error))


def main() -> int:
    app = QApplication(sys.argv)
    app.setQuitOnLastWindowClosed(False)
    try:
        config = load_config()
        if config["general"]["startup_enabled"]:
            set_startup(True)
        visualizer = VisualizerApp(config)
    except Exception as error:
        traceback.print_exc()
        QMessageBox.critical(None, "Sound Visualizer", str(error))
        return 1
    app.aboutToQuit.connect(visualizer.stop)
    visualizer.start()
    if config["general"].get("settings_window_open", False):
        visualizer._show_settings()
    return app.exec()


if __name__ == "__main__":
    raise SystemExit(main())