import torch
import torch.nn.functional as F
import numpy as np
import cv2
from PIL import Image
from config import SAPIENS_LITE_MODELS_PATH

def load_model(task, version):
    try:
        model_path = SAPIENS_LITE_MODELS_PATH[task][version]
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        if torch.cuda.is_available() and torch.cuda.get_device_properties(0).major >= 8:
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True
        model = torch.jit.load(model_path)
        model.eval()
        model.to(device)
        return model, device
    except KeyError as e:
        print(f"Error: Tarea o versión inválida. {e}")
        return None, None

def preprocess_image(image, input_shape):
    img = cv2.resize(image, (input_shape[2], input_shape[1]), interpolation=cv2.INTER_LINEAR).transpose(2, 0, 1)
    img = torch.from_numpy(img)
    img = img[[2, 1, 0], ...].float()
    mean = torch.tensor([123.5, 116.5, 103.5]).view(-1, 1, 1)
    std = torch.tensor([58.5, 57.0, 57.5]).view(-1, 1, 1)
    img = (img - mean) / std
    return img.unsqueeze(0)

def post_process_normal(result, original_shape):
    if result.dim() == 3:
        result = result.unsqueeze(0)
    elif result.dim() == 4:
        pass
    else:
        raise ValueError(f"Unexpected result dimension: {result.dim()}")
    
    seg_logits = F.interpolate(result, size=original_shape, mode="bilinear", align_corners=False).squeeze(0)
    normal_map = seg_logits.float().cpu().numpy().transpose(1, 2, 0)  # H x W x 3
    return normal_map

def visualize_normal(normal_map):
    normal_map_norm = np.linalg.norm(normal_map, axis=-1, keepdims=True)
    normal_map_normalized = normal_map / (normal_map_norm + 1e-5)  # Add a small epsilon to avoid division by zero
    
    normal_map_vis = ((normal_map_normalized + 1) / 2 * 255).astype(np.uint8)
    normal_map_vis = normal_map_vis[:, :, ::-1]  # RGB to BGR
    
    return normal_map_vis

def process_image_or_video(input_data, task='normal', version='sapiens_0.3b'):
    model, device = load_model(task, version)
    if model is None or device is None:
        return None

    input_shape = (3, 1024, 768)

    def process_frame(frame):
        if isinstance(frame, Image.Image):
            frame = np.array(frame)
        
        if frame.shape[2] == 4:  # RGBA
            frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2RGB)
        
        img = preprocess_image(frame, input_shape)
        
        with torch.no_grad():
            result = model(img.to(device))
        
        normal_map = post_process_normal(result, (frame.shape[0], frame.shape[1]))
        normal_image = visualize_normal(normal_map)
        
        return Image.fromarray(cv2.cvtColor(normal_image, cv2.COLOR_BGR2RGB))

    if isinstance(input_data, np.ndarray):  # Video frame
        return process_frame(input_data)
    elif isinstance(input_data, Image.Image):  # Imagen
        return process_frame(input_data)
    else:
        print("Tipo de entrada no soportado. Por favor, proporcione una imagen PIL o un frame de video numpy.")
        return None