import cv2
import numpy as np


class DocumentProcessor:
    """
    Processador especializado para documentos, textos e OCR.
    """
    
    def __init__(self):
        pass
    
    def enhance_text(self, img):
        """
        Melhora imagem para leitura de texto.
        
        Args:
            img: Imagem BGR
            
        Returns:
            Imagem otimizada para texto
        """
        # Se já for grayscale, usar direto
        if len(img.shape) == 2:
            gray = img
        else:
            # Converter para grayscale
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # Denoise
        gray = cv2.fastNlMeansDenoising(gray, None, h=10, templateWindowSize=7, searchWindowSize=21)
        
        # Melhorar contraste com CLAHE
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        gray = clahe.apply(gray)
        
        # Retornar como BGR para consistência
        return cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
    
    def binarize(self, img, method='adaptive'):
        """
        Binariza a imagem para OCR.
        
        Args:
            img: Imagem BGR
            method: Método de binarização ('adaptive', 'otsu', 'simple')
            
        Returns:
            Imagem binarizada
        """
        # Converter para grayscale se necessário
        if len(img.shape) == 3:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        else:
            gray = img
        
        # Denoise primeiro
        gray = cv2.fastNlMeansDenoising(gray, None, h=10)
        
        if method == 'adaptive':
            # Binarização adaptativa - melhor para iluminação não-uniforme
            binary = cv2.adaptiveThreshold(
                gray, 255,
                cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                cv2.THRESH_BINARY,
                11, 2
            )
        elif method == 'otsu':
            # Método de Otsu - automático
            _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        else:
            # Threshold simples
            _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
        
        return binary
    
    def deskew(self, img):
        """
        Corrige inclinação do documento.
        
        Args:
            img: Imagem BGR
            
        Returns:
            Imagem com inclinação corrigida
        """
        # Converter para grayscale
        if len(img.shape) == 3:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        else:
            gray = img
        
        # Binarizar
        binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
        
        # Detectar ângulo de rotação
        coords = np.column_stack(np.where(binary > 0))
        angle = cv2.minAreaRect(coords)[-1]
        
        # Ajustar ângulo
        if angle < -45:
            angle = -(90 + angle)
        else:
            angle = -angle
        
        # Rotacionar imagem
        (h, w) = img.shape[:2]
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, angle, 1.0)
        
        if len(img.shape) == 3:
            rotated = cv2.warpAffine(
                img, M, (w, h),
                flags=cv2.INTER_CUBIC,
                borderMode=cv2.BORDER_REPLICATE
            )
        else:
            rotated = cv2.warpAffine(
                img, M, (w, h),
                flags=cv2.INTER_CUBIC,
                borderMode=cv2.BORDER_REPLICATE
            )
        
        return rotated
    
    def remove_noise(self, img):
        """
        Remove ruído usando morfologia.
        
        Args:
            img: Imagem binária
            
        Returns:
            Imagem limpa
        """
        # Kernel para operações morfológicas
        kernel = np.ones((2, 2), np.uint8)
        
        # Opening: remove ruído pequeno
        opening = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel, iterations=1)
        
        # Closing: fecha buracos pequenos
        closing = cv2.morphologyEx(opening, cv2.MORPH_CLOSE, kernel, iterations=1)
        
        return closing
    
    def enhance_for_ocr(self, img):
        """
        Pipeline completo de otimização para OCR.
        
        Args:
            img: Imagem BGR
            
        Returns:
            Imagem otimizada para OCR
        """
        # Converter para grayscale
        if len(img.shape) == 3:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        else:
            gray = img
        
        # 1. Denoise
        gray = cv2.fastNlMeansDenoising(gray, None, h=10)
        
        # 2. Aumentar contraste
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        gray = clahe.apply(gray)
        
        # 3. Binarizar
        binary = cv2.adaptiveThreshold(
            gray, 255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY,
            11, 2
        )
        
        # 4. Limpar ruído
        binary = self.remove_noise(binary)
        
        # 5. Leve sharpen
        kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
        sharpened = cv2.filter2D(binary, -1, kernel)
        
        return sharpened
    
    def remove_shadows(self, img):
        """
        Remove sombras de documentos fotografados.
        
        Args:
            img: Imagem BGR
            
        Returns:
            Imagem sem sombras
        """
        # Converter para LAB
        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        
        # Aplicar blur para estimar iluminação
        blur = cv2.GaussianBlur(l, (0, 0), sigmaX=3, sigmaY=3)
        
        # Normalizar
        l = cv2.divide(l, blur, scale=255)
        
        # Recombinar
        lab = cv2.merge([l, a, b])
        img = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
        
        return img
    
    def increase_dpi(self, img, target_dpi=300, current_dpi=72):
        """
        Aumenta DPI da imagem para impressão/OCR.
        
        Args:
            img: Imagem BGR
            target_dpi: DPI desejado
            current_dpi: DPI atual
            
        Returns:
            Imagem com DPI aumentado
        """
        scale = target_dpi / current_dpi
        
        width = int(img.shape[1] * scale)
        height = int(img.shape[0] * scale)
        
        # Usar INTER_CUBIC ou INTER_LANCZOS4 para upscaling
        resized = cv2.resize(img, (width, height), interpolation=cv2.INTER_LANCZOS4)
        
        return resized
    
    def detect_text_regions(self, img):
        """
        Detecta regiões de texto na imagem.
        
        Args:
            img: Imagem BGR
            
        Returns:
            Lista de bounding boxes (x, y, w, h)
        """
        # Converter para grayscale
        if len(img.shape) == 3:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        else:
            gray = img
        
        # Binarizar
        binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
        
        # Operações morfológicas para conectar texto
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (20, 5))
        dilated = cv2.dilate(binary, kernel, iterations=2)
        
        # Encontrar contornos
        contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # Extrair bounding boxes
        boxes = []
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            # Filtrar regiões muito pequenas
            if w > 50 and h > 20:
                boxes.append((x, y, w, h))
        
        return boxes
