from collections import deque
import json
import os
import time
import cv2
import numpy as np
import torch
from gal5000.gal_service import Gal5000Camera
from Python.OAK.datasets.gal5000.raw_segformer_service import RawSegformerService

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")

INFERIR = False

cfg_path = r"config.json"

# ---- Carrega config ----
with open(cfg_path, "r") as f:
    config = json.load(f)

MODELO        = config["camera"]         # ex: "gal5000"
MODEL_NAME    = config["model_name"]     # ex: "segformer_b3"
modelo_folder = config["modelo"]         # ex: "dif"

USE_NDVI  = bool(config.get("use_ndvi", False))
CHANNELS  = int(config.get("channels", 4))
RESOLUCAO = config["resolucao"]
W, H      = RESOLUCAO[0], RESOLUCAO[1]
preview_w = RESOLUCAO[0]
preview_h = RESOLUCAO[1]
fps_window = 30

if INFERIR:
    save_path = os.path.join(MODELO, "backup", modelo_folder, MODEL_NAME, f"raw{CHANNELS}")
    ckpt_path = os.path.join(save_path, "best_miou.pt")
    
    # ---- Instancia service do modelo RAW (4 ou 5 canais) ----
    model_svc = RawSegformerService(
        config_path=cfg_path,
        ckpt_path=ckpt_path,
        device=device,
        use_amp=True,
    )

print("[mode] Câmera ao vivo (GAL5000 + SegFormer RAW)")
cam = Gal5000Camera(raw_w=W, raw_h=H, use_auto_exposure=True)

win = "GAL5000 RAW + SegFormer (Q=quit)"
cv2.namedWindow(win, cv2.WINDOW_NORMAL)

tq = deque(maxlen=max(5, int(fps_window)))

with cam:
    #cam.configure_scaling_and_fps(bin_h=2, bin_v=2, dec_h=1, dec_v=1, fps=20.0)
    cam.configure_fps(fps_window)
    cam.start_streaming()
    
    while True:
        t0 = time.time()

        # raw4_base: (4,H,W) float32 0..1 [R,G,IR,B] já redimensionado
        raw4_base, dbg = cam.grab_raw4(out_h=preview_h, out_w=preview_w, timeout_ms=2000)

        # Separa canais básicos
        r  = raw4_base[0]
        g  = raw4_base[1]
        ir = raw4_base[2]
        b  = raw4_base[3]

        if INFERIR:
            # Deixa o service montar a entrada final com 4 ou 5 canais
            # (R,G,IR,B) | (R,G,B,NDVI) | (R,G,B,IR,NDVI)
            raw_input = model_svc.build_raw_input(r, g, ir, b)  # (C,H,W) float32 0..1
            # Inferência + previews via service
            pred_ids, rgb, pred_rgb, overlay, t_inf, t_pvw = model_svc.infer_and_preview(raw_input)
        else:
            rgb = np.stack([r, g, b], axis=0)      # (3,H,W)
            rgb = (rgb * 255.0).clip(0, 255).astype("uint8")
            overlay = np.transpose(rgb, (1, 2, 0))  # (H,W,3) -> formato OpenCV

        overlay = np.ascontiguousarray(overlay)
        if overlay.dtype != np.uint8:
            overlay = overlay.astype(np.uint8)

        tq.append(time.time() - t0)
        fps = 1.0 / (sum(tq) / len(tq))

        status = cam.get_status()
        
        shape = dbg["raw_shape"]
        lat = dbg["latency_s"] * 1000
        tc  = dbg.get("t_capture", 0) * 1000
        tconv = dbg.get("t_convert", 0) * 1000
        t_ae = dbg.get("t_ae", 0) * 1000
        
        lines = [
            f"{preview_w}x{preview_h} | FPS~{fps:.1f} | {shape}",
            f"EX={status['exp_raw']}  G={status['gain_a']}/{status['gain_d']}",
            f"C={CHANNELS}  NDVI={int(USE_NDVI)}",
            f"cap={tc:.1f}ms   ae={t_ae:.1f}ms   conv={tconv:.1f}ms",
        ]
        y = 24
        for line in lines:
            cv2.putText(overlay, line, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
            y += 22

        cv2.imshow(win, cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
        k = cv2.waitKey(1) & 0xFF
        if k in (ord("q"), ord("Q"), 27):
            break

cv2.destroyAllWindows()