rta

#!/usr/bin/env python3
import cProfile
import pstats
import numpy as np
import soundfile as sf
import argparse
from AudioSource import  RealTimeAudioSource, FileAudioSource
import AudioSource
import os

# setup argparse before opening pygame
argparse = argparse.ArgumentParser(description='Audio Visualizer')
argparse.add_argument('--mode', choices=['spl', 'acf'], default="", help='Mode to run the visualizer in')
argparse.add_argument('--source', type=str, help='Use test data instead of real-time audio')
argparse.add_argument('--windowsize', type=int, default=65536, help='Window size for FFT')
argparse.add_argument('--rotate', choices=['true','false','True','False'], default=None)
argparse.add_argument('--profile', action='store_true', help='Profile the code')

args = argparse.parse_args()
import AppMode
if args.rotate:
    args.rotate = args.rotate.lower()
    if args.rotate == 'true':
        AppMode.rotate = True
    elif args.rotate == 'false':
        AppMode.rotate = False

windowsize = int(args.windowsize)
if os.path.exists(args.source):
    audio_source = FileAudioSource(args.source) 
elif args.source == "-l":
    AudioSource.list_audio_devices()
    exit()
else:
    audio_source = RealTimeAudioSource(source=args.source)
next(audio_source) # read a chunk and discard - this is necessary to initialize samplerate


import pygame
pygame.init()
pygame.display.set_caption('Audio Visualizer')

class AudioVisualizer:
    def __init__(self):
        self.spl_mode = AppMode.SPLMode()
        if AudioSource.samplerate is None:
            raise RuntimeError("Samplerate not set")
        self.acf_mode = AppMode.ACFMode(windowsize, AudioSource.samplerate)
        self.current_mode = None
        self.switch_mode(args.mode)

    def switch_mode(self, mode_name):
        previous = self.current_mode
        if mode_name == 'acf':
            print(f'switching modes from {self.current_mode} to acf')
            self.current_mode = self.acf_mode
            self.current_mode.setup_plot()
        elif mode_name == 'spl':
            print(f'switching modes from {self.current_mode} to spl')
            self.current_mode = self.spl_mode
            self.current_mode.setup_plot()
        else:
            if self.current_mode == self.spl_mode:
                print(f'switching modes from {self.current_mode} to acf')
                self.current_mode = self.acf_mode
            else:
                print(f'switching modes from {self.current_mode} to spl')
                self.current_mode = self.spl_mode
        print(f'previous mode: {previous}, current mode: {self.current_mode}')
        self.redraw()

    def redraw(self):
        self.current_mode.setup_plot()

    def process_audio_chunk(self, audio_chunk):
        if self.spl_mode is not None:
            self.spl_mode.process_data(audio_chunk)
        if self.acf_mode is not None:
            self.acf_mode.process_data(audio_chunk)
        self.current_mode.update_plot()

def scan_buttons():
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            raise KeyboardInterrupt
        elif event.type == pygame.KEYDOWN:
            return event.key
    return False

def main():
    visualizer = AudioVisualizer()
    run = True
    while run:
        try:
            button_press = scan_buttons()
            if button_press:
                print('got keypress')
                visualizer.switch_mode(button_press)
            chunk = next(audio_source)
            visualizer.process_audio_chunk(chunk)
            pygame.display.flip()
            pygame.time.wait(10)
        except KeyboardInterrupt:
            run = False

if __name__ == "__main__":
    if args.profile:
        import cProfile, pstats
        profiler = cProfile.Profile()
        profiler.enable()
    main()
    if args.profile:
        profiler.disable()
        with open('profile_output.txt', 'w') as f:
            ps = pstats.Stats(profiler, stream=f)
            ps.strip_dirs().sort_stats('cumulative').print_stats(20)
        print("Profiling data saved to 'profile_output.txt'")