optimizer.py

"""
Ultra Supreme Optimizer - Main optimization engine for image analysis
"""

# IMPORTANT: spaces must be imported BEFORE torch or any CUDA-using library
import spaces
import gc
import logging
from datetime import datetime
from typing import Tuple, Dict, Any, Optional

import torch
import numpy as np
from PIL import Image
from clip_interrogator import Config, Interrogator

from analyzer import UltraSupremeAnalyzer

logger = logging.getLogger(__name__)


class UltraSupremeOptimizer:
    """Main optimizer class for ultra supreme image analysis"""
    
    def __init__(self):
        self.interrogator: Optional[Interrogator] = None
        self.analyzer = UltraSupremeAnalyzer()
        self.usage_count = 0
        self.device = self._get_device()
        self.is_initialized = False

    @staticmethod
    def _get_device() -> str:
        """Determine the best available device for computation"""
        if torch.cuda.is_available():
            return "cuda"
        elif torch.backends.mps.is_available():
            return "mps"
        else:
            return "cpu"

    def initialize_model(self) -> bool:
        """Initialize the CLIP interrogator model"""
        if self.is_initialized:
            return True
        
        try:
            config = Config(
                clip_model_name="ViT-L-14/openai",
                download_cache=True,
                chunk_size=2048,
                quiet=True,
                device=self.device
            )
            
            self.interrogator = Interrogator(config)
            self.is_initialized = True
            
            # Clean up memory after initialization
            if self.device == "cpu":
                gc.collect()
            else:
                torch.cuda.empty_cache()
                
            return True
            
        except Exception as e:
            logger.error(f"Initialization error: {e}")
            return False

    def optimize_image(self, image: Any) -> Optional[Image.Image]:
        """Optimize image for processing"""
        if image is None:
            return None
            
        try:
            # Convert to PIL Image if necessary
            if isinstance(image, np.ndarray):
                image = Image.fromarray(image)
            elif not isinstance(image, Image.Image):
                image = Image.open(image)
            
            # Convert to RGB if necessary
            if image.mode != 'RGB':
                image = image.convert('RGB')
            
            # Resize if too large
            max_size = 768 if self.device != "cpu" else 512
            if image.size[0] > max_size or image.size[1] > max_size:
                image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
            
            return image
            
        except Exception as e:
            logger.error(f"Image optimization error: {e}")
            return None

    @spaces.GPU
    def generate_ultra_supreme_prompt(self, image: Any) -> Tuple[str, str, int, Dict[str, int]]:
        """
        Generate ultra supreme prompt from image
        
        Returns:
            Tuple of (prompt, analysis_info, score, breakdown)
        """
        try:
            # Initialize model if needed
            if not self.is_initialized:
                if not self.initialize_model():
                    return "❌ Model initialization failed.", "Please refresh and try again.", 0, {}
            
            # Validate input
            if image is None:
                return "❌ Please upload an image.", "No image provided.", 0, {}
            
            self.usage_count += 1
            
            # Optimize image
            image = self.optimize_image(image)
            if image is None:
                return "❌ Image processing failed.", "Invalid image format.", 0, {}
            
            start_time = datetime.now()
            
            # ULTRA SUPREME TRIPLE CLIP ANALYSIS
            logger.info("ULTRA SUPREME ANALYSIS - Maximum intelligence deployment")
            
            clip_fast = self.interrogator.interrogate_fast(image)
            clip_classic = self.interrogator.interrogate_classic(image) 
            clip_best = self.interrogator.interrogate(image)
            
            logger.info(f"ULTRA CLIP Results:\nFast: {clip_fast}\nClassic: {clip_classic}\nBest: {clip_best}")
            
            # ULTRA SUPREME ANALYSIS
            ultra_analysis = self.analyzer.ultra_supreme_analysis(clip_fast, clip_classic, clip_best)
            
            # BUILD ULTRA SUPREME FLUX PROMPT
            optimized_prompt = self.analyzer.build_ultra_supreme_prompt(
                ultra_analysis, 
                [clip_fast, clip_classic, clip_best]
            )
            
            # CALCULATE ULTRA SUPREME SCORE
            score, breakdown = self.analyzer.calculate_ultra_supreme_score(optimized_prompt, ultra_analysis)
            
            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()
            
            # Memory cleanup
            if self.device == "cpu":
                gc.collect()
            else:
                torch.cuda.empty_cache()
            
            # Generate analysis report
            analysis_info = self._generate_analysis_report(
                ultra_analysis, clip_fast, clip_classic, clip_best, 
                score, breakdown, duration
            )
            
            return optimized_prompt, analysis_info, score, breakdown
            
        except Exception as e:
            logger.error(f"Ultra supreme generation error: {e}")
            return f"❌ Error: {str(e)}", "Please try with a different image.", 0, {}

    def _generate_analysis_report(self, ultra_analysis: Dict[str, Any], 
                                  clip_fast: str, clip_classic: str, clip_best: str,
                                  score: int, breakdown: Dict[str, int], 
                                  duration: float) -> str:
        """Generate detailed analysis report"""
        
        gpu_status = "⚡ ZeroGPU" if torch.cuda.is_available() else "💻 CPU"
        
        # Format detected elements - Fixed the .title() error by checking for None
        features = ", ".join(ultra_analysis["facial_ultra"]["facial_hair"]) if ultra_analysis["facial_ultra"]["facial_hair"] else "None detected"
        cultural = ", ".join(ultra_analysis["demographic"]["cultural_religious"]) if ultra_analysis["demographic"]["cultural_religious"] else "None detected"
        clothing = ", ".join(ultra_analysis["clothing_accessories"]["eyewear"] + ultra_analysis["clothing_accessories"]["headwear"]) if ultra_analysis["clothing_accessories"]["eyewear"] or ultra_analysis["clothing_accessories"]["headwear"] else "None detected"
        
        # Safe access to potentially None values
        age_category = ultra_analysis["demographic"].get("age_category", "Unspecified")
        if age_category and age_category != "Unspecified":
            age_category = age_category.replace("_", " ").title()
        
        setting_type = ultra_analysis["environmental"].get("setting_type", "Standard")
        if setting_type and setting_type != "Standard":
            setting_type = setting_type.title()
        
        primary_emotion = ultra_analysis["emotional_state"].get("primary_emotion", "Neutral")
        if primary_emotion and primary_emotion != "Neutral":
            primary_emotion = primary_emotion.title()
        
        analysis_info = f"""**🚀 ULTRA SUPREME ANALYSIS COMPLETE**
**Processing:** {gpu_status} • {duration:.1f}s • Triple CLIP Ultra Intelligence  
**Ultra Score:** {score}/100 • Breakdown: Structure({breakdown.get('structure',0)}) Features({breakdown.get('features',0)}) Cultural({breakdown.get('cultural',0)}) Emotional({breakdown.get('emotional',0)}) Technical({breakdown.get('technical',0)})  
**Generation:** #{self.usage_count}  
**🧠 ULTRA DEEP DETECTION:**
- **Age Category:** {age_category} (Confidence: {ultra_analysis["demographic"].get("age_confidence", 0)})
- **Cultural Context:** {cultural}  
- **Facial Features:** {features}
- **Accessories:** {clothing}
- **Setting:** {setting_type}
- **Emotion:** {primary_emotion}
- **Total Features:** {ultra_analysis["intelligence_metrics"]["total_features_detected"]}
**📊 CLIP ANALYSIS SOURCES:**
- **Fast:** {clip_fast[:50]}...
- **Classic:** {clip_classic[:50]}...  
- **Best:** {clip_best[:50]}...
**⚡ ULTRA OPTIMIZATION:** Applied absolute maximum depth analysis with Pariente AI research rules"""
        
        return analysis_info