import numpy as np
import soundfile as sf
from scipy import signal
import gradio as gr
import shutil
import tempfile
from IPython.display import Audio, display

def generate_vinyl_sound(noise_ratio, lowcut, highcut, duration, pop_rate):
    # Parameters
    sample_rate = 44100  # sample rate in Hz
    num_samples = int(duration * sample_rate)

    # Generate pink noise
    pink_noise = np.random.normal(0, 1, num_samples)
    b, a = signal.butter(2, [0.002, 0.4], btype='band')
    pink_noise = signal.lfilter(b, a, pink_noise)

    # Apply band-pass filter to pink noise
    nyquist_rate = 0.5 * sample_rate
    low = lowcut / nyquist_rate
    high = highcut / nyquist_rate
    b, a = signal.butter(1, [low, high], btype='band')
    pink_noise = signal.lfilter(b, a, pink_noise)

    # Generate low-frequency rumble
    rumble_noise = np.random.normal(0, 1, num_samples)
    b, a = signal.butter(2, 0.002, btype='low')
    rumble_noise = signal.lfilter(b, a, rumble_noise)

    # Generate high-frequency hiss
    hiss_noise = np.random.normal(0, 1, num_samples)
    b, a = signal.butter(2, 0.4, btype='high')
    hiss_noise = signal.lfilter(b, a, hiss_noise)

    # Generate pops
    num_pops = int(duration * pop_rate)
    pop_times = np.random.randint(0, num_samples, num_pops)
    pop_data = np.zeros(num_samples)
    pop_data[pop_times] = np.random.normal(0, 1, num_pops) * 0.2  # random loudness

    # Create a simple low-pass filter to make the pops sound more like clicks
    b, a = signal.butter(4, 0.1)
    pop_data = signal.lfilter(b, a, pop_data)

    # Combine the noises and pops
    vinyl_sound = noise_ratio * (pink_noise + 0.05 * rumble_noise + 0.05 * hiss_noise) + (1 - noise_ratio) * pop_data

    # Normalize to between -1 and 1
    vinyl_sound /= np.max(np.abs(vinyl_sound))

    return vinyl_sound.astype(np.float32), sample_rate

def convert_to_wav(data, sample_rate):
    # Normalize to between -1 and 1
    data /= np.max(np.abs(data))

    # Save to a temporary .wav file
    temp_file = tempfile.mktemp(suffix=".wav")
    sf.write(temp_file, data, sample_rate)

    return temp_file

def play_and_download_sound(noise_ratio, lowcut, highcut, duration, pop_rate):
    data, sample_rate = generate_vinyl_sound(noise_ratio, lowcut, highcut, duration, pop_rate)
    temp_file = convert_to_wav(data, sample_rate)
    audio = Audio(temp_file)
    display(audio)

    return audio

iface = gr.Interface(
    fn=play_and_download_sound,
    inputs=[
        gr.inputs.Slider(minimum=0, maximum=1, default=0.001, step=0.001, label="Noise Ratio"),
        gr.inputs.Slider(minimum=20, maximum=20000, default=300, step=10, label="Lowcut Frequency (Hz)"),
        gr.inputs.Slider(minimum=20, maximum=20000, default=5000, step=10, label="Highcut Frequency (Hz)"),
        gr.inputs.Slider(minimum=1, maximum=60, default=10, step=1, label="Duration (seconds)"),
        gr.inputs.Slider(minimum=1, maximum=100, default=10, step=1, label="Pop Rate (pops per second)")
    ],
    outputs=[
        gr.outputs.Audio(type="numpy", label="Vinyl Sound"),
        gr.outputs.Button(label="Download Vinyl Sound", type="button", callback=download_sound)
    ],
    title="Vinyl Sound Generator",
    description="Generate a synthetic vinyl sound using pink noise, rumble, hiss, and pops. Adjust the noise ratio, bandpass frequencies, duration, and pop rate to modify the sound.",
    examples=[
        [0.001, 300, 5000, 10, 10],
        [0.002, 500, 4000, 15, 20],
        [0.005, 200, 6000, 20, 30]
    ]
)

iface.launch()