import gradio as gr
import numpy as np
from scipy.io.wavfile import read, write
import tempfile

def dynamic_vibration_simulation(file):
    try:
        # Read the uploaded WAV file
        sample_rate, data = read(file)
        
        # Debug: Print sample rate and shape
        print(f"Sample Rate: {sample_rate}, Data Shape: {data.shape}")

        # Convert to mono if stereo
        if len(data.shape) > 1:
            data = data.mean(axis=1)
        
        # Normalize the data to range [-1, 1]
        data = data / np.max(np.abs(data))
        
        # Parameters for vibration simulation
        segment_duration = 0.1  # Analyze 0.1-second segments
        segment_length = int(sample_rate * segment_duration)

        # Initialize the low-frequency waveform
        low_freq_wave = []

        for i in range(0, len(data), segment_length):
            segment = data[i:i + segment_length]

            if len(segment) < segment_length:
                break  # Skip incomplete segment at the end

            # Calculate amplitude (RMS value) of the segment
            rms_amplitude = np.sqrt(np.mean(segment**2))
            print(f"RMS Amplitude for segment {i // segment_length}: {rms_amplitude}")

            # Map amplitude to vibration parameters
            if rms_amplitude > 0.5:
                freq = 50  # Strong vibration frequency (50 Hz)
                duration = segment_duration
            elif rms_amplitude > 0.2:
                freq = 30  # Medium vibration frequency (30 Hz)
                duration = segment_duration
            else:
                freq = 0  # Silence (no vibration)
                duration = segment_duration

            # Generate waveform for this segment
            t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
            if freq > 0:
                wave = 0.5 * np.sin(2 * np.pi * freq * t)
            else:
                wave = np.zeros_like(t)  # Silence
            
            # Append the generated wave
            low_freq_wave.extend(wave)
        
        # Convert the waveform to 16-bit PCM
        low_freq_wave = (np.array(low_freq_wave) * 32767).astype(np.int16)
        
        # Save the generated waveform to a temporary file
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
        write(temp_file.name, sample_rate, low_freq_wave)

        print("Output WAV file generated successfully.")
        return temp_file.name
    except Exception as e:
        print(f"Error: {e}")
        return "Error in processing the uploaded WAV file."

# Gradio Interface
interface = gr.Interface(
    fn=dynamic_vibration_simulation,
    inputs=gr.Audio(label="Upload a WAV file", type="filepath"),
    outputs=gr.Audio(label="Dynamic Vibration Audio"),
    title="Dynamic Vibration Simulation",
    description="Upload a WAV file. The vibration simulation dynamically adapts based on the sound's amplitude and variation."
)

interface.launch()