"""
Utility functions for visualization and analysis
"""

import numpy as np
import matplotlib.pyplot as plt
import librosa
import librosa.display

def create_waveform_plot(y, sr, title="Audio Waveform"):
    """
    Create waveform visualization
    
    Args:
        y (np.array): Audio time series
        sr (int): Sample rate
        title (str): Plot title
    
    Returns:
        matplotlib.figure.Figure: Waveform plot
    """
    fig, ax = plt.subplots(figsize=(10, 3))
    librosa.display.waveshow(y, sr=sr, ax=ax, color='#2E86DE')
    ax.set_title(title, fontsize=14, fontweight='bold')
    ax.set_xlabel('Time (seconds)', fontsize=11)
    ax.set_ylabel('Amplitude', fontsize=11)
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    return fig

def create_spectrogram_plot(y, sr, title="Spectrogram"):
    """
    Create spectrogram visualization
    
    Args:
        y (np.array): Audio time series
        sr (int): Sample rate
        title (str): Plot title
    
    Returns:
        matplotlib.figure.Figure: Spectrogram plot
    """
    fig, ax = plt.subplots(figsize=(10, 4))
    D = librosa.amplitude_to_db(np.abs(librosa.stft(y)), ref=np.max)
    img = librosa.display.specshow(D, sr=sr, x_axis='time', y_axis='hz', 
                                   ax=ax, cmap='viridis')
    ax.set_title(title, fontsize=14, fontweight='bold')
    ax.set_xlabel('Time (seconds)', fontsize=11)
    ax.set_ylabel('Frequency (Hz)', fontsize=11)
    fig.colorbar(img, ax=ax, format='%+2.0f dB')
    plt.tight_layout()
    return fig

def create_mel_spectrogram_plot(y, sr, title="Mel Spectrogram"):
    """
    Create mel spectrogram visualization
    
    Args:
        y (np.array): Audio time series
        sr (int): Sample rate
        title (str): Plot title
    
    Returns:
        matplotlib.figure.Figure: Mel spectrogram plot
    """
    fig, ax = plt.subplots(figsize=(10, 4))
    S = librosa.feature.melspectrogram(y=y, sr=sr)
    S_dB = librosa.power_to_db(S, ref=np.max)
    img = librosa.display.specshow(S_dB, sr=sr, x_axis='time', y_axis='mel',
                                   ax=ax, cmap='magma')
    ax.set_title(title, fontsize=14, fontweight='bold')
    ax.set_xlabel('Time (seconds)', fontsize=11)
    ax.set_ylabel('Mel Frequency', fontsize=11)
    fig.colorbar(img, ax=ax, format='%+2.0f dB')
    plt.tight_layout()
    return fig

def format_probability_text(prob_dict, top_k=None):
    """
    Format probability dictionary as text with progress bars
    
    Args:
        prob_dict (dict): Dictionary of emotion: probability
        top_k (int): Show only top K emotions (None for all)
    
    Returns:
        str: Formatted text
    """
    # Sort by probability
    sorted_probs = sorted(prob_dict.items(), key=lambda x: x[1], reverse=True)
    
    if top_k:
        sorted_probs = sorted_probs[:top_k]
    
    text = ""
    for emotion, prob in sorted_probs:
        # Create progress bar
        bar_length = int(prob * 30)
        bar = '█' * bar_length + '░' * (30 - bar_length)
        
        text += f"**{emotion.capitalize()}**: {bar} {prob*100:.2f}%\n"
    
    return text

def get_emotion_emoji(emotion):
    """
    Get emoji for emotion
    
    Args:
        emotion (str): Emotion name
    
    Returns:
        str: Emoji character
    """
    emoji_map = {
        'angry': '😠',
        'calm': '😌',
        'disgust': '🤢',
        'fearful': '😨',
        'happy': '😊',
        'neutral': '😐',
        'sad': '😢',
        'surprised': '😲'
    }
    return emoji_map.get(emotion.lower(), '🎭')