"""
Image noise models and denoising methods.

Image noise is fundamentally different from 1-D time-series noise:
spatial correlation matters, the noise is often signal-dependent (Poisson
shot noise in photon-limited imaging), and the human visual system is the
ultimate judge of what denoising "should" look like.

Noise models
------------
* gaussian: read noise, thermal noise in CMOS
* salt_pepper: dead/hot pixels, transmission errors
* speckle: coherent imaging (ultrasound, SAR, laser)
* shot_noise: photon counting (Poisson, variance = mean)

Denoising
---------
* denoise_median: salt-and-pepper removal
* denoise_bilateral: edge-preserving Gaussian smoothing
"""

import numpy as np
from scipy.ndimage import median_filter
from scipy.signal import convolve2d


def add_gaussian(img: np.ndarray, sigma: float = 0.05,
                 seed: int | None = None) -> np.ndarray:
    """Add Gaussian white noise to an image.

    Parameters
    ----------
    img : ndarray, values in [0, 1]
        Input image.
    sigma : float
        Noise standard deviation.
    seed : int or None
        RNG seed.

    Returns
    -------
    noisy : ndarray
        Image with additive Gaussian noise, clipped to [0, 1].
    """
    rng = np.random.default_rng(seed)
    noise = rng.normal(0.0, sigma, size=img.shape)
    return np.clip(img + noise, 0.0, 1.0)


def add_salt_pepper(img: np.ndarray, prob: float = 0.02,
                    seed: int | None = None) -> np.ndarray:
    """Add salt-and-pepper noise (random dead/hot pixels).

    Parameters
    ----------
    img : ndarray, values in [0, 1]
    prob : float
        Probability of a pixel being corrupted (half salt, half pepper).
    seed : int or None

    Returns
    -------
    noisy : ndarray
    """
    rng = np.random.default_rng(seed)
    noisy = img.copy()
    mask = rng.random(size=img.shape) < prob
    salt_mask = mask & (rng.random(size=img.shape) < 0.5)
    pepper_mask = mask & ~salt_mask
    noisy[salt_mask] = 1.0
    noisy[pepper_mask] = 0.0
    return noisy


def add_speckle(img: np.ndarray, sigma: float = 0.1,
                seed: int | None = None) -> np.ndarray:
    """Add multiplicative speckle noise.

    Speckle noise multiplies the signal: noisy = img * (1 + N(0, sigma)).
    Characteristic of coherent imaging (ultrasound, SAR, laser).

    Parameters
    ----------
    img : ndarray, values in [0, 1]
    sigma : float
        Noise standard deviation (relative).
    seed : int or None

    Returns
    -------
    noisy : ndarray
    """
    rng = np.random.default_rng(seed)
    noise = rng.normal(0.0, sigma, size=img.shape)
    return np.clip(img * (1.0 + noise), 0.0, 1.0)


def add_shot_noise(img: np.ndarray, peak_photons: float = 100.0,
                   seed: int | None = None) -> np.ndarray:
    """Add Poisson (photon shot) noise.

    Models the fundamental quantum limit of photon detection.  Each pixel
    value is treated as the expected photon count scaled to [0, 1], and a
    Poisson draw replaces it.  Variance equals the mean; darker regions
    have lower absolute noise but higher noise relative to signal.

    Parameters
    ----------
    img : ndarray, values in [0, 1]
    peak_photons : float
        Expected photon count at full brightness (img=1.0).  Lower values
        = noisier image (fewer photons collected).
    seed : int or None

    Returns
    -------
    noisy : ndarray
    """
    rng = np.random.default_rng(seed)
    photons = img * peak_photons
    noisy_photons = rng.poisson(photons)
    return np.clip(noisy_photons / peak_photons, 0.0, 1.0)


def denoise_median(img: np.ndarray, size: int = 3) -> np.ndarray:
    """Median filter, optimal for salt-and-pepper noise.

    Parameters
    ----------
    img : ndarray
    size : int
        Filter window size (must be odd).

    Returns
    -------
    filtered : ndarray
    """
    return median_filter(img, size=size)


def denoise_bilateral(img: np.ndarray, sigma_spatial: float = 3.0,
                      sigma_intensity: float = 0.1) -> np.ndarray:
    """Edge-preserving bilateral filter.

    Weights pixels by both spatial distance (Gaussian) and intensity
    difference (Gaussian, so edges are preserved).  This is a simplified
    implementation for educational use; production code should use the
    fast approximations in OpenCV or scikit-image.

    Parameters
    ----------
    img : ndarray
    sigma_spatial : float
        Spatial Gaussian sigma (pixels).
    sigma_intensity : float
        Intensity Gaussian sigma (same units as image, typically [0,1]).

    Returns
    -------
    filtered : ndarray
    """
    radius = int(np.ceil(3 * sigma_spatial))
    rows, cols = img.shape
    result = np.zeros_like(img)

    # Spatial Gaussian kernel.
    y, x = np.ogrid[-radius:radius+1, -radius:radius+1]
    g_spatial = np.exp(-(x**2 + y**2) / (2 * sigma_spatial**2))

    for i in range(rows):
        # Bounds for the neighbourhood.
        r_min = max(0, i - radius)
        r_max = min(rows, i + radius + 1)
        for j in range(cols):
            c_min = max(0, j - radius)
            c_max = min(cols, j + radius + 1)

            patch = img[r_min:r_max, c_min:c_max]
            g_intensity = np.exp(-(patch - img[i, j])**2 / (2 * sigma_intensity**2))

            # Align the spatial kernel to the patch.
            dy = r_min - (i - radius)
            dx = c_min - (j - radius)
            g_s = g_spatial[dy:dy+patch.shape[0], dx:dx+patch.shape[1]]

            weights = g_s * g_intensity
            result[i, j] = np.sum(weights * patch) / np.sum(weights)

    return result