import gradio as gr
import numpy as np
import librosa

# Function to load audio using librosa
def load_audio(file):
    audio, sr = librosa.load(file, sr=None)  # Load the audio file
    return audio, sr

# Function to get a relevant audio segment based on onset detection
def get_segment(audio, sr):
    onset_env = librosa.onset.onset_strength(y=audio, sr=sr)
    onset_frames = librosa.onset.onset_detect(onset_envelope=onset_env, sr=sr, backtrack=True)

    if len(onset_frames) == 0:
        return audio[:2048], sr  # Return the first 2048 samples if no onsets

    # Calculate energy over time
    energy = np.array([np.sum(np.abs(audio[i:i + 2048]**2)) for i in range(0, len(audio), 2048)])
    
    # Threshold for musical relevance
    energy_threshold = np.mean(energy) + np.std(energy)
    relevant_onsets = [f for f in onset_frames if f < len(energy) and energy[f] > energy_threshold]

    start_frame = relevant_onsets[0] if relevant_onsets else onset_frames[0]
    start_sample = librosa.frames_to_samples(start_frame)

    segment_length = sr  # Length of segment in samples (1 second)
    end_sample = min(start_sample + segment_length, len(audio))

    return audio[start_sample:end_sample], sr

# Function to extend music by adding silence
def extend_music(file, added_minutes):
    audio, sr = load_audio(file)
    segment, sr = get_segment(audio, sr)

    # Calculate additional samples
    additional_samples = int(added_minutes * 60 * sr)
    extended_audio = np.concatenate([segment, np.zeros(additional_samples)])  # Add silence

    # Normalize audio to the range of [-1, 1]
    extended_audio = extended_audio / np.max(np.abs(extended_audio)) if np.max(np.abs(extended_audio)) > 0 else extended_audio

    # Return the audio as a NumPy array (float32) and the sample rate as an integer
    return extended_audio.astype(np.float32), sr  # Ensure it's a float32 NumPy array

# Gradio UI setup
with gr.Blocks() as app:
    gr.Markdown("# Audio Extender")
    
    audio_input = gr.Audio(type="filepath", label="Upload Audio File")
    added_minutes = gr.Slider(minimum=0, maximum=10, label="Additional Minutes")
    audio_output = gr.Audio(type="numpy", label="Extended Audio")
    
    submit_button = gr.Button("Extend Audio")
    submit_button.click(extend_music, inputs=[audio_input, added_minutes], outputs=audio_output)

# Launch the app with a public link
app.launch(share=True)