import gradio as gr
import subprocess
import os
import tempfile
import librosa
import librosa.display
import matplotlib.pyplot as plt
import numpy as np
import scipy.ndimage
from pathlib import Path
import logging
import warnings
import shutil
import requests
from spleeter.separator import Separator
from spleeter import SpleeterError

# Set matplotlib backend for web display
plt.switch_backend('Agg')
warnings.filterwarnings('ignore')

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

class AudioAnalyzer:
    def __init__(self):
        self.temp_dir = Path(tempfile.mkdtemp())
        self.spleeter_model_dir = self.temp_dir / "pretrained_models"
        self.spleeter_model_dir.mkdir(exist_ok=True)
        logger.info(f"Temporary directory: {self.temp_dir}")

    def cleanup(self):
        """Clean up temporary directory."""
        if self.temp_dir.exists():
            shutil.rmtree(self.temp_dir)
            logger.info(f"Cleaned up temporary directory: {self.temp_dir}")

    def download_spleeter_model(self, model_url="https://github.com/deezer/spleeter/releases/download/v1.4.0/2stems.tar.gz", progress=gr.Progress()):
        """Download Spleeter pretrained model, handling redirects."""
        model_path = self.spleeter_model_dir / "2stems.tar.gz"
        if model_path.exists():
            logger.info("Spleeter model already downloaded.")
            return str(model_path)

        progress(0.1, desc="Downloading Spleeter model...")
        try:
            response = requests.get(model_url, stream=True, allow_redirects=True)
            if response.status_code != 200:
                raise Exception(f"Failed to download model: HTTP {response.status_code}")

            with open(model_path, "wb") as f:
                total_size = int(response.headers.get('content-length', 0))
                downloaded = 0
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        downloaded += len(chunk)
                        if total_size > 0:
                            progress(downloaded / total_size, desc="Downloading Spleeter model...")
            logger.info(f"Spleeter model downloaded to: {model_path}")
            progress(1.0, desc="Model download complete!")
            return str(model_path)
        except Exception as e:
            logger.error(f"Error downloading Spleeter model: {e}")
            return None

    def download_youtube_audio(self, video_url: str, progress=gr.Progress()):
        """Download audio from YouTube video using yt-dlp."""
        if not video_url:
            return None, "Please provide a YouTube URL"

        progress(0.1, desc="Initializing download...")
        output_dir = self.temp_dir / "downloaded_audio"
        output_dir.mkdir(exist_ok=True)
        output_file = output_dir / "audio.mp3"

        command = [
            "yt-dlp",
            "-x",
            "--audio-format", "mp3",
            "-o", str(output_file),
            "--no-playlist",
            "--restrict-filenames",
            video_url
        ]

        try:
            progress(0.3, desc="Downloading audio...")
            result = subprocess.run(command, check=True, capture_output=True, text=True)
            progress(1.0, desc="Download complete!")
            return str(output_file), f"Successfully downloaded audio: {output_file.name}"
        except FileNotFoundError:
            return None, "yt-dlp not found. Please install it: pip install yt-dlp"
        except subprocess.CalledProcessError as e:
            return None, f"Download failed: {e.stderr}"
        except Exception as e:
            return None, f"Unexpected error: {str(e)}"

    def extract_basic_features(self, audio_path: str, sr=16000, progress=gr.Progress()):
        """Extract basic audio features and create visualizations."""
        if not audio_path or not Path(audio_path).exists():
            return None, None, "Invalid or missing audio file"

        try:
            progress(0.1, desc="Loading audio...")
            y, sr = librosa.load(audio_path, sr=sr)
            duration = librosa.get_duration(y=y, sr=sr)

            max_duration = 60
            if duration > max_duration:
                y = y[:int(sr * max_duration)]
                duration = max_duration

            progress(0.3, desc="Computing features...")
            features = {
                'duration': duration,
                'sample_rate': sr,
                'samples': len(y),
                'tempo': librosa.beat.beat_track(y=y, sr=sr)[0],
                'mfcc': librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13),
                'spectral_centroid': librosa.feature.spectral_centroid(y=y, sr=sr)[0],
                'spectral_rolloff': librosa.feature.spectral_rolloff(y=y, sr=sr)[0],
                'zero_crossing_rate': librosa.feature.zero_crossing_rate(y)[0]
            }

            progress(0.5, desc="Computing mel spectrogram...")
            hop_length = 512
            S_mel = librosa.feature.melspectrogram(y=y, sr=sr, hop_length=hop_length, n_mels=80)
            S_dB = librosa.power_to_db(S_mel, ref=np.max)

            progress(0.8, desc="Creating visualizations...")
            fig, axes = plt.subplots(2, 2, figsize=(15, 10))

            time_axis = librosa.frames_to_time(range(len(y)), sr=sr)
            axes[0, 0].plot(time_axis, y)
            axes[0, 0].set_title('Waveform')
            axes[0, 0].set_xlabel('Time (s)')
            axes[0, 0].set_ylabel('Amplitude')

            librosa.display.specshow(S_dB, sr=sr, hop_length=hop_length, 
                                    x_axis='time', y_axis='mel', ax=axes[0, 1])
            axes[0, 1].set_title('Mel Spectrogram')

            librosa.display.specshow(features['mfcc'], sr=sr, x_axis='time', ax=axes[1, 0])
            axes[1, 0].set_title('MFCC')

            times = librosa.frames_to_time(range(len(features['spectral_centroid'])), sr=sr, hop_length=hop_length)
            axes[1, 1].plot(times, features['spectral_centroid'], label='Spectral Centroid')
            axes[1, 1].plot(times, features['spectral_rolloff'], label='Spectral Rolloff')
            axes[1, 1].set_title('Spectral Features')
            axes[1, 1].set_xlabel('Time (s)')
            axes[1, 1].legend()

            plt.tight_layout()
            plot_path = self.temp_dir / f"basic_features_{np.random.randint(10000)}.png"
            plt.savefig(plot_path, dpi=150, bbox_inches='tight')
            plt.close()

            summary = f"""
**Audio Summary:**
- Duration: {duration:.2f} seconds
- Sample Rate: {sr} Hz
- Estimated Tempo: {features['tempo']:.1f} BPM
- Number of Samples: {len(y):,}

**Feature Shapes:**
- MFCC: {features['mfcc'].shape}
- Spectral Centroid: {features['spectral_centroid'].shape}
- Spectral Rolloff: {features['spectral_rolloff'].shape}
- Zero Crossing Rate: {features['zero_crossing_rate'].shape}
            """

            progress(1.0, desc="Analysis complete!")
            return str(plot_path), summary, None

        except Exception as e:
            return None, None, f"Error processing audio: {str(e)}"

    def extract_chroma_features(self, audio_path: str, sr=16000, progress=gr.Progress()):
        """Extract and visualize enhanced chroma features."""
        if not audio_path or not Path(audio_path).exists():
            return None, None, "Invalid or missing audio file"

        try:
            progress(0.1, desc="Loading audio...")
            y, sr = librosa.load(audio_path, sr=sr)
            max_duration = 30
            if len(y) > sr * max_duration:
                y = y[:int(sr * max_duration)]

            progress(0.3, desc="Computing chroma variants...")
            chroma_orig = librosa.feature.chroma_cqt(y=y, sr=sr)
            y_harm = librosa.effects.harmonic(y=y, margin=8)
            chroma_harm = librosa.feature.chroma_cqt(y=y_harm, sr=sr)
            chroma_filter = np.minimum(chroma_harm,
                                     librosa.decompose.nn_filter(chroma_harm,
                                                                aggregate=np.median,
                                                                metric='cosine'))
            chroma_smooth = scipy.ndimage.median_filter(chroma_filter, size=(1, 9))
            chroma_stft = librosa.feature.chroma_stft(y=y, sr=sr)
            chroma_cens = librosa.feature.chroma_cens(y=y, sr=sr)

            progress(0.8, desc="Creating visualizations...")
            fig, axes = plt.subplots(3, 2, figsize=(15, 12))
            axes = axes.flatten()

            for i, (chroma, title) in enumerate([
                (chroma_orig, 'Original Chroma (CQT)'),
                (chroma_harm, 'Harmonic Chroma'),
                (chroma_filter, 'Non-local Filtered'),
                (chroma_smooth, 'Median Filtered'),
                (chroma_stft, 'Chroma (STFT)'),
                (chroma_cens, 'CENS Features')
            ]):
                librosa.display.specshow(chroma, y_axis='chroma', x_axis='time', ax=axes[i])
                axes[i].set_title(title)

            plt.tight_layout()
            plot_path = self.temp_dir / f"chroma_features_{np.random.randint(10000)}.png"
            plt.savefig(plot_path, dpi=150, bbox_inches='tight')
            plt.close()

            summary = "Chroma feature analysis complete! Visualizations show different chroma extraction methods for harmonic analysis."
            progress(1.0, desc="Chroma analysis complete!")
            return str(plot_path), summary, None

        except Exception as e:
            return None, None, f"Error processing chroma features: {str(e)}"

    def generate_patches(self, audio_path: str, sr=16000, patch_duration=5.0, hop_duration=1.0, progress=gr.Progress()):
        """Generate fixed-duration patches for transformer input."""
        if not audio_path or not Path(audio_path).exists():
            return None, None, "Invalid or missing audio file"

        try:
            progress(0.1, desc="Loading audio...")
            y, sr = librosa.load(audio_path, sr=sr)

            progress(0.3, desc="Computing mel spectrogram...")
            hop_length = 512
            S_mel = librosa.feature.melspectrogram(y=y, sr=sr, hop_length=hop_length, n_mels=80)
            S_dB = librosa.power_to_db(S_mel, ref=np.max)

            progress(0.5, desc="Generating patches...")
            patch_frames = librosa.time_to_frames(patch_duration, sr=sr, hop_length=hop_length)
            hop_frames = librosa.time_to_frames(hop_duration, sr=sr, hop_length=hop_length)
            patches = librosa.util.frame(S_dB, frame_length=patch_frames, hop_length=hop_frames)

            progress(0.8, desc="Creating visualizations...")
            num_patches_to_show = min(6, patches.shape[-1])
            fig, axes = plt.subplots(2, 3, figsize=(18, 8))
            axes = axes.flatten()

            for i in range(num_patches_to_show):
                librosa.display.specshow(patches[..., i], y_axis='mel', x_axis='time', 
                                       ax=axes[i], sr=sr, hop_length=hop_length)
                axes[i].set_title(f'Patch {i+1}')

            for i in range(num_patches_to_show, len(axes)):
                axes[i].set_visible(False)

            plt.tight_layout()
            plot_path = self.temp_dir / f"patches_{np.random.randint(10000)}.png"
            plt.savefig(plot_path, dpi=150, bbox_inches='tight')
            plt.close()

            summary = f"""
**Patch Generation Summary:**
- Total patches generated: {patches.shape[-1]}
- Patch duration: {patch_duration} seconds
- Hop duration: {hop_duration} seconds
- Patch shape (mels, time, patches): {patches.shape}
- Each patch covers {patch_frames} time frames
            """

            progress(1.0, desc="Patch generation complete!")
            return str(plot_path), summary, None

        except Exception as e:
            return None, None, f"Error generating patches: {str(e)}"

    def separate_audio(self, audio_path: str, progress=gr.Progress()):
        """Separate audio into vocals and accompaniment using Spleeter."""
        if not audio_path or not Path(audio_path).exists():
            return None, None, "Invalid or missing audio file"

        try:
            progress(0.1, desc="Preparing Spleeter model...")
            model_path = self.download_spleeter_model(progress=progress)
            if not model_path:
                return None, None, "Failed to download Spleeter model"

            separator = Separator('spleeter:2stems', pretrained_model_dir=str(self.spleeter_model_dir))
            output_dir = self.temp_dir / "separated_audio"
            output_dir.mkdir(exist_ok=True)

            progress(0.5, desc="Separating audio...")
            separator.separate_to_file(audio_path, output_dir)

            vocals_path = output_dir / Path(audio_path).stem / "vocals.wav"
            accomp_path = output_dir / Path(audio_path).stem / "accompaniment.wav"

            if not vocals_path.exists() or not accomp_path.exists():
                return None, None, "Separation failed: output files not found"

            progress(1.0, desc="Separation complete!")
            return str(vocals_path), str(accomp_path), None

        except SpleeterError as e:
            return None, None, f"Spleeter error: {str(e)}"
        except Exception as e:
            return None, None, f"Error during separation: {str(e)}"

def create_gradio_interface():
    """Create Gradio interface for audio analysis and separation."""
    analyzer = AudioAnalyzer()

    with gr.Blocks(title="🎵 Audio Analysis & Separation Suite", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🎵 Audio Analysis & Separation Suite

        Analyze audio from YouTube videos or uploaded files. Extract features, generate transformer patches, or separate vocals and accompaniment using Spleeter.

        **Features:**
        - 📊 **Basic Features**: Waveform, Mel Spectrogram, MFCC, Spectral Analysis, Tempo
        - 🎼 **Chroma Features**: Harmonic content analysis with multiple methods
        - 🧩 **Transformer Patches**: Fixed-duration patches for deep learning
        - 🎤 **Source Separation**: Separate vocals and accompaniment with Spleeter

        **Requirements**: `yt-dlp` (`pip install yt-dlp`), `spleeter` (`pip install spleeter`), `ffmpeg` (system install).
        """)

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📁 Audio Input")
                with gr.Group():
                    gr.Markdown("**Download from YouTube:**")
                    youtube_url = gr.Textbox(
                        label="YouTube URL",
                        placeholder="https://www.youtube.com/watch?v=...",
                        info="Paste a YouTube video URL to extract audio"
                    )
                    download_btn = gr.Button("📥 Download Audio", variant="primary")
                    download_status = gr.Textbox(label="Download Status", interactive=False)

                with gr.Group():
                    gr.Markdown("**Or upload audio file:**")
                    audio_file = gr.Audio(
                        label="Upload Audio File",
                        type="filepath",
                        info="Supported formats: MP3, WAV, FLAC, etc."
                    )

            with gr.Column(scale=2):
                gr.Markdown("### 🔍 Analysis & Separation Results")
                with gr.Tabs():
                    with gr.Tab("📊 Basic Features"):
                        basic_plot = gr.Image(label="Feature Visualizations")
                        basic_summary = gr.Markdown(label="Feature Summary")
                        basic_btn = gr.Button("🔍 Analyze Basic Features", variant="secondary")

                    with gr.Tab("🎼 Chroma Features"):
                        chroma_plot = gr.Image(label="Chroma Visualizations")
                        chroma_summary = gr.Markdown(label="Chroma Summary")
                        chroma_btn = gr.Button("🎼 Analyze Chroma Features", variant="secondary")

                    with gr.Tab("🧩 Transformer Patches"):
                        with gr.Row():
                            patch_duration = gr.Slider(
                                label="Patch Duration (seconds)",
                                minimum=1.0, maximum=10.0, value=5.0, step=0.5,
                                info="Duration of each patch"
                            )
                            hop_duration = gr.Slider(
                                label="Hop Duration (seconds)",
                                minimum=0.1, maximum=5.0, value=1.0, step=0.1,
                                info="Time between patch starts"
                            )
                        patches_plot = gr.Image(label="Generated Patches")
                        patches_summary = gr.Markdown(label="Patch Summary")
                        patches_btn = gr.Button("🧩 Generate Patches", variant="secondary")

                    with gr.Tab("🎤 Source Separation"):
                        vocals_output = gr.Audio(label="Vocals", type="filepath")
                        accomp_output = gr.Audio(label="Accompaniment", type="filepath")
                        separate_btn = gr.Button("🎤 Separate Vocals & Accompaniment", variant="secondary")

                error_output = gr.Textbox(label="Error Messages", interactive=False)

        gr.Markdown("""
        ### ℹ️ Usage Tips
        - **Processing Limits**: 60s for basic features, 30s for chroma, full length for separation
        - **YouTube**: Ensure URLs are valid and respect terms of service
        - **Spleeter**: Requires `ffmpeg` (install via `apt`, `brew`, or download)
        - **Visualizations**: High-quality, suitable for research
        - **Storage**: Temporary files are cleaned up on interface close
        """)

        # Event handlers
        download_btn.click(
            fn=analyzer.download_youtube_audio,
            inputs=[youtube_url],
            outputs=[audio_file, download_status]
        )

        basic_btn.click(
            fn=analyzer.extract_basic_features,
            inputs=[audio_file],
            outputs=[basic_plot, basic_summary, error_output]
        )

        chroma_btn.click(
            fn=analyzer.extract_chroma_features,
            inputs=[audio_file],
            outputs=[chroma_plot, chroma_summary, error_output]
        )

        patches_btn.click(
            fn=analyzer.generate_patches,
            inputs=[audio_file, patch_duration, hop_duration],
            outputs=[patches_plot, patches_summary, error_output]
        )

        separate_btn.click(
            fn=analyzer.separate_audio,
            inputs=[audio_file],
            outputs=[vocals_output, accomp_output, error_output]
        )

        audio_file.change(
            fn=analyzer.extract_basic_features,
            inputs=[audio_file],
            outputs=[basic_plot, basic_summary, error_output]
        )

        demo.unload(fn=analyzer.cleanup)

    return demo

if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch()