"""
OralGuard AI - Prediction Utilities for Hugging Face Spaces
Handles image preprocessing, prediction, and ensemble logic
"""

import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import numpy as np
import time
import streamlit as st

class OralGuardPredictor:
    """
    Handles predictions for OralGuard AI models with on-demand loading
    """
    def __init__(self, model_loader, img_size=224):
        self.model_loader = model_loader
        self.img_size = img_size
        self.device = model_loader.device
        self.class_names = ['Normal', 'OSCC']
        
        # Define transforms
        self.transform = transforms.Compose([
            transforms.Resize((img_size, img_size)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])
    
    def preprocess_image(self, image):
        """
        Preprocess a PIL Image for model input
        
        Args:
            image: PIL Image
            
        Returns:
            Preprocessed tensor
        """
        # Convert to RGB if necessary
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # Apply transforms
        img_tensor = self.transform(image)
        
        # Add batch dimension
        img_tensor = img_tensor.unsqueeze(0)
        
        return img_tensor.to(self.device)
    
    def predict_single_model(self, image, model_type, magnification):
        """
        Get prediction from a single model (loads on-demand)
        
        Args:
            image: PIL Image
            model_type: 'efficientnet', 'vgg16', or 'vit'
            magnification: '100x' or '400x'
            
        Returns:
            Dictionary with prediction results
        """
        start_time = time.time()
        
        # Load model (downloads if needed)
        model = self.model_loader.load_model(model_type, magnification)
        
        # Preprocess image
        img_tensor = self.preprocess_image(image)
        
        # Make prediction
        with torch.no_grad():
            outputs = model(img_tensor)
            probabilities = F.softmax(outputs, dim=1)
            confidence, predicted = torch.max(probabilities, 1)
            
        inference_time = time.time() - start_time
        
        # Convert to numpy
        probs = probabilities.cpu().numpy()[0]
        pred_class = predicted.cpu().item()
        conf_score = confidence.cpu().item()
        
        return {
            'prediction': self.class_names[pred_class],
            'confidence': float(conf_score),
            'probabilities': {
                'Normal': float(probs[0]),
                'OSCC': float(probs[1])
            },
            'inference_time': inference_time,
            'model': f"{model_type}_{magnification}"
        }
    
    def predict_ensemble(self, image, magnification, weights=None):
        """
        Get ensemble prediction from all three models
        
        Args:
            image: PIL Image
            magnification: '100x' or '400x'
            weights: Optional weights for models [efficientnet, vgg16, vit]
            
        Returns:
            Dictionary with ensemble prediction results
        """
        start_time = time.time()
        
        # Default equal weights
        if weights is None:
            weights = [1/3, 1/3, 1/3]
        
        # Get predictions from all models
        model_types = ['efficientnet', 'vgg16', 'vit']
        all_probs = []
        individual_results = {}
        
        progress_bar = st.progress(0)
        status_text = st.empty()
        
        for idx, model_type in enumerate(model_types):
            status_text.text(f"Loading and running {model_type}...")
            
            result = self.predict_single_model(image, model_type, magnification)
            all_probs.append([
                result['probabilities']['Normal'],
                result['probabilities']['OSCC']
            ])
            individual_results[model_type] = result
            
            progress_bar.progress((idx + 1) / len(model_types))
        
        progress_bar.empty()
        status_text.empty()
        
        # Weighted average of probabilities
        all_probs = np.array(all_probs)
        weights = np.array(weights).reshape(-1, 1)
        ensemble_probs = (all_probs * weights).sum(axis=0)
        
        # Get prediction
        pred_class = np.argmax(ensemble_probs)
        confidence = float(ensemble_probs[pred_class])
        
        inference_time = time.time() - start_time
        
        return {
            'prediction': self.class_names[pred_class],
            'confidence': confidence,
            'probabilities': {
                'Normal': float(ensemble_probs[0]),
                'OSCC': float(ensemble_probs[1])
            },
            'inference_time': inference_time,
            'model': f'ensemble_{magnification}',
            'individual_results': individual_results,
            'ensemble_weights': {
                'efficientnet': float(weights[0][0]),
                'vgg16': float(weights[1][0]),
                'vit': float(weights[2][0])
            }
        }
    
    def predict_batch(self, images, model_type, magnification):
        """
        Predict on a batch of images
        
        Args:
            images: List of PIL Images
            model_type: 'efficientnet', 'vgg16', 'vit', or 'ensemble'
            magnification: '100x' or '400x'
            
        Returns:
            List of prediction results
        """
        results = []
        
        progress_bar = st.progress(0)
        status_text = st.empty()
        
        for idx, image in enumerate(images):
            status_text.text(f"Processing image {idx+1}/{len(images)}...")
            
            try:
                if model_type == 'ensemble':
                    # Use silent version to avoid nested progress bars
                    result = self.predict_ensemble_silent(image, magnification)
                else:
                    result = self.predict_single_model(image, model_type, magnification)
                
                result['image_index'] = idx
                result['status'] = 'success'
                results.append(result)
                
            except Exception as e:
                error_msg = str(e)
                st.error(f"❌ Error processing image {idx+1}: {error_msg}")
                results.append({
                    'image_index': idx,
                    'status': 'error',
                    'error': error_msg,
                    'prediction': 'Error',
                    'confidence': 0.0
                })
            
            progress_bar.progress((idx + 1) / len(images))
        
        progress_bar.empty()
        status_text.empty()
        
        # Show summary
        successful = sum(1 for r in results if r.get('status') == 'success')
        if successful < len(results):
            st.warning(f"⚠️ {len(results) - successful} images failed to process")
        
        return results
    
    def predict_ensemble_silent(self, image, magnification, weights=None):
        """
        Get ensemble prediction without UI elements (for batch processing)
        
        Args:
            image: PIL Image
            magnification: '100x' or '400x'
            weights: Optional weights for models [efficientnet, vgg16, vit]
            
        Returns:
            Dictionary with ensemble prediction results
        """
        start_time = time.time()
        
        # Default equal weights
        if weights is None:
            weights = [1/3, 1/3, 1/3]
        
        # Get predictions from all models
        model_types = ['efficientnet', 'vgg16', 'vit']
        all_probs = []
        individual_results = {}
        
        for model_type in model_types:
            result = self.predict_single_model(image, model_type, magnification)
            all_probs.append([
                result['probabilities']['Normal'],
                result['probabilities']['OSCC']
            ])
            individual_results[model_type] = result
        
        # Weighted average of probabilities
        all_probs = np.array(all_probs)
        weights = np.array(weights).reshape(-1, 1)
        ensemble_probs = (all_probs * weights).sum(axis=0)
        
        # Get prediction
        pred_class = np.argmax(ensemble_probs)
        confidence = float(ensemble_probs[pred_class])
        
        inference_time = time.time() - start_time
        
        return {
            'prediction': self.class_names[pred_class],
            'confidence': confidence,
            'probabilities': {
                'Normal': float(ensemble_probs[0]),
                'OSCC': float(ensemble_probs[1])
            },
            'inference_time': inference_time,
            'model': f'ensemble_{magnification}',
            'individual_results': individual_results,
            'ensemble_weights': {
                'efficientnet': float(weights[0][0]),
                'vgg16': float(weights[1][0]),
                'vit': float(weights[2][0])
            }
        }
    
    def get_risk_level(self, prediction, confidence):
        """
        Determine risk level based on prediction and confidence
        
        Args:
            prediction: 'Normal' or 'OSCC'
            confidence: confidence score (0-1)
            
        Returns:
            Risk level and color
        """
        if prediction == 'Normal':
            if confidence >= 0.9:
                return 'Low Risk', 'green'
            elif confidence >= 0.7:
                return 'Low-Medium Risk', 'lightgreen'
            else:
                return 'Uncertain', 'orange'
        else:  # OSCC
            if confidence >= 0.9:
                return 'High Risk', 'red'
            elif confidence >= 0.7:
                return 'Medium-High Risk', 'orange'
            else:
                return 'Uncertain', 'orange'
    
    def format_result_for_display(self, result):
        """
        Format prediction result for display
        
        Args:
            result: Prediction result dictionary
            
        Returns:
            Formatted dictionary
        """
        risk_level, risk_color = self.get_risk_level(
            result['prediction'],
            result['confidence']
        )
        
        formatted = {
            'Prediction': result['prediction'],
            'Confidence': f"{result['confidence']*100:.2f}%",
            'Risk Level': risk_level,
            'Risk Color': risk_color,
            'Probability (Normal)': f"{result['probabilities']['Normal']*100:.2f}%",
            'Probability (OSCC)': f"{result['probabilities']['OSCC']*100:.2f}%",
            'Inference Time': f"{result['inference_time']:.3f}s",
            'Model Used': result['model']
        }
        
        # Add individual results for ensemble
        if 'individual_results' in result:
            formatted['Individual Models'] = {}
            for model, res in result['individual_results'].items():
                formatted['Individual Models'][model] = {
                    'prediction': res['prediction'],
                    'confidence': f"{res['confidence']*100:.2f}%"
                }
        
        return formatted