"""
Metrics Manager
Coordinates all IQA metrics and provides a unified interface.
"""

import os
from typing import Dict, List, Optional, Any
from pathlib import Path
import logging

from .deqa_metric import DeQAMetric
from .traditional_metrics import TraditionalMetrics
from .pyiqa_metrics import PyIQAMetrics

logger = logging.getLogger(__name__)


class MetricsManager:
    """Manages all IQA metrics and provides unified scoring interface."""
    
    def __init__(self, enable_deqa: bool = True, enable_traditional: bool = True, 
                 enable_pyiqa: bool = True, pyiqa_selected_metrics: Optional[List[str]] = None):
        """
        Initialize the metrics manager.
        
        Args:
            enable_deqa: Whether to enable DeQA metric
            enable_traditional: Whether to enable traditional metrics
            enable_pyiqa: Whether to enable PyIQA metrics
        """
        self.metrics = {}
        
        # Initialize DeQA metric
        if enable_deqa:
            try:
                self.metrics['deqa'] = DeQAMetric()
                logger.info("DeQA metric initialized successfully")
            except Exception as e:
                logger.warning(f"Failed to initialize DeQA metric: {e}")
                self.metrics['deqa'] = None
        
        # Initialize traditional metrics
        if enable_traditional:
            try:
                self.metrics['traditional'] = TraditionalMetrics()
                logger.info("Traditional metrics initialized successfully")
            except Exception as e:
                logger.warning(f"Failed to initialize traditional metrics: {e}")
                self.metrics['traditional'] = None
        
        # Initialize PyIQA metrics
        if enable_pyiqa:
            try:
                self.metrics['pyiqa'] = PyIQAMetrics(selected_metrics=pyiqa_selected_metrics)
                logger.info("PyIQA metrics initialized successfully")
                
                # Preload priority metrics to avoid download delays
                logger.info("Preloading priority PyIQA metrics...")
                preload_results = self.metrics['pyiqa'].preload_priority_metrics()
                logger.info(f"Preloaded {sum(preload_results.values())}/{len(preload_results)} priority metrics")
                
            except Exception as e:
                logger.warning(f"Failed to initialize PyIQA metrics: {e}")
                self.metrics['pyiqa'] = None
    
    def get_available_metrics(self) -> List[str]:
        """Get list of available metric types."""
        return list(self.metrics.keys())
    
    def calculate_all_metrics(self, image_path: str) -> Dict[str, Any]:
        """
        Calculate all available metrics for an image.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Dictionary containing all calculated metrics
        """
        results = {}
        
        # Calculate DeQA score
        if 'deqa' in self.metrics and self.metrics['deqa'] is not None:
            try:
                deqa_score = self.metrics['deqa'].score_single_image(image_path)
                if deqa_score is not None:
                    results['deqa_score'] = deqa_score
            except Exception as e:
                logger.error(f"Error calculating DeQA score: {e}")
        
        # Calculate traditional metrics
        if 'traditional' in self.metrics and self.metrics['traditional'] is not None:
            try:
                traditional_metrics = self.metrics['traditional'].calculate_all_metrics(image_path)
                results.update(traditional_metrics)
            except Exception as e:
                logger.error(f"Error calculating traditional metrics: {e}")
        
        # Calculate PyIQA metrics (all available metrics)
        if 'pyiqa' in self.metrics and self.metrics['pyiqa'] is not None:
            try:
                # Calculate all available PyIQA metrics
                pyiqa_metrics = self.metrics['pyiqa'].calculate_all_pyiqa_metrics(image_path)
                results.update(pyiqa_metrics)
                
                # Also keep individual MUSIQ score for compatibility
                if 'pyiqa_musiq' in pyiqa_metrics:
                    results['musiq_score'] = pyiqa_metrics['pyiqa_musiq']
                    
            except Exception as e:
                logger.error(f"Error calculating PyIQA metrics: {e}")
        
        # Add file information
        try:
            file_size = os.path.getsize(image_path)
            results['file_size_bytes'] = file_size
            results['file_size_mb'] = round(file_size / (1024 * 1024), 2)
            results['file_path'] = image_path
        except Exception as e:
            logger.error(f"Error getting file information: {e}")
        
        return results
    
    def calculate_metrics_batch(self, image_paths: List[str]) -> Dict[str, Dict[str, Any]]:
        """
        Calculate metrics for multiple images.
        
        Args:
            image_paths: List of image file paths
            
        Returns:
            Dictionary mapping image names to their metrics
        """
        results = {}
        
        for image_path in image_paths:
            try:
                image_name = Path(image_path).name
                results[image_name] = self.calculate_all_metrics(image_path)
            except Exception as e:
                logger.error(f"Error processing {image_path}: {e}")
                results[Path(image_path).name] = {'error': str(e)}
        
        return results
    
    def get_metric_info(self) -> Dict[str, Dict[str, Any]]:
        """Get information about all available metrics."""
        info = {}
        
        if 'deqa' in self.metrics and self.metrics['deqa'] is not None:
            info['deqa'] = {
                'name': self.metrics['deqa'].get_metric_name(),
                'description': self.metrics['deqa'].get_description(),
                'score_range': self.metrics['deqa'].get_score_range(),
                'type': 'learning_based'
            }
        
        if 'traditional' in self.metrics and self.metrics['traditional'] is not None:
            info['traditional'] = {
                'name': 'Traditional IQA Metrics',
                'description': 'Classical image quality assessment metrics',
                'metrics': self.metrics['traditional'].get_metric_names(),
                'descriptions': self.metrics['traditional'].get_metric_descriptions(),
                'type': 'statistical'
            }
        
        if 'pyiqa' in self.metrics and self.metrics['pyiqa'] is not None:
            info['pyiqa'] = {
                'name': 'PyIQA Framework Metrics',
                'description': 'Advanced IQA metrics from PyIQA framework',
                'available_metrics': self.metrics['pyiqa'].get_available_metrics(),
                'musiq_info': self.metrics['pyiqa'].get_musiq_info(),
                'type': 'learning_based'
            }
        
        return info

    def get_selected_pyiqa_metrics(self) -> Optional[List[str]]:
        """Return the explicitly selected PyIQA metrics, if any."""
        if 'pyiqa' in self.metrics and self.metrics['pyiqa'] is not None:
            try:
                return getattr(self.metrics['pyiqa'], 'selected_metrics', None)
            except Exception:
                return None
        return None
    
    def get_deqa_score_only(self, image_path: str) -> Optional[float]:
        """Get only the DeQA score for an image."""
        if 'deqa' in self.metrics and self.metrics['deqa'] is not None:
            return self.metrics['deqa'].score_single_image(image_path)
        return None
    
    def get_deqa_scores_batch(self, image_paths: List[str]) -> List[Optional[float]]:
        """Get DeQA scores for multiple images."""
        if 'deqa' in self.metrics and self.metrics['deqa'] is not None:
            return self.metrics['deqa'].score_multiple_images(image_paths)
        return [None] * len(image_paths)
    
    def get_deqa_scorer(self):
        """Get the DeQA scorer object for batch processing."""
        if 'deqa' in self.metrics and self.metrics['deqa'] is not None:
            return self.metrics['deqa']
        return None
    
    def get_traditional_metrics_only(self, image_path: str) -> Dict[str, float]:
        """Get only traditional metrics for an image."""
        if 'traditional' in self.metrics and self.metrics['traditional'] is not None:
            return self.metrics['traditional'].calculate_all_metrics(image_path)
        return {}
    
    def get_musiq_score_only(self, image_path: str) -> Optional[float]:
        """Get only the MUSIQ score for an image."""
        if 'pyiqa' in self.metrics and self.metrics['pyiqa'] is not None:
            return self.metrics['pyiqa'].calculate_musiq_score(image_path)
        return None
    
    def get_musiq_scores_batch(self, image_paths: List[str]) -> List[Optional[float]]:
        """Get MUSIQ scores for multiple images."""
        if 'pyiqa' in self.metrics and self.metrics['pyiqa'] is not None:
            return self.metrics['pyiqa'].calculate_musiq_scores_batch(image_paths)
        return [None] * len(image_paths)
    
    def get_pyiqa_metric_score(self, metric_name: str, image_path: str, 
                              reference_path: Optional[str] = None) -> Optional[float]:
        """Get score for a specific PyIQA metric."""
        if 'pyiqa' in self.metrics and self.metrics['pyiqa'] is not None:
            return self.metrics['pyiqa'].calculate_metric_score(metric_name, image_path, reference_path)
        return None
    
    def export_metrics_summary(self, results: Dict[str, Dict[str, Any]]) -> Dict[str, Any]:
        """
        Generate a summary of all metrics across all images.
        
        Args:
            results: Results from calculate_metrics_batch
            
        Returns:
            Summary statistics for all metrics
        """
        summary = {
            'total_images': len(results),
            'successful_images': 0,
            'failed_images': 0,
            'metrics_summary': {}
        }
        
        # Count successful and failed images
        for image_name, image_results in results.items():
            if 'error' in image_results:
                summary['failed_images'] += 1
            else:
                summary['successful_images'] += 1
        
        # Calculate summary statistics for each metric
        if summary['successful_images'] > 0:
            # Get all metric names from successful results
            all_metrics = set()
            for image_results in results.values():
                if 'error' not in image_results:
                    all_metrics.update(image_results.keys())
            
            # Remove file info metrics
            file_info_metrics = {'file_size_bytes', 'file_size_mb', 'file_path'}
            metrics_to_analyze = all_metrics - file_info_metrics
            
            for metric in metrics_to_analyze:
                values = []
                for image_results in results.values():
                    if 'error' not in image_results and metric in image_results:
                        values.append(image_results[metric])
                
                if values:
                    summary['metrics_summary'][metric] = {
                        'count': len(values),
                        'mean': sum(values) / len(values),
                        'min': min(values),
                        'max': max(values),
                        'std': self._calculate_std(values)
                    }
        
        return summary
    
    def _calculate_std(self, values: List[float]) -> float:
        """Calculate standard deviation."""
        if len(values) < 2:
            return 0.0
        
        mean = sum(values) / len(values)
        variance = sum((x - mean) ** 2 for x in values) / (len(values) - 1)
        return variance ** 0.5
