"""
data.py — CIFAR-10 with synthetic source / target domain split.

Source domain  = standard augmentation, subset of classes
Target domain  = Raman-style channel noise + jitter, ALL classes
Self-sup proxy = flattened image as pseudo-spectrum, random masking
"""

import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms

# ── Transforms ─────────────────────────────────────────────────────────────────

MEAN = (0.4914, 0.4822, 0.4465)
STD  = (0.2023, 0.1994, 0.2010)

TRAIN_TRANSFORM = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(MEAN, STD),
])

EVAL_TRANSFORM = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(MEAN, STD),
])

# ── Raman-style augmentation ──────────────────────────────────────────────────

class RamanStyleAugmentation:
    """
    Proxy for inter-device spectral variance in Raman spectroscopy.
    Simulates:
      • Gain variation  (instrument intensity scaling)
      • Channel bias   (per-detector offset drift)
      • Additive Gaussian noise (photon shot noise)
    """

    def __init__(self, noise_std=0.1, gain_range=(0.8, 1.2), bias_std=0.05):
        self.noise_std = noise_std
        self.gain_lo, self.gain_hi = gain_range
        self.bias_std = bias_std

    def __call__(self, img):
        # img: Tensor [C, H, W]
        gain = torch.empty(3).uniform_(self.gain_lo, self.gain_hi).view(3, 1, 1)
        bias = torch.empty(3).normal_(0, self.bias_std).view(3, 1, 1)
        noise = torch.randn_like(img) * self.noise_std
        img = img * gain + bias + noise
        return img.clamp(0.0, 1.0)

# ── Dataset wrappers ───────────────────────────────────────────────────────────

class CIFAR10Source(Dataset):
    """
    Source domain: clean images, only classes < num_source_classes.
    Domain label = 0 (source).
    """

    def __init__(self, indices, transform, num_classes=5):
        self.base = datasets.CIFAR10(
            root='/opt/datasets', train=True, download=False, transform=transform
        )
        self.indices = indices
        self.num_classes = num_classes
        # Filter to source-class subset (first num_classes labels)
        self._filtered = [
            i for i in indices if self.base.targets[i] < num_classes
        ]

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

    def __getitem__(self, idx):
        real = self._filtered[idx]
        img, label = self.base[real]
        return img, label, 0   # domain_label = 0 (source)

class CIFAR10Target(Dataset):
    """
    Target domain: Raman-style augmentation applied inline during __getitem__,
    ALL classes exposed (no class filtering). Domain label = 1 (target).

    Unlike CIFAR10Source, this class applies the noise model directly in
    __getitem__ rather than via a separate transform pipeline, so that the
    augmentation is baked into the dataset contract and each access returns
    a freshly-noised sample.
    """

    def __init__(self, indices, noise_std=0.1):
        self.base = datasets.CIFAR10(
            root='/opt/datasets', train=True, download=False, transform=None
        )
        self.indices = indices
        self.noise_std = noise_std
        # Cache mean/std as buffers for normalization (computed once)
        self._mean = torch.tensor(MEAN).view(3, 1, 1)
        self._std  = torch.tensor(STD).view(3, 1, 1)

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

    def __getitem__(self, idx):
        real = self.indices[idx]
        img, label = self.base[real]          # PIL [H, W, C] uint8

        # ToTensor + normalize inline
        img = torch.from_numpy(
            np.array(img).transpose(2, 0, 1)
        ).float() / 255.0                     # [3, 32, 32], [0, 1]
        img = (img - self._mean) / self._std  # normalize

        # Raman-style channel perturbations (fresh draw each access)
        gain = (
            torch.empty(3, 1, 1)
            .uniform_(0.8, 1.2)
            .expand_as(img)
        )
        bias = torch.empty(3, 1, 1).normal_(0, 0.05).expand_as(img)
        noise = torch.randn_like(img) * self.noise_std
        img = img * gain + bias + noise
        img = img.clamp(0.0, 1.0)

        # Re-normalize back to distribution with original stats for model input
        img = img * self._std + self._mean

        return img, label, 1   # domain_label = 1 (target)

class CIFAR10Test(Dataset):
    """
    Test set — clean images (ToTensor + Normalize), all 10 classes.
    Domain label = 0 (treated as source-like for evaluation).
    """

    def __init__(self, transform=None):
        self.base = datasets.CIFAR10(
            root='/opt/datasets', train=False, download=False, transform=None
        )
        self.transform = transform
        self._mean = torch.tensor(MEAN).view(3, 1, 1)
        self._std  = torch.tensor(STD).view(3, 1, 1)

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

    def __getitem__(self, idx):
        img, label = self.base[idx]          # PIL [H, W, C] uint8

        # Manual ToTensor + Normalize (avoids torchvision path for clarity)
        img = torch.from_numpy(
            np.array(img).transpose(2, 0, 1)
        ).float() / 255.0                     # [3, 32, 32], [0, 1]
        img = (img - self._mean) / self._std  # normalize

        if self.transform is not None:
            img = self.transform(img)

        return img, label, 0

# ── DataLoader factory ────────────────────────────────────────────────────────

def get_dataloaders(config):
    """
    Returns dict with keys: 'source_train', 'source_val', 'target_train',
    'target_val', 'test'.

    Source domain = first `source_ratio` fraction of classes (clean).
    Target domain = all classes with Raman-style inline augmentation.
    """
    bs = config.batch_size
    nw = config.num_workers
    pm = config.pin_memory

    num_classes = config.num_classes
    num_source_classes = int(num_classes * config.source_ratio)

    # ── Base datasets ──────────────────────────────────────────
    full_train = datasets.CIFAR10(
        root='/opt/datasets', train=True, download=False, transform=None
    )
    all_indices = list(range(len(full_train)))

    # Stratified split: 80 % train, 20 % val  (same split for both domains)
    np.random.seed(config.seed)
    np.random.shuffle(all_indices)
    split = int(len(all_indices) * 0.8)
    train_idx, val_idx = all_indices[:split], all_indices[split:]

    # Source train / val  (clean, class-filtered)
    src_train_ds = CIFAR10Source(train_idx, TRAIN_TRANSFORM, num_source_classes)
    src_val_ds   = CIFAR10Source(val_idx,   EVAL_TRANSFORM,   num_source_classes)

    # Target train / val  (Raman noise inline, no class filter)
    tgt_train_ds = CIFAR10Target(train_idx, noise_std=config.target_noise_std)
    tgt_val_ds   = CIFAR10Target(val_idx,   noise_std=config.target_noise_std)

    # Test: clean, all classes
    test_ds = CIFAR10Test(transform=EVAL_TRANSFORM)

    loader_cfg = dict(
        batch_size=bs, num_workers=nw, pin_memory=pm,
        shuffle=True, drop_last=True
    )
    loader_cfg_val = dict(
        batch_size=bs, num_workers=nw, pin_memory=pm,
        shuffle=False, drop_last=False
    )

    return {
        'source_train': DataLoader(src_train_ds, **loader_cfg),
        'source_val':   DataLoader(src_val_ds,   **loader_cfg_val),
        'target_train': DataLoader(tgt_train_ds, **loader_cfg),
        'target_val':   DataLoader(tgt_val_ds,   **loader_cfg_val),
        'test':         DataLoader(test_ds,      **loader_cfg_val),
    }

# ── Self-supervised collate (masked spectrum reconstruction) ─────────────────

def selfsup_mask_collate(batch, mask_ratio=0.4):
    """
    Converts a batch into masked-spectrum format.
    Flattens image [B,3,32,32] → pseudo-spectrum [B, 3072].
    Masks mask_ratio fraction with zeros (proxy for masked Raman wavenumbers).
    Returns (masked_spectra, targets, labels).
    """
    images_list, labels_list = zip(*[(img, l) for img, l, _ in batch])
    images = torch.stack(images_list)            # [B, 3, 32, 32]
    labels = torch.tensor(labels_list, dtype=torch.long)

    B = images.size(0)
    x_flat = images.view(B, -1)                  # [B, 3072]

    mask = torch.rand(B, 3072) < mask_ratio       # Bernoulli mask
    masked = x_flat.clone()
    masked[mask] = 0.0                           # [MASK] = zero

    return masked, x_flat.detach().clone(), labels