data_loaders.py

"""
data_loaders.py
---------------
Datasets e DataLoaders para treino/validação/teste de classificação de nódulos pulmonares.
- LunaDataset: carrega patches 2.5D (.npy) a partir de um CSV com colunas npy_path e class_target.
- prepare_loaders: divide os dados em hold-out (15% teste) + K-Fold nos 85% restantes, com
  amostragem balanceada no treino e transformações por tipo de modelo (ResNet vs VMamba).
"""

import os
import torch
import pandas as pd
import numpy as np
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
from sklearn.model_selection import GroupKFold, GroupShuffleSplit
import albumentations as A
from albumentations.pytorch import ToTensorV2

# ---------------------------------------------------------------------------
# Caminhos: CSV com patches 2.5D (saída do pre_processing.py).
# No Google Colab, defina os.environ["PROJECT_PATH"] = PROJECT_PATH na 1ª célula
# para que todos os resultados sejam salvos na pasta do projeto (Drive).
# ---------------------------------------------------------------------------
BASE_DIR = os.environ.get("PROJECT_PATH", os.getcwd())
DATA_DIR = os.path.join(BASE_DIR, "data")
CSV_25D = os.path.join(DATA_DIR, "processed_all_data_25d.csv")
PROCESSED_CSV = CSV_25D
BATCH_SIZE = 16


class LunaDataset(Dataset):
    """
    Dataset PyTorch para patches 2.5D de nódulos (LUNA16 + LIDC).
    Espera um DataFrame com colunas: npy_path (caminho do .npy), class_target (0=benigno, 1=maligno).
    Retorna (imagem [C,H,W], label float) com imagem em 3 canais 224x224.
    """

    def __init__(self, dataframe, transforms=None):
        self.df = dataframe
        self.transforms = transforms

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

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        npy_path = row['npy_path']
        label = float(row['class_target'])

        # Carrega patch do disco; fallback para zeros se arquivo ausente ou erro
        if not os.path.exists(npy_path):
            image = np.zeros((224, 224, 3), dtype=np.float32)
        else:
            try:
                image = np.load(npy_path).astype(np.float32)
            except Exception:
                image = np.zeros((224, 224, 3), dtype=np.float32)

        # Garante 3 canais (H, W, 3) para compatibilidade com redes que esperam RGB
        if len(image.shape) == 2:
            image = np.expand_dims(image, axis=2)
        if image.shape[2] == 1:
            image = np.repeat(image, 3, axis=2)

        # Normalização HU: janela pulmonar [-1000, 400] -> [0, 1] (só se ainda estiver em HU)
        if image.max() > 1.0:
            min_hu, max_hu = -1000.0, 400.0
            image = np.clip(image, min_hu, max_hu)
            image = (image - min_hu) / (max_hu - min_hu)
            image = image.astype(np.float32)

        # Aplica augmentations (Albumentations) ou apenas converte para tensor (C, H, W)
        if self.transforms:
            augmented = self.transforms(image=image)
            image = augmented['image']
        else:
            image = torch.from_numpy(image).permute(2, 0, 1)

        return image, torch.tensor(label, dtype=torch.float32)

def get_transforms(split='train', model_type='resnet'):
    """
    Retorna pipeline de transformações (Albumentations) para treino ou validação.
    - split: 'train' (flips, rotação, affine, coarse dropout) ou 'val'/'test' (só normalização).
    - model_type: 'resnet' usa normalização ImageNet; 'vmamba_tiny' usa média 0 e std 1 (entrada [0,1]).
    """
    # Normalização: ResNet foi pré-treinado no ImageNet; VMamba usa entrada já em [0,1]
    if model_type == 'resnet':
        norm = A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
    else:
        norm = A.Normalize(mean=(0, 0, 0), std=(1, 1, 1))

    # Pipeline: treino com data augmentation; val/test só normaliza e converte para tensor
    if split == 'train':
        return A.Compose([
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),
            A.Rotate(limit=20, p=0.5),
            A.Affine(scale=(0.9, 1.1), translate_percent=(0.1, 0.1), p=0.5),
            A.OneOf([
                # CORREÇÃO: Sintaxe moderna do CoarseDropout (sem fill_value explícito se padrão)
                A.CoarseDropout(
                    num_holes_range=(1, 8),
                    hole_height_range=(8, 16),
                    hole_width_range=(8, 16),
                    # fill_value removido pois o padrão já é 0 e algumas versões reclamam
                    p=0.5
                ),
            ], p=0.3),
            norm, 
            ToTensorV2()
        ])
    else:
        return A.Compose([
            norm, 
            ToTensorV2()
        ])

def make_balanced_sampler(df):
    """
    Cria WeightedRandomSampler para balancear as classes no treino.
    Amostras da classe minoritária têm maior peso, reduzindo viés quando há desbalanceamento.
    Retorna None se houver menos de 2 classes.
    """
    targets = df['class_target'].values
    class_counts = np.bincount(targets.astype(int))

    if len(class_counts) < 2:
        return None

    class_weights = 1.0 / class_counts
    sample_weights = class_weights[targets.astype(int)]

    sampler = WeightedRandomSampler(
        weights=sample_weights, num_samples=len(sample_weights), replacement=True
    )
    return sampler


def split_holdout(df, test_size=0.15):
    """
    Separa um conjunto de teste (hold-out) por paciente (seriesuid).
    Garante que nenhum paciente do treino/val apareça no teste (evita data leakage).
    """
    gss = GroupShuffleSplit(n_splits=1, test_size=test_size, random_state=42)
    groups = df['seriesuid']
    
    train_val_idx, test_idx = next(gss.split(df, groups=groups))
    
    train_val_df = df.iloc[train_val_idx].reset_index(drop=True)
    test_df = df.iloc[test_idx].reset_index(drop=True)
    
    return train_val_df, test_df

def prepare_loaders(csv_path=PROCESSED_CSV, fold=0, n_splits=5, batch_size=BATCH_SIZE, model_type='resnet', mode='train_val'):
    """
    Prepara DataLoaders para o pipeline de treino/avaliação.

    - csv_path: CSV com npy_path e class_target (ex.: processed_all_data_25d.csv).
    - fold: índice do fold (0 a n_splits-1) quando mode='train_val'.
    - n_splits: número de folds no K-Fold (GroupKFold por seriesuid).
    - model_type: 'resnet' ou 'vmamba_tiny' (define normalização e transforms).
    - mode: 'train_val' -> (train_loader, val_loader) para o fold dado;
            'test' -> test_loader do hold-out (15% dos pacientes).

    O hold-out é fixo (15%); o K-Fold é aplicado apenas nos 85% restantes.
    """
    if not os.path.exists(csv_path):
        raise FileNotFoundError(f"CSV não encontrado: {csv_path}")

    df = pd.read_csv(csv_path)
    
    # 1. Separação Rigorosa do Hold-out (15%)
    train_val_df, test_df = split_holdout(df, test_size=0.15)
    
    if mode == 'test':
        print(f"🔒 Preparando Dataloader de TESTE FINAL (Hold-out): {len(test_df)} imagens.")
        test_ds = LunaDataset(test_df, transforms=get_transforms('val', model_type))
        test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
        return test_loader

    # 2. Se for treino, aplica K-Fold nos 85% restantes
    gkf = GroupKFold(n_splits=n_splits)
    groups = train_val_df['seriesuid']
    
    folds = list(gkf.split(train_val_df, train_val_df['class_target'], groups))
    train_idx, val_idx = folds[fold]
    
    train_fold = train_val_df.iloc[train_idx].reset_index(drop=True)
    val_fold = train_val_df.iloc[val_idx].reset_index(drop=True)
    
    print(f"📊 [K-Fold {fold+1}/{n_splits}] Train: {len(train_fold)} | Val: {len(val_fold)} (Total Pool: {len(train_val_df)})")
    
    # Sampler e Loaders
    train_sampler = make_balanced_sampler(train_fold)
    
    train_ds = LunaDataset(train_fold, transforms=get_transforms('train', model_type))
    val_ds = LunaDataset(val_fold, transforms=get_transforms('val', model_type))
    
    if train_sampler:
        print("⚖️ Sampler Balanceado ativado.")
        train_loader = DataLoader(train_ds, batch_size=batch_size, sampler=train_sampler, shuffle=False, num_workers=2, pin_memory=True)
    else:
        train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True)
    
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
    
    return train_loader, val_loader