# infer.py (batendo com o treino: RAW->uint8 HWC, LetterBox, forward direto, NMS, masks, unletterbox)
import cv2
import torch
import numpy as np
from pathlib import Path
import yaml

from ultralytics import YOLO
from ultralytics.utils import ops

# ============================
# CONFIG
# ============================
with Path("data.yaml").open("r", encoding="utf-8") as f:
    cfg = yaml.safe_load(f) or {}

CHANNELS = int(cfg.get("channels", 4))          # 4 ou 5
SCALE    = cfg.get("scale", "m")
IMG_SZ   = int(cfg.get("size", 800))

names_dict  = cfg.get("names", {})
colors_dict = cfg.get("colors", {})

CLASS_NAMES = [names_dict[k] for k in sorted(names_dict.keys())] if names_dict else []
CLASS_COLORS = [
    tuple(map(int, colors_dict[k].strip("()").split(",")))
    for k in sorted(colors_dict.keys())
] if colors_dict else []

# Ajuste conforme seu run
WEIGHTS_PATH = f"C:/ZendionInc/agrobot_base/Python/yolov8-seg/runs/segment/ch_{CHANNELS}/sc_{SCALE}/sz_{IMG_SZ}/weed_detector_segformer/weights/best.pt"
TEST_DIR     = Path(r"C:\ZendionInc\agrobot_base\Python\yolov8-seg\images\val")

# ============================
# RAW settings (iguais do treino)
# ============================
RAW_DTYPE = np.float32
RAW_ALREADY_01 = True
RAW_MAX_VAL = 1.0 if RAW_ALREADY_01 else 65535.0  # se algum dia usar uint16

# ============================
# Letterbox (multi-canal)
# ============================
def letterbox_multi(
    img: np.ndarray,
    new_shape: int = 640,
    color: int = 114
):
    """
    img: HWC uint8 (H,W,C)
    Retorna:
      img_lb: letterboxed HWC uint8 (new_shape, new_shape, C)
      ratio: (r, r)
      pad: (dw, dh) em pixels (float)
      new_unpad: (w_new, h_new)
    """
    assert img.ndim == 3, "img deve ser HWC"
    h0, w0 = img.shape[:2]

    # scale ratio (new / old)
    r = min(new_shape / h0, new_shape / w0)
    w_new = int(round(w0 * r))
    h_new = int(round(h0 * r))

    # compute padding
    dw = (new_shape - w_new) / 2
    dh = (new_shape - h_new) / 2

    # resize
    if (w0, h0) != (w_new, h_new):
        img = cv2.resize(img, (w_new, h_new), interpolation=cv2.INTER_LINEAR)

    # pad (top, bottom, left, right)
    top    = int(round(dh - 0.1))
    bottom = int(round(dh + 0.1))
    left   = int(round(dw - 0.1))
    right  = int(round(dw + 0.1))

    # padding constante para qualquer número de canais
    img_lb = np.pad(
        img,
        pad_width=((top, bottom), (left, right), (0, 0)),
        mode="constant",
        constant_values=color
    ).astype(img.dtype, copy=False)

    ratio = (r, r)
    pad = (dw, dh)
    return img_lb, ratio, pad, (w_new, h_new), (top, bottom, left, right)

def unletterbox_mask_to_original(
    masks_lb: torch.Tensor,          # (N, Hlb, Wlb) bool/0-1
    orig_hw: tuple,                  # (h0, w0)
    new_unpad_wh: tuple,             # (w_new, h_new)
    pads_tblr: tuple                 # (top, bottom, left, right)
):
    """
    Remove padding e volta masks para tamanho original.
    """
    h0, w0 = orig_hw
    w_new, h_new = new_unpad_wh
    top, bottom, left, right = pads_tblr

    # crop padding
    masks_crop = masks_lb[:, top:top + h_new, left:left + w_new]  # (N, h_new, w_new)

    # resize para original com interpolate (torch)
    masks_crop = masks_crop.unsqueeze(1).float()  # (N,1,h_new,w_new)
    masks_up = torch.nn.functional.interpolate(
        masks_crop, size=(h0, w0), mode="bilinear", align_corners=False
    )
    return (masks_up[:, 0] > 0.5)  # (N,h0,w0) bool

# ============================
# RAW loading (igual treino)
# ============================
def load_raw_as_uint8_hwc(img_path: Path, canais: int = CHANNELS) -> np.ndarray:
    """
    Lê RAW (4 canais físicos R,G,IR,B) como CHW no arquivo, normaliza fixo e
    retorna HWC uint8 com C=4 ou C=5 (com NDVI em 0..255).
    """
    # usa o preview só pra pegar H,W (igual treino)
    preview = cv2.imread(str(img_path), cv2.IMREAD_UNCHANGED)
    if preview is None:
        raise FileNotFoundError(f"Não consegui ler preview: {img_path}")
    h0, w0 = preview.shape[:2]

    raw_path = img_path.with_suffix(".raw")
    if not raw_path.exists():
        raise FileNotFoundError(f"RAW não encontrado: {raw_path}")

    num_raw_channels = 4
    arr_flat = np.fromfile(str(raw_path), dtype=RAW_DTYPE)
    expected_size = num_raw_channels * h0 * w0
    if arr_flat.size != expected_size:
        raise ValueError(f"RAW size errado {raw_path}: esperado {expected_size}, veio {arr_flat.size}")

    raw_chw = arr_flat.reshape(num_raw_channels, h0, w0)  # (4,H,W)
    arr = np.transpose(raw_chw, (1, 2, 0)).astype(np.float32, copy=False)  # (H,W,4)

    # normalização fixa
    if RAW_ALREADY_01:
        arr_norm = arr
    else:
        arr_norm = arr / float(RAW_MAX_VAL)
    arr_norm = np.clip(arr_norm, 0.0, 1.0)

    # NDVI opcional (5ch)
    if canais == 5:
        R  = arr_norm[..., 0]
        IR = arr_norm[..., 2]
        eps = 1e-6
        ndvi = (IR - R) / (IR + R + eps)           # [-1,1]
        ndvi = np.clip(ndvi, -1.0, 1.0)
        ndvi01 = (ndvi + 1.0) / 2.0                # [0,1]
        arr_norm = np.concatenate([arr_norm, ndvi01[..., None]], axis=-1)  # (H,W,5)
    elif canais != 4:
        raise ValueError(f"canais deve ser 4 ou 5, veio {canais}")

    im_uint8 = (arr_norm * 255.0).round().clip(0, 255).astype(np.uint8)
    return im_uint8  # HWC uint8

def make_bgr_preview_from_raw_uint8(raw_hwc_uint8: np.ndarray) -> np.ndarray:
    """
    Preview BGR (OpenCV) a partir de HWC uint8:
    assume ordem [R,G,IR,B,(NDVI)].
    Usa B=canal 3, G=canal 1, R=canal 0.
    """
    R = raw_hwc_uint8[..., 0]
    G = raw_hwc_uint8[..., 1]
    B = raw_hwc_uint8[..., 3]
    return np.dstack([B, G, R])

# ============================
# Drawing
# ============================
def draw_result(vis_bgr: np.ndarray, boxes_xyxy: np.ndarray, cls_ids: np.ndarray, confs: np.ndarray, masks: np.ndarray | None):
    out = vis_bgr.copy()

    # máscaras (se houver)
    if masks is not None:
        for i in range(masks.shape[0]):
            cls = int(cls_ids[i])
            color = CLASS_COLORS[cls] if cls < len(CLASS_COLORS) else (0, 255, 0)
            m = masks[i].astype(bool)
            overlay = np.array(color, dtype=np.float32)
            out[m] = (out[m] * 0.4 + overlay * 0.6).astype(np.uint8)

    # boxes + labels
    for i in range(len(boxes_xyxy)):
        x1, y1, x2, y2 = boxes_xyxy[i].astype(int)
        cls = int(cls_ids[i])
        conf = float(confs[i])
        color = CLASS_COLORS[cls] if cls < len(CLASS_COLORS) else (0, 255, 0)
        name = CLASS_NAMES[cls] if cls < len(CLASS_NAMES) else str(cls)

        cv2.rectangle(out, (x1, y1), (x2, y2), color, 2)
        txt = f"{name} {conf:.2f}"
        cv2.putText(out, txt, (x1, max(0, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2, cv2.LINE_AA)

    return out

# ============================
# Inference core
# ============================
@torch.no_grad()
def infer_one(model: YOLO, raw_uint8_hwc: np.ndarray, conf_thres=0.25, iou_thres=0.7, max_det=300):
    """
    raw_uint8_hwc: HWC uint8 com C=4 ou 5
    Retorna boxes/masks já no tamanho original (h0,w0)
    """
    device = next(model.model.parameters()).device

    h0, w0 = raw_uint8_hwc.shape[:2]

    # 1) LetterBox igual treino (mantém proporção)
    lb_img, ratio, pad, new_unpad_wh, pads_tblr = letterbox_multi(raw_uint8_hwc, new_shape=IMG_SZ, color=114)

    # 2) Tensor (B,C,H,W) float32 0..1
    x = torch.from_numpy(lb_img).to(device)                   # (H,W,C) uint8
    x = x.permute(2, 0, 1).contiguous().float() / 255.0       # (C,H,W) float
    x = x.unsqueeze(0)                                        # (1,C,H,W)

    # 3) Forward direto
    out = model.model(x)

    # 4) separa preds/proto
    if isinstance(out, (list, tuple)):
        preds_raw = out[0]
        proto = out[1] if len(out) > 1 else None
    else:
        preds_raw = out
        proto = None

    if preds_raw.ndim == 2:
        preds_raw = preds_raw.unsqueeze(0)

    nc = len(model.names) if hasattr(model, "names") else len(CLASS_NAMES)
    det_list = ops.non_max_suppression(
        preds_raw,
        conf_thres=conf_thres,
        iou_thres=iou_thres,
        classes=None,
        agnostic=False,
        max_det=max_det,
        nc=nc
    )
    det = det_list[0]
    if det is None or len(det) == 0:
        return None

    # det: [x1,y1,x2,y2,conf,cls, ...mask coeffs]
    boxes_lb = det[:, :4]
    confs = det[:, 4]
    cls_ids = det[:, 5].to(torch.int64)

    # 5) Máscaras no espaço letterboxed
    masks_orig = None
    if proto is not None and det.shape[1] > 6:
        # proto pode vir com batch
        proto_t = proto
        while isinstance(proto_t, (list, tuple)):
            proto_t = proto_t[0] if len(proto_t) else None
            if proto_t is None:
                break

        if isinstance(proto_t, torch.Tensor):
            if proto_t.ndim == 4:
                proto_t = proto_t[0]  # (C,Hm,Wm)
            proto_t = proto_t.to(det.device)

            mask_coeffs = det[:, 6:].to(det.device)

            # garante C_proto == C_mask
            c_proto = proto_t.shape[0]
            c_mask = mask_coeffs.shape[1]
            if c_proto != c_mask:
                if c_proto > c_mask:
                    proto_t = proto_t[:c_mask]
                else:
                    mask_coeffs = mask_coeffs[:, :c_proto]

            # gera máscaras no tamanho do input (letterboxed)
            masks_lb = ops.process_mask(proto_t, mask_coeffs, boxes_lb, x.shape[2:], upsample=True)  # (N,Hlb,Wlb) bool
            # desfaz letterbox -> original
            masks_orig_t = unletterbox_mask_to_original(masks_lb, (h0, w0), new_unpad_wh, pads_tblr)
            masks_orig = masks_orig_t.detach().cpu().numpy().astype(np.uint8)  # 0/1

    # 6) Boxes: letterbox -> original (usa ratio_pad)
    # ratio_pad esperado: (ratio, pad) onde pad é (dw, dh)
    ratio_pad = (ratio, pad)
    boxes_scaled = ops.scale_boxes(x.shape[2:], boxes_lb.clone(), (h0, w0), ratio_pad=ratio_pad)
    boxes_xyxy = boxes_scaled.detach().cpu().numpy()

    return {
        "boxes": boxes_xyxy,
        "confs": confs.detach().cpu().numpy(),
        "cls": cls_ids.detach().cpu().numpy(),
        "masks": masks_orig
    }

# ============================
# Main loop viewer
# ============================
def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Usando dispositivo: {device}")

    model = YOLO(WEIGHTS_PATH)
    model.to(device)
    model.fuse()

    # sanity check do patch
    first_conv = model.model.model[0].conv
    print("First conv weight shape:", tuple(first_conv.weight.shape))  # deve ser [48, C, 3, 3]

    img_paths = sorted([p for p in TEST_DIR.iterdir() if p.suffix.lower() in [".jpg", ".jpeg", ".png"]])
    if not img_paths:
        print(f"Nenhuma imagem encontrada em {TEST_DIR}")
        return

    idx = 0
    n = len(img_paths)

    CONF = 0.45
    IOU = 0.7

    cv2.namedWindow(f"YOLO {CHANNELS}ch Viewer", cv2.WINDOW_NORMAL)

    while True:
        img_path = img_paths[idx]
        print("\n======================================")
        print(f"Processando: {img_path}")

        try:
            raw_uint8 = load_raw_as_uint8_hwc(img_path, canais=CHANNELS)  # HWC uint8
            preview_bgr = make_bgr_preview_from_raw_uint8(raw_uint8)

            res = infer_one(model, raw_uint8, conf_thres=CONF, iou_thres=IOU)
            if res is None:
                vis = preview_bgr.copy()
                cv2.putText(vis, "SEM DETECCOES", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,0,255), 2, cv2.LINE_AA)
            else:
                vis = draw_result(
                    preview_bgr,
                    boxes_xyxy=res["boxes"],
                    cls_ids=res["cls"],
                    confs=res["confs"],
                    masks=res["masks"]
                )

            txt = f"[{idx+1}/{n}] {img_path.name} | A/D navega | +/- conf {CONF:.2f} | Q/ESC sai"
            cv2.putText(vis, txt, (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 2, cv2.LINE_AA)
            cv2.imshow(f"YOLO {CHANNELS}ch Viewer", vis)

        except Exception as e:
            print(f"Erro ao processar {img_path}: {e}")
            blank = np.zeros((480, 900, 3), dtype=np.uint8)
            cv2.putText(blank, f"Erro em {img_path.name}", (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,0,255), 2)
            cv2.putText(blank, str(e)[:120], (10, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,0,255), 2)
            cv2.imshow(f"YOLO {CHANNELS}ch Viewer", blank)

        k = cv2.waitKey(0) & 0xFF
        if k in (ord("q"), 27):  # q ou ESC
            break
        elif k == ord("a"):
            idx = (idx - 1) % n
        elif k == ord("d"):
            idx = (idx + 1) % n
        elif k in (ord("+"), ord("=")):
            CONF = min(0.99, CONF + 0.05)
        elif k in (ord("-"), ord("_")):
            CONF = max(0.01, CONF - 0.05)

    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()