"""
AI-powered background removal and manipulation.

This module provides capabilities for removing and replacing video backgrounds
using deep learning models.
"""

import torch
import numpy as np
from typing import List, Dict, Tuple, Optional, Union
import logging
from pathlib import Path
import cv2

logger = logging.getLogger(__name__)

class BackgroundRemover:
    """Background removal using deep learning."""

    def __init__(self, model_name: str = "u2net", device: str = "auto"):
        """
        Initialize background remover.

        Args:
            model_name: Name of segmentation model
            device: Device to run inference on ("cpu", "cuda", or "auto")
        """
        self.model_name = model_name

        # Determine device
        if device == "auto":
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        # Load model
        try:
            if model_name == "u2net":
                self.model = torch.hub.load("mateuszbuda/u2net", "u2net")
                self.model.to(self.device)
                self.model.eval()
            else:
                raise ValueError(f"Unsupported model: {model_name}")
        except Exception as e:
            logger.error(f"Failed to load model {model_name}: {e}")
            raise

    def remove_background(self, frame: np.ndarray,
                         threshold: float = 0.5) -> Tuple[np.ndarray, np.ndarray]:
        """
        Remove background from frame.

        Args:
            frame: Input frame (RGB)
            threshold: Segmentation threshold

        Returns:
            Tuple of (frame with transparent background, alpha mask)
        """
        # Preprocess frame
        input_size = 320
        img = cv2.resize(frame, (input_size, input_size))
        img = img / 255.0
        img = img.transpose((2, 0, 1))
        img = torch.from_numpy(img).float().unsqueeze(0).to(self.device)

        # Run inference
        with torch.no_grad():
            pred = self.model(img)
            pred = torch.sigmoid(pred[0]).squeeze()
            mask = pred.cpu().numpy()

        # Resize mask to original size
        mask = cv2.resize(mask, (frame.shape[1], frame.shape[0]))

        # Apply threshold
        mask = (mask > threshold).astype(np.float32)

        # Create RGBA output
        rgba = np.concatenate([frame, mask[..., None] * 255], axis=-1)

        return rgba, mask

    def replace_background(self, frame: np.ndarray, background: np.ndarray,
                          threshold: float = 0.5) -> np.ndarray:
        """
        Replace background with new image.

        Args:
            frame: Input frame (RGB)
            background: Background image (RGB)
            threshold: Segmentation threshold

        Returns:
            Frame with replaced background
        """
        # Remove original background
        _, mask = self.remove_background(frame, threshold)

        # Resize background if needed
        if background.shape[:2] != frame.shape[:2]:
            background = cv2.resize(background, (frame.shape[1], frame.shape[0]))

        # Combine foreground and new background
        mask = mask[..., None]
        result = frame * mask + background * (1 - mask)

        return result.astype(np.uint8)

class BackgroundBlur:
    """Background blur effect using segmentation."""

    def __init__(self, model_name: str = "u2net", device: str = "auto"):
        """
        Initialize background blur.

        Args:
            model_name: Name of segmentation model
            device: Device to run inference on ("cpu", "cuda", or "auto")
        """
        self.remover = BackgroundRemover(model_name, device)

    def blur_background(self, frame: np.ndarray, radius: float = 15,
                       threshold: float = 0.5) -> np.ndarray:
        """
        Blur video background.

        Args:
            frame: Input frame (RGB)
            radius: Blur radius
            threshold: Segmentation threshold

        Returns:
            Frame with blurred background
        """
        # Get segmentation mask
        _, mask = self.remover.remove_background(frame, threshold)
        mask = mask[..., None]

        # Create blurred version of frame
        blurred = cv2.GaussianBlur(frame, (0, 0), radius)

        # Combine original foreground with blurred background
        result = frame * mask + blurred * (1 - mask)

        return result.astype(np.uint8)

class BackgroundMotionDetector:
    """Background motion detection using segmentation."""

    def __init__(self, model_name: str = "u2net", device: str = "auto",
                 history_length: int = 5):
        """
        Initialize motion detector.

        Args:
            model_name: Name of segmentation model
            device: Device to run inference on ("cpu", "cuda", or "auto")
            history_length: Number of frames to keep in history
        """
        self.remover = BackgroundRemover(model_name, device)
        self.history_length = history_length
        self.frame_history = []
        self.mask_history = []

    def detect_motion(self, frame: np.ndarray, threshold: float = 0.5) -> np.ndarray:
        """
        Detect motion in background.

        Args:
            frame: Input frame (RGB)
            threshold: Segmentation threshold

        Returns:
            Motion mask (0-1 float array)
        """
        # Get segmentation mask
        _, mask = self.remover.remove_background(frame, threshold)

        # Update history
        self.frame_history.append(frame)
        self.mask_history.append(mask)

        if len(self.frame_history) > self.history_length:
            self.frame_history.pop(0)
            self.mask_history.pop(0)

        if len(self.frame_history) < 2:
            return np.zeros_like(mask)

        # Compute frame differences
        motion = np.zeros_like(mask)
        for i in range(len(self.frame_history) - 1):
            # Get background regions
            bg_mask1 = 1 - self.mask_history[i]
            bg_mask2 = 1 - self.mask_history[i + 1]

            # Compute difference in background regions
            diff = np.abs(self.frame_history[i + 1] - self.frame_history[i])
            diff = np.mean(diff, axis=2)
            diff = (diff > 25).astype(np.float32)  # Threshold difference

            # Add to motion mask for background regions
            motion += diff * bg_mask1 * bg_mask2

        # Normalize motion mask
        motion = np.clip(motion, 0, 1)

        return motion
