import torch
from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
    def __init__(self, data, labels, transform=None):
        self.data = data
        self.labels = labels
        self.transform = transform

    def __len__(self):
        """데이터셋 크기 반환"""
        return len(self.data)

    def __getitem__(self, idx):
        """인덱스로 샘플 하나를 반환"""
        sample = self.data[idx]
        label = self.labels[idx]
        if self.transform:
            sample = self.transform(sample)
        return sample, label

# 사용 예시
X = torch.randn(1000, 10)     # 1000개 샘플, 10차원
y = torch.randint(0, 3, (1000,))  # 3 클래스

dataset = CustomDataset(X, y)
print(f"데이터셋 크기: {len(dataset)}")
sample, label = dataset[0]
print(f"샘플 shape: {sample.shape}, 레이블: {label}")

dataloader = DataLoader(
    dataset,
    batch_size=32,       # 미니배치 크기
    shuffle=True,        # 에포크마다 데이터 섞기
    num_workers=4,       # 병렬 데이터 로딩 (CPU 코어 수에 맞게)
    drop_last=True,      # 마지막 불완전 배치 제거
    pin_memory=True,     # GPU 전송 속도 향상
)

# 배치 순회
for batch_idx, (batch_data, batch_labels) in enumerate(dataloader):
    print(f"배치 {batch_idx}: data={batch_data.shape}, labels={batch_labels.shape}")
    if batch_idx >= 2:
        break
# 배치 0: data=torch.Size([32, 10]), labels=torch.Size([32])

import torchvision
import torchvision.transforms as transforms

# 이미지 변환 파이프라인
transform = transforms.Compose([
    transforms.Resize(32),
    transforms.ToTensor(),                          # PIL → Tensor, [0, 1] 정규화
    transforms.Normalize((0.5,), (0.5,)),           # [-1, 1] 정규화
])

# MNIST 다운로드 및 로드
train_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=True,
    download=True,
    transform=transform,
)

test_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=False,
    transform=transform,
)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 첫 배치 확인
images, labels = next(iter(train_loader))
print(f"이미지 배치: {images.shape}")   # (64, 1, 32, 32)
print(f"레이블 배치: {labels.shape}")    # (64,)

train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),   # 좌우 반전
    transforms.RandomRotation(10),             # ±10도 회전
    transforms.ColorJitter(brightness=0.2),    # 밝기 변화
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),  # ImageNet 통계
])

# 검증/테스트에는 증강 없이 정규화만
val_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

from torch.utils.data import random_split

full_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True)

# 80% 학습, 20% 검증
train_size = int(0.8 * len(full_dataset))
val_size = len(full_dataset) - train_size
train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size])

print(f"학습: {len(train_dataset)}, 검증: {len(val_dataset)}")

import pandas as pd

class CSVDataset(Dataset):
    """CSV 파일을 읽어 텐서로 변환하는 데이터셋"""
    def __init__(self, csv_path, target_column):
        df = pd.read_csv(csv_path)
        self.features = torch.FloatTensor(df.drop(columns=[target_column]).values)
        self.labels = torch.LongTensor(df[target_column].values)

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

    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]

# 사용
# dataset = CSVDataset('data/train.csv', target_column='label')
# loader = DataLoader(dataset, batch_size=32, shuffle=True)