"""
Noise augmentation transforms.

Provides noise effects for document image augmentation:
- GaussianNoise: Adds Gaussian noise to simulate sensor noise
- SaltPepper: Adds salt and pepper noise for impulse noise effects
"""

from typing import Any

import numpy as np

from shared.augmentation.base import AugmentationResult, BaseAugmentation


class GaussianNoise(BaseAugmentation):
    """
    Adds Gaussian noise to the image.

    Simulates sensor noise from cameras or scanners.
    Document-safe with conservative default parameters.

    Parameters:
        mean: Mean of the Gaussian noise (default: 0).
        std: Standard deviation, can be int or (min, max) tuple (default: (5, 15)).
    """

    name = "gaussian_noise"
    affects_geometry = False

    def _validate_params(self) -> None:
        std = self.params.get("std", (5, 15))
        if isinstance(std, (int, float)):
            if std < 0:
                raise ValueError("std must be non-negative")
        elif isinstance(std, tuple):
            if len(std) != 2 or std[0] < 0 or std[1] < std[0]:
                raise ValueError("std tuple must be (min, max) with min <= max >= 0")

    def apply(
        self,
        image: np.ndarray,
        bboxes: np.ndarray | None = None,
        rng: np.random.Generator | None = None,
    ) -> AugmentationResult:
        rng = rng or np.random.default_rng()

        mean = self.params.get("mean", 0)
        std = self.params.get("std", (5, 15))

        if isinstance(std, tuple):
            std = rng.uniform(std[0], std[1])

        # Generate noise
        noise = rng.normal(mean, std, image.shape).astype(np.float32)

        # Apply noise
        noisy = image.astype(np.float32) + noise
        noisy = np.clip(noisy, 0, 255).astype(np.uint8)

        return AugmentationResult(
            image=noisy,
            bboxes=bboxes.copy() if bboxes is not None else None,
            metadata={"applied_std": std},
        )

    def get_preview_params(self) -> dict[str, Any]:
        return {"mean": 0, "std": 15}


class SaltPepper(BaseAugmentation):
    """
    Adds salt and pepper (impulse) noise to the image.

    Simulates defects from damaged sensors or transmission errors.
    Very sparse by default to preserve document readability.

    Parameters:
        amount: Proportion of pixels to affect, can be float or (min, max) tuple.
                Default: (0.001, 0.005) for very sparse noise.
        salt_vs_pepper: Ratio of salt to pepper (default: 0.5 for equal amounts).
    """

    name = "salt_pepper"
    affects_geometry = False

    def _validate_params(self) -> None:
        amount = self.params.get("amount", (0.001, 0.005))
        if isinstance(amount, (int, float)):
            if not (0 <= amount <= 1):
                raise ValueError("amount must be between 0 and 1")
        elif isinstance(amount, tuple):
            if len(amount) != 2 or not (0 <= amount[0] <= amount[1] <= 1):
                raise ValueError("amount tuple must be (min, max) in range [0, 1]")

    def apply(
        self,
        image: np.ndarray,
        bboxes: np.ndarray | None = None,
        rng: np.random.Generator | None = None,
    ) -> AugmentationResult:
        rng = rng or np.random.default_rng()

        amount = self.params.get("amount", (0.001, 0.005))
        salt_vs_pepper = self.params.get("salt_vs_pepper", 0.5)

        if isinstance(amount, tuple):
            amount = rng.uniform(amount[0], amount[1])

        # Copy image
        output = image.copy()
        h, w = image.shape[:2]
        total_pixels = h * w

        # Calculate number of salt and pepper pixels
        num_salt = int(total_pixels * amount * salt_vs_pepper)
        num_pepper = int(total_pixels * amount * (1 - salt_vs_pepper))

        # Add salt (white pixels)
        if num_salt > 0:
            salt_coords = (
                rng.integers(0, h, num_salt),
                rng.integers(0, w, num_salt),
            )
            output[salt_coords] = 255

        # Add pepper (black pixels)
        if num_pepper > 0:
            pepper_coords = (
                rng.integers(0, h, num_pepper),
                rng.integers(0, w, num_pepper),
            )
            output[pepper_coords] = 0

        return AugmentationResult(
            image=output,
            bboxes=bboxes.copy() if bboxes is not None else None,
            metadata={"applied_amount": amount},
        )

    def get_preview_params(self) -> dict[str, Any]:
        return {"amount": 0.01, "salt_vs_pepper": 0.5}
