import gradio as gr
import torch
import torchaudio
import numpy as np
from vocos import Vocos
import logging
from datetime import datetime
import matplotlib.pyplot as plt
from io import BytesIO
import base64

# Configure detailed logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# Initialize Vocos model
device = "cuda" if torch.cuda.is_available() else "cpu"
logger.info("="*60)
logger.info("APPLICATION STARTUP")
logger.info("="*60)
logger.info(f"PyTorch version: {torch.__version__}")
logger.info(f"Torchaudio version: {torchaudio.__version__}")
logger.info(f"Device: {device}")
logger.info(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    logger.info(f"CUDA device: {torch.cuda.get_device_name(0)}")

logger.info("Loading Vocos model from 'charactr/vocos-mel-24khz'...")
vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz").to(device)
logger.info("✅ Model loaded successfully!")
logger.info("="*60)

def calculate_snr(original, reconstructed):
    """Calculate Signal-to-Noise Ratio"""
    noise = original - reconstructed
    signal_power = np.mean(original ** 2)
    noise_power = np.mean(noise ** 2)
    if noise_power == 0:
        return float('inf')
    return 10 * np.log10(signal_power / noise_power)

def plot_waveforms(original, reconstructed, sample_rate=24000):
    """Create comparison plot of original vs reconstructed waveforms"""
    fig, axes = plt.subplots(3, 1, figsize=(12, 8))
    
    # Time axis
    time_orig = np.arange(len(original)) / sample_rate
    time_recon = np.arange(len(reconstructed)) / sample_rate
    
    # Original waveform
    axes[0].plot(time_orig, original, color='#0080FF', alpha=0.7, linewidth=0.5)
    axes[0].set_title('Original Waveform', fontsize=12, fontweight='bold')
    axes[0].set_ylabel('Amplitude')
    axes[0].grid(True, alpha=0.3)
    axes[0].set_xlim(0, max(time_orig[-1], time_recon[-1]))
    
    # Reconstructed waveform
    axes[1].plot(time_recon, reconstructed, color='#FF4500', alpha=0.7, linewidth=0.5)
    axes[1].set_title('Reconstructed Waveform', fontsize=12, fontweight='bold')
    axes[1].set_ylabel('Amplitude')
    axes[1].grid(True, alpha=0.3)
    axes[1].set_xlim(0, max(time_orig[-1], time_recon[-1]))
    
    # Overlay comparison (first 2 seconds)
    max_samples = min(len(original), len(reconstructed), sample_rate * 2)
    axes[2].plot(time_orig[:max_samples], original[:max_samples], 
                 color='#0080FF', alpha=0.6, label='Original', linewidth=0.8)
    axes[2].plot(time_recon[:max_samples], reconstructed[:max_samples], 
                 color='#FF4500', alpha=0.6, label='Reconstructed', linewidth=0.8)
    axes[2].set_title('Overlay Comparison (First 2 seconds)', fontsize=12, fontweight='bold')
    axes[2].set_xlabel('Time (s)')
    axes[2].set_ylabel('Amplitude')
    axes[2].legend(loc='upper right')
    axes[2].grid(True, alpha=0.3)
    
    plt.tight_layout()
    return fig

def plot_spectrograms(original, reconstructed, sample_rate=24000):
    """Create spectrogram comparison"""
    fig, axes = plt.subplots(2, 1, figsize=(12, 8))
    
    # Original spectrogram
    axes[0].specgram(original, Fs=sample_rate, cmap='viridis', NFFT=1024)
    axes[0].set_title('Original Spectrogram', fontsize=12, fontweight='bold')
    axes[0].set_ylabel('Frequency (Hz)')
    
    # Reconstructed spectrogram
    axes[1].specgram(reconstructed, Fs=sample_rate, cmap='viridis', NFFT=1024)
    axes[1].set_title('Reconstructed Spectrogram', fontsize=12, fontweight='bold')
    axes[1].set_xlabel('Time (s)')
    axes[1].set_ylabel('Frequency (Hz)')
    
    plt.tight_layout()
    return fig

def process_audio(audio_input, enable_noise_reduction, volume_boost, trim_silence):
    """
    Process uploaded audio through Vocos reconstruction with optional enhancements
    
    Args:
        audio_input: tuple of (sample_rate, audio_data) from Gradio
        enable_noise_reduction: bool to apply noise reduction
        volume_boost: float for volume adjustment (dB)
        trim_silence: bool to trim leading/trailing silence
    
    Returns:
        tuple: (sample_rate, reconstructed_audio), stats_text, waveform_plot, spectrogram_plot
    """
    logger.info("\n" + "="*60)
    logger.info("NEW AUDIO PROCESSING REQUEST")
    logger.info("="*60)
    
    if audio_input is None:
        logger.warning("❌ No audio input provided")
        return None, "⚠️ No audio provided", None, None
    
    try:
        # Unpack input
        sample_rate, audio_data = audio_input
        logger.info(f"📥 Input received:")
        logger.info(f"   - Sample rate: {sample_rate} Hz")
        logger.info(f"   - Audio shape: {audio_data.shape}")
        logger.info(f"   - Audio dtype: {audio_data.dtype}")
        logger.info(f"   - Audio range: [{audio_data.min():.4f}, {audio_data.max():.4f}]")
        
        # Convert to float32 and normalize to [-1, 1]
        original_dtype = audio_data.dtype
        if audio_data.dtype == np.int16:
            logger.info("🔄 Converting from int16 to float32")
            audio_data = audio_data.astype(np.float32) / 32768.0
        elif audio_data.dtype == np.int32:
            logger.info("🔄 Converting from int32 to float32")
            audio_data = audio_data.astype(np.float32) / 2147483648.0
        else:
            logger.info("🔄 Converting to float32")
            audio_data = audio_data.astype(np.float32)
        
        logger.info(f"   - Normalized range: [{audio_data.min():.4f}, {audio_data.max():.4f}]")
        
        # Handle stereo to mono conversion
        if len(audio_data.shape) > 1:
            logger.info(f"🔄 Converting stereo to mono (shape: {audio_data.shape})")
            audio_data = audio_data.mean(axis=1)
            logger.info(f"   - New shape: {audio_data.shape}")
        else:
            logger.info("✅ Audio is already mono")
        
        # Trim silence if requested
        if trim_silence:
            logger.info("✂️ Trimming silence...")
            threshold = 0.01
            non_silent = np.where(np.abs(audio_data) > threshold)[0]
            if len(non_silent) > 0:
                audio_data = audio_data[non_silent[0]:non_silent[-1]+1]
                logger.info(f"   - Trimmed to {len(audio_data)} samples")
        
        # Convert to torch tensor
        waveform = torch.from_numpy(audio_data).float().unsqueeze(0)  # (1, N)
        logger.info(f"🔄 Converted to torch tensor: {waveform.shape}")
        logger.info(f"   - Duration: {waveform.shape[1] / sample_rate:.2f} seconds")
        
        # Store original for comparison
        original_waveform = waveform.clone()
        original_sample_rate = sample_rate
        
        # Resample to 24kHz if needed
        if sample_rate != 24000:
            logger.info(f"🔄 Resampling from {sample_rate} Hz to 24000 Hz")
            resampler = torchaudio.transforms.Resample(
                orig_freq=sample_rate, 
                new_freq=24000
            )
            waveform = resampler(waveform)
            original_waveform = resampler(original_waveform)
            logger.info(f"   - Resampled shape: {waveform.shape}")
            logger.info(f"   - New duration: {waveform.shape[1] / 24000:.2f} seconds")
        else:
            logger.info("✅ Sample rate already 24000 Hz, no resampling needed")
        
        # Move to device
        logger.info(f"🔄 Moving tensor to {device}")
        waveform = waveform.to(device)
        
        # Run Vocos reconstruction
        logger.info("🎵 Running Vocos reconstruction...")
        start_time = datetime.now()
        with torch.inference_mode():
            reconstructed = vocos(waveform)
        end_time = datetime.now()
        processing_time = (end_time - start_time).total_seconds()
        
        logger.info(f"✅ Reconstruction complete!")
        logger.info(f"   - Processing time: {processing_time:.3f} seconds")
        logger.info(f"   - Output shape: {reconstructed.shape}")
        logger.info(f"   - Output range: [{reconstructed.min():.4f}, {reconstructed.max():.4f}]")
        
        # Convert back to numpy
        output_audio = reconstructed.squeeze(0).cpu().numpy()
        
        # Apply noise reduction if enabled
        if enable_noise_reduction:
            logger.info("🔇 Applying noise reduction...")
            # Simple spectral gating
            threshold = np.percentile(np.abs(output_audio), 5)
            mask = np.abs(output_audio) > threshold
            output_audio = output_audio * mask
            logger.info(f"   - Noise threshold: {threshold:.6f}")
        
        # Apply volume boost if requested
        if volume_boost != 0:
            logger.info(f"🔊 Applying volume boost: {volume_boost:.1f} dB")
            gain = 10 ** (volume_boost / 20)
            output_audio = output_audio * gain
            # Prevent clipping
            max_val = np.abs(output_audio).max()
            if max_val > 1.0:
                output_audio = output_audio / max_val
                logger.info(f"   - Normalized to prevent clipping (peak: {max_val:.3f})")
        
        logger.info(f"🔄 Final output: {output_audio.shape}")
        logger.info(f"   - Output dtype: {output_audio.dtype}")
        
        # Calculate quality metrics
        original_np = original_waveform.squeeze(0).cpu().numpy()
        min_len = min(len(original_np), len(output_audio))
        
        original_energy = np.mean(original_np[:min_len] ** 2)
        reconstructed_energy = np.mean(output_audio[:min_len] ** 2)
        snr = calculate_snr(original_np[:min_len], output_audio[:min_len])
        correlation = np.corrcoef(original_np[:min_len], output_audio[:min_len])[0, 1]
        
        logger.info(f"📊 Quality metrics:")
        logger.info(f"   - Original energy: {original_energy:.6f}")
        logger.info(f"   - Reconstructed energy: {reconstructed_energy:.6f}")
        logger.info(f"   - Energy ratio: {reconstructed_energy/original_energy:.4f}")
        logger.info(f"   - SNR: {snr:.2f} dB")
        logger.info(f"   - Correlation: {correlation:.4f}")
        
        # Create statistics text
        stats_text = f"""
### 📊 Processing Statistics

**Input Information:**
- Sample Rate: {original_sample_rate} Hz
- Duration: {len(original_np) / 24000:.2f} seconds
- Channels: Mono
- Samples: {len(original_np):,}

**Processing Options:**
- Noise Reduction: {'✅ Enabled' if enable_noise_reduction else '❌ Disabled'}
- Volume Boost: {volume_boost:+.1f} dB
- Trim Silence: {'✅ Enabled' if trim_silence else '❌ Disabled'}

**Quality Metrics:**
- Processing Time: {processing_time:.3f} seconds
- Signal-to-Noise Ratio: {snr:.2f} dB
- Correlation: {correlation:.4f}
- Energy Ratio: {reconstructed_energy/original_energy:.4f}

**Output Information:**
- Sample Rate: 24000 Hz
- Samples: {len(output_audio):,}
- Peak Amplitude: {np.abs(output_audio).max():.4f}
"""
        
        # Generate plots
        logger.info("📈 Generating visualizations...")
        waveform_plot = plot_waveforms(original_np, output_audio, 24000)
        spectrogram_plot = plot_spectrograms(original_np, output_audio, 24000)
        
        logger.info("="*60)
        logger.info("✅ PROCESSING COMPLETE")
        logger.info("="*60 + "\n")
        
        return (24000, output_audio), stats_text, waveform_plot, spectrogram_plot
    
    except Exception as e:
        logger.error("="*60)
        logger.error("❌ ERROR DURING PROCESSING")
        logger.error("="*60)
        logger.error(f"Error type: {type(e).__name__}")
        logger.error(f"Error message: {str(e)}")
        logger.exception("Full traceback:")
        logger.error("="*60 + "\n")
        raise gr.Error(f"Failed to process audio: {str(e)}")

# Create Gradio interface
with gr.Blocks(title="Vocos Audio Reconstruction", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎵 Vocos Audio Reconstruction Studio
    
    Upload an audio file to hear it reconstructed through the **Vocos neural vocoder** with advanced processing options.
    
    **Features:**
    - 🎯 High-quality neural audio reconstruction
    - 🔇 Optional noise reduction
    - 🔊 Volume boost control
    - ✂️ Automatic silence trimming
    - 📊 Detailed quality metrics
    - 📈 Visual waveform & spectrogram analysis
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📥 Input")
            audio_input = gr.Audio(
                label="Upload Audio File or Record",
                type="numpy",
                sources=["upload", "microphone"]
            )
            
            gr.Markdown("### ⚙️ Processing Options")
            enable_noise_reduction = gr.Checkbox(
                label="Enable Noise Reduction",
                value=False,
                info="Apply spectral gating to reduce background noise"
            )
            
            volume_boost = gr.Slider(
                minimum=-20,
                maximum=20,
                value=0,
                step=1,
                label="Volume Boost (dB)",
                info="Adjust output volume (-20 to +20 dB)"
            )
            
            trim_silence = gr.Checkbox(
                label="Trim Silence",
                value=False,
                info="Remove leading and trailing silence"
            )
            
            process_btn = gr.Button("🔄 Reconstruct Audio", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Output")
            audio_output = gr.Audio(
                label="Reconstructed Audio",
                type="numpy"
            )
            
            stats_output = gr.Markdown("*Process audio to see statistics*")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 📈 Waveform Comparison")
            waveform_plot = gr.Plot(label="Waveform Analysis")
        
        with gr.Column():
            gr.Markdown("### 🎼 Spectrogram Comparison")
            spectrogram_plot = gr.Plot(label="Spectrogram Analysis")
    
    gr.Markdown(f"""
    ---
    ### ℹ️ Technical Information
    
    **Model Details:**
    - Model: [Vocos Mel-24kHz](https://huggingface.co/charactr/vocos-mel-24khz)
    - Architecture: Neural vocoder with mel-spectrogram backbone
    - Target Sample Rate: 24 kHz
    
    **System Information:**
    - Device: **{device.upper()}**
    - PyTorch: {torch.__version__}
    - Torchaudio: {torchaudio.__version__}
    
    **Supported Formats:**
    - Input: WAV, MP3, FLAC, OGG, M4A (any format supported by your browser)
    - Output: WAV at 24 kHz
    
    ### 📝 Quality Metrics Explained
    - **SNR (Signal-to-Noise Ratio):** Higher is better (>20 dB is good)
    - **Correlation:** Closer to 1.0 means higher similarity
    - **Energy Ratio:** Closer to 1.0 means similar loudness
    
    ### 💡 Tips
    - For best results, use clear audio recordings
    - Enable noise reduction for recordings with background noise
    - Use volume boost for quiet recordings
    - Check the visualizations to compare quality
    """)
    
    # Set up event handlers
    process_btn.click(
        fn=process_audio,
        inputs=[audio_input, enable_noise_reduction, volume_boost, trim_silence],
        outputs=[audio_output, stats_output, waveform_plot, spectrogram_plot]
    )
    
    # Auto-process on upload with default settings
    audio_input.change(
        fn=lambda x: process_audio(x, False, 0, False),
        inputs=audio_input,
        outputs=[audio_output, stats_output, waveform_plot, spectrogram_plot]
    )
    
    # Examples
    gr.Markdown("### 🎯 Quick Start Examples")
    gr.Examples(
        examples=[
            ["Example: Upload your audio file above"],
        ],
        inputs=audio_input,
        label="Try these settings:"
    )

# Launch the app
if __name__ == "__main__":
    logger.info("🚀 Launching Gradio interface...")
    demo.launch()
    logger.info("👋 Application shutdown")