"""
PyIQA Metrics Implementation
Integrates with the PyIQA framework for various image quality assessment metrics.
"""

import torch
import pyiqa
from typing import Dict, List, Optional, Union, Any
import logging
from pathlib import Path

logger = logging.getLogger(__name__)


class PyIQAMetrics:
    """PyIQA framework integration for various IQA metrics."""
    
    def __init__(self, device: Optional[str] = None, selected_metrics: Optional[List[str]] = None):
        """
        Initialize PyIQA metrics.
        
        Args:
            device: Device to use ('cuda', 'cpu', or None for auto-detection)
        """
        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
        self.device_torch = torch.device(self.device)
        self.metrics = {}
        self.available_metrics = pyiqa.list_models()
        # Optional user-selected subset; if provided, we will prefer/limit to these
        self.selected_metrics = selected_metrics
        
        logger.info(f"PyIQA initialized on device: {self.device}")
        logger.info(f"Available metrics: {len(self.available_metrics)}")
    
    def get_available_metrics(self) -> List[str]:
        """Get list of all available PyIQA metrics."""
        return self.available_metrics
    
    def create_metric(self, metric_name: str, **kwargs) -> Any:
        """
        Create a specific PyIQA metric.
        
        Args:
            metric_name: Name of the metric to create
            **kwargs: Additional arguments for metric creation
            
        Returns:
            PyIQA metric instance
        """
        try:
            if metric_name not in self.available_metrics:
                logger.warning(f"Metric '{metric_name}' not available in PyIQA")
                return None
            
            # Create metric with device specification
            metric = pyiqa.create_metric(metric_name, device=self.device, **kwargs)
            logger.info(f"Created PyIQA metric: {metric_name}")
            return metric
            
        except Exception as e:
            logger.error(f"Failed to create metric '{metric_name}': {e}")
            return None
    
    def calculate_musiq_score(self, image_path: str) -> Optional[float]:
        """
        Calculate MUSIQ score for a single image.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            MUSIQ score or None if failed
        """
        try:
            # Create MUSIQ metric
            musiq_metric = self.create_metric('musiq')
            if musiq_metric is None:
                return None
            
            # Calculate score (MUSIQ is no-reference, so only one image needed)
            score = musiq_metric(image_path)
            
            # Convert to float if it's a tensor
            if hasattr(score, 'item'):
                score = float(score.item())
            elif hasattr(score, 'tolist'):
                score = float(score.tolist()[0])
            else:
                score = float(score)
            
            logger.debug(f"MUSIQ score for {image_path}: {score}")
            return score
            
        except Exception as e:
            logger.error(f"Error calculating MUSIQ score for {image_path}: {e}")
            return None
    
    def calculate_musiq_scores_batch(self, image_paths: List[str]) -> List[Optional[float]]:
        """
        Calculate MUSIQ scores for multiple images.
        
        Args:
            image_paths: List of image file paths
            
        Returns:
            List of MUSIQ scores or None for failed images
        """
        try:
            # Create MUSIQ metric once
            musiq_metric = self.create_metric('musiq')
            if musiq_metric is None:
                return [None] * len(image_paths)
            
            scores = []
            for image_path in image_paths:
                try:
                    score = musiq_metric(image_path)
                    
                    # Convert to float if it's a tensor
                    if hasattr(score, 'item'):
                        score = float(score.item())
                    elif hasattr(score, 'tolist'):
                        score = float(score.tolist()[0])
                    else:
                        score = float(score)
                    
                    scores.append(score)
                    logger.debug(f"MUSIQ score for {image_path}: {score}")
                    
                except Exception as e:
                    logger.error(f"Error calculating MUSIQ score for {image_path}: {e}")
                    scores.append(None)
            
            return scores
            
        except Exception as e:
            logger.error(f"Error in batch MUSIQ calculation: {e}")
            return [None] * len(image_paths)
    
    def calculate_metric_score(self, metric_name: str, image_path: str, 
                             reference_path: Optional[str] = None) -> Optional[float]:
        """
        Calculate score for a specific metric.
        
        Args:
            metric_name: Name of the PyIQA metric
            image_path: Path to the image file
            reference_path: Path to reference image (for full-reference metrics)
            
        Returns:
            Metric score or None if failed
        """
        try:
            # Create metric
            metric = self.create_metric(metric_name)
            if metric is None:
                return None
            
            # Determine FR/NR with explicit lists first, then fallback to metric attribute
            name_l = metric_name.lower()
            explicit_fr = {
                'ssim', 'ms_ssim', 'fsim', 'vif', 'dists', 'psnr', 'lpips', 'fid', 'vmaf'
            }
            explicit_nr = {
                'brisque','niqe','piqe','nrqm','nima','paq2piq','dbcnn','hyperiqa',
                'musiq','topiq_nr','clipiqa+_vitl14_512','maniqa','ahiq','unique','uranker'
            }
            if name_l in explicit_fr:
                is_full_reference = True
            elif name_l in explicit_nr:
                is_full_reference = False
            else:
                try:
                    if hasattr(metric, 'metric_mode'):
                        is_full_reference = str(getattr(metric, 'metric_mode')).upper() == 'FR'
                    else:
                        is_full_reference = name_l in explicit_fr
                except Exception:
                    is_full_reference = name_l in explicit_fr
            
            if is_full_reference and reference_path is None:
                logger.warning(f"Metric '{metric_name}' requires reference image")
                return None
            
            # Calculate score
            if is_full_reference:
                score = metric(image_path, reference_path)
            else:
                score = metric(image_path)
            
            # Convert to float if it's a tensor
            if hasattr(score, 'item'):
                score = float(score.item())
            elif hasattr(score, 'tolist'):
                score = float(score.tolist()[0])
            else:
                score = float(score)
            
            logger.debug(f"{metric_name} score for {image_path}: {score}")
            return score
            
        except Exception as e:
            logger.error(f"Error calculating {metric_name} score for {image_path}: {e}")
            return None
    
    def get_metric_info(self, metric_name: str) -> Dict[str, Any]:
        """
        Get information about a specific metric.
        
        Args:
            metric_name: Name of the metric
            
        Returns:
            Dictionary containing metric information
        """
        try:
            metric = self.create_metric(metric_name)
            if metric is None:
                return {}
            
            info = {
                'name': metric_name,
                'available': True,
                'device': self.device
            }
            
            # Get metric properties if available
            if hasattr(metric, 'lower_better'):
                info['lower_better'] = metric.lower_better
            
            if hasattr(metric, 'metric_mode'):
                info['mode'] = metric.metric_mode
            
            # Determine if it's full-reference or no-reference
            if metric_name.lower() in ['psnr', 'ssim', 'lpips', 'fid', 'vif', 'vmaf']:
                info['type'] = 'full_reference'
                info['description'] = f'Full-reference {metric_name.upper()} metric'
            else:
                info['type'] = 'no_reference'
                info['description'] = f'No-reference {metric_name.upper()} metric'
            
            return info
            
        except Exception as e:
            logger.error(f"Error getting info for metric '{metric_name}': {e}")
            return {
                'name': metric_name,
                'available': False,
                'error': str(e)
            }
    
    def get_musiq_info(self) -> Dict[str, Any]:
        """Get specific information about MUSIQ metric."""
        return self.get_metric_info('musiq')
    
    def get_metric_names_by_type(self) -> Dict[str, List[str]]:
        """Get metrics categorized by type."""
        full_reference = []
        no_reference = []
        
        for metric_name in self.available_metrics:
            if metric_name.lower() in ['psnr', 'ssim', 'lpips', 'fid', 'vif', 'vmaf']:
                full_reference.append(metric_name)
            else:
                no_reference.append(metric_name)
        
        return {
            'full_reference': full_reference,
            'no_reference': no_reference
        }
    
    def test_metric(self, metric_name: str, test_image_path: str, 
                   reference_path: Optional[str] = None) -> bool:
        """
        Test if a metric works with given images.
        
        Args:
            metric_name: Name of the metric to test
            test_image_path: Path to test image
            reference_path: Path to reference image (if needed)
            
        Returns:
            True if metric works, False otherwise
        """
        try:
            score = self.calculate_metric_score(metric_name, test_image_path, reference_path)
            return score is not None
        except Exception as e:
            logger.error(f"Test failed for metric '{metric_name}': {e}")
            return False
    
    def calculate_all_pyiqa_metrics(self, image_path: str) -> Dict[str, float]:
        """
        Calculate all available PyIQA metrics for an image.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Dictionary containing all PyIQA metric scores
        """
        results = {}
        
        # Define metrics to skip (require reference images or have download issues)
        skip_metrics = {
            'psnr', 'ssim', 'lpips', 'fid', 'vif', 'vmaf', 'brisque_matlab', 'ckdn',
            'clipiqa', 'clipiqa_plus', 'clipiqa_plus_plus', 'clipiqa_plus_plus_plus'
        }
        
        # If the user provided an explicit subset, use that set (filtered for availability)
        if self.selected_metrics:
            metric_iter = [m for m in self.selected_metrics if m in self.available_metrics and m not in skip_metrics]
        else:
            # Define priority metrics that are known to work well
            priority_metrics = [
                'musiq', 'niqe', 'brisque', 'piqe', 'paq2piq', 'hyperiqa', 'dbcnn',
                'nima', 'koncept512', 'koncept512_plus', 'koncept512_plus_plus'
            ]
            # Start with priority, then others (excluding problematic ones)
            remaining = [
                m for m in self.available_metrics
                if m not in skip_metrics and m not in priority_metrics
            ]
            metric_iter = [m for m in priority_metrics if m in self.available_metrics] + remaining
        
        # Calculate along the chosen iterator
        for metric_name in metric_iter:
            try:
                score = self.calculate_metric_score(metric_name, image_path)
                if score is not None:
                    results[f'pyiqa_{metric_name}'] = score
            except Exception as e:
                logger.debug(f"Failed to calculate {metric_name}: {e}")
                continue
        
        logger.info(f"Successfully calculated {len(results)} PyIQA metrics")
        return results
    
    def preload_priority_metrics(self) -> Dict[str, bool]:
        """
        Preload priority metrics to avoid download delays during analysis.
        
        Returns:
            Dictionary mapping metric names to success status
        """
        priority_metrics = [
            'musiq', 'niqe', 'brisque', 'piqe', 'paq2piq', 'hyperiqa', 'dbcnn',
            'nima', 'koncept512', 'koncept512_plus', 'koncept512_plus_plus'
        ]
        
        results = {}
        logger.info("Preloading priority PyIQA metrics...")
        
        for metric_name in priority_metrics:
            if metric_name in self.available_metrics:
                try:
                    metric = self.create_metric(metric_name)
                    if metric is not None:
                        results[metric_name] = True
                        logger.info(f"✓ Preloaded {metric_name}")
                    else:
                        results[metric_name] = False
                        logger.warning(f"✗ Failed to preload {metric_name}")
                except Exception as e:
                    results[metric_name] = False
                    logger.warning(f"✗ Failed to preload {metric_name}: {e}")
            else:
                results[metric_name] = False
                logger.warning(f"✗ {metric_name} not available")
        
        success_count = sum(results.values())
        logger.info(f"Preloaded {success_count}/{len(priority_metrics)} priority metrics")
        return results
