import numpy as np
import cv2
from PIL import Image
import random

def generate_mask(image, thresh_val=None, min_size=9, crop_top=0, crop_bottom=0, crop_left=0, crop_right=0, max_size=52, typical_unit_size=40):
    """
    Generate a mask from the input image using optimized OpenCV functions.
    Replaces skimage dependency for better performance in Cloud Run.
    """
    if image is None:
        return None, None, None
    
    # 1. Conversión eficiente a NumPy
    if isinstance(image, Image.Image):
        if image.mode != 'L':
            image = image.convert('L')
        img = np.asarray(image)
    elif isinstance(image, np.ndarray):
        if len(image.shape) == 3:
            img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            img = image
    else:
        return None, None, None
    
    h, w = img.shape
    
    # 2. Recorte (Slicing)
    y1 = max(0, crop_top)
    y2 = min(h, h - crop_bottom)
    x1 = max(0, crop_left)
    x2 = min(w, w - crop_right)
    
    if y2 <= y1 or x2 <= x1:
        return None, None, None
    
    # Usamos slicing directo. OpenCV maneja la memoria eficientemente.
    img_crop = img[y1:y2, x1:x2]
    
    # 3. Binarización
    # Suavizar para mejorar el cálculo del umbral, sin perder la imagen original
    img_blur = cv2.GaussianBlur(img_crop, (5, 5), 0)

    if thresh_val is None or thresh_val == "" or str(thresh_val).lower() in ["none", "auto"]:
        # img_blur para Otsu
        otsu_thresh, bw = cv2.threshold(img_blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        print(f"Umbral calculado automáticamente (Otsu): {otsu_thresh}")
    else:
        try:
            thresh_int = int(thresh_val)
        except ValueError:
            print(f"Advertencia: '{thresh_val}' no es un número válido. Usando Otsu por defecto.")
            otsu_thresh, bw = cv2.threshold(img_blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
            thresh_int = otsu_thresh # Solo para el print
        else:
            _, bw = cv2.threshold(img_crop, thresh_int, 255, cv2.THRESH_BINARY)
            print(f"Umbral manual aplicado: {thresh_int}")
    
    # Fallback automático si no hay blancos (check rápido con np.any)
    # if not np.any(bw):
    #     _, bw = cv2.threshold(img_crop, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    
    # 4. Etiquetado optimizado (Connected Components)
    # Reemplaza a skimage.measure.label y regionprops
    # stats es una matriz donde cada fila es [x, y, w, h, area]
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(bw, connectivity=8)
    
    # 5. Generación de Máscara Vectorizada (Sin bucles lentos)
    # stats[:, 4] es la columna de áreas. El índice 0 es el fondo.
    valid_indices = np.where(stats[:, 4] >= min_size)[0]
    valid_indices = valid_indices[valid_indices != 0] # Quitar fondo
    
    # Crear Look-Up Table (LUT) para generar la máscara final rápido
    lut = np.zeros(num_labels, dtype=np.uint8)
    lut[valid_indices] = 255
    mask = lut[labels] # Aplicación vectorial de la LUT
    
    # Crear versión color para visualización
    mask_color = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
    mask_color[mask > 0] = [0, 255, 0] # Verde (formato BGR para OpenCV interno, luego PIL lo interpreta)
    
    # Empaquetamos datos (compatible con la estructura anterior pero con extras de OpenCV)
    processed_data = {
        "mask": mask,
        "labels": labels,
        "stats": stats,          # Nuevo: Stats de OpenCV (x,y,w,h,area)
        "centroids": centroids,  # Nuevo: Centroides de OpenCV
        "min_size": min_size,
        "max_size": max_size,
        "typical_unit_size": typical_unit_size,
        "img_crop": img_crop
    }
    
    # Convertir a PIL para el return final
    mask_img = Image.fromarray(mask_color)
    img_crop_pil = Image.fromarray(img_crop)
    
    return mask_img, processed_data, img_crop_pil


def visualize_structures(
    input_data, 
    min_size=1, 
    visualization_mode="contours",
    max_size=1000,
    typical_unit_size=40
):
    """
    Analyze and visualize structures using optimized OpenCV calls.
    """
    # 1. Extracción de datos
    stats = None
    centroids = None
    
    if isinstance(input_data, dict):
        if input_data is None:
            return None, "Debes procesar una imagen primero"
        mask = input_data["mask"]
        img_crop = input_data.get("img_crop", mask.copy())
        # Intentamos recuperar stats ya calculados para ahorrar tiempo
        stats = input_data.get("stats")
        centroids = input_data.get("centroids")
        min_size = input_data.get("min_size", min_size)
        max_size = input_data.get("max_size", max_size)
        typical_unit_size = input_data.get("typical_unit_size", typical_unit_size)
    elif isinstance(input_data, Image.Image):
        mask = np.array(input_data.convert("L"))
        mask = (mask > 0).astype(np.uint8) * 255
        img_crop = mask
    else:
        return None, "Input data type not supported"
    
    # Imagen base a color
    if len(img_crop.shape) == 2:
        img_color = cv2.cvtColor(img_crop, cv2.COLOR_GRAY2RGB)
    else:
        img_color = img_crop.copy()

    # 2. Si no tenemos stats (vino de una imagen directa), calcularlos
    if stats is None:
        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, connectivity=8)

    # 3. Filtrado vectorial de estructuras válidas
    # stats[:, 4] es el área.
    areas = stats[:, 4]
    valid_mask = (areas >= min_size)
    valid_mask[0] = False # Ignorar el fondo (label 0)
    valid_indices = np.where(valid_mask)[0]
    
    count = len(valid_indices)
    
    if count == 0:
        return Image.fromarray(img_color), "No se detectaron estructuras."

    total_units = 0
    structure_counter = 0

    # Lógica específica por modo de visualización
    
    # -- MODO 1: Bounding Boxes (Muy rápido, usa stats directos) --
    if visualization_mode == "bounding_boxes":
        for idx in valid_indices:
            structure_counter += 1
            x, y, w, h = stats[idx, 0:4]
            area = areas[idx]
            
            # Calcular unidades
            units = 1 if area <= max_size else max(1, round(area / typical_unit_size))
            total_units += units
            
            label_text = str(structure_counter) if units == 1 else f"{structure_counter}({units}u)"
            
            cv2.rectangle(img_color, (x, y), (x+w, y+h), (0, 255, 0), 2)
            cv2.putText(img_color, label_text, (x, y-5), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 1)

    # -- MODO 2: Connected Components (Coloreado por labels) --
    elif visualization_mode == "connected_components":
        # Generar colores aleatorios para la LUT
        # Necesitamos colores para todos los labels (incluso los inválidos, para simplificar indexing)
        colors = np.random.randint(0, 255, size=(len(stats), 3), dtype=np.uint8)
        colors[0] = [0, 0, 0] # Fondo negro
        
        # Crear imagen de etiquetas coloreada
        # Recuperamos labels si no los tenemos (raro, pero posible si input fue Image)
        if 'labels' not in locals():
             _, labels, _, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
             
        # Pintar
        colored_labels = colors[labels]
        
        # Iterar solo para poner texto y contar
        for idx in valid_indices:
            structure_counter += 1
            area = areas[idx]
            cx, cy = int(centroids[idx][0]), int(centroids[idx][1])
            
            units = 1 if area <= max_size else max(1, round(area / typical_unit_size))
            total_units += units
            
            label_text = str(structure_counter) if units == 1 else f"{structure_counter}({units}u)"
            cv2.putText(colored_labels, label_text, (cx-10, cy), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
            
        # Mezclar con original
        img_color = cv2.addWeighted(img_color, 0.3, colored_labels, 0.7, 0)

    # -- MODO 3: Contornos / Relleno / Elipses / Círculos --
    else:
        # Encontrar contornos es rápido en la máscara binaria ya limpia
        # Usamos la máscara en lugar de 'labels' porque findContours espera imagen de 8-bits
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # Filtramos contornos manualmente por si acaso (aunque la máscara ya debería estar limpia)
        valid_contours = [c for c in contours if cv2.contourArea(c) >= min_size]
        
        # Recalculamos count basado en contornos encontrados
        count = len(valid_contours)
        structure_counter = 0
        
        for i, cnt in enumerate(valid_contours):
            structure_counter += 1
            area = cv2.contourArea(cnt)
            
            units = 1 if area <= max_size else max(1, round(area / typical_unit_size))
            total_units += units
            
            # Centroide para texto
            M = cv2.moments(cnt)
            if M["m00"] != 0:
                cx = int(M["m10"] / M["m00"])
                cy = int(M["m01"] / M["m00"])
            else:
                x,y,w,h = cv2.boundingRect(cnt)
                cx, cy = x + w//2, y + h//2

            label_text = f"{structure_counter}" if units == 1 else f"{structure_counter}({units}u)"

            if visualization_mode == "contours":
                cv2.drawContours(img_color, [cnt], -1, (0, 255, 0), 2)
                cv2.putText(img_color, label_text, (cx-10, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 1)
                
            elif visualization_mode == "filled":
                color = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255))
                cv2.drawContours(img_color, [cnt], -1, color, -1)
                cv2.putText(img_color, label_text, (cx-10, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
                
            elif visualization_mode == "ellipses":
                if len(cnt) >= 5:
                    ellipse = cv2.fitEllipse(cnt)
                    cv2.ellipse(img_color, ellipse, (0, 255, 0), 2)
                    cv2.putText(img_color, label_text, (cx-10, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 1)
                else:
                    # Fallback a rectangulo si no alcanza para elipse
                    x,y,w,h = cv2.boundingRect(cnt)
                    cv2.rectangle(img_color, (x, y), (x+w, y+h), (0, 255, 0), 2)
                    
            else: # Default: Circles / Points
                cv2.circle(img_color, (cx, cy), 8, (0, 0, 255), 2)
                cv2.putText(img_color, label_text, (cx-10, cy-15), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1)

    return Image.fromarray(img_color), f"Estructuras: {count} | Unidades estimadas: {total_units}"