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, rpm, rumble_level, hiss_level):
    # Parameters
    sample_rate = 44100  # sample rate in Hz
    num_samples = int(duration * sample_rate)

    # Calculate wow and flutter based on RPM (more realistic variations)
    wow_freq = rpm / 60 * 0.1  # Example: 0.1 Hz variation for 33 1/3 RPM
    flutter_freq = np.random.uniform(1, 10)  # Flutter frequency in Hz
    wow_flutter = 0.001 * np.sin(2 * np.pi * wow_freq * np.arange(num_samples) / sample_rate) + \
                  0.0005 * np.sin(2 * np.pi * flutter_freq * np.arange(num_samples) / sample_rate)

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

    # Apply band-pass filter to pink noise (optimized using sosfilt)
    nyquist_rate = 0.5 * sample_rate
    low = lowcut / nyquist_rate
    high = highcut / nyquist_rate
    sos = signal.butter(1, [low, high], btype='band', output='sos')
    pink_noise = signal.sosfilt(sos, pink_noise)

    # Generate low-frequency rumble (optimized using sosfilt, controllable level)
    rumble_noise = np.random.normal(0, 1, num_samples)
    sos = signal.butter(2, 0.005, btype='low', output='sos')  # Lower cutoff for more rumble
    rumble_noise = signal.sosfilt(sos, rumble_noise) * rumble_level  # Adjust amplitude

    # Generate high-frequency hiss (optimized using sosfilt, controllable level)
    hiss_noise = np.random.normal(0, 1, num_samples)
    sos = signal.butter(2, 0.4, btype='high', output='sos')
    hiss_noise = signal.sosfilt(sos, hiss_noise) * hiss_level  # Adjust amplitude

    # Generate pops with varying loudness and frequency
    num_pops = int(duration * pop_rate)
    pop_times = np.random.randint(0, num_samples, num_pops)
    pop_data = np.zeros(num_samples)
    pop_amplitudes = np.random.uniform(0.05, 0.2, num_pops)
    pop_data[pop_times] = pop_amplitudes

    # Apply pop filter (optimized by applying to the entire signal at once)
    pop_filter_freq = np.random.uniform(0.05, 0.2)
    b, a = signal.butter(4, pop_filter_freq)
    pop_data = signal.lfilter(b, a, pop_data)

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

    # Apply wow and flutter
    vinyl_sound = vinyl_sound * (1 + wow_flutter)

    # 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, rpm, rumble_level, hiss_level):
    data, sample_rate = generate_vinyl_sound(noise_ratio, lowcut, highcut, duration, pop_rate, rpm, rumble_level, hiss_level)
    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)"),
        gr.inputs.Number(default=33.33, label="RPM"),
        gr.inputs.Slider(minimum=0, maximum=10, default=0.1, step=0.01, label="Rumble Level"),
        gr.inputs.Slider(minimum=0, maximum=10, default=0.05, step=0.01, label="Hiss Level")
    ],
    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, pop rate, RPM, rumble level, and hiss level to modify the sound.",
    examples=[
        [0.0005, 300, 5000, 30, 10, 33.33, 0.1, 0.05],  # Example for 33 1/3 RPM
        [0.001, 500, 4000, 30, 50, 45.0, 0.15, 0.08],      # Example for 45 RPM
        [0.0008, 400, 6000, 20, 15, 78.0, 0.2, 0.1]      # Example for 78 RPM
    ]
)

iface.launch()