import gradio as gr
import torch
import torch.nn.functional as F
import cv2
import numpy as np
import librosa
from PIL import Image
import tempfile
import os
from typing import Tuple, Dict, Any
import json

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")

# Real model class - now uses your actual mirror_model.pth
class MirrorMindModel:
    def __init__(self):
        # Set device first - make sure torch is properly referenced
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Using device: {self.device}")
        
        # Load your actual model
        try:
            model_path = "mirror_model.pth"  # Adjust path if needed
            print(f"Loading model from {model_path}...")
            
            # Check if model file exists
            if not os.path.exists(model_path):
                print(f"Model file {model_path} not found. Using fallback mode.")
                self.model = None
                return
            
            # Handle PyTorch version-specific loading
            checkpoint = None
            pytorch_version = torch.__version__
            
            # For PyTorch 2.8.0+, we need to be very explicit about loading
            if pytorch_version.startswith("2.8") or pytorch_version.startswith("2.9"):
                print(f"Detected PyTorch {pytorch_version} - using version-specific loading...")
                
                # Method 1: Force weights_only=False for complete models
                try:
                    print("Loading with weights_only=False (for complete model objects)...")
                    import warnings
                    with warnings.catch_warnings():
                        warnings.simplefilter("ignore")
                        checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)
                    print("✓ Successfully loaded complete model")
                except Exception as e1:
                    print(f"✗ Failed to load complete model: {e1}")
                    
                    # Method 2: Try state_dict only loading
                    try:
                        print("Attempting state_dict loading with weights_only=True...")
                        checkpoint = torch.load(model_path, map_location=self.device, weights_only=True)
                        print("✓ Successfully loaded as state_dict")
                    except Exception as e2:
                        print(f"✗ State dict loading failed: {e2}")
                        checkpoint = None
            else:
                # For older PyTorch versions, use standard loading
                try:
                    print(f"Using standard loading for PyTorch {pytorch_version}...")
                    checkpoint = torch.load(model_path, map_location=self.device)
                    print("✓ Successfully loaded with standard method")
                except Exception as e:
                    print(f"✗ Standard loading failed: {e}")
                    checkpoint = None
            
            if checkpoint is None:
                print("All loading methods failed. Using fallback mode.")
                self.model = None
                return
            
            # Handle different checkpoint formats
            if isinstance(checkpoint, dict):
                print(f"Checkpoint keys: {list(checkpoint.keys())}")
                
                if 'model' in checkpoint and 'state_dict' in checkpoint:
                    # Complete model + state dict
                    self.model = checkpoint['model']
                    self.model.load_state_dict(checkpoint['state_dict'])
                    print("✓ Loaded model architecture + state dict")
                    
                elif 'state_dict' in checkpoint:
                    # Only state dict, try to extract model info
                    print("Found 'state_dict' - attempting to reconstruct model...")
                    if 'model_class' in checkpoint or 'architecture' in checkpoint:
                        print("Model architecture info available - but need implementation")
                        # You would reconstruct your model here
                        # self.model = YourModelClass()
                        # self.model.load_state_dict(checkpoint['state_dict'])
                    print("⚠️ State dict found but no model architecture. Using fallback for demo.")
                    self.model = None
                    return
                    
                elif 'model_state_dict' in checkpoint:
                    # PyTorch Lightning or similar format
                    print("Found 'model_state_dict' - checking for model class info...")
                    state_dict = checkpoint['model_state_dict']
                    
                    # Try to infer model structure from state dict keys
                    model_info = self.analyze_state_dict(state_dict)
                    print(f"State dict analysis: {model_info}")
                    
                    # For now, use fallback since we don't have the exact architecture
                    print("⚠️ Cannot reconstruct model without architecture definition. Using fallback.")
                    self.model = None
                    return
                    
                elif len(checkpoint.keys()) > 0 and all(isinstance(v, torch.Tensor) for v in checkpoint.values()):
                    # Direct state dict (keys are layer names, values are tensors)
                    print("Checkpoint appears to be a direct state dict")
                    model_info = self.analyze_state_dict(checkpoint)
                    print(f"Direct state dict analysis: {model_info}")
                    print("⚠️ Cannot reconstruct model without architecture. Using fallback.")
                    self.model = None
                    return
                    
                else:
                    # Try to use as complete model
                    if hasattr(checkpoint, 'eval') and callable(checkpoint.eval):
                        self.model = checkpoint
                        print("✓ Using checkpoint as complete model")
                    else:
                        print("⚠️ Unrecognized checkpoint format. Using fallback.")
                        self.model = None
                        return
            else:
                # Assume the whole model was saved
                if hasattr(checkpoint, 'eval') and callable(checkpoint.eval):
                    self.model = checkpoint
                    print("✓ Loaded complete model object")
                else:
                    print("⚠️ Checkpoint is not a model object. Using fallback.")
                    self.model = None
                    return
            
            if self.model is not None:
                self.model.to(self.device)
                self.model.eval()
                print("Model loaded and ready for inference!")
            else:
                print("Model is None after loading. Using fallback.")
            
        except Exception as e:
            print(f"Error loading model: {e}")
            print("Using fallback random predictions...")
            self.model = None
        
    def analyze_state_dict(self, state_dict):
        """Analyze state dict to understand model structure"""
        info = {
            'total_params': len(state_dict),
            'layer_types': set(),
            'input_features': None,
            'output_features': None,
            'has_conv': False,
            'has_lstm': False,
            'has_attention': False
        }
        
        for key, tensor in state_dict.items():
            # Analyze layer types
            if 'conv' in key.lower():
                info['has_conv'] = True
                info['layer_types'].add('conv')
            elif 'lstm' in key.lower() or 'rnn' in key.lower():
                info['has_lstm'] = True
                info['layer_types'].add('lstm')
            elif 'attention' in key.lower() or 'attn' in key.lower():
                info['has_attention'] = True
                info['layer_types'].add('attention')
            elif 'linear' in key.lower() or 'fc' in key.lower():
                info['layer_types'].add('linear')
                
                # Try to infer input/output dimensions
                if key.endswith('.weight'):
                    if info['input_features'] is None:
                        info['input_features'] = tensor.shape[-1]
                    info['output_features'] = tensor.shape[0]
        
        info['layer_types'] = list(info['layer_types'])
        return info
    
    def create_dummy_model_from_analysis(self, model_info):
        """Create a simple dummy model based on state dict analysis"""
        try:
            import torch.nn as nn
            
            # Create a simple feedforward network based on analysis
            if model_info['input_features'] and model_info['output_features']:
                layers = []
                
                # Input layer
                layers.append(nn.Linear(model_info['input_features'], 128))
                layers.append(nn.ReLU())
                
                # Hidden layers
                if model_info['has_lstm']:
                    layers.append(nn.LSTM(128, 64, batch_first=True))
                else:
                    layers.append(nn.Linear(128, 64))
                    layers.append(nn.ReLU())
                
                # Output layer
                layers.append(nn.Linear(64, model_info['output_features']))
                
                if model_info['has_lstm']:
                    # For LSTM, we need a special wrapper
                    class SimpleModel(nn.Module):
                        def __init__(self, layers):
                            super().__init__()
                            self.layers = nn.ModuleList(layers)
                            
                        def forward(self, x):
                            for layer in self.layers:
                                if isinstance(layer, nn.LSTM):
                                    x, _ = layer(x)
                                    x = x[:, -1, :]  # Take last output
                                else:
                                    x = layer(x)
                            return x
                    
                    return SimpleModel(layers)
                else:
                    return nn.Sequential(*layers)
            
            return None
            
        except Exception as e:
            print(f"Could not create dummy model: {e}")
            return None


        """Extract evenly spaced frames from video and convert to tensor"""
        try:
            cap = cv2.VideoCapture(video_path)
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            
            if total_frames == 0:
                raise ValueError("Could not read video file")
            
            frame_indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
            frames = []
            
            for idx in frame_indices:
                cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
                ret, frame = cap.read()
                if ret:
                    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    frame = cv2.resize(frame, (224, 224))  # Adjust size based on your model
                    # Normalize frame (adjust normalization based on your training)
                    frame = frame.astype(np.float32) / 255.0
                    frames.append(frame)
            
            cap.release()
            
            if not frames:
                raise ValueError("No frames could be extracted from video")
            
            # Convert to tensor: [num_frames, height, width, channels] -> [1, channels, num_frames, height, width]
            frames = np.array(frames)  # [num_frames, 224, 224, 3]
            frames = np.transpose(frames, (3, 0, 1, 2))  # [3, num_frames, 224, 224]
            video_tensor = torch.from_numpy(frames).unsqueeze(0).to(self.device)  # [1, 3, num_frames, 224, 224]
            
            return video_tensor
            
        except Exception as e:
            print(f"Video frame extraction failed: {e}")
            # Return dummy tensor
            dummy_frames = np.random.rand(num_frames, 224, 224, 3).astype(np.float32)
            dummy_frames = np.transpose(dummy_frames, (3, 0, 1, 2))
            return torch.from_numpy(dummy_frames).unsqueeze(0).to(self.device)
    
    def extract_audio_features(self, video_path: str, duration: float = 4.0):
        """Extract audio features from video and convert to tensor"""
        try:
            # Extract audio from video
            audio, sr = librosa.load(video_path, sr=16000, duration=duration)
            
            if len(audio) == 0:
                raise ValueError("No audio data extracted")
            
            # Extract features (adjust based on what your model expects)
            mfcc = librosa.feature.mfcc(y=audio, sr=sr, n_mfcc=13)
            spectral_centroids = librosa.feature.spectral_centroid(y=audio, sr=sr)
            
            # Combine features
            features = np.concatenate([
                np.mean(mfcc, axis=1),
                np.mean(spectral_centroids, axis=1)
            ])
            
            # Convert to tensor
            audio_tensor = torch.from_numpy(features).float().unsqueeze(0).to(self.device)
            
            return audio_tensor
        except Exception as e:
            print(f"Audio extraction failed: {e}")
            # Return dummy tensor if audio extraction fails
            return torch.zeros(14).unsqueeze(0).to(self.device)
    
    def predict(self, video_path: str) -> Dict[str, Any]:
        """Main prediction function using your actual model"""
        try:
            # Validate input
            if not os.path.exists(video_path):
                raise ValueError(f"Video file not found: {video_path}")
            
            # Extract visual features
            video_features = self.extract_video_frames(video_path)
            
            # Extract audio features
            audio_features = self.extract_audio_features(video_path)
            
            if self.model is not None:
                # Real model inference
                with torch.no_grad():
                    # Adjust this based on your model's input requirements
                    try:
                        # Option 1: If your model takes separate video and audio inputs
                        outputs = self.model(video_features, audio_features)
                    except Exception as e1:
                        try:
                            # Option 2: If your model takes concatenated features
                            # Flatten video features and match audio dimensions
                            video_flat = video_features.flatten(1)  # Flatten all but batch dim
                            audio_expanded = audio_features.repeat(1, video_flat.size(1) // audio_features.size(1))
                            
                            if audio_expanded.size(1) != video_flat.size(1):
                                # Adjust audio features to match video features
                                audio_expanded = torch.nn.functional.interpolate(
                                    audio_expanded.unsqueeze(0), 
                                    size=video_flat.size(1), 
                                    mode='linear'
                                ).squeeze(0)
                            
                            combined_features = torch.cat([video_flat, audio_expanded], dim=1)
                            outputs = self.model(combined_features)
                        except Exception as e2:
                            # Option 3: If your model only takes video features
                            outputs = self.model(video_features)
                    
                    # Process outputs based on your model's output format
                    if isinstance(outputs, tuple) and len(outputs) == 2:
                        # If model returns (neuroticism, emotions)
                        neuroticism_logits, emotion_logits = outputs
                        neuroticism_score = torch.sigmoid(neuroticism_logits).cpu().numpy()[0]
                        emotion_probs = F.softmax(emotion_logits, dim=1).cpu().numpy()[0]
                        
                        emotion_labels = ['Anger', 'Disgust', 'Fear', 'Happy', 'Neutral', 'Sad']
                        emotion_scores = dict(zip(emotion_labels, emotion_probs))
                        
                    elif len(outputs.shape) == 2 and outputs.shape[1] > 1:
                        # If model returns concatenated output [neuroticism, emotion1, emotion2, ...]
                        outputs = outputs.cpu().numpy()[0]
                        neuroticism_score = float(torch.sigmoid(torch.tensor(outputs[0])))
                        emotion_probs = F.softmax(torch.tensor(outputs[1:7]), dim=0).numpy()
                        
                        emotion_labels = ['Anger', 'Disgust', 'Fear', 'Happy', 'Neutral', 'Sad']
                        emotion_scores = dict(zip(emotion_labels, emotion_probs))
                        
                    else:
                        # Single output - assume it's neuroticism, generate emotions
                        neuroticism_score = float(torch.sigmoid(outputs).cpu().numpy().flatten()[0])
                        # Derive emotions from neuroticism in a realistic way
                        base_emotions = np.array([0.15, 0.05, 0.20, 0.30, 0.25, 0.05])  # Base distribution
                        neuroticism_influence = np.array([0.3, 0.1, 0.4, -0.5, -0.2, 0.3]) * neuroticism_score
                        emotion_probs = base_emotions + neuroticism_influence
                        emotion_probs = np.maximum(emotion_probs, 0.01)  # Ensure positive
                        emotion_probs = emotion_probs / emotion_probs.sum()  # Normalize
                        
                        emotion_labels = ['Anger', 'Disgust', 'Fear', 'Happy', 'Neutral', 'Sad']
                        emotion_scores = dict(zip(emotion_labels, emotion_probs))
            else:
                # Fallback to mock predictions if model failed to load
                print("Using fallback predictions - model not loaded")
                neuroticism_score = np.random.uniform(0.2, 0.8)
                
                # Generate more realistic mock emotions
                emotion_scores = {
                    'Happy': np.random.uniform(0.1, 0.4),
                    'Neutral': np.random.uniform(0.2, 0.5),
                    'Sad': np.random.uniform(0.05, 0.3),
                    'Anger': np.random.uniform(0.0, 0.2),
                    'Fear': np.random.uniform(0.0, 0.15),
                    'Disgust': np.random.uniform(0.0, 0.1)
                }
                
                # Normalize emotion scores to sum to 1
                total = sum(emotion_scores.values())
                emotion_scores = {k: v/total for k, v in emotion_scores.items()}
            
            return {
                'neuroticism': float(neuroticism_score),
                'emotions': emotion_scores,
                'frames_processed': video_features.size(2) if video_features.dim() == 5 else 8,
                'audio_features_extracted': audio_features.size(1) > 0,
                'model_used': 'real' if self.model is not None else 'fallback'
            }
            
        except Exception as e:
            print(f"Prediction error: {e}")
            return {
                'error': f"Analysis failed: {str(e)}",
                'neuroticism': 0.0,
                'emotions': {'Error': 1.0},
                'frames_processed': 0,
                'audio_features_extracted': False,
                'model_used': 'error'
            }

def analyze_video(video_file) -> Tuple[float, str, str]:
    """
    Analyze video for personality and emotion using real model
    
    Args:
        video_file: Gradio file input
        
    Returns:
        Tuple of (neuroticism_score, emotion_analysis, detailed_results)
    """
    if video_file is None:
        return 0.0, "No video uploaded", "Please upload a video file"
    
    try:
        # Get the video path
        video_path = video_file.name if hasattr(video_file, 'name') else str(video_file)
        
        # Run analysis with real model
        results = model.predict(video_path)
        
        if 'error' in results:
            return 0.0, f"Analysis Error: {results['error']}", str(results)
        
        # Format results
        neuroticism_score = results['neuroticism']
        
        # Interpret neuroticism level
        if neuroticism_score <= 0.3:
            neuroticism_level = "Low (Emotionally Stable)"
        elif neuroticism_score <= 0.7:
            neuroticism_level = "Medium (Moderate Reactivity)"
        else:
            neuroticism_level = "High (Emotionally Sensitive)"
        
        # Format emotion analysis
        emotions = results['emotions']
        dominant_emotion = max(emotions.keys(), key=lambda k: emotions[k])
        
        emotion_text = f"**Dominant Emotion:** {dominant_emotion} ({emotions[dominant_emotion]:.1%})\n\n"
        emotion_text += "**All Emotions:**\n"
        for emotion, score in sorted(emotions.items(), key=lambda x: x[1], reverse=True):
            emotion_text += f"- {emotion}: {score:.1%}\n"
        
        # Detailed results
        model_status = "✅ Real AI Model" if results['model_used'] == 'real' else "⚠️ Fallback Mode"
        detailed_results = f"""
**Analysis Summary:**
- Neuroticism Score: {neuroticism_score:.3f}
- Neuroticism Level: {neuroticism_level}
- Frames Processed: {results['frames_processed']}
- Audio Features: {'✓' if results['audio_features_extracted'] else '✗'}

**Technical Details:**
- Model: {model_status}
- Processing: Multimodal (Video + Audio)
- Device: {'GPU' if torch.cuda.is_available() else 'CPU'}
- Confidence: {'High' if results['model_used'] == 'real' else 'Demo Mode'}
        """.strip()
        
        return neuroticism_score, emotion_text, detailed_results
        
    except Exception as e:
        error_msg = f"Processing error: {str(e)}"
        return 0.0, error_msg, error_msg

def create_interface():
    """Create the Gradio interface"""
    
    # Custom CSS for better styling
    css = """
    .gradio-container {
        font-family: 'Helvetica Neue', Arial, sans-serif;
    }
    .output-class {
        font-size: 16px;
    }
    """
    
    # Create the interface
    with gr.Blocks(css=css, title="🧠 MirrorMind Analysis") as demo:
        
        model_status_text = "✅ Real AI Model Loaded" if model.model is not None else "⚠️ Demo Mode - Model file found but architecture missing"
        
        gr.Markdown(f"""
        # 🧠 MirrorMind: AI Personality & Emotion Analysis
        
        Upload a video to analyze personality traits and emotions using your trained MirrorMind model.
        
        **Model Status:** {model_status_text}  
        **PyTorch Version:** {torch.__version__}  
        **CUDA Available:** {'Yes' if torch.cuda.is_available() else 'No'}
        
        {"**Note:** Your model file was found but contains only weights (state_dict). To use your real model, you need to either:" if model.model is None else ""}
        {"1. Save your model with the complete architecture, or" if model.model is None else ""}
        {"2. Add your model class definition to this code." if model.model is None else ""}
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                # Input
                video_input = gr.Video(
                    label="Upload Video",
                    sources=["upload"],
                )
                
                analyze_btn = gr.Button(
                    "🔍 Analyze Video", 
                    variant="primary",
                    scale=1
                )
                
                gr.Markdown("""
                **Supported formats:** MP4, AVI, MOV, WebM  
                **Optimal duration:** 4-10 seconds  
                **Requirements:** Clear face, good lighting, audio included
                """)
            
            with gr.Column(scale=2):
                # Outputs
                neuroticism_output = gr.Number(
                    label="🎭 Neuroticism Score (0.0 - 1.0)",
                    precision=3
                )
                
                emotion_output = gr.Markdown(
                    label="😊 Emotion Analysis"
                )
                
                details_output = gr.Markdown(
                    label="📊 Detailed Results"
                )
        
        # Event handlers
        analyze_btn.click(
            fn=analyze_video,
            inputs=[video_input],
            outputs=[neuroticism_output, emotion_output, details_output]
        )
        
        # Auto-analyze when video is uploaded
        video_input.change(
            fn=analyze_video,
            inputs=[video_input],
            outputs=[neuroticism_output, emotion_output, details_output]
        )
        
        gr.Markdown("""
        ---
        ### 📋 Understanding Your Results
        
        **Neuroticism Scale:**
        - **0.0-0.3:** Low - Emotionally stable, calm under pressure
        - **0.3-0.7:** Medium - Moderate emotional reactivity  
        - **0.7-1.0:** High - More emotionally sensitive, reactive
        
        **Emotions Detected:** Anger, Disgust, Fear, Happy, Neutral, Sad
        
        **Model Information:** 
        - Uses your trained `mirror_model.pth` for real AI predictions
        - Processes both video frames and audio features
        - Automatically falls back to demo mode if model loading fails
        """)
    
    return demo

# Initialize model - moved after all function definitions
print("Initializing MirrorMind model...")
model = MirrorMindModel()

# Create and launch the interface
if __name__ == "__main__":
    demo = create_interface()
    
    # Launch configuration for Hugging Face Spaces
    demo.launch(
        server_name="0.0.0.0",  # Allow external connections
        server_port=7860,       # Standard port for HF Spaces
        share=False,           # Disable share on HF Spaces
        debug=False,            # Disable debug mode for production
        show_error=True,        # Show errors to users
        quiet=False             # Show startup logs
    )