import gradio as gr
import torch
import torch.nn as nn
import numpy as np
import librosa
import cv2
import re
from transformers import Wav2Vec2Processor, Wav2Vec2Model, AutoTokenizer, AutoModel
from torchvision import models
import tempfile
import os
from huggingface_hub import hf_hub_download

# Configuration
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
SAMPLE_RATE = 16000
TEXT_MAX_LEN = 64
LABELS = ["angry", "happy", "neutral", "sad"]

# Load processors
processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# Model Architecture (same as training)
class ResNetVideoEncoder(nn.Module):
    def __init__(self, out_dim=768):
        super().__init__()
        base = models.resnet18(pretrained=False)
        self.backbone = nn.Sequential(*list(base.children())[:-1])
        self.proj = nn.Linear(512, out_dim)

    def forward(self, x):
        B, C, T, H, W = x.shape
        feats = []
        for t in range(T):
            ft = self.backbone(x[:, :, t])
            feats.append(ft.squeeze(-1).squeeze(-1))
        feats = torch.stack(feats, dim=1).mean(1)
        return self.proj(feats)

def mean_pool(x, mask):
    mask = mask[:, :x.size(1)]
    mask = mask.unsqueeze(-1).float()
    return (x * mask).sum(1) / mask.sum(1).clamp(min=1e-6)

class HBF(nn.Module):
    def __init__(self, d=768, n_layers=6):
        super().__init__()
        self.proj_a = nn.ModuleList([nn.Linear(d, d) for _ in range(n_layers)])
        self.proj_t = nn.ModuleList([nn.Linear(d, d) for _ in range(n_layers)])
        self.proj_v = nn.ModuleList([nn.Linear(d, d) for _ in range(n_layers)])
        self.fwd1 = nn.ModuleList([nn.Linear(3*d, d) for _ in range(n_layers)])
        self.fwd2 = nn.ModuleList([nn.Linear(d, d) for _ in range(n_layers)])
        self.drop = nn.Dropout(0.1)
        self.act1, self.act2 = nn.GELU(), nn.Tanh()
        self.n = n_layers

    def forward(self, a, t, v):
        v_prev = None
        for i in range(self.n):
            va = self.act2(self.drop(self.proj_a[i](a)))
            vt = self.act2(self.drop(self.proj_t[i](t)))
            vv = self.act2(self.drop(self.proj_v[i](v)))
            cat = torch.cat([va, vt, vv] if v_prev is None else [va, vt, v_prev], -1)
            x = self.act1(self.fwd1[i](cat))
            v_prev = self.fwd2[i](x)
        return v_prev

class AVVideoModel(nn.Module):
    def __init__(self, num_classes, n_layers=6):
        super().__init__()
        self.a_enc = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base-960h")
        self.t_enc = AutoModel.from_pretrained("bert-base-uncased")
        self.v_enc = ResNetVideoEncoder()
        self.hbf = HBF(n_layers=n_layers)
        self.fc = nn.Linear(768, num_classes)
        self.fc_audio = nn.Linear(768, num_classes)
        self.fc_text = nn.Linear(768, num_classes)
        self.fc_video = nn.Linear(768, num_classes)

    def forward(self, audio, audio_mask, text_ids, text_mask, video):
        a_out = self.a_enc(audio, attention_mask=audio_mask, return_dict=True)
        t_out = self.t_enc(input_ids=text_ids, attention_mask=text_mask, return_dict=True)
        
        a_pool = mean_pool(a_out.last_hidden_state, audio_mask)
        t_pool = mean_pool(t_out.last_hidden_state, text_mask)
        v_pool = self.v_enc(video)
        
        a_pool = torch.nan_to_num(a_pool, nan=0.0, posinf=1e4, neginf=-1e4)
        t_pool = torch.nan_to_num(t_pool, nan=0.0, posinf=1e4, neginf=-1e4)
        v_pool = torch.nan_to_num(v_pool, nan=0.0, posinf=1e4, neginf=-1e4)
        
        a_logits = self.fc_audio(a_pool)
        t_logits = self.fc_text(t_pool)
        v_logits = self.fc_video(v_pool)
        
        fused = self.hbf(a_pool, t_pool, v_pool)
        fused = torch.nan_to_num(fused, nan=0.0, posinf=1e4, neginf=-1e4)
        fused_logits = self.fc(fused)
        
        return fused_logits, a_logits, t_logits, v_logits

# Load model
# Load model
model = AVVideoModel(num_classes=len(LABELS)).to(DEVICE)

# Download model from Hugging Face Model Hub
try:
    model_path = hf_hub_download(
        repo_id="your-username/emotion-model",  # CHANGE THIS
        filename="model_weights.pth"            # YOUR FILE NAME
    )

    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
    model.eval()
    print("✅ Model loaded from Hugging Face")

except Exception as e:
    print(f"❌ Failed to load model: {e}")

def extract_video_frames(video_path, max_frames=8, resize=(224, 224)):
    """Extract frames from video file"""
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None
    
    frames = []
    while len(frames) < max_frames:
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame = cv2.resize(frame, resize)
        frames.append(frame)
    
    cap.release()
    
    if len(frames) == 0:
        return None
    
    # Pad if needed
    while len(frames) < max_frames:
        frames.append(frames[-1])
    
    frames = np.array(frames[:max_frames], dtype=np.uint8)
    return frames

def preprocess_inputs(audio_path, text, video_path):
    """Preprocess all three modalities"""
    
    # Audio
    wav, _ = librosa.load(audio_path, sr=SAMPLE_RATE)
    audio_inputs = processor(wav, sampling_rate=SAMPLE_RATE, return_tensors="pt")
    audio_values = audio_inputs.input_values.to(DEVICE)
    audio_mask = torch.ones_like(audio_values).to(DEVICE)
    
    # Text
    text_clean = re.sub(r"[^a-zA-Z0-9\s]", "", text.lower())
    text_inputs = tokenizer(
        text_clean,
        truncation=True,
        padding="max_length",
        max_length=TEXT_MAX_LEN,
        return_tensors="pt"
    )
    text_ids = text_inputs.input_ids.to(DEVICE)
    text_mask = text_inputs.attention_mask.to(DEVICE)
    
    # Video
    frames = extract_video_frames(video_path)
    if frames is None:
        raise ValueError("Could not extract frames from video")
    
    frames_tensor = torch.tensor(frames).permute(0, 3, 1, 2).float() / 255.0
    frames_tensor = frames_tensor.unsqueeze(0).permute(0, 2, 1, 3, 4).to(DEVICE)
    
    return audio_values, audio_mask, text_ids, text_mask, frames_tensor

def predict_emotion(audio_file, text_input, video_file):
    """Main prediction function"""
    
    if audio_file is None or video_file is None or not text_input.strip():
        return "Please provide all three inputs: audio, text, and video", None, None, None, None
    
    try:
        # Preprocess
        audio, audio_mask, text_ids, text_mask, video = preprocess_inputs(
            audio_file, text_input, video_file
        )
        
        # Inference
        with torch.no_grad():
            fused_logits, a_logits, t_logits, v_logits = model(
                audio, audio_mask, text_ids, text_mask, video
            )
        
        # Get probabilities
        fused_probs = torch.softmax(fused_logits, dim=1)[0].cpu().numpy()
        audio_probs = torch.softmax(a_logits, dim=1)[0].cpu().numpy()
        text_probs = torch.softmax(t_logits, dim=1)[0].cpu().numpy()
        video_probs = torch.softmax(v_logits, dim=1)[0].cpu().numpy()
        
        # Format results
        fused_result = {LABELS[i]: float(fused_probs[i]) for i in range(len(LABELS))}
        audio_result = {LABELS[i]: float(audio_probs[i]) for i in range(len(LABELS))}
        text_result = {LABELS[i]: float(text_probs[i]) for i in range(len(LABELS))}
        video_result = {LABELS[i]: float(video_probs[i]) for i in range(len(LABELS))}
        
        predicted_emotion = LABELS[fused_probs.argmax()]
        confidence = float(fused_probs.max())
        
        result_text = f"🎯 **Predicted Emotion: {predicted_emotion.upper()}**\n\n**Confidence: {confidence:.2%}**"
        
        return result_text, fused_result, audio_result, text_result, video_result
        
    except Exception as e:
        return f"Error: {str(e)}", None, None, None, None

# Gradio Interface
with gr.Blocks(title="Multimodal Emotion Recognition", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎭 Multimodal Emotion Recognition
        
        This system predicts emotions using **Audio**, **Text**, and **Video** inputs simultaneously.
        
        ### How to use:
        1. Upload an audio file (WAV, MP3)
        2. Enter the transcript or spoken text
        3. Upload a video file (MP4, AVI)
        4. Click "Predict Emotion"
        
        The model will analyze all three modalities and provide predictions.
        """
    )
    
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(type="filepath", label="🎤 Audio Input")
            text_input = gr.Textbox(
                label="📝 Text Transcript",
                placeholder="Enter what was said in the audio/video...",
                lines=3
            )
            video_input = gr.Video(label="🎥 Video Input")
            predict_btn = gr.Button("🚀 Predict Emotion", variant="primary", size="lg")
        
        with gr.Column():
            result_text = gr.Markdown(label="Result")
            
            with gr.Accordion("📊 Detailed Predictions", open=True):
                fused_output = gr.Label(label="🔗 Fused Prediction", num_top_classes=4)
                audio_output = gr.Label(label="🎤 Audio-only Prediction", num_top_classes=4)
                text_output = gr.Label(label="📝 Text-only Prediction", num_top_classes=4)
                video_output = gr.Label(label="🎥 Video-only Prediction", num_top_classes=4)
    
    predict_btn.click(
        fn=predict_emotion,
        inputs=[audio_input, text_input, video_input],
        outputs=[result_text, fused_output, audio_output, text_output, video_output]
    )
    
    gr.Markdown(
        """
        ---
        ### 📌 Notes:
        - Supported emotions: **Angry, Happy, Neutral, Sad**
        - Model uses Wav2Vec2 (audio), BERT (text), and ResNet18 (video)
        - Best results with clear audio, accurate transcripts, and visible faces
        """
    )

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