"""
DeQA Metric Implementation
Deep Quality Assessment metric using the DeQA model.
"""

import torch
from transformers import AutoModelForCausalLM
from PIL import Image
from typing import List, Union, Optional
import logging

logger = logging.getLogger(__name__)


class DeQAMetric:
    """DeQA model wrapper for image quality scoring."""
    
    def __init__(self, model_name: str = "zhiyuanyou/DeQA-Score-Mix3"):
        """
        Initialize the DeQA metric.
        
        Args:
            model_name: HuggingFace model name for DeQA
        """
        self.model_name = model_name
        self.model = None
        self._is_loaded = False
        
    def load_model(self) -> None:
        """Load the DeQA scoring model."""
        if self._is_loaded:
            return
            
        logger.info("Loading DeQA model...")
        try:
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                trust_remote_code=True,
                attn_implementation="eager",
                torch_dtype=torch.float16,
                device_map="auto",
            )
            self._is_loaded = True
            logger.info("DeQA model loaded successfully!")
        except Exception as e:
            logger.error(f"Failed to load DeQA model: {e}")
            raise
    
    def score_single_image(self, image_path: str) -> Optional[float]:
        """
        Score a single image using DeQA model.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            DeQA score (0-5 scale) or None if failed
        """
        if not self._is_loaded:
            self.load_model()
        
        try:
            image = Image.open(image_path)
            scores = self.model.score([image])
            
            # Convert tensor to float
            if hasattr(scores, 'item'):
                return float(scores.item())
            elif hasattr(scores, 'tolist'):
                return float(scores.tolist()[0])
            else:
                return float(scores[0])
                
        except Exception as e:
            logger.error(f"Error scoring image {image_path}: {e}")
            return None
    
    def score_multiple_images(self, image_paths: List[str]) -> List[Optional[float]]:
        """
        Score multiple images using DeQA model.
        
        Args:
            image_paths: List of image file paths
            
        Returns:
            List of DeQA scores (0-5 scale) or None for failed images
        """
        if not self._is_loaded:
            self.load_model()
        
        try:
            images = [Image.open(path) for path in image_paths]
            scores = self.model.score(images)
            
            # Convert tensor to list of floats
            if hasattr(scores, 'tolist'):
                return [float(score) for score in scores.tolist()]
            else:
                return [float(score) for score in scores]
                
        except Exception as e:
            logger.error(f"Error scoring multiple images: {e}")
            return [None] * len(image_paths)
    
    def score_multiple_images_from_pil(self, images: List[Image.Image]) -> List[Optional[float]]:
        """
        Score multiple PIL images directly (for batch processing).
        
        Args:
            images: List of PIL Image objects
            
        Returns:
            List of DeQA scores (0-5 scale) or None for failed images
        """
        if not self._is_loaded:
            self.load_model()
        
        try:
            # Filter out None images
            valid_images = [img for img in images if img is not None]
            
            if not valid_images:
                return [None] * len(images)
            
            # Score images
            scores = self.model.score(valid_images)
            
            # Convert scores to list of floats and map back to original positions
            result_scores = []
            valid_idx = 0
            
            for img in images:
                if img is None:
                    result_scores.append(None)
                else:
                    try:
                        score = scores[valid_idx]
                        if hasattr(score, 'item'):
                            result_scores.append(float(score.item()))
                        elif hasattr(score, 'tolist'):
                            result_scores.append(float(score.tolist()))
                        else:
                            result_scores.append(float(score))
                        valid_idx += 1
                    except Exception as e:
                        logger.warning(f"Failed to convert score: {e}")
                        result_scores.append(None)
                        valid_idx += 1
            
            return result_scores
            
        except Exception as e:
            logger.error(f"Error scoring PIL images: {e}")
            return [None] * len(images)
    
    def get_metric_name(self) -> str:
        """Get the name of this metric."""
        return "DeQA"
    
    def get_score_range(self) -> tuple:
        """Get the score range for this metric."""
        return (0.0, 5.0)
    
    def get_description(self) -> str:
        """Get description of this metric."""
        return "Deep Quality Assessment using DeQA-Score-Mix3 model"
