import torch
import torch.nn as nn
from torch.utils.data import Dataset
import h5py
import numpy as np
import os
import glob
import pickle

class RobotDiffusionDataset(Dataset):
    def __init__(self, 
                 dataset_dir, 
                 pred_horizon=16, 
                 obs_horizon=2, 
                 action_horizon=8,
                 camera_names=['r_vis', 'top', 'front'],
                 normalization_type='gaussian'): 
        """
        Args:
            dataset_dir: 存放 episode_*.hdf5 的文件夹路径
            pred_horizon: 预测未来动作的长度 (Tp)
            obs_horizon: 历史观测长度 (To)
            action_horizon: 执行动作长度 (Ta) - 在Dataset中主要影响Evaluation，这里作为参数保留
        """
        self.dataset_dir = dataset_dir
        self.pred_horizon = pred_horizon
        self.obs_horizon = obs_horizon
        self.action_horizon = action_horizon
        self.camera_names = camera_names
        self.normalization_type = normalization_type
        # 1. 扫描所有HDF5文件并建立索引
        # 格式: [(file_path, episode_length), ...]
        self.episode_files = sorted(glob.glob(os.path.join(dataset_dir, 'episode_*.hdf5')))
        self.indices = []
        
        print(f"Found {len(self.episode_files)} episodes. Building index...")
        
        for file_path in self.episode_files:
            with h5py.File(file_path, 'r') as f:
                # 获取该 episode 的长度 (例如 700)
                l = f['action'].shape[0]
                # 保存每个有效 step 的索引信息
                # (file_path, episode_length, current_step_index)
                for i in range(l):
                    self.indices.append((file_path, l, i))
        
        # 2. 统计数据
        with open(os.path.join(dataset_dir, 'data_stats.pkl'), 'rb') as f:
            self.stats = pickle.load(f)

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

    def __getitem__(self, idx):
        file_path, episode_len, start_ts = self.indices[idx]
        
        # -----------------------------
        # 1. 打开文件
        # -----------------------------
        # 注意: 在 __getitem__ 中打开文件对多进程 DataLoader 更友好
        # 如果追求极致IO性能，可以考虑使用 h5py 的 swmr 模式或内存缓存
        with h5py.File(file_path, 'r') as root:
            
            # -----------------------------
            # 2. 处理 Action (Prediction Target)
            # -----------------------------
            # 目标: 获取 [t, t + pred_horizon] 的动作
            action_start = start_ts
            action_end = min(start_ts + self.pred_horizon, episode_len)
            
            actions = root['action'][action_start:action_end] # shape: (T_subset, 16)
            
            # Padding: 如果剩余动作不足 pred_horizon，复制最后一步
            if len(actions) < self.pred_horizon:
                pad_len = self.pred_horizon - len(actions)
                last_action = actions[-1]
                # 重复最后一行
                pad_content = np.repeat(last_action[np.newaxis, :], pad_len, axis=0)
                actions = np.concatenate([actions, pad_content], axis=0)
            
            # 归一化 Action
            if self.stats:
                actions = self._normalize_data(actions, self.stats['action'])

            # -----------------------------
            # 3. 处理 Observations (History)
            # -----------------------------
            # 目标: 获取 [t - obs_horizon + 1, t + 1] 的观测
            # 索引逻辑:
            # 如果 obs_horizon=2, current_ts=0 -> indices=[0, 0] (Padding)
            # 如果 obs_horizon=2, current_ts=5 -> indices=[4, 5]
            
            start_idx_raw = start_ts - (self.obs_horizon - 1)
            start_idx = max(start_idx_raw, 0)
            end_idx = start_ts + 1
            pad_len = max(0, -start_idx_raw)

            # Qpos
            qpos_data = root['observations/qpos']
            qpos_val = qpos_data[start_idx:end_idx]
            
            if pad_len > 0:
                first_frame = qpos_val[0]
                padding = np.repeat(first_frame[np.newaxis, :], pad_len, axis=0)
                qpos_val = np.concatenate([padding, qpos_val], axis=0)
                
            if self.stats:
                qpos_val = self._normalize_data(qpos_val, self.stats['qpos'])

            # Images
            image_dict = {}
            for cam_name in self.camera_names:
                img_dset = root['observations']['images'][cam_name]
                imgs_np = img_dset[start_idx:end_idx]  # (T, H, W, C)
                
                if pad_len > 0:
                    first_frame = imgs_np[0]
                    padding = np.repeat(first_frame[np.newaxis, ...], pad_len, axis=0)
                    imgs_np = np.concatenate([padding, imgs_np], axis=0)
                
                # 转换为 Tensor: (T, H, W, C) -> (T, C, H, W)
                imgs_tensor = torch.from_numpy(imgs_np).float() / 255.0
                imgs_tensor = torch.einsum('thwc->tchw', imgs_tensor)
                image_dict[cam_name] = imgs_tensor

            # ==============================
            # 3. 组装 Batch
            # ==============================
            data_batch = {
                'action': torch.from_numpy(actions).float(),
                'qpos': torch.from_numpy(qpos_val).float(),
            }
            for cam_name, img_tensor in image_dict.items():
                data_batch[f'image_{cam_name}'] = img_tensor

            return data_batch

    def _normalize_data(self, data, stats):
            if self.normalization_type == 'min_max':
                # 之前的逻辑: [-1, 1]
                min_val = stats['min']
                max_val = stats['max']
                data = (data - min_val) / (max_val - min_val + 1e-8)
                return data * 2 - 1
                
            elif self.normalization_type == 'gaussian':
                # 新逻辑: Mean/Std
                mean = stats['mean']
                std = stats['std']
                # (data - mean) / std
                # 这里的 data 是 numpy array
                return (data - mean) / (std + 1e-8)