"""
Data Loader Module for Verril-Learn

Handles MNIST dataset loading with intentional data poisoning
to simulate adversarial attacks or privacy errors.

Poison Mechanism:
- Takes a percentage of digit '7' samples
- Maliciously relabels them as '1'
- This simulates a label-flipping attack
"""

import torch
from torch.utils.data import DataLoader, Subset, Dataset
from torchvision import datasets, transforms
import numpy as np
from typing import Tuple, List, Optional


class PoisonedDataset(Dataset):
    """
    A wrapper dataset that applies label poisoning to selected indices.
    
    This allows us to track which samples are poisoned while maintaining
    efficient data loading through PyTorch's DataLoader.
    """
    
    def __init__(
        self,
        base_dataset: Dataset,
        poison_indices: List[int],
        source_label: int = 7,
        target_label: int = 1
    ):
        """
        Args:
            base_dataset: The original MNIST dataset
            poison_indices: List of indices to poison
            source_label: Original label to flip (default: 7)
            target_label: Label to flip to (default: 1)
        """
        self.base_dataset = base_dataset
        self.poison_indices = set(poison_indices)
        self.source_label = source_label
        self.target_label = target_label
    
    def __len__(self) -> int:
        return len(self.base_dataset)
    
    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
        image, label = self.base_dataset[idx]
        
        # If this index is poisoned and has the source label, flip it
        if idx in self.poison_indices and label == self.source_label:
            label = self.target_label
        
        return image, label


def get_mnist_transforms() -> transforms.Compose:
    """
    Standard MNIST preprocessing transforms.
    
    Returns:
        transforms.Compose: Composition of ToTensor and Normalize
    """
    return transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))  # MNIST mean and std
    ])


def get_clean_mnist(
    data_dir: str = './data',
    batch_size: int = 64,
    num_workers: int = 0
) -> Tuple[DataLoader, DataLoader]:
    """
    Load clean (unpoisoned) MNIST dataset.
    
    Args:
        data_dir: Directory to download/load MNIST
        batch_size: Batch size for DataLoaders
        num_workers: Number of workers for data loading
    
    Returns:
        Tuple of (train_loader, test_loader)
    """
    transform = get_mnist_transforms()
    
    train_dataset = datasets.MNIST(
        root=data_dir,
        train=True,
        download=True,
        transform=transform
    )
    
    test_dataset = datasets.MNIST(
        root=data_dir,
        train=False,
        download=True,
        transform=transform
    )
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers
    )
    
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers
    )
    
    return train_loader, test_loader


def get_poisoned_mnist(
    poison_ratio: float = 0.1,
    source_label: int = 7,
    target_label: int = 1,
    data_dir: str = './data',
    batch_size: int = 64,
    num_workers: int = 0,
    random_seed: Optional[int] = 42
) -> Tuple[DataLoader, DataLoader, List[int]]:
    """
    Load MNIST with label poisoning attack.
    
    This function simulates a data poisoning attack by flipping labels
    of a specified percentage of samples from source_label to target_label.
    
    Attack Description:
        - Identifies all samples with source_label (default: digit '7')
        - Randomly selects poison_ratio% of these samples
        - Changes their labels to target_label (default: digit '1')
        - Returns poisoned train loader and clean test loader
    
    Args:
        poison_ratio: Fraction of source_label samples to poison (0.0 to 1.0)
        source_label: The digit label to attack (default: 7)
        target_label: The label to flip to (default: 1)
        data_dir: Directory to download/load MNIST
        batch_size: Batch size for DataLoaders
        num_workers: Number of workers for data loading
        random_seed: Random seed for reproducibility
    
    Returns:
        Tuple containing:
            - train_loader: DataLoader with poisoned training data
            - test_loader: DataLoader with clean test data
            - poison_indices: List of indices that were poisoned
    
    Example:
        >>> train_loader, test_loader, poison_idx = get_poisoned_mnist(
        ...     poison_ratio=0.1,
        ...     source_label=7,
        ...     target_label=1
        ... )
        >>> print(f"Poisoned {len(poison_idx)} samples")
    """
    if not 0.0 <= poison_ratio <= 1.0:
        raise ValueError(f"poison_ratio must be between 0 and 1, got {poison_ratio}")
    
    if random_seed is not None:
        np.random.seed(random_seed)
        torch.manual_seed(random_seed)
    
    transform = get_mnist_transforms()
    
    # Load base datasets
    train_dataset = datasets.MNIST(
        root=data_dir,
        train=True,
        download=True,
        transform=transform
    )
    
    test_dataset = datasets.MNIST(
        root=data_dir,
        train=False,
        download=True,
        transform=transform
    )
    
    # Find all indices with the source label
    source_indices = []
    for idx in range(len(train_dataset)):
        _, label = train_dataset[idx]
        if label == source_label:
            source_indices.append(idx)
    
    print(f"[Verril-Learn] Found {len(source_indices)} samples with label {source_label}")
    
    # Randomly select indices to poison
    num_poison = int(len(source_indices) * poison_ratio)
    poison_indices = np.random.choice(
        source_indices,
        size=num_poison,
        replace=False
    ).tolist()
    
    print(f"[Verril-Learn] Poisoning {num_poison} samples ({poison_ratio*100:.1f}%)")
    print(f"[Verril-Learn] Label flip: {source_label} → {target_label}")
    
    # Create poisoned dataset wrapper
    poisoned_dataset = PoisonedDataset(
        base_dataset=train_dataset,
        poison_indices=poison_indices,
        source_label=source_label,
        target_label=target_label
    )
    
    # Create DataLoaders
    train_loader = DataLoader(
        poisoned_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers
    )
    
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers
    )
    
    return train_loader, test_loader, poison_indices


def split_poisoned_and_retain(
    train_dataset: Dataset,
    poison_indices: List[int],
    batch_size: int = 64,
    num_workers: int = 0
) -> Tuple[DataLoader, DataLoader]:
    """
    Split the training data into poisoned and retain (clean) subsets.
    
    This is useful for the unlearning algorithm which needs:
    - poisoned_loader: Data to forget (maximize loss on)
    - retain_loader: Data to remember (maintain performance on)
    
    Args:
        train_dataset: The full training dataset
        poison_indices: List of poisoned sample indices
        batch_size: Batch size for DataLoaders
        num_workers: Number of workers for data loading
    
    Returns:
        Tuple of (poisoned_loader, retain_loader)
    """
    poison_set = set(poison_indices)
    all_indices = set(range(len(train_dataset)))
    retain_indices = list(all_indices - poison_set)
    
    poisoned_subset = Subset(train_dataset, poison_indices)
    retain_subset = Subset(train_dataset, retain_indices)
    
    poisoned_loader = DataLoader(
        poisoned_subset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers
    )
    
    retain_loader = DataLoader(
        retain_subset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers
    )
    
    print(f"[Verril-Learn] Split data: {len(poison_indices)} poisoned, {len(retain_indices)} retain")
    
    return poisoned_loader, retain_loader


def verify_poisoning(
    dataloader: DataLoader,
    poison_indices: List[int],
    source_label: int = 7,
    target_label: int = 1,
    num_samples: int = 10
) -> None:
    """
    Verify that poisoning was applied correctly by checking sample labels.
    
    Args:
        dataloader: The poisoned DataLoader
        poison_indices: List of poisoned indices
        source_label: Original label
        target_label: Poisoned label
        num_samples: Number of samples to verify
    """
    print(f"\n[Verril-Learn] Verifying poisoning...")
    print(f"Checking first {num_samples} poisoned samples:")
    
    dataset = dataloader.dataset
    checked = 0
    
    for idx in poison_indices[:num_samples]:
        _, label = dataset[idx]
        expected = target_label
        status = "✓" if label == expected else "✗"
        print(f"  Index {idx}: label={label} (expected {expected}) {status}")
        checked += 1
    
    print(f"Verification complete for {checked} samples.\n")
