import paho.mqtt.client as mqtt
import numpy as np
import cv2
import json

def on_connect(client, userdata, flags, rc):
    print("Connected with result code " + str(rc))
    client.subscribe("kinect/depth")

def on_message(client, userdata, msg):
    depth_pixels = json.loads(msg.payload)
    depth_values = [pixel['Depth'] for pixel in depth_pixels]  # Extraia os valores de profundidade
    heatmap, obstacles = process_depth_data(depth_values)
    send_obstacles(client, obstacles)

def process_depth_data(depth_values):
    # Converta os valores de profundidade para um array numpy
    depth_array = np.array(depth_values, dtype=np.uint16)
    
    # Redimensione o array para a forma correta, se necessário
    depth_array = depth_array.reshape((480, 640))  # Exemplo de resolução 480x640, ajuste conforme necessário
    
    # Gere o mapa de calor (heatmap)
    heatmap = cv2.applyColorMap(cv2.convertScaleAbs(depth_array, alpha=0.03), cv2.COLORMAP_JET)
    
    # Segmente a imagem para encontrar obstáculos
    obstacles = segment_image_by_color(heatmap, depth_array)
    
    return heatmap, obstacles

def segment_image_by_color(heatmap, depth_array):
    # Defina a cor alvo e o limiar
    target_color = np.array([0, 0, 255])  # Exemplo: vermelho puro
    threshold = 50  # Ajuste conforme necessário
    
    # Converta o mapa de calor para o espaço de cores HSV
    hsv_heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2HSV)
    target_hsv = cv2.cvtColor(np.uint8([[target_color]]), cv2.COLOR_BGR2HSV)[0][0]
    
    # Defina os limites inferior e superior para a segmentação
    lower_bound = np.array([target_hsv[0] - threshold, 100, 100])
    upper_bound = np.array([target_hsv[0] + threshold, 255, 255])
    
    # Crie uma máscara para a cor alvo
    mask = cv2.inRange(hsv_heatmap, lower_bound, upper_bound)
    
    # Encontre contornos na máscara
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    obstacles = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        obstacle = calculate_obstacle_properties(x, y, w, h, depth_array)
        obstacles.append(obstacle)
    
    return obstacles

def calculate_obstacle_properties(x, y, width, height, depth_array):
    # Calcular a distância média do obstáculo
    obstacle_region = depth_array[y:y+height, x:x+width]
    distance = np.mean(obstacle_region)
    
    # Calcular a relação pixels para milímetros
    pixel_to_mm = calculate_pixel_to_mm_ratio(distance)
    
    # Converter dimensões de pixels para milímetros
    width_mm = width * pixel_to_mm
    height_mm = height * pixel_to_mm
    
    obstacle = {
        "X": x,
        "Y": y,
        "Width": width,
        "Height": height,
        "Distance_mm": distance,
        "Width_mm": width_mm,
        "Height_mm": height_mm
    }
    
    return obstacle

def calculate_pixel_to_mm_ratio(distance):
    # Exemplo de cálculo simplificado da relação pixels para milímetros
    # A relação exata pode variar e pode precisar de ajustes
    # Aqui assumimos uma relação linear simplificada
    focal_length = 580  # Focal length do Kinect em pixels
    sensor_width_mm = 70  # Largura do sensor em milímetros

    pixel_to_mm = (distance / focal_length) * (sensor_width_mm / 640) * 10
    return pixel_to_mm

def send_obstacles(client, obstacles):
    message = json.dumps(obstacles)
    client.publish("kinect/obstacles", message)

client = mqtt.Client()
client.on_connect = on_connect
client.on_message = on_message

client.connect("localhost", 1883, 60)  # Substitua pelo seu broker

client.loop_forever()