import gradio as gr
import torch
import numpy as np
from transformers import AutoModel, Wav2Vec2Processor
import librosa
import scipy.stats as stats
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from sklearn.cluster import KMeans
from scipy.signal import find_peaks
import warnings
warnings.filterwarnings('ignore')

# Configurar o modelo
MODEL_NAME = "marcosremar2/wavlm-large-deploy"

print("Carregando modelo...")
try:
    processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
    model = AutoModel.from_pretrained(MODEL_NAME)
    model.eval()
    print("Modelo carregado com sucesso!")
except Exception as e:
    print(f"Erro ao carregar modelo: {e}")
    processor = None
    model = None

def analyze_prosodic_features(features):
    """Analisa características prosódicas nas features"""
    # Features relacionadas à prosódia geralmente estão nas primeiras dimensões
    prosodic_dims = features[:, :128]  # Primeiras 128 dimensões
    
    # Análise de variabilidade (relacionada ao ritmo e stress)
    variance = np.var(prosodic_dims, axis=0)
    high_variance_dims = np.where(variance > np.percentile(variance, 90))[0]
    
    # Análise de tendências temporais (relacionada à entonação)
    trends = []
    for dim in range(min(32, prosodic_dims.shape[1])):
        slope, _, r_value, _, _ = stats.linregress(range(len(prosodic_dims)), prosodic_dims[:, dim])
        trends.append((dim, slope, r_value**2))
    
    # Ordenar por força da tendência
    trends.sort(key=lambda x: abs(x[1]), reverse=True)
    
    return {
        'high_variance_dims': high_variance_dims[:10],
        'variance_stats': {
            'mean': np.mean(variance),
            'std': np.std(variance),
            'max': np.max(variance)
        },
        'strongest_trends': trends[:5]
    }

def analyze_phonetic_features(features):
    """Analisa características fonéticas"""
    # Features fonéticas geralmente estão nas dimensões médias
    phonetic_dims = features[:, 128:512]  # Dimensões 128-512
    
    # Análise de padrões de ativação
    activation_patterns = np.mean(phonetic_dims, axis=0)
    highly_active = np.where(activation_patterns > np.percentile(activation_patterns, 85))[0]
    
    # Análise de correlações entre dimensões (indicativo de co-articulação)
    correlations = np.corrcoef(phonetic_dims.T)
    high_corr_pairs = []
    
    for i in range(min(50, correlations.shape[0])):
        for j in range(i+1, min(50, correlations.shape[1])):
            if abs(correlations[i, j]) > 0.7:
                high_corr_pairs.append((i, j, correlations[i, j]))
    
    return {
        'highly_active_dims': highly_active[:15],
        'activation_stats': {
            'mean': np.mean(activation_patterns),
            'std': np.std(activation_patterns),
            'active_ratio': len(highly_active) / len(activation_patterns)
        },
        'high_correlations': sorted(high_corr_pairs, key=lambda x: abs(x[2]), reverse=True)[:5]
    }

def analyze_temporal_features(features):
    """Analisa características temporais e de duração"""
    # Análise da evolução temporal das features
    frame_energy = np.mean(features**2, axis=1)  # Energia por frame
    
    # Detectar segmentos de alta/baixa energia (aproximação de segmentação)
    energy_threshold = np.mean(frame_energy)
    high_energy_frames = np.where(frame_energy > energy_threshold)[0]
    
    # Análise de transições (mudanças abruptas indicam fronteiras fonéticas)
    feature_diff = np.diff(features, axis=0)
    transition_strength = np.mean(feature_diff**2, axis=1)
    strong_transitions = np.where(transition_strength > np.percentile(transition_strength, 80))[0]
    
    return {
        'energy_analysis': {
            'mean_energy': np.mean(frame_energy),
            'energy_variance': np.var(frame_energy),
            'high_energy_ratio': len(high_energy_frames) / len(frame_energy)
        },
        'transition_analysis': {
            'num_strong_transitions': len(strong_transitions),
            'avg_transition_strength': np.mean(transition_strength),
            'transition_density': len(strong_transitions) / len(features)
        }
    }

def analyze_speaker_features(features):
    """Analisa características do falante"""
    # Features de speaker geralmente estão nas últimas dimensões
    speaker_dims = features[:, 512:]  # Dimensões finais
    
    # Consistência das features de speaker (devem ser relativamente estáveis)
    speaker_stability = np.std(speaker_dims, axis=0)
    stable_dims = np.where(speaker_stability < np.percentile(speaker_stability, 30))[0]
    
    # Média geral das características do speaker
    speaker_profile = np.mean(speaker_dims, axis=0)
    
    return {
        'stability_analysis': {
            'num_stable_dims': len(stable_dims),
            'mean_stability': np.mean(speaker_stability),
            'stability_ratio': len(stable_dims) / speaker_dims.shape[1]
        },
        'speaker_profile_stats': {
            'profile_mean': np.mean(speaker_profile),
            'profile_std': np.std(speaker_profile),
            'dominant_features': np.where(speaker_profile > np.percentile(speaker_profile, 75))[0][:10]
        }
    }

def detect_phoneme_boundaries(features, audio_duration):
    """Detecta fronteiras de fonemas usando mudanças nas features"""
    # Calcular diferenças entre frames consecutivos
    feature_diff = np.diff(features, axis=0)
    change_magnitude = np.mean(feature_diff**2, axis=1)
    
    # Encontrar picos (possíveis fronteiras de fonemas)
    peaks, _ = find_peaks(change_magnitude, 
                         height=np.percentile(change_magnitude, 70),
                         distance=3)  # Mínimo 3 frames entre picos
    
    # Converter frames para timestamps
    frame_duration = audio_duration / len(features)
    phoneme_boundaries = peaks * frame_duration
    
    # Estimar fonemas (simplificado baseado em energia e mudanças espectrais)
    phoneme_types = []
    for i in range(len(phoneme_boundaries)):
        if i < len(peaks):
            frame_idx = peaks[i]
            # Análise simplificada das características espectrais
            spectral_features = features[frame_idx, 128:256]  # Features fonéticas
            energy = np.mean(features[frame_idx]**2)
            
            # Classificação básica baseada em energia e padrões espectrais
            if energy > np.percentile([np.mean(features[j]**2) for j in range(len(features))], 80):
                if np.mean(spectral_features[:32]) > np.mean(spectral_features[32:]):
                    phoneme_types.append("VOGAL")
                else:
                    phoneme_types.append("CONSOANTE_FORTE")
            else:
                phoneme_types.append("CONSOANTE_FRACA")
    
    return phoneme_boundaries, phoneme_types

def detect_word_boundaries(features, audio_duration):
    """Detecta fronteiras de palavras usando pausas e mudanças prosódicas"""
    # Energia por frame
    frame_energy = np.mean(features**2, axis=1)
    
    # Detectar pausas (baixa energia)
    energy_threshold = np.percentile(frame_energy, 20)
    low_energy_frames = frame_energy < energy_threshold
    
    # Encontrar segmentos contínuos de baixa energia (pausas)
    pause_starts = []
    pause_ends = []
    in_pause = False
    
    for i, is_low in enumerate(low_energy_frames):
        if is_low and not in_pause:
            pause_starts.append(i)
            in_pause = True
        elif not is_low and in_pause:
            pause_ends.append(i)
            in_pause = False
    
    # Converter para timestamps
    frame_duration = audio_duration / len(features)
    word_boundaries = []
    
    for start, end in zip(pause_starts, pause_ends):
        if (end - start) * frame_duration > 0.1:  # Pausas > 100ms
            word_boundaries.append((start * frame_duration, end * frame_duration))
    
    return word_boundaries

def analyze_speaker_changes(features, audio_duration):
    """Analisa mudanças de falante usando clustering das features de speaker"""
    speaker_features = features[:, 512:]  # Features de falante
    
    if speaker_features.shape[1] < 10:
        return [], []
    
    # Suavizar features para reduzir ruído
    window_size = 5
    smoothed_features = np.array([
        np.mean(speaker_features[max(0, i-window_size):i+window_size+1], axis=0)
        for i in range(len(speaker_features))
    ])
    
    # Clustering para identificar diferentes falantes
    try:
        n_speakers = min(3, len(smoothed_features) // 10)  # Máximo 3 falantes
        if n_speakers < 2:
            return [], []
            
        kmeans = KMeans(n_clusters=n_speakers, random_state=42, n_init=10)
        speaker_labels = kmeans.fit_predict(smoothed_features)
        
        # Detectar mudanças de falante
        speaker_changes = []
        current_speaker = speaker_labels[0]
        
        for i, label in enumerate(speaker_labels[1:], 1):
            if label != current_speaker:
                timestamp = i * (audio_duration / len(features))
                speaker_changes.append((timestamp, current_speaker, label))
                current_speaker = label
        
        return speaker_changes, speaker_labels
    except:
        return [], []

def create_temporal_visualization(features, audio, sr, audio_duration):
    """Cria visualizações temporais das features"""
    fig, axes = plt.subplots(4, 1, figsize=(15, 12))
    
    # 1. Waveform com energia
    time_axis = np.linspace(0, audio_duration, len(audio))
    axes[0].plot(time_axis, audio, alpha=0.7, color='blue', linewidth=0.5)
    axes[0].set_title('Forma de Onda do Áudio', fontsize=12, fontweight='bold')
    axes[0].set_ylabel('Amplitude')
    axes[0].grid(True, alpha=0.3)
    
    # 2. Energia por frame
    frame_energy = np.mean(features**2, axis=1)
    frame_time = np.linspace(0, audio_duration, len(frame_energy))
    axes[1].plot(frame_time, frame_energy, color='red', linewidth=2)
    axes[1].set_title('Energia por Frame (Detecção de Pausas)', fontsize=12, fontweight='bold')
    axes[1].set_ylabel('Energia')
    axes[1].grid(True, alpha=0.3)
    
    # 3. Features Prosódicas (primeiras 16 dimensões)
    prosodic_features = features[:, :16]
    im1 = axes[2].imshow(prosodic_features.T, aspect='auto', cmap='viridis',
                        extent=[0, audio_duration, 0, 16])
    axes[2].set_title('Features Prosódicas (Entonação, Ritmo)', fontsize=12, fontweight='bold')
    axes[2].set_ylabel('Dimensão')
    plt.colorbar(im1, ax=axes[2], label='Valor da Feature')
    
    # 4. Features Fonéticas (dimensões 128-144)
    phonetic_features = features[:, 128:144]
    im2 = axes[3].imshow(phonetic_features.T, aspect='auto', cmap='plasma',
                        extent=[0, audio_duration, 0, 16])
    axes[3].set_title('Features Fonéticas (Fonemas)', fontsize=12, fontweight='bold')
    axes[3].set_xlabel('Tempo (s)')
    axes[3].set_ylabel('Dimensão')
    plt.colorbar(im2, ax=axes[3], label='Valor da Feature')
    
    plt.tight_layout()
    return fig

def create_segmentation_visualization(features, audio_duration, phoneme_boundaries, 
                                    word_boundaries, speaker_changes):
    """Cria visualização da segmentação temporal"""
    fig, ax = plt.subplots(1, 1, figsize=(15, 8))
    
    # Features de energia como base
    frame_energy = np.mean(features**2, axis=1)
    frame_time = np.linspace(0, audio_duration, len(frame_energy))
    ax.plot(frame_time, frame_energy, color='black', alpha=0.7, linewidth=1)
    
    # Marcar fronteiras de fonemas
    for i, boundary in enumerate(phoneme_boundaries):
        ax.axvline(x=boundary, color='blue', linestyle='--', alpha=0.7)
        if i < len(phoneme_boundaries):
            ax.text(boundary, max(frame_energy) * 0.9, f'F{i+1}', 
                   rotation=90, fontsize=8, color='blue')
    
    # Marcar pausas (fronteiras de palavras)
    for start, end in word_boundaries:
        rect = patches.Rectangle((start, 0), end-start, max(frame_energy),
                               linewidth=0, facecolor='yellow', alpha=0.3)
        ax.add_patch(rect)
    
    # Marcar mudanças de falante
    for timestamp, old_speaker, new_speaker in speaker_changes:
        ax.axvline(x=timestamp, color='red', linestyle='-', linewidth=3, alpha=0.8)
        ax.text(timestamp, max(frame_energy) * 0.7, f'S{old_speaker}→S{new_speaker}',
               rotation=90, fontsize=10, color='red', fontweight='bold')
    
    ax.set_title('Segmentação Temporal: Fonemas, Palavras e Falantes', 
                fontsize=14, fontweight='bold')
    ax.set_xlabel('Tempo (s)')
    ax.set_ylabel('Energia')
    ax.grid(True, alpha=0.3)
    
    # Legenda
    from matplotlib.lines import Line2D
    legend_elements = [
        Line2D([0], [0], color='blue', linestyle='--', label='Fronteiras de Fonemas'),
        Line2D([0], [0], color='yellow', linewidth=10, alpha=0.3, label='Pausas (Palavras)'),
        Line2D([0], [0], color='red', linewidth=3, label='Mudanças de Falante')
    ]
    ax.legend(handles=legend_elements, loc='upper right')
    
    plt.tight_layout()
    return fig

def process_audio(audio_file):
    """Processa o arquivo de áudio com análise completa e visualizações"""
    if audio_file is None:
        return "❌ Por favor, carregue um arquivo de áudio.", None, None
    
    if processor is None or model is None:
        return "❌ Erro: Modelo não foi carregado corretamente.", None, None
    
    try:
        # Carregar áudio
        audio, sr = librosa.load(audio_file, sr=16000)
        
        # Limitar duração
        max_duration = 30  # segundos
        if len(audio) > max_duration * sr:
            audio = audio[:max_duration * sr]
        
        # Processar com o modelo
        inputs = processor(audio, sampling_rate=16000, return_tensors="pt", padding=True)
        
        with torch.no_grad():
            outputs = model(inputs.input_values)
            hidden_states = outputs.last_hidden_state
        
        # Converter para numpy para análise
        features = hidden_states.squeeze(0).cpu().numpy()
        duration = len(audio) / sr
        
        # Análises temporais detalhadas
        phoneme_boundaries, phoneme_types = detect_phoneme_boundaries(features, duration)
        word_boundaries = detect_word_boundaries(features, duration)
        speaker_changes, speaker_labels = analyze_speaker_changes(features, duration)
        
        # Criar visualizações
        temporal_viz = create_temporal_visualization(features, audio, sr, duration)
        segmentation_viz = create_segmentation_visualization(
            features, duration, phoneme_boundaries, word_boundaries, speaker_changes)
        
        # Análise detalhada de fonemas
        phoneme_analysis = ""
        if len(phoneme_boundaries) > 0:
            phoneme_analysis = "\n## 🗣️ **Análise Temporal de Fonemas**\n"
            for i, (boundary, ptype) in enumerate(zip(phoneme_boundaries, phoneme_types)):
                phoneme_analysis += f"- **Fonema {i+1}**: {ptype} aos {boundary:.2f}s\n"
        
        # Análise de palavras/pausas
        word_analysis = ""
        if len(word_boundaries) > 0:
            word_analysis = "\n## 📝 **Análise de Palavras/Pausas**\n"
            for i, (start, end) in enumerate(word_boundaries):
                duration_pause = end - start
                word_analysis += f"- **Pausa {i+1}**: {start:.2f}s - {end:.2f}s (duração: {duration_pause:.2f}s)\n"
        
        # Análise de falantes
        speaker_analysis = ""
        if len(speaker_changes) > 0:
            speaker_analysis = "\n## 👥 **Diarização de Falantes**\n"
            current_speaker = 0
            speaker_analysis += f"- **Início**: Falante {current_speaker}\n"
            for timestamp, old_speaker, new_speaker in speaker_changes:
                speaker_analysis += f"- **{timestamp:.2f}s**: Mudança de Falante {old_speaker} → Falante {new_speaker}\n"
        else:
            speaker_analysis = "\n## 👤 **Análise de Falante**\n- Apenas um falante detectado no áudio\n"
        
        # Estatísticas gerais
        num_frames = features.shape[0]
        frame_rate = num_frames / duration
        
        result = f"""
# 🎵 Análise Completa WavLM-Large com Visualizações

## 📊 **Informações Básicas**
- **Duração**: {duration:.2f} segundos
- **Frames extraídos**: {num_frames}
- **Taxa de frames**: {frame_rate:.1f} frames/segundo
- **Resolução temporal**: {duration/num_frames:.3f}s por frame

## 🔍 **Resumo da Segmentação**
- **Fonemas detectados**: {len(phoneme_boundaries)}
- **Pausas detectadas**: {len(word_boundaries)}
- **Mudanças de falante**: {len(speaker_changes)}
- **Qualidade da análise**: {'✅ Excelente' if num_frames > 100 else '⚠️ Limitada (áudio curto)'}

{phoneme_analysis}

{word_analysis}

{speaker_analysis}

## 📈 **Interpretação das Visualizações**

**Gráfico 1 - Temporal Features:**
- Mostra evolução das características ao longo do tempo
- Prosódia (entonação) e fonemas em tempo real

**Gráfico 2 - Segmentação:**
- Linhas azuis: fronteiras de fonemas
- Áreas amarelas: pausas entre palavras
- Linhas vermelhas: mudanças de falante

## 🎯 **Aplicações Práticas**
- **Transcrição temporal**: Use os timestamps para sincronizar texto
- **Análise prosódica**: Veja padrões de entonação
- **Diarização**: Identifique quem fala quando
- **Segmentação**: Encontre fronteiras naturais da fala
        """
        
        return result, temporal_viz, segmentation_viz
        
    except Exception as e:
        return f"❌ Erro ao processar: {str(e)}", None, None

# Interface Gradio com visualizações
iface = gr.Interface(
    fn=process_audio,
    inputs=gr.Audio(type="filepath", label="📁 Carregar Arquivo de Áudio"),
    outputs=[
        gr.Textbox(label="📊 Análise Temporal Detalhada", lines=25),
        gr.Plot(label="📈 Visualização Temporal das Features"),
        gr.Plot(label="🎯 Segmentação: Fonemas, Palavras e Falantes")
    ],
    title="🎵 WavLM-Large: Análise Temporal Avançada",
    description="""
    **Análise completa com timestamps e visualizações**
    
    ✨ **Novidades desta versão:**
    • 🗣️ **Timestamps de fonemas** - Quando cada som foi pronunciado
    • 📝 **Detecção de pausas** - Fronteiras de palavras/frases  
    • 👥 **Diarização de falantes** - Quem fala quando
    • 📈 **Visualizações temporais** - Gráficos das características
    • 🎯 **Segmentação visual** - Mapa temporal completo
    
    Carregue um arquivo de áudio para análise temporal completa!
    """,
    examples=None,
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch()