"""
Traditional IQA Metrics Implementation
Uses pyiqa library for various image quality assessment metrics.
"""

import cv2
import numpy as np
from typing import Dict, Optional, Tuple
import logging

logger = logging.getLogger(__name__)


class TraditionalMetrics:
    """Traditional IQA metrics using pyiqa and OpenCV."""
    
    def __init__(self):
        """Initialize the traditional metrics calculator."""
        self.metrics = {}
        
    def calculate_all_metrics(self, image_path: str) -> Dict[str, float]:
        """
        Calculate all traditional IQA metrics for an image.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Dictionary containing all calculated metrics
        """
        try:
            # Load image
            image = cv2.imread(image_path)
            if image is None:
                logger.error(f"Failed to load image: {image_path}")
                return {}
            
            # Convert BGR to RGB
            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            # Calculate metrics
            metrics = {}
            
            # Basic statistical metrics
            metrics.update(self._calculate_statistical_metrics(image_rgb))
            
            # Sharpness metrics
            metrics.update(self._calculate_sharpness_metrics(image_rgb))
            
            # Color metrics
            metrics.update(self._calculate_color_metrics(image_rgb))
            
            # Noise metrics
            metrics.update(self._calculate_noise_metrics(image_rgb))
            
            # Edge metrics
            metrics.update(self._calculate_edge_metrics(image_rgb))
            
            return metrics
            
        except Exception as e:
            logger.error(f"Error calculating metrics for {image_path}: {e}")
            return {}
    
    def _calculate_statistical_metrics(self, image: np.ndarray) -> Dict[str, float]:
        """Calculate basic statistical metrics."""
        metrics = {}
        
        # Convert to grayscale for some metrics
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray = image
        
        # Brightness (mean intensity)
        metrics['brightness'] = float(np.mean(gray))
        
        # Contrast (standard deviation)
        metrics['contrast'] = float(np.std(gray))
        
        # Entropy (information content)
        metrics['entropy'] = self._calculate_entropy(gray)
        
        # Variance
        metrics['variance'] = float(np.var(gray))
        
        return metrics
    
    def _calculate_sharpness_metrics(self, image: np.ndarray) -> Dict[str, float]:
        """Calculate sharpness-related metrics."""
        metrics = {}
        
        # Convert to grayscale
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray = image
        
        # Laplacian variance (sharpness)
        laplacian = cv2.Laplacian(gray, cv2.CV_64F)
        metrics['sharpness_laplacian'] = float(np.var(laplacian))
        
        # Sobel gradient magnitude
        sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        gradient_magnitude = np.sqrt(sobelx**2 + sobely**2)
        metrics['sharpness_sobel'] = float(np.mean(gradient_magnitude))
        
        # Tenengrad (gradient-based sharpness)
        metrics['sharpness_tenengrad'] = float(np.sum(gradient_magnitude))
        
        return metrics
    
    def _calculate_color_metrics(self, image: np.ndarray) -> Dict[str, float]:
        """Calculate color-related metrics."""
        metrics = {}
        
        if len(image.shape) != 3:
            return metrics
        
        # Colorfulness (saturation)
        hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
        metrics['colorfulness'] = float(np.mean(hsv[:, :, 1]))  # Saturation channel
        
        # Color variance
        metrics['color_variance'] = float(np.var(image))
        
        # RGB channel statistics
        for i, channel in enumerate(['red', 'green', 'blue']):
            metrics[f'{channel}_mean'] = float(np.mean(image[:, :, i]))
            metrics[f'{channel}_std'] = float(np.std(image[:, :, i]))
        
        return metrics
    
    def _calculate_noise_metrics(self, image: np.ndarray) -> Dict[str, float]:
        """Calculate noise-related metrics."""
        metrics = {}
        
        # Convert to grayscale
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray = image
        
        # Median absolute deviation (noise estimation)
        median = np.median(gray)
        mad = np.median(np.abs(gray - median))
        metrics['noise_mad'] = float(mad)
        
        # Noise level using high-frequency components
        kernel = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]])
        high_freq = cv2.filter2D(gray, cv2.CV_64F, kernel)
        metrics['noise_high_freq'] = float(np.std(high_freq))
        
        return metrics
    
    def _calculate_edge_metrics(self, image: np.ndarray) -> Dict[str, float]:
        """Calculate edge-related metrics."""
        metrics = {}
        
        # Convert to grayscale
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray = image
        
        # Canny edge detection
        edges = cv2.Canny(gray, 100, 200)
        edge_density = np.sum(edges > 0) / (edges.shape[0] * edges.shape[1])
        metrics['edge_density'] = float(edge_density)
        
        # Edge strength
        sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        edge_strength = np.sqrt(sobelx**2 + sobely**2)
        metrics['edge_strength'] = float(np.mean(edge_strength))
        
        return metrics
    
    def _calculate_entropy(self, image: np.ndarray) -> float:
        """Calculate image entropy."""
        hist = cv2.calcHist([image], [0], None, [256], [0, 256])
        hist = hist / hist.sum()
        hist = hist[hist > 0]  # Remove zero probabilities
        entropy = -np.sum(hist * np.log2(hist))
        return float(entropy)
    
    def get_metric_names(self) -> list:
        """Get list of all available metric names."""
        return [
            'brightness', 'contrast', 'entropy', 'variance',
            'sharpness_laplacian', 'sharpness_sobel', 'sharpness_tenengrad',
            'colorfulness', 'color_variance', 'red_mean', 'green_mean', 'blue_mean',
            'red_std', 'green_std', 'blue_std',
            'noise_mad', 'noise_high_freq',
            'edge_density', 'edge_strength'
        ]
    
    def get_metric_descriptions(self) -> Dict[str, str]:
        """Get descriptions for all metrics."""
        return {
            'brightness': 'Average pixel intensity (0-255)',
            'contrast': 'Standard deviation of pixel values',
            'entropy': 'Information content of the image',
            'variance': 'Variance of pixel values',
            'sharpness_laplacian': 'Sharpness using Laplacian operator',
            'sharpness_sobel': 'Sharpness using Sobel gradient',
            'sharpness_tenengrad': 'Sharpness using gradient magnitude',
            'colorfulness': 'Color saturation level',
            'color_variance': 'Overall color variance',
            'red_mean': 'Red channel mean value',
            'green_mean': 'Green channel mean value',
            'blue_mean': 'Blue channel mean value',
            'red_std': 'Red channel standard deviation',
            'green_std': 'Green channel standard deviation',
            'blue_std': 'Blue channel standard deviation',
            'noise_mad': 'Noise level using median absolute deviation',
            'noise_high_freq': 'Noise level using high-frequency components',
            'edge_density': 'Percentage of edge pixels',
            'edge_strength': 'Average edge strength'
        }
