# ══════════════════════════════════════════════════════════════════════════════
# PYTORCH — Transforms & DataLoaders
# ══════════════════════════════════════════════════════════════════════════════
def get_pytorch_transforms(img_size: int = 150):
    """
    Retourne (train_transform, val_transform).
    Augmentation scène-aware pour le dataset Intel (6 classes naturelles RGB).
    """
    from torchvision import transforms

    # Statistiques ImageNet — optimal pour images naturelles RGB 3 canaux
    MEAN = [0.485, 0.456, 0.406]
    STD  = [0.229, 0.224, 0.225]

    train_transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomVerticalFlip(p=0.1),
        transforms.RandomRotation(degrees=40),
        transforms.ColorJitter(
            brightness=0.3, contrast=0.2, saturation=0.1, hue=0.05
        ),
        transforms.RandomGrayscale(p=0.05),
        transforms.ToTensor(),
        transforms.Normalize(MEAN, STD),      # ← normalisation ici, pas dans le modèle
        transforms.RandomErasing(p=0.15, scale=(0.02, 0.15)),
    ])

    val_transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.ToTensor(),
        transforms.Normalize(MEAN, STD),      # ← même normalisation en val/test
    ])

    return train_transform, val_transform


def get_pytorch_loaders(
    train_dir: str,
    test_dir:  str,
    img_size:  int = 150,
    batch_size: int = 64,
):
    from torch.utils.data import DataLoader
    from torchvision import datasets

    train_tf, val_tf = get_pytorch_transforms(img_size)

    train_loader = DataLoader(
        datasets.ImageFolder(train_dir, transform=train_tf),
        batch_size=batch_size, shuffle=True,
        num_workers=2, pin_memory=True,
    )
    test_loader = DataLoader(
        datasets.ImageFolder(test_dir, transform=val_tf),
        batch_size=batch_size, shuffle=False,
        num_workers=2, pin_memory=True,
    )
    return train_loader, test_loader


# ══════════════════════════════════════════════════════════════════════════════
# TENSORFLOW — Dataset pipeline
# ══════════════════════════════════════════════════════════════════════════════
def get_tf_datasets(
    train_dir: str,
    test_dir: str,
    img_size: int = 228,
    batch_size: int = 64,
):
    import tensorflow as tf

    # Same preprocessing as in the notebook
    norm_layer = tf.keras.layers.Rescaling(1.0 / 255.0)

    # ── Raw loading ──────────────────────────────────────────────────────────
    train_ds = tf.keras.utils.image_dataset_from_directory(
        train_dir,
        seed=123,
        image_size=(img_size, img_size),
        batch_size=batch_size,
        shuffle=True,
        label_mode="int",
    )

    test_ds = tf.keras.utils.image_dataset_from_directory(
        test_dir,
        seed=123,
        image_size=(img_size, img_size),
        batch_size=batch_size,
        shuffle=False,
        label_mode="int",
    )

    # ── Normalization only ───────────────────────────────────────────────────
    train_ds = train_ds.map(
        lambda x, y: (norm_layer(x), y),
        num_parallel_calls=tf.data.AUTOTUNE
    )

    test_ds = test_ds.map(
        lambda x, y: (norm_layer(x), y),
        num_parallel_calls=tf.data.AUTOTUNE
    )

    # ── Performance ──────────────────────────────────────────────────────────
    train_ds = train_ds.prefetch(tf.data.AUTOTUNE)
    test_ds = test_ds.prefetch(tf.data.AUTOTUNE)

    return train_ds, test_ds

# ══════════════════════════════════════════════════════════════════════════════
# INFÉRENCE — Preprocessing image unique (Flask / production)
# ══════════════════════════════════════════════════════════════════════════════
def preprocess_image_pytorch(pil_img, img_size: int = 150):
    """Prépare une image PIL pour l'inférence PyTorch. Retourne (1,3,H,W)."""
    import torch
    from torchvision import transforms

    tf = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])
    return tf(pil_img).unsqueeze(0)


def preprocess_image_tf(pil_img, img_size: int = 150):
    """
    Prépare une image PIL pour l'inférence TensorFlow. Retourne (1,H,W,3).
    Normalisation identique au pipeline val/test : ÷255 → [0,1].
    """
    import numpy as np
    arr = np.array(pil_img.resize((img_size, img_size)), dtype=np.float32)
    arr = arr / 255.0              # ← même normalisation que normalize_only()
    return np.expand_dims(arr, 0)  # (1, H, W, 3)