import cv2
import numpy as np
import matplotlib.pyplot as plt

# Carregar a imagem
image_path = "mapa_calor.jpg"
image = cv2.imread(image_path)

# Converter a imagem de BGR para RGB
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# Definir intervalos de cores em HSV para segmentação
lower_blue = np.array([100, 150, 0])
upper_blue = np.array([140, 255, 255])
lower_green = np.array([40, 52, 72])
upper_green = np.array([80, 255, 255])
lower_red1 = np.array([0, 70, 50])
upper_red1 = np.array([10, 255, 255])
lower_red2 = np.array([170, 70, 50])
upper_red2 = np.array([180, 255, 255])

# Converter a imagem para HSV
image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

# Criar máscaras para cada cor
mask_blue = cv2.inRange(image_hsv, lower_blue, upper_blue)
mask_green = cv2.inRange(image_hsv, lower_green, upper_green)
mask_red1 = cv2.inRange(image_hsv, lower_red1, upper_red1)
mask_red2 = cv2.inRange(image_hsv, lower_red2, upper_red2)
mask_red = cv2.bitwise_or(mask_red1, mask_red2)

# Função para calcular a distância com base na cor
def get_distance_from_color(color):
    blue_dist = 500
    red_dist = 4000
    dist = blue_dist + (red_dist - blue_dist) * (color[0] / 255)
    return dist

# Função para processar a máscara e encontrar o maior contorno
def get_largest_contour(mask):
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        largest_contour = max(contours, key=cv2.contourArea)
        return largest_contour
    return None

# Encontrar o maior contorno em cada máscara
contour_blue = get_largest_contour(mask_blue)
contour_green = get_largest_contour(mask_green)
contour_red = get_largest_contour(mask_red)

# Desenhar caixas delimitadoras e calcular distâncias
output_image = image_rgb.copy()
object_details = []

for contour, color in zip([contour_blue, contour_green, contour_red], [(0, 0, 255), (0, 255, 0), (255, 0, 0)]):
    if contour is not None:
        x, y, w, h = cv2.boundingRect(contour)
        color_center = image_rgb[y + h // 2, x + w // 2]
        distance = get_distance_from_color(color_center)
        object_details.append((x, y, w, h, distance))
        cv2.rectangle(output_image, (x, y), (x + w, y + h), (255, 255, 255), 2)
        cv2.putText(output_image, f'D: {distance:.0f}mm', (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)

# Mostrar a imagem processada
plt.figure(figsize=(10, 10))
plt.imshow(output_image)
plt.axis('off')
plt.show()