import librosa
import librosa.display
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import butter, filtfilt
import sounddevice as sd
from tkinter import Tk
from tkinter.filedialog import askopenfilename
import threading
from matplotlib.widgets import Cursor

# Function to apply a band-pass filter
def bandpass_filter(data, lowcut, highcut, fs, order=5):
    nyquist = 0.5 * fs
    low = lowcut / nyquist
    high = highcut / nyquist
    b, a = butter(order, [low, high], btype='band')
    y = filtfilt(b, a, data)
    return y

# Function to detect snare hits
def detect_snare_hits(audio_file):
    # Load the audio file
    y, sr = librosa.load(audio_file, sr=None)

    # Apply a band-pass filter to focus on snare frequencies
    y_filtered = bandpass_filter(y, 1000, 5000, sr)

    # Detect onsets (potential snare hits)
    onset_frames = librosa.onset.onset_detect(y=y_filtered, sr=sr, backtrack=True)
    onset_times = librosa.frames_to_time(onset_frames, sr=sr)

    return y, sr, onset_times

# Function to visualize only snare hits with zoom functionality
def visualize_snare_hits(y, sr, onset_times, output_file=None):
    fig, ax = plt.subplots(figsize=(14, 5))
    
    # Plot snare hits only
    ax.plot(onset_times, np.zeros_like(onset_times), 'go', label='Snare Hit')

    # Add cursor for interactivity
    cursor = Cursor(ax, useblit=True, color='red', linewidth=1)

    def on_click(event):
        if event.inaxes == ax:
            time = event.xdata
            closest_time = min(onset_times, key=lambda x: abs(x - time))
            ax.annotate(f'Snare Hit\nTime: {closest_time:.2f}s', xy=(closest_time, 0), 
                        xytext=(closest_time + 0.5, 0.1),
                        arrowprops=dict(facecolor='black', shrink=0.05))
            plt.draw()

    fig.canvas.mpl_connect('button_press_event', on_click)

    ax.set_title('Detected Snare Hits')
    ax.set_xlabel('Time (s)')
    ax.set_ylabel('Amplitude')
    ax.legend(loc="upper right")

    # Save the plot if an output file is specified, otherwise show it
    if output_file:
        plt.savefig(output_file)
    else:
        plt.show()

# Function to play/pause the audio
def play_audio(y, sr):
    global audio_playing
    play_pos = 0

    def callback(outdata, frames, time, status):
        nonlocal play_pos
        if audio_playing:
            if play_pos + frames > len(y):
                frames = len(y) - play_pos
                outdata[:frames] = y[play_pos:play_pos+frames].reshape(-1, 1)
                outdata[frames:] = 0
            else:
                outdata[:] = y[play_pos:play_pos+frames].reshape(-1, 1)
            play_pos += frames
        else:
            outdata[:] = np.zeros((frames, 1))
    
    with sd.OutputStream(samplerate=sr, channels=1, callback=callback):
        while audio_playing and play_pos < len(y):
            sd.sleep(100)

# Function to toggle play/pause
def toggle_play_pause():
    global audio_playing
    audio_playing = not audio_playing

# Main function to run the app
def main():
    global audio_playing

    # Open a file dialog to select the audio file
    Tk().withdraw()  # Prevent the root window from appearing
    audio_file = askopenfilename(filetypes=[("Mp3", "*.mp3"), ("All files", "*.*")])

    if not audio_file:
        print("No file selected.")
        return

    # Detect snare hits
    y, sr, onset_times = detect_snare_hits(audio_file)

    # Visualize the snare hits
    visualize_snare_hits(y, sr, onset_times)

    # Play/Pause functionality
    audio_playing = False
    play_thread = threading.Thread(target=play_audio, args=(y, sr))
    play_thread.start()

    while True:
        command = input("Enter 'play' to play, 'pause' to pause, or 'exit' to exit: ").strip().lower()
        if command == 'play':
            toggle_play_pause()
        elif command == 'pause':
            toggle_play_pause()
        elif command == 'exit':
            audio_playing = False
            play_thread.join()
            break
        else:
            print("Invalid command. Please enter 'play', 'pause', or 'exit'.")

# Run the app
if __name__ == '__main__':
    main()