# custom_channel_dataset.py

from pathlib import Path
import cv2
import numpy as np
from ultralytics.data.dataset import YOLODataset
from ultralytics.data.augment import Compose, LetterBox, Format

# 🔧 AJUSTE AQUI PRO SEU RAW REAL:
RAW_ALREADY_01 = True
RAW_DTYPE = np.float32   # ou np.uint16, conforme seu arquivo .raw
RAW_MAX_VAL = 1.0 if RAW_ALREADY_01 else 65535.0


class CustomChannelSegDataset(YOLODataset):
    """
    Dataset de segmentação para imagens 4-canais (R, G, IR, B).

    Convenção:
      - O train.txt / val.txt continuam apontando para as IMAGENS RGB normais,
        ex: data/images/train/20260126_121341_900.jpg

      - Ao lado de cada imagem, você terá um .raw
        ex: data/images/train/20260126_121341_900.raw

        Esse .raw deve ser (H, W, 4) float32 em [0,1],
        já com os canais [R, G, IR, B] na ordem que combinamos.
    """

    def __init__(self, img_path, data, canais, imgsz, augment, hyp, **kwargs):
        self.canais = canais

        data = dict(data)
        data["channels"] = canais

        super().__init__(
            img_path=img_path,
            data=data,
            imgsz=imgsz,
            augment=augment,
            hyp=hyp,
            task="segment",
            **kwargs,
        )

        if len(self.im_files):
            p0 = self.im_files[0]
            preview0 = cv2.imread(p0, cv2.IMREAD_UNCHANGED)
            if preview0 is None:
                raise FileNotFoundError(f"Não consegui ler preview: {p0}")
            self._fixed_hw = preview0.shape[:2]  # (h0, w0)
        else:
            self._fixed_hw = None

    def load_image(self, i):
        """
        Lê a imagem a partir do .raw e retorna:
        im    : np.ndarray uint8 com shape (H, W, C) onde C = 4 ou 5
        shapes: (h0, w0) tamanho original
        """
        im_path = self.im_files[i]  # caminho listado no train/val.txt

        # 1) Usa o JPG/PNG só pra descobrir H, W
        #preview = cv2.imread(im_path, cv2.IMREAD_UNCHANGED)
        #if preview is None:
        #    raise FileNotFoundError(f"Não consegui ler preview: {im_path}")
        #h0, w0 = preview.shape[:2]
        h0, w0 = self._fixed_hw

        # 2) Caminho do RAW correspondente
        p = Path(im_path)
        raw_path = p.with_suffix(".raw")  # troca extensão por .raw
        if not raw_path.exists():
            raise FileNotFoundError(f"RAW 4ch não encontrado: {raw_path}")

        # 3) Carrega o RAW como CHW (4, H, W)
        #    IMPORTANTE: o arquivo físico sempre tem 4 canais (R, G, IR, B)
        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"Tamanho inesperado no RAW {raw_path}: "
                f"esperado {expected_size}, veio {arr_flat.size}"
            )

        raw_chw = arr_flat.reshape(num_raw_channels, h0, w0)  # (4, H, W)

        # 4) Converte para HWC float32 (H, W, 4)
        arr = np.transpose(raw_chw, (1, 2, 0)).astype(np.float32)  # (H, W, 4)

        # 5) Normaliza para 0..1 de forma consistente
        arr = arr.astype(np.float32, copy=False)
        if RAW_ALREADY_01:
            arr_norm = arr
        else:
            # Normalização por escala fixa (ex.: 65535 p/ uint16)
            arr_norm = arr / float(RAW_MAX_VAL)
        # Segurança: clamp
        arr_norm = np.clip(arr_norm, 0.0, 1.0)

        # 6) Se self.canais == 5, calculamos NDVI e empilhamos como 5º canal
        #    Ordem final: [R, G, IR, B, NDVI]
        if self.canais == 5:
            # índices assumidos: 0=R, 1=G, 2=IR, 3=B
            R  = arr_norm[..., 0]
            IR = arr_norm[..., 2]

            # NDVI bruto
            eps = 1e-6
            ndvi = (IR - R) / (IR + R + eps)   # faixa ~[-1, 1]

            # clampa e reescala pra [0, 1] (fica alinhado com outros canais)
            ndvi = np.clip(ndvi, -1.0, 1.0)
            ndvi01 = (ndvi + 1.0) / 2.0

            # empilha como 5º canal
            arr_norm = np.concatenate(
                [arr_norm, ndvi01[..., None]],
                axis=-1
            )  # (H, W, 5)

        elif self.canais == 4:
            # usa só os 4 canais originais
            pass
        else:
            raise ValueError(f"self.canais deve ser 4 ou 5, veio {self.canais}.")

        # 7) Converte de 0..1 para 0..255 uint8 pro YOLO
        im = np.clip(arr_norm * 255.0, 0, 255).astype(np.uint8)

        #if i == 0:
        #    print("im dtype:", im.dtype, "min/max:", im.min(), im.max(), "shape:", im.shape)

        h, w = im.shape[:2]
        return im, (h0, w0), (h, w)

    def build_targets(self, batch):
        # Garante que as masks e labels aceitem o formato custom
        return super().build_targets(batch)

    def build_transforms_disabled(self, hyp=None):
        # Pega o tamanho da imagem (imgsz)
        new_shape = getattr(self, 'imgsz', 800)
        
        return Compose([
            # 1. Redimensiona a imagem 4-canais e as masks
            LetterBox(new_shape=new_shape),
            # 2. Converte tudo para o formato que o Segment Head entende
            Format(
                bbox_format="xywh",      # Formato padrão do YOLO (centro_x, centro_y, largura, altura)
                normalize=True,          # Normaliza bboxes para 0-1
                return_mask=True,        # ESSENCIAL para segmentação
                batch_idx=True,          # Necessário para o DataLoader do Trainer
                mask_ratio=4,            # Padrão do YOLOv8-seg
                mask_overlap=True,       # Padrão do YOLOv8-seg
            )
        ])