"""Custom augmentation effects operating on the RGBA signature layer.

These act on the placed signature *before* it is blended onto the document,
so they affect only the ink itself (stroke geometry, ink density) rather
than the surrounding paper.
"""

from __future__ import annotations

import cv2
import numpy as np

from src.utils.transform_utils import rotate_and_scale_rgba, smooth_noise_field


def stroke_morphology(rgba: np.ndarray, strength: float, thicken: bool) -> np.ndarray:
    """Thicken or thin ink strokes via morphological dilation/erosion of the alpha channel."""
    ksize = max(1, int(round(1 + strength * 5)))
    ksize = ksize + 1 if ksize % 2 == 0 else ksize
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ksize, ksize))
    alpha = rgba[:, :, 3]
    op = cv2.dilate if thicken else cv2.erode
    new_alpha = op(alpha, kernel, iterations=1)
    out = rgba.copy()
    out[:, :, 3] = new_alpha
    return out


def motion_blur_signature(rgba: np.ndarray, strength: float, angle_deg: float) -> np.ndarray:
    """Directional (hand-shake style) motion blur applied to the signature stamp."""
    length = max(3, int(round(3 + strength * 16)))
    kernel = np.zeros((length, length), dtype=np.float32)
    kernel[length // 2, :] = 1.0
    center = (length / 2 - 0.5, length / 2 - 0.5)
    rot = cv2.getRotationMatrix2D(center, angle_deg, 1.0)
    kernel = cv2.warpAffine(kernel, rot, (length, length))
    kernel_sum = kernel.sum()
    if kernel_sum > 1e-6:
        kernel /= kernel_sum

    out = rgba.astype(np.float32)
    for c in range(4):
        out[:, :, c] = cv2.filter2D(out[:, :, c], -1, kernel, borderType=cv2.BORDER_CONSTANT)
    return np.clip(out, 0, 255).astype(np.uint8)


def gaussian_blur_signature(rgba: np.ndarray, sigma: float) -> np.ndarray:
    """Isotropic blur of the signature stamp (softens edges, simulates scan blur/feathering)."""
    if sigma <= 0:
        return rgba
    blurred = cv2.GaussianBlur(rgba.astype(np.float32), (0, 0), sigmaX=sigma)
    return np.clip(blurred, 0, 255).astype(np.uint8)


def ink_fading(rgba: np.ndarray, strength: float, rng: np.random.Generator) -> np.ndarray:
    """Fade ink unevenly, simulating a pen running low on ink mid-stroke."""
    h, w = rgba.shape[:2]
    if h < 2 or w < 2:
        return rgba
    field = smooth_noise_field((h, w), rng, blur_frac=0.15)
    multiplier = 1.0 - strength * (1.0 - field) * 0.9
    multiplier = np.clip(multiplier, 0.05, 1.0)
    out = rgba.copy()
    out[:, :, 3] = np.clip(out[:, :, 3].astype(np.float32) * multiplier, 0, 255).astype(np.uint8)
    return out


def ink_bleed(rgba: np.ndarray, strength: float) -> np.ndarray:
    """Simulate ink bleeding/feathering into paper fibers around stroke edges."""
    ksize = max(1, int(round(1 + strength * 4)))
    ksize = ksize + 1 if ksize % 2 == 0 else ksize
    alpha = rgba[:, :, 3].astype(np.float32)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ksize, ksize))
    dilated = cv2.dilate(alpha, kernel, iterations=1)
    halo = cv2.GaussianBlur(dilated, (0, 0), sigmaX=max(ksize * 0.6, 0.8))
    bled = np.maximum(alpha, halo * (0.35 + 0.35 * strength))
    out = rgba.copy()
    out[:, :, 3] = np.clip(bled, 0, 255).astype(np.uint8)
    return out


def _fit_to_canvas(rgba: np.ndarray, canvas_h: int, canvas_w: int) -> np.ndarray:
    """Center-crop/pad ``rgba`` back to ``(canvas_h, canvas_w)``.

    ``rotate_and_scale_rgba`` auto-expands its output to avoid clipping,
    which is correct for a standalone stamp but not here: these jitter
    effects run on the full-canvas placement layer, and every downstream
    stage (blending, geometry detection) assumes that layer stays exactly
    canvas-sized.
    """
    h, w = rgba.shape[:2]
    if h == canvas_h and w == canvas_w:
        return rgba
    out = np.zeros((canvas_h, canvas_w, 4), dtype=rgba.dtype)
    src_x1, src_y1 = max((w - canvas_w) // 2, 0), max((h - canvas_h) // 2, 0)
    src = rgba[src_y1:src_y1 + min(h, canvas_h), src_x1:src_x1 + min(w, canvas_w)]
    sh, sw = src.shape[:2]
    dst_x1, dst_y1 = max((canvas_w - w) // 2, 0), max((canvas_h - h) // 2, 0)
    out[dst_y1:dst_y1 + sh, dst_x1:dst_x1 + sw] = src
    return out


def jitter_rotation(rgba: np.ndarray, extra_deg: float) -> np.ndarray:
    """Apply a small extra rotation on top of the base placement rotation."""
    h, w = rgba.shape[:2]
    return _fit_to_canvas(rotate_and_scale_rgba(rgba, extra_deg, 1.0), h, w)