"""Post-processing effects for rendered images."""

import logging
from typing import Tuple, Optional

import torch
import torch.nn.functional as F

logger = logging.getLogger(__name__)


def apply_bloom(
    images: torch.Tensor,
    intensity: float = 0.5,
    threshold: float = 0.8,
    radius: int = 20
) -> torch.Tensor:
    """
    Apply bloom/glow effect to bright areas.

    Args:
        images: Input images [N, H, W, C] or [H, W, C]
        intensity: Bloom intensity
        threshold: Brightness threshold
        radius: Blur radius for bloom

    Returns:
        Images with bloom applied
    """
    was_3d = len(images.shape) == 3
    if was_3d:
        images = images.unsqueeze(0)

    device = images.device
    N, H, W, C = images.shape

    output = []
    for frame in images:
        # Extract bright areas
        brightness = frame.mean(dim=-1, keepdim=True)
        bright_mask = (brightness > threshold).float()
        bright_areas = frame * bright_mask

        # Blur the bright areas
        # Convert to NCHW for convolution
        bright_nchw = bright_areas.permute(2, 0, 1).unsqueeze(0)  # [1, C, H, W]

        # Create Gaussian kernel
        kernel_size = radius * 2 + 1
        sigma = radius / 3
        kernel_1d = _gaussian_kernel_1d(kernel_size, sigma, device)

        # Separable blur
        padded = F.pad(bright_nchw, (radius, radius, radius, radius), mode='reflect')

        # Horizontal blur
        blurred_h = F.conv2d(
            padded,
            kernel_1d.view(1, 1, 1, -1).expand(C, 1, 1, -1),
            groups=C
        )

        # Vertical blur
        blurred = F.conv2d(
            blurred_h,
            kernel_1d.view(1, 1, -1, 1).expand(C, 1, -1, 1),
            groups=C
        )

        # Convert back to HWC
        bloom = blurred.squeeze(0).permute(1, 2, 0)

        # Composite
        result = frame + bloom * intensity

        output.append(result.clamp(0, 1))

    result = torch.stack(output)

    if was_3d:
        result = result.squeeze(0)

    return result


def apply_depth_of_field(
    images: torch.Tensor,
    depth_map: torch.Tensor,
    focal_distance: float = 0.5,
    aperture: float = 0.1,
    max_blur: int = 20
) -> torch.Tensor:
    """
    Apply depth of field blur effect.

    Args:
        images: Input images [N, H, W, C]
        depth_map: Depth values [N, H, W, 1] (0 = close, 1 = far)
        focal_distance: Distance in focus (0-1)
        aperture: Aperture size (larger = more blur)
        max_blur: Maximum blur radius

    Returns:
        Images with DoF applied
    """
    was_3d = len(images.shape) == 3
    if was_3d:
        images = images.unsqueeze(0)
        depth_map = depth_map.unsqueeze(0)

    device = images.device
    N, H, W, C = images.shape

    output = []
    for i in range(N):
        frame = images[i]
        depth = depth_map[i]

        # Compute blur amount based on distance from focal plane
        blur_amount = torch.abs(depth - focal_distance) * aperture * max_blur
        blur_amount = blur_amount.squeeze(-1)  # [H, W]

        # For simplicity, use uniform blur based on average
        # (true DoF would need per-pixel blur which is expensive)
        avg_blur = int(blur_amount.mean().item())

        if avg_blur > 0:
            # Convert for convolution
            frame_nchw = frame.permute(2, 0, 1).unsqueeze(0)

            # Gaussian blur
            kernel_size = avg_blur * 2 + 1
            sigma = avg_blur / 3
            kernel = _gaussian_kernel_2d(kernel_size, sigma, device)

            padded = F.pad(frame_nchw, (avg_blur, avg_blur, avg_blur, avg_blur), mode='reflect')
            blurred = F.conv2d(padded, kernel.unsqueeze(0).unsqueeze(0).expand(C, 1, -1, -1), groups=C)

            blurred_hwc = blurred.squeeze(0).permute(1, 2, 0)

            # Blend based on blur amount
            blend = (blur_amount / max_blur).unsqueeze(-1).clamp(0, 1)
            result = frame * (1 - blend) + blurred_hwc * blend
        else:
            result = frame

        output.append(result)

    result = torch.stack(output)

    if was_3d:
        result = result.squeeze(0)

    return result


def apply_motion_blur(
    images: torch.Tensor,
    intensity: float = 0.5
) -> torch.Tensor:
    """
    Apply simple motion blur between frames.

    Args:
        images: Video frames [N, H, W, C]
        intensity: Blur intensity (0-1)

    Returns:
        Motion blurred frames
    """
    if len(images.shape) != 4 or images.shape[0] < 2:
        return images

    N, H, W, C = images.shape
    output = images.clone()

    for i in range(1, N):
        # Blend with previous frame
        prev = images[i - 1]
        curr = images[i]
        output[i] = curr * (1 - intensity) + prev * intensity

    return output


def apply_outline(
    images: torch.Tensor,
    width: float = 1.0,
    color: Tuple[float, float, float] = (0.0, 0.0, 0.0),
    threshold: float = 0.1
) -> torch.Tensor:
    """
    Apply edge outline effect.

    Args:
        images: Input images [N, H, W, C] or [H, W, C]
        width: Outline width
        color: Outline color
        threshold: Edge detection threshold

    Returns:
        Images with outlines
    """
    was_3d = len(images.shape) == 3
    if was_3d:
        images = images.unsqueeze(0)

    device = images.device
    N, H, W, C = images.shape

    # Sobel kernels
    sobel_x = torch.tensor([
        [-1, 0, 1],
        [-2, 0, 2],
        [-1, 0, 1]
    ], dtype=images.dtype, device=device)

    sobel_y = torch.tensor([
        [-1, -2, -1],
        [0, 0, 0],
        [1, 2, 1]
    ], dtype=images.dtype, device=device)

    output = []
    for frame in images:
        # Convert to grayscale
        gray = frame.mean(dim=-1)

        # Pad for convolution
        gray_4d = gray.unsqueeze(0).unsqueeze(0)  # [1, 1, H, W]
        padded = F.pad(gray_4d, (1, 1, 1, 1), mode='replicate')

        # Apply Sobel
        gx = F.conv2d(padded, sobel_x.view(1, 1, 3, 3))
        gy = F.conv2d(padded, sobel_y.view(1, 1, 3, 3))

        # Edge magnitude
        edges = torch.sqrt(gx ** 2 + gy ** 2).squeeze()

        # Dilate edges for width
        if width > 1:
            dilation = int(width)
            edges_4d = edges.unsqueeze(0).unsqueeze(0)
            kernel = torch.ones(1, 1, dilation * 2 + 1, dilation * 2 + 1, device=device)
            edges_dilated = F.conv2d(
                F.pad(edges_4d, (dilation, dilation, dilation, dilation), mode='replicate'),
                kernel
            ).squeeze()
            edges = (edges_dilated > 0).float()
        else:
            edges = (edges > threshold * width).float()

        # Apply outline color
        outline_mask = edges.unsqueeze(-1)
        outline_color = torch.tensor(color, dtype=frame.dtype, device=device)
        result = frame * (1 - outline_mask) + outline_color.view(1, 1, 3) * outline_mask

        output.append(result)

    result = torch.stack(output)

    if was_3d:
        result = result.squeeze(0)

    return result


def apply_vignette(
    images: torch.Tensor,
    intensity: float = 0.5,
    radius: float = 0.8
) -> torch.Tensor:
    """
    Apply vignette effect (darkening at edges).

    Args:
        images: Input images [N, H, W, C] or [H, W, C]
        intensity: Vignette intensity
        radius: Vignette radius (1 = no effect at center)

    Returns:
        Images with vignette
    """
    was_3d = len(images.shape) == 3
    if was_3d:
        images = images.unsqueeze(0)

    device = images.device
    N, H, W, C = images.shape

    # Create vignette mask
    y, x = torch.meshgrid(
        torch.linspace(-1, 1, H, device=device),
        torch.linspace(-1, 1, W, device=device),
        indexing='ij'
    )

    dist = torch.sqrt(x ** 2 + y ** 2)
    vignette = 1 - torch.clamp((dist - radius) / (1 - radius), 0, 1) * intensity
    vignette = vignette.unsqueeze(-1)  # [H, W, 1]

    output = images * vignette.unsqueeze(0)

    if was_3d:
        output = output.squeeze(0)

    return output


def _gaussian_kernel_1d(size: int, sigma: float, device: torch.device) -> torch.Tensor:
    """Create 1D Gaussian kernel."""
    x = torch.arange(size, dtype=torch.float32, device=device) - size // 2
    kernel = torch.exp(-x ** 2 / (2 * sigma ** 2))
    return kernel / kernel.sum()


def _gaussian_kernel_2d(size: int, sigma: float, device: torch.device) -> torch.Tensor:
    """Create 2D Gaussian kernel."""
    kernel_1d = _gaussian_kernel_1d(size, sigma, device)
    kernel_2d = kernel_1d.unsqueeze(1) @ kernel_1d.unsqueeze(0)
    return kernel_2d