import json
import os
import depthai as dai
import numpy as np
import cv2
import time
from utils import converter_mask_ids_para_rgb, carregar_labelmap_completo

# ⚙️ Configurações
with open("config.json", "r") as f:
    config = json.load(f)
MODELO = config["camera"]
MODEL_NAME = config["model_name"]
RESOLUCAO = config["resolucao"]
SHAVES = config["shaves"]
ROI_INICIO = 0.0
ROI_TAMANHO = 1.0
MAIN_CLASS_NAME = config["main_class_name"]
model_to_use = config["model_to_use"]
labelmap_path = os.path.join(MODELO, "dataset", "labelmap.txt")
model_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME)
model_name = ""
if model_to_use == "geral":
    model_name = f"{MODEL_NAME}_best.pth"
elif model_to_use == "main_class":
    MAIN_CLASS_NAME = config["main_class_name"]
    model_name = f"{MODEL_NAME}_best_f1_{MAIN_CLASS_NAME}.pth"
elif model_to_use == "es":
    ES_CLASSES_NAME = config["es_classes"]
    model_name = f"{MODEL_NAME}_best_es_{ES_CLASSES_NAME}.pth"
else:
    model_name = f"{MODEL_NAME}_best.pth"

blob_path = os.path.join(model_path, f"{model_name.replace(".pth", "")}_openvino_2022.1_{SHAVES}shave.blob")

# Carregar mapa de cores
_, colormap_rgb, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path)
NUM_CLASSES = len(classes)
IGNORE_ID = ignore_rgb[0]

# Criar pipeline
pipeline = dai.Pipeline()

# Câmera
cam = pipeline.createColorCamera()
cam.setBoardSocket(dai.CameraBoardSocket.CAM_A)
cam.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
cam.setInterleaved(False)
cam.setColorOrder(dai.ColorCameraProperties.ColorOrder.RGB)
cam.setFps(30)

# ImageManip (redimensiona para a entrada da rede)
manip = pipeline.createImageManip()

y1 = 1.0 - (ROI_INICIO + ROI_TAMANHO)
y2 = 1.0 - ROI_INICIO

manip.initialConfig.setCropRect(0.0, y1, 1.0, y2)
manip.initialConfig.setResize(RESOLUCAO[0], RESOLUCAO[1])
manip.initialConfig.setFrameType(dai.RawImgFrame.Type.RGB888p)
manip.initialConfig.setKeepAspectRatio(False)
cam.video.link(manip.inputImage)

# Neural network
nn = pipeline.createNeuralNetwork()
nn.setBlobPath(blob_path)
manip.out.link(nn.input)

# Saída RGB para overlay (sem redimensionar)
xout_rgb = pipeline.createXLinkOut()
xout_rgb.setStreamName("rgb")
cam.video.link(xout_rgb.input)

# Saída NN
xout_nn = pipeline.createXLinkOut()
xout_nn.setStreamName("nn")
nn.out.link(xout_nn.input)

# E garanta que o video stream é 16:9
#cam.setVideoSize(384, 384)  # ou 1280x720

# Rodar pipeline
with dai.Device(pipeline) as device:
    rgb_queue = device.getOutputQueue("rgb", maxSize=1, blocking=False)
    nn_queue = device.getOutputQueue("nn", maxSize=1, blocking=False)

    print("Rodando inferência na OAK... Pressione 'q' para sair.")
    prev_time = time.time()

    H, W = RESOLUCAO[1], RESOLUCAO[0]

    in_rgb = rgb_queue.get()
    frame  = in_rgb.getCvFrame()
    frame_h, frame_w = frame.shape[:2]

    y_start = int(y1 * frame_h)
    y_end   = int(y2 * frame_h)
    roi_h   = y_end - y_start
    roi_w   = frame_w

    pred_ids = np.empty((H, W), dtype=np.uint8)
    pred_rgb = np.empty((H, W, 3), dtype=np.uint8)
    overlay = np.empty((H, W, 3), dtype=np.uint8)

    lut = np.zeros((256, 3), dtype=np.uint8)
    for i, color in enumerate(colormap_rgb):
        lut[i] = color
    lut[IGNORE_ID] = (255, 255, 255)

    while True:
        in_rgb = rgb_queue.get()
        in_nn = nn_queue.get()

        # Inferência - saída é um vetor flat [num_classes * H * W]
        out_raw = in_nn.getFirstLayerFp16()
        arr16 = np.frombuffer(np.asarray(out_raw, dtype=np.float16), dtype=np.float16)
        arr16 = arr16.reshape(NUM_CLASSES, H, W)
        pred_ids = arr16.argmax(axis=0).astype(np.uint8, copy=False)

        pred_rgb[:] = lut[pred_ids]

        frame = in_rgb.getCvFrame()
        cv2.resize(frame, (W, H), interpolation=cv2.INTER_AREA, dst=overlay)
        cv2.addWeighted(overlay, 0.4, pred_rgb, 0.6, 0, dst=overlay)
        
        # FPS
        now = time.time()
        fps_inst = 1.0 / (now - prev_time)
        prev_time = now
        fps = 0.9 * fps + 0.1 * fps_inst if 'fps' in locals() else fps_inst

        cv2.putText(overlay, f"FPS: {fps:.1f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
        #cv2.imshow("Segmentacao - OAK (on-board)", cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
        cv2.imshow("Segmentacao - OAK (on-board)", overlay)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cv2.destroyAllWindows()