#!/usr/bin/env python3
"""
预训练数据加载器
支持三个任务的数据加载
"""

import pickle
from pathlib import Path
from typing import Dict, List, Optional
import torch
from torch.utils.data import Dataset
from PIL import Image
import torchvision.transforms as transforms


class PretrainDataset(Dataset):
    """
    预训练数据集
    
    Args:
        data_file: pretrain_data.pkl路径
        split: 'train', 'val', 或 'test'
        task: 'task1', 'task2', 'task3', 或 'all'
        transform: 图像变换
    """
    
    def __init__(
        self,
        data_file: str,
        split: str = "train",
        task: str = "all",
        transform: Optional[transforms.Compose] = None
    ):
        self.split = split
        self.task = task
        self.transform = transform or self.default_transform()
        
        # 加载数据
        with open(data_file, "rb") as f:
            all_data = pickle.load(f)
        
        split_data = all_data[split]
        
        # 根据任务选择数据
        self.samples = []
        if task == "all":
            self.samples.extend(split_data["task1_environment"])
            self.samples.extend(split_data["task2_accident_detection"])
            self.samples.extend(split_data["task3_sequence_prediction"])
        elif task == "task1":
            self.samples = split_data["task1_environment"]
        elif task == "task2":
            self.samples = split_data["task2_accident_detection"]
        elif task == "task3":
            self.samples = split_data["task3_sequence_prediction"]
        else:
            raise ValueError(f"未知任务: {task}")
        
        print(f"加载 {split} 集, 任务 {task}: {len(self.samples)} 样本")
    
    def default_transform(self):
        """默认图像变换"""
        return transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        sample = self.samples[idx]
        task_type = sample["task"]
        
        if task_type in ["environment", "accident_detection"]:
            # 单帧任务
            image = Image.open(sample["image_path"]).convert("RGB")
            image = self.transform(image)
            
            return {
                "task": task_type,
                "image": image,
                "label": sample["label"],
                "metadata": sample["metadata"]
            }
        
        elif task_type == "sequence_prediction":
            # 序列任务
            images = []
            for img_path in sample["image_sequence"]:
                img = Image.open(img_path).convert("RGB")
                img = self.transform(img)
                images.append(img)
            
            images = torch.stack(images)  # [T, C, H, W]
            
            return {
                "task": task_type,
                "image_sequence": images,
                "label": sample["label"],
                "metadata": sample["metadata"]
            }
        
        else:
            raise ValueError(f"未知任务类型: {task_type}")


def collate_fn(batch):
    """
    自定义collate函数，处理不同任务的batch
    """
    # 按任务分组
    single_frame_batch = []
    sequence_batch = []
    
    for item in batch:
        if item["task"] in ["environment", "accident_detection"]:
            single_frame_batch.append(item)
        elif item["task"] == "sequence_prediction":
            sequence_batch.append(item)
    
    result = {}
    
    # 处理单帧任务
    if single_frame_batch:
        result["single_frame"] = {
            "task": [x["task"] for x in single_frame_batch],
            "images": torch.stack([x["image"] for x in single_frame_batch]),
            "labels": [x["label"] for x in single_frame_batch],
            "metadata": [x["metadata"] for x in single_frame_batch]
        }
    
    # 处理序列任务（需要padding到相同长度）
    if sequence_batch:
        max_len = max(x["image_sequence"].shape[0] for x in sequence_batch)
        
        padded_sequences = []
        masks = []
        
        for item in sequence_batch:
            seq = item["image_sequence"]
            seq_len = seq.shape[0]
            
            # Padding
            if seq_len < max_len:
                padding = torch.zeros(max_len - seq_len, *seq.shape[1:])
                seq = torch.cat([seq, padding], dim=0)
            
            # Mask (1=有效, 0=padding)
            mask = torch.ones(max_len)
            mask[seq_len:] = 0
            
            padded_sequences.append(seq)
            masks.append(mask)
        
        result["sequence"] = {
            "task": [x["task"] for x in sequence_batch],
            "sequences": torch.stack(padded_sequences),  # [B, T, C, H, W]
            "masks": torch.stack(masks),  # [B, T]
            "labels": [x["label"] for x in sequence_batch],
            "metadata": [x["metadata"] for x in sequence_batch]
        }
    
    return result


# ============ 使用示例 ============
if __name__ == "__main__":
    from torch.utils.data import DataLoader
    
    # 数据路径
    data_file = "PROJECT_ROOT/data/dataset/pretrain/train/pretrain_data.pkl"
    
    # 创建数据集
    train_dataset = PretrainDataset(
        data_file=data_file,
        split="train",
        task="all"  # 或 "task1", "task2", "task3"
    )
    
    # 创建DataLoader
    train_loader = DataLoader(
        train_dataset,
        batch_size=8,
        shuffle=True,
        num_workers=4,
        collate_fn=collate_fn
    )
    
    # 测试
    print("\n测试DataLoader:")
    for batch in train_loader:
        print(f"Batch keys: {batch.keys()}")
        
        if "single_frame" in batch:
            sf = batch["single_frame"]
            print(f"  单帧任务: {len(sf['images'])} 样本")
            print(f"    图像shape: {sf['images'].shape}")
            print(f"    标签示例: {sf['labels'][0]}")
        
        if "sequence" in batch:
            seq = batch["sequence"]
            print(f"  序列任务: {len(seq['sequences'])} 样本")
            print(f"    序列shape: {seq['sequences'].shape}")
            print(f"    Mask shape: {seq['masks'].shape}")
            print(f"    标签示例: {seq['labels'][0]}")
        
        break  # 只测试一个batch
    
    print("\n✅ 数据加载器测试通过!")