import cv2
import numpy as np
import warnings


class SuperResolution:
    """
    Processador de super-resolução para upscaling de imagens.
    
    Nota: Para usar modelos de IA como Real-ESRGAN, você precisa baixar
    os pesos do modelo separadamente. Esta implementação usa upscaling
    de alta qualidade do OpenCV como fallback.
    """
    
    def __init__(self, use_ai=False, model_path=None):
        """
        Inicializa o processador de super-resolução.
        
        Args:
            use_ai: Tentar usar modelos de IA (Real-ESRGAN)
            model_path: Caminho para os pesos do modelo
        """
        self.use_ai = use_ai
        self.model_path = model_path
        self.ai_model = None
        
        if use_ai:
            self._init_ai_model()
    
    def _init_ai_model(self):
        """Inicializa modelo de IA se disponível."""
        try:
            # Tentar importar Real-ESRGAN
            from basicsr.archs.rrdbnet_arch import RRDBNet
            import torch
            
            # Configurar modelo
            self.ai_model = RRDBNet(
                num_in_ch=3,
                num_out_ch=3,
                num_feat=64,
                num_block=23,
                num_grow_ch=32,
                scale=4
            )
            
            # Carregar pesos se disponível
            if self.model_path and torch.cuda.is_available():
                self.ai_model.load_state_dict(torch.load(self.model_path))
                self.ai_model.eval()
                self.ai_model = self.ai_model.cuda()
            else:
                warnings.warn(
                    "Modelo de IA não configurado. Usando upscaling tradicional."
                )
                self.ai_model = None
                
        except ImportError:
            warnings.warn(
                "Dependências de IA não encontradas. Instale: pip install basicsr realesrgan"
            )
            self.ai_model = None
    
    def upscale(self, img, scale=2):
        """
        Aumenta resolução da imagem.
        
        Args:
            img: Imagem BGR
            scale: Fator de escala (2, 4, 8)
            
        Returns:
            Imagem com resolução aumentada
        """
        if self.ai_model is not None:
            return self._upscale_with_ai(img, scale)
        else:
            return self._upscale_traditional(img, scale)
    
    def _upscale_with_ai(self, img, scale):
        """
        Upscale usando modelo de IA (Real-ESRGAN).
        
        Args:
            img: Imagem BGR
            scale: Fator de escala
            
        Returns:
            Imagem upscaled
        """
        try:
            import torch
            
            # Converter para tensor
            img_tensor = torch.from_numpy(
                np.transpose(img, (2, 0, 1))
            ).float().div(255).unsqueeze(0)
            
            if torch.cuda.is_available():
                img_tensor = img_tensor.cuda()
            
            # Inferência
            with torch.no_grad():
                output = self.ai_model(img_tensor)
            
            # Converter de volta para numpy
            output = output.squeeze().cpu().numpy()
            output = np.transpose(output, (1, 2, 0))
            output = (output * 255).clip(0, 255).astype(np.uint8)
            
            return output
            
        except Exception as e:
            warnings.warn(f"Erro no upscale com IA: {e}. Usando método tradicional.")
            return self._upscale_traditional(img, scale)
    
    def _upscale_traditional(self, img, scale):
        """
        Upscale usando interpolação de alta qualidade.
        
        Args:
            img: Imagem BGR
            scale: Fator de escala
            
        Returns:
            Imagem upscaled
        """
        height, width = img.shape[:2]
        new_width = int(width * scale)
        new_height = int(height * scale)
        
        # LANCZOS4 é a melhor interpolação para upscaling
        upscaled = cv2.resize(
            img,
            (new_width, new_height),
            interpolation=cv2.INTER_LANCZOS4
        )
        
        # Aplicar unsharp mask para recuperar detalhes
        upscaled = self._unsharp_mask(upscaled, sigma=1.0, strength=0.5)
        
        return upscaled
    
    def _unsharp_mask(self, img, sigma=1.0, strength=0.5):
        """
        Aplica unsharp mask para melhorar nitidez.
        
        Args:
            img: Imagem BGR
            sigma: Tamanho do blur
            strength: Intensidade
            
        Returns:
            Imagem com sharpen aplicado
        """
        blurred = cv2.GaussianBlur(img, (0, 0), sigma)
        sharpened = cv2.addWeighted(img, 1.0 + strength, blurred, -strength, 0)
        return sharpened
    
    def upscale_smart(self, img, target_width=None, target_height=None):
        """
        Upscale inteligente para dimensão alvo.
        
        Args:
            img: Imagem BGR
            target_width: Largura alvo (opcional)
            target_height: Altura alvo (opcional)
            
        Returns:
            Imagem redimensionada
        """
        height, width = img.shape[:2]
        
        if target_width and target_height:
            new_width = target_width
            new_height = target_height
        elif target_width:
            ratio = target_width / width
            new_width = target_width
            new_height = int(height * ratio)
        elif target_height:
            ratio = target_height / height
            new_height = target_height
            new_width = int(width * ratio)
        else:
            return img
        
        # Calcular escala
        scale = max(new_width / width, new_height / height)
        
        if scale > 1:
            # Upscaling
            return self.upscale(img, scale=int(np.ceil(scale)))
        else:
            # Downscaling - usar INTER_AREA
            return cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA)
    
    def multi_scale_enhance(self, img):
        """
        Melhoria multi-escala para preservar detalhes.
        
        Args:
            img: Imagem BGR
            
        Returns:
            Imagem melhorada
        """
        # Criar pirâmide gaussiana
        pyramid = [img]
        for i in range(3):
            pyramid.append(cv2.pyrDown(pyramid[-1]))
        
        # Reconstruir com detalhes preservados
        result = pyramid[-1]
        for i in range(len(pyramid) - 2, -1, -1):
            result = cv2.pyrUp(result)
            # Ajustar tamanho se necessário
            if result.shape[:2] != pyramid[i].shape[:2]:
                result = cv2.resize(result, (pyramid[i].shape[1], pyramid[i].shape[0]))
            # Mesclar com detalhes
            result = cv2.addWeighted(result, 0.6, pyramid[i], 0.4, 0)
        
        return result
    
    def enhance_texture(self, img):
        """
        Realça texturas após upscaling.
        
        Args:
            img: Imagem BGR
            
        Returns:
            Imagem com texturas realçadas
        """
        # Converter para LAB
        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        
        # Realçar detalhes finos no canal L
        kernel = np.array([
            [0, -1, 0],
            [-1, 5, -1],
            [0, -1, 0]
        ]) / 5.0
        
        l_enhanced = cv2.filter2D(l, -1, kernel)
        
        # Recombinar
        lab_enhanced = cv2.merge([l_enhanced, a, b])
        result = cv2.cvtColor(lab_enhanced, cv2.COLOR_LAB2BGR)
        
        return result
