import numpy as np
from skimage.metrics import peak_signal_noise_ratio, structural_similarity, mean_squared_error, normalized_root_mse
import imageio.v3 as imageio
from typing import Dict, List
from .core import naturalize_rgb_image


__all__ = ["compute_all_metrics", "compare_images"]


def _compute_nf(image: np.ndarray) -> float:
    """Compute the mean naturalness factor across RGB channels for an image."""
    _, nfs, _, _, _ = naturalize_rgb_image(image)
    return float(np.mean(nfs))


def compute_all_metrics(real: np.ndarray, synthetic: np.ndarray, verbose = False) -> Dict[str, float]:
    """
    Compute all image quality metrics between a reference and a synthetic image.

    Args:
        real:       Reference (real) image as a numpy array.
        synthetic:  Comparison (synthetic) image as a numpy array.
        verbose:    Whether to show status logs.

    Returns:
        Dictionary containing MSE, PSNR, SSIM, NRMSE, and dNf values.
        dNf is the delta naturalness factor (synthetic Nf minus real Nf).
    """
    if real.shape != synthetic.shape:
        raise ValueError(
            f"Images must have the same shape. Got {real.shape} and {synthetic.shape}."
        )

    if verbose: print("Computing naturalness for first image...")
    nf_real = _compute_nf(real)
    if verbose: print("Computing naturalness for second image...")
    nf_synthetic = _compute_nf(synthetic)
    
    if verbose: print("Calculating remaining metrics...")
    results = {
        "MSE":   mean_squared_error(real, synthetic) / (255 ** 2),
        "PSNR":  peak_signal_noise_ratio(real, synthetic),
        "SSIM":  structural_similarity(real, synthetic, channel_axis=-1),
        "NRMSE": normalized_root_mse(real, synthetic),
        "Nf_real":      nf_real,
        "Nf_synthetic": nf_synthetic,
        "dNf":          nf_synthetic - nf_real,
    }

    for key, value in results.items():
        if key == "PSNR":
            print(f"  {key:<12} {value:.3f} dB")
        else:
            print(f"  {key:<12} {value:.4f}")

    return results


def compare_images(real: str, images: List[str], verbose = False) -> Dict[str, Dict[str, float]]:
    """
    Compare a reference image against a list of synthetic images,
    printing a summary and returning all metrics.

    Args:
        real:   Path to a reference (real) image.
        images: List of image paths.
        verbose: Whether to show status logs

    Returns:
        Dictionary mapping each image name to its computed metrics.
    """
    results = {}
    if verbose: print(f"Reading image {real}")
    real_img = imageio.imread(real)

    for img_str in images:
        if verbose: print(f"Reading image {img_str}")
        img = imageio.imread(img_str)

        if img.shape != real_img.shape:
            print(f"Skipping '{img_str}': shape mismatch ({img.shape} vs {real_img.shape}).")
            continue

        metrics = compute_all_metrics(real_img, img, verbose)
        results[img_str] = metrics

        print(
            f"{img_str}: "
            f"PSNR={metrics['PSNR']:.2f} dB, "
            f"MSE={metrics['MSE']:.6f}, "
            f"SSIM={metrics['SSIM']:.4f}, "
            f"NRMSE={metrics['NRMSE']:.4f}, "
            f"Nf_real={metrics['Nf_real']:.4f}, "
            f"Nf_synthetic={metrics['Nf_synthetic']:.4f}, "
            f"dNf={metrics['dNf']:+.4f}"
        )

    return results