import gradio as gr
import librosa
import numpy as np
import soundfile as sf

def reverse_segments_dynamic(audio_path):
    # Load audio with original sampling rate
    y, sr = librosa.load(audio_path, sr=None)

    # Get a dynamic tempo estimate
    onset_env = librosa.onset.onset_strength(y=y, sr=sr)
    tempo, beat_frames = librosa.beat.beat_track(y=y, sr=sr, onset_envelope=onset_env, trim=False)

    # Calculate a full tempogram (BPM over time)
    oenv = librosa.onset.onset_strength(y=y, sr=sr)
    tempogram = librosa.feature.tempogram(onset_envelope=oenv, sr=sr)
    bpm_changes = librosa.beat.tempo(onset_envelope=oenv, sr=sr, aggregate=None)  # Array of BPM values over time

    # Convert beat positions to sample indices
    beat_samples = librosa.frames_to_samples(beat_frames)

    # Reverse the entire audio first
    y_reversed = y[::-1]

    # Cut into segments based on detected beats
    segments = [y_reversed[beat_samples[i]:beat_samples[i+1]] for i in range(len(beat_samples)-1)]

    # Reverse each segment back to forward
    processed_audio = np.concatenate([segment[::-1] for segment in segments])

    # Save the processed file
    output_path = "output.wav"
    sf.write(output_path, processed_audio, sr)

    return output_path

iface = gr.Interface(
    fn=reverse_segments_dynamic,
    inputs=gr.Audio(type="filepath"),
    outputs=gr.Audio(),
    title="Dynamic BPM Beat-Reversed Music",
    description="Detects BPM changes over time and reverses music accordingly, keeping beats aligned."
)

iface.launch()