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
from typing import Tuple

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)
    return temp_file, temp_file

iface = gr.Interface(
    fn=play_and_download_sound,
    inputs=[
        gr.inputs.Slider(minimum=0, maximum=0.005, default=0.0005, step=0.0001, 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=600, default=30, step=1, label="Duration (seconds)"),
        gr.inputs.Slider(minimum=1, maximum=180, default=10, step=1, label="Pop Rate (pops per second)")
    ],
    outputs=[
        gr.outputs.Audio(type="numpy", label="Vinyl Sound"),
        gr.outputs.HTML("<button onclick='downloadSound()'>Download Vinyl Sound</button>"
                        "<script>function downloadSound(){"
                        "var link = document.createElement('a');"
                        "link.href = '/download';"
                        "link.download = 'vinyl_sound.wav';"
                        "link.click();}</script>")
    ],
    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.0005, 300, 5000, 30, 10],
        [0.001, 500, 4000, 30, 50],
        [0.002, 200, 6000, 30, 100]
    ]
)

iface.launch()