import cv2
import numpy as np
from skimage import exposure


class ImageFilters:
    """
    Filtros e melhorias de imagem usando OpenCV.
    """
    
    def __init__(self, denoise_strength=10, contrast_clip=3.0):
        """
        Inicializa os filtros.
        
        Args:
            denoise_strength: Força do denoise (0-30)
            contrast_clip: Intensidade do CLAHE (1.0-5.0)
        """
        self.denoise_strength = denoise_strength
        self.contrast_clip = contrast_clip
    
    def denoise(self, img, strength=None):
        """
        Remove ruído da imagem.
        
        Args:
            img: Imagem BGR
            strength: Força do denoise (sobrescreve padrão)
            
        Returns:
            Imagem com ruído reduzido
        """
        if strength is None:
            strength = self.denoise_strength
        
        # Converter para o tipo correto se necessário
        if img.dtype != np.uint8:
            img = img.astype(np.uint8)
        
        if len(img.shape) == 3:
            # Imagem colorida
            return cv2.fastNlMeansDenoisingColored(
                img, None,
                h=strength,
                hColor=strength,
                templateWindowSize=7,
                searchWindowSize=21
            )
        else:
            # Imagem em escala de cinza
            return cv2.fastNlMeansDenoising(
                img, None,
                h=strength,
                templateWindowSize=7,
                searchWindowSize=21
            )
    
    def enhance_contrast(self, img):
        """
        Melhora o contraste usando CLAHE (Contrast Limited Adaptive Histogram Equalization).
        
        Args:
            img: Imagem BGR
            
        Returns:
            Imagem com contraste melhorado
        """
        # Converter para LAB
        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        
        # Aplicar CLAHE no canal L
        clahe = cv2.createCLAHE(
            clipLimit=self.contrast_clip,
            tileGridSize=(8, 8)
        )
        cl = clahe.apply(l)
        
        # Recombinar canais
        limg = cv2.merge((cl, a, b))
        
        # Converter de volta para BGR
        return cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
    
    def sharpen(self, img, strength=1.0):
        """
        Aplica sharpening na imagem.
        
        Args:
            img: Imagem BGR
            strength: Intensidade do sharpen (0.5-2.0)
            
        Returns:
            Imagem com sharpen aplicado
        """
        # Kernel de sharpen
        kernel = np.array([
            [-1, -1, -1],
            [-1,  9, -1],
            [-1, -1, -1]
        ]) * strength / 9
        
        kernel[1, 1] = 1 + (8 * strength / 9)
        
        return cv2.filter2D(img, -1, kernel)
    
    def unsharp_mask(self, img, sigma=1.0, strength=1.5):
        """
        Aplica unsharp mask para sharpening mais refinado.
        
        Args:
            img: Imagem BGR
            sigma: Tamanho do blur
            strength: Intensidade do efeito
            
        Returns:
            Imagem com unsharp mask aplicado
        """
        # Criar versão borrada
        blurred = cv2.GaussianBlur(img, (0, 0), sigma)
        
        # Sharpened = Original + (Original - Blurred) * strength
        sharpened = cv2.addWeighted(img, 1.0 + strength, blurred, -strength, 0)
        
        return sharpened
    
    def enhance_dark_ui(self, img):
        """
        Otimiza imagens de interfaces escuras (WhatsApp Dark, etc).
        
        Args:
            img: Imagem BGR
            
        Returns:
            Imagem otimizada
        """
        # Converter para LAB
        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        
        # Ajuste de gamma para realçar áreas escuras
        gamma = 0.8
        l = self._adjust_gamma(l, gamma)
        
        # CLAHE mais agressivo
        clahe = cv2.createCLAHE(clipLimit=4.0, tileGridSize=(8, 8))
        l = clahe.apply(l)
        
        # Recombinar
        limg = cv2.merge((l, a, b))
        img = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
        
        # Redução seletiva de ruído (menor que normal para não perder detalhes)
        img = cv2.fastNlMeansDenoisingColored(img, None, 5, 5, 7, 21)
        
        return img
    
    def adjust_brightness_contrast(self, img, brightness=0, contrast=0):
        """
        Ajusta brilho e contraste.
        
        Args:
            img: Imagem BGR
            brightness: Ajuste de brilho (-100 a 100)
            contrast: Ajuste de contraste (-100 a 100)
            
        Returns:
            Imagem ajustada
        """
        if brightness != 0:
            if brightness > 0:
                shadow = brightness
                highlight = 255
            else:
                shadow = 0
                highlight = 255 + brightness
            alpha_b = (highlight - shadow) / 255
            gamma_b = shadow
            
            img = cv2.addWeighted(img, alpha_b, img, 0, gamma_b)
        
        if contrast != 0:
            f = 131 * (contrast + 127) / (127 * (131 - contrast))
            alpha_c = f
            gamma_c = 127 * (1 - f)
            
            img = cv2.addWeighted(img, alpha_c, img, 0, gamma_c)
        
        return img
    
    def _adjust_gamma(self, img, gamma=1.0):
        """
        Ajusta gamma da imagem.
        
        Args:
            img: Imagem em escala de cinza
            gamma: Valor de gamma
            
        Returns:
            Imagem com gamma ajustado
        """
        inv_gamma = 1.0 / gamma
        table = np.array([
            ((i / 255.0) ** inv_gamma) * 255
            for i in range(256)
        ]).astype("uint8")
        
        return cv2.LUT(img, table)
    
    def enhance_edges(self, img):
        """
        Realça bordas na imagem.
        
        Args:
            img: Imagem BGR
            
        Returns:
            Imagem com bordas realçadas
        """
        # Converter para escala de cinza
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # Detectar bordas
        edges = cv2.Canny(gray, 100, 200)
        
        # Dilatar bordas
        kernel = np.ones((2, 2), np.uint8)
        edges = cv2.dilate(edges, kernel, iterations=1)
        
        # Converter bordas para BGR
        edges_bgr = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
        
        # Adicionar bordas à imagem original
        enhanced = cv2.addWeighted(img, 0.9, edges_bgr, 0.1, 0)
        
        return enhanced
    
    def auto_white_balance(self, img):
        """
        Balanço automático de branco.
        
        Args:
            img: Imagem BGR
            
        Returns:
            Imagem com balanço de branco corrigido
        """
        result = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        avg_a = np.average(result[:, :, 1])
        avg_b = np.average(result[:, :, 2])
        
        result[:, :, 1] = result[:, :, 1] - ((avg_a - 128) * (result[:, :, 0] / 255.0) * 1.1)
        result[:, :, 2] = result[:, :, 2] - ((avg_b - 128) * (result[:, :, 0] / 255.0) * 1.1)
        
        result = cv2.cvtColor(result, cv2.COLOR_LAB2BGR)
        
        return result
