import os
from PIL import Image
import numpy as np
import cv2

import torch
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
from torchvision import transforms
from sklearn.model_selection import train_test_split

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]


class CLAHETransform:
    """Apply CLAHE to the L channel of LAB color space.

    Operates on PIL Images; returns a PIL Image.
    """

    def __init__(self, clip_limit: float = 2.0, tile_grid_size: tuple = (8, 8)):
        self.clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)

    def __call__(self, img: Image.Image) -> Image.Image:
        img_np = np.array(img)  # RGB uint8
        lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
        lab[:, :, 0] = self.clahe.apply(lab[:, :, 0])
        result = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
        return Image.fromarray(result)


class DRDataset(Dataset):
    """Diabetic Retinopathy dataset from flat class-folder layout."""

    def __init__(self, image_paths: list, labels: list, transform=None):
        self.image_paths = image_paths
        self.labels = labels
        self.transform = transform

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int):
        img = Image.open(self.image_paths[idx]).convert('RGB')
        label = self.labels[idx]
        if self.transform is not None:
            img = self.transform(img)
        return img, label


def split_dataset(data_root: str, seed: int = 42):
    """Scan class folders 0–4 and return stratified 70/15/15 splits.

    Returns:
        (train_paths, train_labels, val_paths, val_labels, test_paths, test_labels)
    """
    all_paths, all_labels = [], []
    for class_idx in range(5):
        class_dir = os.path.join(data_root, str(class_idx))
        for fname in sorted(os.listdir(class_dir)):
            if fname.lower().endswith(('.jpg', '.jpeg', '.png')):
                all_paths.append(os.path.join(class_dir, fname))
                all_labels.append(class_idx)

    # 70% train, 30% temp
    train_paths, temp_paths, train_labels, temp_labels = train_test_split(
        all_paths, all_labels, test_size=0.30, stratify=all_labels, random_state=seed
    )
    # 50/50 of temp -> 15% val, 15% test
    val_paths, test_paths, val_labels, test_labels = train_test_split(
        temp_paths, temp_labels, test_size=0.50, stratify=temp_labels, random_state=seed
    )

    print(f"Dataset split — train: {len(train_paths)}, val: {len(val_paths)}, test: {len(test_paths)}")
    return train_paths, train_labels, val_paths, val_labels, test_paths, test_labels


def build_transforms(use_clahe: bool, training: bool, img_size: int = 224):
    """Build a transform pipeline for training or eval.

    Training path uses RandomResizedCrop for scale/position augmentation.
    Eval path uses a deterministic Resize to preserve comparability.
    CLAHE is applied before spatial transforms at a slightly larger size
    so the crop has room to operate.
    """
    ops = []
    if training:
        if use_clahe:
            # Resize larger so RandomResizedCrop still sees enough context after CLAHE
            ops.append(transforms.Resize((img_size + 32, img_size + 32)))
            ops.append(CLAHETransform(clip_limit=2.0, tile_grid_size=(8, 8)))
        ops += [
            transforms.RandomResizedCrop(img_size, scale=(0.8, 1.0), ratio=(0.9, 1.1)),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomVerticalFlip(p=0.5),
            transforms.RandomRotation(degrees=30),
            transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.02),
            transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 1.0)),
        ]
    else:
        ops.append(transforms.Resize((img_size, img_size)))
        if use_clahe:
            ops.append(CLAHETransform(clip_limit=2.0, tile_grid_size=(8, 8)))
    ops += [
        transforms.ToTensor(),
        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
    ]
    return transforms.Compose(ops)


def build_dataloaders(
    train_paths, train_labels,
    val_paths, val_labels,
    test_paths, test_labels,
    use_clahe: bool,
    batch_size: int = 32,
    num_workers: int = 4,
    img_size: int = 224,
):
    """Create DataLoaders for all three splits.

    The training DataLoader uses WeightedRandomSampler so that each class
    is sampled at approximately equal frequency, counteracting class imbalance.
    """
    train_ds = DRDataset(train_paths, train_labels, transform=build_transforms(use_clahe, training=True, img_size=img_size))
    val_ds = DRDataset(val_paths, val_labels, transform=build_transforms(use_clahe, training=False, img_size=img_size))
    test_ds = DRDataset(test_paths, test_labels, transform=build_transforms(use_clahe, training=False, img_size=img_size))

    # Per-sample weights: sqrt-inverse-frequency for moderate class balancing.
    # Full inverse (1/count) causes over-prediction of minorities; sqrt gives a
    # softer boost that retains enough majority-class signal.
    labels_arr = np.array(train_labels)
    class_counts = np.bincount(labels_arr, minlength=5)
    class_sample_weights = 1.0 / np.sqrt(class_counts.astype(float))
    sample_weights = class_sample_weights[labels_arr]
    sampler = WeightedRandomSampler(
        weights=torch.from_numpy(sample_weights).float(),
        num_samples=len(train_labels),
        replacement=True,
    )

    train_loader = DataLoader(
        train_ds, batch_size=batch_size, sampler=sampler,
        num_workers=num_workers, pin_memory=True,
        persistent_workers=True, prefetch_factor=4,
    )
    val_loader = DataLoader(
        val_ds, batch_size=batch_size * 2, shuffle=False,
        num_workers=num_workers, pin_memory=True,
        persistent_workers=True, prefetch_factor=4,
    )
    test_loader = DataLoader(
        test_ds, batch_size=batch_size * 2, shuffle=False,
        num_workers=num_workers, pin_memory=True,
        persistent_workers=True, prefetch_factor=4,
    )
    return train_loader, val_loader, test_loader