import os
import cv2
import numpy as np

# ========= CONFIG =========
RAW_PATH  = r"dataset/ds/raws/20260115_152748_007.raw"      # ajuste se precisar
W = 2592                                     # largura do RAW (pixels do mosaico)
H = 2056                                     # altura do RAW
UPSCALE = 2                                  # aumenta preview (2x fica bom)

# overlay
ALPHA = 0.45                                 # transparência da máscara
MIN_IR = 15                                  # ignora pixels muito escuros no IR (ruído)
MIN_G  = 20                                  # ignora pixels muito escuros no G (ruído)

# ========= RAW decode =========
def read_raw_mosaic(path, w, h):
    raw = np.fromfile(path, dtype=np.uint8)
    if raw.size != w * h:
        raise RuntimeError(f"RAW size mismatch: got {raw.size}, expected {w*h}. "
                           f"Confira W/H.")
    return raw.reshape(h, w)

def split_4ch(raw):
    # 2x2 pattern:
    # [R, G]
    # [IR,B]
    R  = raw[0::2, 0::2]
    G  = raw[0::2, 1::2]
    IR = raw[1::2, 0::2]
    B  = raw[1::2, 1::2]
    return R, G, IR, B

def norm8(x, p_lo=2, p_hi=98):
    lo = np.percentile(x, p_lo)
    hi = np.percentile(x, p_hi)
    if hi <= lo + 1:
        return x.astype(np.uint8)
    y = (x.astype(np.float32) - lo) * (255.0 / (hi - lo))
    return np.clip(y, 0, 255).astype(np.uint8)

def make_rgb_preview(R, G, B, upscale=2):
    Rn, Gn, Bn = norm8(R), norm8(G), norm8(B)
    bgr = np.dstack([Bn, Gn, Rn])  # OpenCV = BGR
    if upscale != 1:
        bgr = cv2.resize(bgr, (bgr.shape[1]*upscale, bgr.shape[0]*upscale), interpolation=cv2.INTER_NEAREST)
    return bgr

# ========= Simple spectral classifier =========
def classify_cane_weed(G, IR, thr_ratio, thr_ir_bias):
    """
    Retorna mask_cane, mask_weed em resolução H/2 x W/2.

    Padrões:
    - ERVA: ratio = G/(IR+1) maior
    - CANA: IR relativamente maior + ratio menor

    thr_ratio: limiar principal de G/IR
    thr_ir_bias: adicional: favorece CANA quando IR está alto
    """
    Gf  = G.astype(np.float32)
    IRf = IR.astype(np.float32)

    ratio = Gf / (IRf + 1.0)

    valid = (Gf >= MIN_G) & (IRf >= MIN_IR)

    # regra: erva se ratio > thr_ratio
    weed = valid & (ratio >= thr_ratio)

    # cana: ratio baixo OU IR alto (bias)
    # IR alto relativo: IR > (G - thr_ir_bias) ajuda puxar cana
    cane = valid & (ratio < thr_ratio)
    

    # resolve conflitos: se cair em ambos, usa ratio como desempate
    both = weed & cane
    if np.any(both):
        # se ratio alto -> weed, senão -> cane
        weed[both] = ratio[both] >= thr_ratio
        cane[both] = ~weed[both]

    # pixels válidos mas não classificados: decide pelo ratio
    undec = valid & ~(weed | cane)
    if np.any(undec):
        weed[undec] = ratio[undec] >= thr_ratio
        cane[undec] = ~weed[undec]

    return cane, weed, ratio, valid

def morph_cleanup(mask, k=3):
    if k <= 1:
        return mask
    ker = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
    m = mask.astype(np.uint8) * 255
    m = cv2.medianBlur(m, 3)
    m = cv2.morphologyEx(m, cv2.MORPH_OPEN, ker, iterations=1)
    m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, ker, iterations=1)
    return m > 0

def overlay_classes(bgr, cane_mask, weed_mask, upscale=2):
    # sobe masks pro tamanho do preview
    h2, w2 = cane_mask.shape
    if upscale != 1:
        cane = cv2.resize(cane_mask.astype(np.uint8)*255, (w2*upscale, h2*upscale), interpolation=cv2.INTER_NEAREST)
        weed = cv2.resize(weed_mask.astype(np.uint8)*255, (w2*upscale, h2*upscale), interpolation=cv2.INTER_NEAREST)
    else:
        cane = cane_mask.astype(np.uint8)*255
        weed = weed_mask.astype(np.uint8)*255

    out = bgr.copy()

    # cores (BGR): cana=azul, erva=verde
    cane_col = np.zeros_like(out)
    cane_col[:, :, 0] = cane  # Blue

    weed_col = np.zeros_like(out)
    weed_col[:, :, 1] = weed  # Green

    # combina overlays
    mask_any = (cane > 0) | (weed > 0)
    overlay = np.clip(cane_col + weed_col, 0, 255).astype(np.uint8)

    out[mask_any] = (out[mask_any].astype(np.float32) * (1 - ALPHA) + overlay[mask_any].astype(np.float32) * ALPHA).astype(np.uint8)
    return out

def main():
    raw = read_raw_mosaic(RAW_PATH, W, H)
    R, G, IR, B = split_4ch(raw)

    base = make_rgb_preview(R, G, B, upscale=UPSCALE)

    cv2.namedWindow("overlay", cv2.WINDOW_NORMAL)
    cv2.namedWindow("debug", cv2.WINDOW_NORMAL)

    # sliders
    # ratio em escala 0..300 -> 0.00..3.00
    cv2.createTrackbar("thr_ratio x100", "overlay", 270, 500, lambda v: None)  # 2.70 inicial
    cv2.createTrackbar("ir_bias", "overlay", 5, 100, lambda v: None)  # 5 inicial
    cv2.createTrackbar("morph_k", "overlay", 5, 21, lambda v: None)            # 5 inicial

    while True:
        thr_ratio = cv2.getTrackbarPos("thr_ratio x100", "overlay") / 100.0
        thr_ir_bias = float(cv2.getTrackbarPos("ir_bias", "overlay"))
        mk = cv2.getTrackbarPos("morph_k", "overlay")
        if mk % 2 == 0:
            mk += 1

        cane, weed, ratio, valid = classify_cane_weed(G, IR, thr_ratio, thr_ir_bias)

        cane2 = morph_cleanup(cane, k=mk)
        weed2 = morph_cleanup(weed, k=mk)

        out = overlay_classes(base, cane2, weed2, upscale=UPSCALE)

        # debug views
        ratio_vis = norm8(ratio, 2, 98)
        if UPSCALE != 1:
            ratio_vis = cv2.resize(ratio_vis, (ratio_vis.shape[1]*UPSCALE, ratio_vis.shape[0]*UPSCALE), interpolation=cv2.INTER_NEAREST)

        # desenha texto rápido
        txt = f"thr_ratio={thr_ratio:.2f}  ir_bias={thr_ir_bias:.0f}  morph_k={mk}"
        cv2.putText(out, txt, (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,0,0), 3, cv2.LINE_AA)
        cv2.putText(out, txt, (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,255,255), 2, cv2.LINE_AA)

        cv2.imshow("overlay", out)
        cv2.imshow("debug", ratio_vis)

        k = cv2.waitKey(10) & 0xFF
        if k in (ord('q'), 27):
            break

    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()