"""
Quality metrics for image compression evaluation.

Provides PSNR, SSIM, and other standard image quality metrics.
"""

from __future__ import annotations
import numpy as np
from typing import Dict, Tuple
try:
    from skimage.metrics import structural_similarity as ssim_func
    from skimage.metrics import peak_signal_noise_ratio as psnr_func
    SKIMAGE_AVAILABLE = True
except ImportError:
    SKIMAGE_AVAILABLE = False
    ssim_func = None
    psnr_func = None


def psnr(original: np.ndarray, compressed: np.ndarray, data_range: float = 255.0) -> float:
    """
    Calculate Peak Signal-to-Noise Ratio.
    
    Higher is better. Typical values:
    - 30-40 dB: Good quality
    - 40-50 dB: Excellent quality
    - >50 dB: Near-perfect
    
    Args:
        original: Original image array
        compressed: Compressed/reconstructed image array
        data_range: Maximum possible pixel value (255 for uint8)
        
    Returns:
        PSNR in decibels
    """
    if psnr_func is None:
        raise ImportError("scikit-image is required for PSNR. Install with: pip install sempress[image]")
    return float(psnr_func(original, compressed, data_range=data_range))


def ssim(original: np.ndarray, compressed: np.ndarray, 
         data_range: float = 255.0, multichannel: bool = None) -> float:
    """
    Calculate Structural Similarity Index.
    
    Ranges from -1 to 1, where 1 is perfect similarity.
    Typical values:
    - 0.9-1.0: Excellent quality
    - 0.8-0.9: Good quality
    - <0.8: Visible degradation
    
    Args:
        original: Original image array
        compressed: Compressed/reconstructed image array
        data_range: Maximum possible pixel value
        multichannel: Whether to treat as color image. Auto-detected if None.
        
    Returns:
        SSIM score (0 to 1)
    """
    if ssim_func is None:
        raise ImportError("scikit-image is required for SSIM. Install with: pip install sempress[image]")
    if multichannel is None:
        multichannel = original.ndim == 3 and original.shape[2] > 1

    # Use channel_axis parameter for newer scikit-image versions
    if multichannel:
        return float(ssim_func(original, compressed, data_range=data_range,
                              channel_axis=2))
    else:
        return float(ssim_func(original, compressed, data_range=data_range))


def mse(original: np.ndarray, compressed: np.ndarray) -> float:
    """Calculate Mean Squared Error."""
    return float(np.mean((original.astype(np.float32) - compressed.astype(np.float32)) ** 2))


def mae(original: np.ndarray, compressed: np.ndarray) -> float:
    """Calculate Mean Absolute Error."""
    return float(np.mean(np.abs(original.astype(np.float32) - compressed.astype(np.float32))))


def compression_ratio(original_size: int, compressed_size: int) -> float:
    """Calculate compression ratio."""
    return float(original_size / compressed_size) if compressed_size > 0 else 0.0


def bits_per_pixel(compressed_size: int, image_shape: Tuple[int, ...]) -> float:
    """
    Calculate bits per pixel (BPP).
    
    Lower is better. Typical values:
    - JPEG: 0.5-2.0 BPP
    - PNG: 4-8 BPP (lossless)
    - Sempress: Target 0.3-1.5 BPP
    """
    h, w = image_shape[:2]
    total_pixels = h * w
    total_bits = compressed_size * 8
    return float(total_bits / total_pixels)


def eval_image(original: np.ndarray, reconstructed: np.ndarray, 
               compressed_size: int) -> Dict[str, float]:
    """
    Comprehensive image quality evaluation.
    
    Args:
        original: Original image array
        reconstructed: Reconstructed image array
        compressed_size: Size of compressed blob in bytes
        
    Returns:
        Dictionary of quality metrics
    """
    # Calculate file sizes
    original_size = original.nbytes
    
    # Quality metrics
    psnr_val = psnr(original, reconstructed)
    ssim_val = ssim(original, reconstructed)
    mse_val = mse(original, reconstructed)
    mae_val = mae(original, reconstructed)
    
    # Compression metrics
    comp_ratio = compression_ratio(original_size, compressed_size)
    bpp = bits_per_pixel(compressed_size, original.shape)
    
    return {
        "psnr_db": psnr_val,
        "ssim": ssim_val,
        "mse": mse_val,
        "mae": mae_val,
        "compression_ratio": comp_ratio,
        "bits_per_pixel": bpp,
        "original_size_bytes": original_size,
        "compressed_size_bytes": compressed_size,
        "size_reduction_percent": (1 - compressed_size / original_size) * 100
    }


def compare_codecs(original: np.ndarray, compressed_versions: Dict[str, Tuple[np.ndarray, int]]) -> Dict:
    """
    Compare multiple compression codecs.
    
    Args:
        original: Original image
        compressed_versions: Dict mapping codec name to (reconstructed_image, compressed_size)
        
    Returns:
        Comparison table
    """
    results = {}
    
    for codec_name, (recon, comp_size) in compressed_versions.items():
        results[codec_name] = eval_image(original, recon, comp_size)
    
    return results


def quality_grade(psnr_val: float, ssim_val: float) -> str:
    """
    Assign quality grade based on PSNR and SSIM.
    
    Returns:
        Grade: A+ (excellent), A (very good), B (good), C (fair), D (poor)
    """
    if psnr_val >= 45 and ssim_val >= 0.98:
        return "A+"
    elif psnr_val >= 40 and ssim_val >= 0.95:
        return "A"
    elif psnr_val >= 35 and ssim_val >= 0.90:
        return "B"
    elif psnr_val >= 30 and ssim_val >= 0.85:
        return "C"
    else:
        return "D"
