from PIL import Image
import glob
import os
import cv2
import numpy as np
import torch
from torch.utils.data import Dataset

from utils import carregar_labelmap_completo, compute_roi_indices, resize_keep_width

# ----------------------------
# Dataset
# ----------------------------
class ROISegDataset(Dataset):
    def __init__(self, root, out_dir, zona_inicio, faixa_atuacao, input_w=384, min_input_h=96, labelmap_path="labelmap.txt",
                 mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):

        self.out_dir = out_dir
        self.img_paths = sorted(glob.glob(os.path.join(root, "images", "*")))
        self.msk_paths = sorted(glob.glob(os.path.join(root, "masks", "*")))
        assert len(self.img_paths) == len(self.msk_paths) and len(self.img_paths) > 0, \
            f"Nenhuma imagem/máscara encontrada em {root}"

        self.zona_inicio = zona_inicio
        self.faixa_atuacao = faixa_atuacao
        self.input_w = input_w
        self.min_input_h = min_input_h
        self.mean = np.array(mean, dtype=np.float32).reshape(1, 1, 3)
        self.std = np.array(std, dtype=np.float32).reshape(1, 1, 3)

        _, _, self.classes, self.ignore_rgb = carregar_labelmap_completo(labelmap_path)
        self.ignore_id = self.ignore_rgb[0]

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, idx):
        img_rgb = np.array(Image.open(self.img_paths[idx]).convert("RGB"))
        msk_grayscale = np.array(Image.open(self.msk_paths[idx]).convert("L"))

        H, W = img_rgb.shape[:2]
        y_fim, y_inicio = compute_roi_indices(H, self.zona_inicio, self.faixa_atuacao)

        img_roi = img_rgb[y_fim:y_inicio, 0:W]
        msk_roi = msk_grayscale[y_fim:y_inicio, 0:W]

        img_in = resize_keep_width(img_roi, self.input_w, self.min_input_h, cv2.INTER_AREA)
        msk_ids = resize_keep_width(msk_roi, self.input_w, self.min_input_h, cv2.INTER_NEAREST)

        valores_validos = list(range(len(self.classes))) + [255]
        msk_ids[np.isin(msk_ids, valores_validos, invert=True)] = self.ignore_id  # converte inválidos em ignore

        valores_unicos = np.unique(msk_ids)
        if np.all(msk_ids == 255):
            raise ValueError(f"Máscara {self.msk_paths[idx]} está só com valor de ignore (255)")

        img_f = img_in.astype(np.float32) / 255.0
        img_f = (img_f - self.mean) / self.std
        img_chw = np.transpose(img_f, (2, 0, 1))

        if idx == 0:
            cv2.imwrite(os.path.join(self.out_dir, "debug_roi_input.png"), img_roi)
            cv2.imwrite(os.path.join(self.out_dir, "debug_roi_mask.png"), msk_roi)

        return torch.from_numpy(img_chw).float(), torch.from_numpy(msk_ids.astype(np.int64))