data_loader_batch.py

import cv2
import torch
import numpy as np
import scipy.io as scio
import random
from concurrent.futures import ThreadPoolExecutor, as_completed
import os
import h5py

# 预导入deformation函数以提高性能
try:
    from deformation import elastic_deformation, homography_deformation, turbulence_deformation, visualize_all_deformations
    DEFORMATION_AVAILABLE = True
except ImportError:
    DEFORMATION_AVAILABLE = False
    print("Warning: deformation.py not available, deformation will be disabled")


def _choose_deform_degree(degree):
    """
    Map a user-friendly degree to a concrete severity label.
    degree can be: 'mild' | 'medium' | 'strong' | 'random'.
    
    Paper-aligned parameter levels:
    | Level  | Elastic α | Homography MPO | Turbulence D/r0 |
    | ------ | --------- | -------------- | --------------- |
    | Mild   | 200       | 1 pixel        | 1               |
    | Medium | 600       | 2 pixels       | 3               |
    | Strong | 1000      | 4 pixels       | 5               |
    """
    if degree is None:
        return 'random'
    if isinstance(degree, str):
        d = degree.lower().strip()
        if d in ['mild', 'medium', 'strong', 'random']:
            return d
        return 'random'
    # numeric severity -> map to level
    try:
        val = float(degree)
        if val <= 1.0:
            return 'mild'
        elif val <= 3.0:
            return 'medium'
        else:
            return 'strong'
    except Exception:
        return 'random'


def _apply_lr_deformation(lr_hsi, deform_type='mixed', deform_degree='mild', seed=None, dataset=None):
    """
    Apply deformation ONLY on LR-HSI (H, W, C) numpy array.
    
    Uses paper-aligned parameter levels from deformation.py:
    | Level  | Elastic α | Elastic σ | Homography MPO | Turbulence D/r₀ |
    | ------ | --------- | --------- | -------------- | --------------- |
    | Mild   | 300       | 20        | 5%             | 1.0             |
    | Medium | 600       | 15        | 10%            | 2.0             |
    | Strong | 1000      | 15        | 20%            | 3.0             |
    
    Args:
        lr_hsi: Input LR-HSI image (H, W, C)
        deform_type: 'elastic' | 'homography' | 'turbulence' | 'mixed' | 'none'
        deform_degree: 'mild' | 'medium' | 'strong' | 'random'
        seed: Random seed (if None and dataset provided, uses dataset-specific seed)
        dataset: Dataset name for automatic seed selection (e.g., 'PaviaU', 'Botswana')
    
    Notes:
      - This function intentionally does NOT touch ref or HR-MSI.
      - If dataset is provided, uses dataset-specific fixed seed for reproducibility.
    """
    if lr_hsi is None:
        return None

    deform_type = (deform_type or 'mixed').lower().strip()
    degree = _choose_deform_degree(deform_degree)

    # Handle 'none' type
    if deform_type == 'none':
        return lr_hsi

    # Handle random level
    if degree == 'random':
        rng = np.random.RandomState(seed if seed is not None else None)
        degree = rng.choice(['mild', 'medium', 'strong'])

    # Choose deformation type
    if deform_type in ['elastic', 'homography', 'turbulence']:
        chosen = deform_type
    else:
        # mixed / unknown -> random pick
        rng = np.random.RandomState(seed if seed is not None else None)
        chosen = rng.choice(['elastic', 'homography', 'turbulence'])

    # Ensure float for deformation ops
    lr_hsi = lr_hsi.astype(np.float32, copy=False)

    if not DEFORMATION_AVAILABLE:
        return lr_hsi

    # Use the new unified interface from deformation.py with dataset support

    from deformation import apply_deformation
    # If dataset is provided, use it for automatic seed selection
    # Otherwise, use the provided seed
    return apply_deformation(lr_hsi, deform_type=chosen, level=degree, seed=seed, dataset=dataset)



class DatasetBatchSampler:
    """
    数据集batch采样器，用于Flow-based方法，支持批量采样和预采样
    """
    def __init__(self, root='data', dataset=None, size=128, n_select_bands=5, scale_ratio=4,
                 gaussian_kernel=5, sigma=2, batch_size=4, n_batches=1000,
                 # LR deformation controls
                 with_lr_defor=True,
                 lr_deform_type='',       # 'elastic' | 'homography' | 'turbulence' | 'mixed' | 'none'
                 lr_deform_degree='',    # 'mild' | 'medium' | 'strong' | 'random' | float
                 lr_deform_seed=None,
                 save_vis=True,
                 vis_dir='vis'):
        self.root = root
        self.dataset = dataset
        self.size = size
        self.n_select_bands = n_select_bands
        self.scale_ratio = scale_ratio
        self.gaussian_kernel = gaussian_kernel
        self.sigma = sigma
        self.batch_size = batch_size
        self.n_batches = n_batches

        # LR deformation parameters
        self.with_lr_defor = with_lr_defor
        self.lr_deform_type = lr_deform_type
        self.lr_deform_degree = lr_deform_degree
        self.lr_deform_seed = lr_deform_seed
        self.save_vis = save_vis
        self.vis_dir = vis_dir

        # 加载数据
        self._load_data()

        # 准备训练和测试数据
        self._prepare_train_test_data()

        if self.save_vis:
            self._save_initial_visualizations()

        # 预采样训练数据
        self._pre_sample_batches()

        print(f"Dataset {dataset} loaded. Train region: {self.train_region_shape}, Batch size: {batch_size}, Pre-sampled batches: {n_batches}")

    def _load_data(self):
        """加载原始数据"""
        if self.dataset == 'Pavia':
            img = scio.loadmat(self.root + '/' + 'Pavia.mat')['pavia']*1.0
        elif self.dataset == 'PaviaU':
            img = scio.loadmat(self.root + '/' + 'PaviaU.mat')['paviaU']*1.0
        elif self.dataset == 'Washington':
            img = scio.loadmat(self.root + '/' + 'Washington.mat')['Washington_DC']*1.0
        elif self.dataset == 'Houston_HSI':
            img = scio.loadmat(self.root + '/' + 'Houston_HSI.mat')['Houston_HSI']*1.0
        elif self.dataset == 'Salinas_corrected':
            img = scio.loadmat(self.root + '/' + 'Salinas_corrected.mat')['salinas_corrected']*1.0
        elif self.dataset == "Urban":
            img = scio.loadmat(self.root + '/Urban.mat')['Y']*1.0
            img = np.reshape(img, (162, 307, 307))
            img = np.transpose(img,(2,1,0))
        elif self.dataset == "IndianP":
            self.size = 64
            img = scio.loadmat(self.root +'/Indian_pines_corrected.mat')
            img = img['indian_pines_corrected'] * 1.0
            img = img[:144,:144,:]
        elif self.dataset == 'Botswana':
            img = scio.loadmat(self.root +'/Botswana.mat')
            img = img['Botswana'] * 1.0
        elif self.dataset == 'Xuzhou':
            mat = scio.loadmat(self.root + '/' + 'xuzhou.mat')
            img = mat['xuzhou']
        elif self.dataset == 'Chikusei':
            mat = h5py.File(self.root + '/' + 'Chikusei.mat')
            img = np.transpose(mat['chikusei'])
        elif self.dataset == 'Xiongan':
            mat = h5py.File(self.root + '/' + 'xiongan.mat')
            img = np.transpose(mat['XiongAn'])


        print(img.shape)
        # 最大最小值归一化到0-1范围 (适合神经网络)
        max_val = np.max(img)
        min_val = np.min(img)
        img = (img - min_val) / (max_val - min_val + 1e-8)

        # 存储归一化参数，用于后续反归一化
        self.data_min = min_val
        self.data_max = max_val
        self.data_range = max_val - min_val

        self.img = img
        self.n_bands = img.shape[2]

        # 预加载并预处理光谱响应函数 (SRF)
        self._load_spectral_response()

        # 预计算优化参数
        self.kernel_size = min(self.gaussian_kernel, 5)  # 限制高斯核大小
        self.target_size = (self.size//self.scale_ratio, self.size//self.scale_ratio)

    def _load_spectral_response(self):
        """预加载光谱响应函数，避免每次采样时的重复加载"""
        try:
            srf = scio.loadmat('data/srf/gf1_srf.mat')
            R = srf['srf']
            R = np.transpose(R, (1, 0))  # [n_bands, 5]
            R = self.sample_srf(R, self.n_bands)  # [n_bands, 5]
            R = self.normalize(R)  # [n_bands, 5]
            self.srf_matrix = R.astype(np.float32)  # 预转换为float32
            print("SRF matrix loaded and preprocessed.")
        except Exception as e:
            print(f"Warning: Failed to load SRF matrix: {e}")
            # 如果加载失败，使用单位矩阵作为fallback
            self.srf_matrix = np.eye(self.n_bands, 5, dtype=np.float32)

    def _prepare_train_test_data(self):
        """准备训练和测试数据区域
        
        新流程：
        1. HR-HSI (ref) - 原始高分辨率HSI，作为训练目标
        2. 对整个 ref 做 deformation -> ref_deform
        3. ref_deform 高斯模糊+下采样 -> LR-HSI_deform (训练输入)
        4. ref 生成 HR-MSI (不变形，作为条件输入)
        """
        # throwing up the edge
        w_edge = self.img.shape[0]//self.scale_ratio*self.scale_ratio-self.img.shape[0]
        h_edge = self.img.shape[1]//self.scale_ratio*self.scale_ratio-self.img.shape[1]
        w_edge = -1 if w_edge==0 else w_edge
        h_edge = -1 if h_edge==0 else h_edge
        img = self.img[:w_edge, :h_edge, :]

        width, height, n_bands = img.shape
        self.full_width = width
        self.full_height = height

        # 根据数据集设置测试区域
        if self.dataset == 'IndianP':
            w_str, w_end = 0, 64
            h_str, h_end = 0, 64
        elif self.dataset == 'Botswana':
            w_str, w_end = 561, 561 + 128
            h_str, h_end = 75, 75 + 128
        elif self.dataset == 'PaviaU':
            w_str, w_end = 210, 210 + 128
            h_str, h_end = 105, 233
        elif self.dataset == 'Pavia':
            w_str, w_end = 240, 240 + 128
            h_str, h_end = 416, 416 + 128
        else:
            # 默认居中裁剪
            w_str = (width - self.size) // 2
            h_str = (height - self.size) // 2
            w_end = w_str + self.size
            h_end = h_str + self.size

        # 测试数据 (不做deformation)
        self.test_ref = img[w_str:w_end, h_str:h_end, :].copy()
        test_lr = cv2.GaussianBlur(self.test_ref, (self.gaussian_kernel, self.gaussian_kernel), self.sigma)
        test_lr = cv2.resize(test_lr, (self.size//self.scale_ratio, self.size//self.scale_ratio))
        self.test_lr = test_lr

        # 生成测试MSI (使用预加载的光谱响应函数)
        self.test_hr = np.matmul(self.test_ref.reshape(self.test_ref.shape[0] * self.test_ref.shape[1], -1),
                               self.srf_matrix).reshape(self.test_ref.shape[0], self.test_ref.shape[1], -1)

        # 测试区域mask：标记测试区域像素（用于训练时屏蔽）
        test_mask = np.zeros_like(img[:, :, 0])
        test_mask[w_str:w_end, h_str:h_end] = 1  # 测试区域标记为1
        self.test_mask = test_mask

        # ========== 训练数据准备 ==========
        # 数据流程 (重要！):
        #   - ref (HR-HSI): 训练目标，不做畸变
        #   - MSI (HR-MSI): 从原始 ref 生成，不做畸变
        #   - LR-HSI: 先对 ref 做畸变，再做高斯模糊+下采样
        # 
        # 这样训练时：
        #   输入: 畸变的 LR-HSI + 未畸变的 MSI
        #   目标: 未畸变的 ref (HR-HSI)
        #   任务: 学习从畸变输入恢复清晰的目标
        
        # 1. train_img_ref: HR-HSI 原始图 (训练目标) - 测试区域mask为0
        #    【重要】这是训练目标，不做畸变！
        self.train_img_ref = (img * (1 - test_mask)[:, :, np.newaxis]).astype(np.float32)
        self.train_region_shape = self.train_img_ref.shape

        # 2. 生成 HR-MSI (整图) - 从原始 ref 生成，不做 deformation
        #    【重要】MSI 也不做畸变，作为条件输入
        train_msi_flat = np.dot(self.train_img_ref.reshape(-1, n_bands), self.srf_matrix)
        self.train_img_msi = train_msi_flat.reshape(width, height, -1).astype(np.float32)
        
        # 3. 对整图做 deformation，生成 ref_deform
        #    【重要】畸变只用于生成 LR-HSI，不影响 ref 和 MSI
        if self.with_lr_defor:
            print(f"Applying {self.lr_deform_type} deformation (degree: {self.lr_deform_degree}) to generate LR-HSI...")
            print(f"  Using dataset-specific seed for: {self.dataset}")
            self.train_img_deform = _apply_lr_deformation(
                self.train_img_ref.copy(),  # 使用 copy()，不影响原始 ref
                deform_type=self.lr_deform_type,
                deform_degree=self.lr_deform_degree,
                seed=self.lr_deform_seed,
                dataset=self.dataset  # 使用数据集名称自动选择种子
            )
            if self.train_img_deform is None:
                return 
        else:
            self.train_img_deform = self.train_img_ref.copy()
        
        # 4. 从 deformed ref 生成 LR-HSI (整图高斯模糊+下采样)
        #    【重要】只有 LR-HSI 是畸变的！
        lr_width = width // self.scale_ratio
        lr_height = height // self.scale_ratio
        self.train_img_lr = cv2.GaussianBlur(
            self.train_img_deform,  # 从畸变图生成
            (self.kernel_size, self.kernel_size), 
            self.sigma
        )
        self.train_img_lr = cv2.resize(
            self.train_img_lr, 
            (lr_height, lr_width),  # cv2.resize expects (width, height)
            interpolation=cv2.INTER_LINEAR
        ).astype(np.float32)
        
        print(f"Train data shapes: ref={self.train_img_ref.shape} (clean), lr={self.train_img_lr.shape} (deformed), msi={self.train_img_msi.shape} (clean)")

        # 计算可以采样的位置：覆盖整个图像的所有可能patch
        # 由于测试区域像素已被mask为0，不会泄露数据
        self.valid_positions = []
        for i in range(0, width - self.size + 1, self.size // 2):  # i是width方向
            for j in range(0, height - self.size + 1, self.size // 2):  # j是height方向
                self.valid_positions.append((i, j))

        print(f"Valid training positions: {len(self.valid_positions)}")

    def _save_initial_visualizations(self):
        try:
            os.makedirs(os.path.join(self.vis_dir, str(self.dataset)), exist_ok=True)
            orig = self.test_ref.astype(np.float32, copy=False)
            seed = self.lr_deform_seed
            if DEFORMATION_AVAILABLE:
                # 使用 apply_deformation 接口，与训练数据使用相同的参数
                from deformation import apply_deformation, get_deformation_params
                
                # 获取当前配置的 deformation 参数
                params = get_deformation_params(self.lr_deform_degree)
                
                # 使用统一接口生成三种类型的变形图像
                elastic_img = apply_deformation(orig.copy(), deform_type='elastic', 
                                                level=self.lr_deform_degree, seed=seed, dataset=self.dataset)
                homography_img = apply_deformation(orig.copy(), deform_type='homography', 
                                                   level=self.lr_deform_degree, seed=seed, dataset=self.dataset)
                turbulence_img = apply_deformation(orig.copy(), deform_type='turbulence', 
                                                   level=self.lr_deform_degree, seed=seed, dataset=self.dataset)
                
                # 保存时在文件名中包含畸变程度信息
                save_path = os.path.join(self.vis_dir, str(self.dataset), 
                                         f'initial_deformations_{self.lr_deform_degree}.png')
                visualize_all_deformations(orig, elastic_img, homography_img, turbulence_img, 
                                           dataset=self.dataset, save_path=save_path)
                print(f"  Saved initial deformations visualization: {save_path}")
                print(f"    Using level '{self.lr_deform_degree}': α={params['elastic_alpha']}, "
                      f"MPO={params['homography_mpo']}%, D/r₀={params['turbulence_dr0']}")

            def to_uint8(img):
                x = img
                if x.dtype != np.uint8:
                    x = (x * 255.0).clip(0, 255).astype(np.uint8)
                return x

            try:
                from deformation import get_rgb_bands, extract_rgb_from_hsi
                h, w, c = orig.shape
                r, g, b = get_rgb_bands(self.dataset, c)
                ref_rgb = extract_rgb_from_hsi(orig, r, g, b)
                ref_rgb = to_uint8(ref_rgb)
                ref_path = os.path.join(self.vis_dir, str(self.dataset), 'test_ref_rgb.png')
                cv2.imwrite(ref_path, cv2.cvtColor(ref_rgb, cv2.COLOR_RGB2BGR))

                lr = self.test_lr.astype(np.float32, copy=False)
                r2, g2, b2 = get_rgb_bands(self.dataset, lr.shape[2])
                lr_rgb = extract_rgb_from_hsi(lr, r2, g2, b2)
                lr_rgb = to_uint8(lr_rgb)
                lr_path = os.path.join(self.vis_dir, str(self.dataset), 'test_lr_rgb.png')
                cv2.imwrite(lr_path, cv2.cvtColor(lr_rgb, cv2.COLOR_RGB2BGR))

                hr = self.test_hr.astype(np.float32, copy=False)
                r3, g3, b3 = 0, min(1, hr.shape[2]-1), min(2, hr.shape[2]-1)
                hr_rgb = np.stack([hr[:, :, r3], hr[:, :, g3], hr[:, :, b3]], axis=2)
                hr_rgb = to_uint8(hr_rgb)
                hr_path = os.path.join(self.vis_dir, str(self.dataset), 'test_hr_msi_rgb.png')
                cv2.imwrite(hr_path, cv2.cvtColor(hr_rgb, cv2.COLOR_RGB2BGR))
            except Exception:
                pass
        except Exception as e:
            print(f"Warning: Failed to save initial visualizations: {e}")

    def _sample_single_batch(self, args):
        """采样单个batch - 从预计算的整图中采样patch
        
        返回: (LR-HSI_deform, HR-HSI_ref, HR-MSI)
        - LR-HSI_deform: 从 deformed 整图下采样后裁剪的 patch
        - HR-HSI_ref: 原始 ref 的 patch (训练目标)
        - HR-MSI: 从原始 ref 生成的 MSI patch (条件输入)
        """
        batch_idx, batch_positions = args
        
        lr_size = self.size // self.scale_ratio

        # 预分配numpy数组以提高性能
        batch_hsi_hr = np.empty((self.batch_size, self.size, self.size, self.n_bands), dtype=np.float32)
        batch_hsi_lr = np.empty((self.batch_size, lr_size, lr_size, self.n_bands), dtype=np.float32)
        batch_msi_hr = np.empty((self.batch_size, self.size, self.size, 5), dtype=np.float32)

        for sample_idx, (i, j) in enumerate(batch_positions):
            # 从预计算的整图中采样对应位置的patch
            # i, j 是 HR 图像的坐标
            
            # 1. HR-HSI ref (训练目标)
            patch_hr = self.train_img_ref[i:i+self.size, j:j+self.size, :]
            
            # 2. LR-HSI (从整图LR采样对应位置)
            # LR图像坐标需要按 scale_ratio 缩放
            lr_i = i // self.scale_ratio
            lr_j = j // self.scale_ratio
            patch_lr = self.train_img_lr[lr_i:lr_i+lr_size, lr_j:lr_j+lr_size, :]
            
            # 3. HR-MSI (条件输入)
            patch_msi = self.train_img_msi[i:i+self.size, j:j+self.size, :]

            batch_hsi_hr[sample_idx] = patch_hr
            batch_hsi_lr[sample_idx] = patch_lr
            batch_msi_hr[sample_idx] = patch_msi

        # 转换为tensor - 直接pin_memory
        batch_hsi_hr_tensor = torch.from_numpy(batch_hsi_hr).permute(0, 3, 1, 2).pin_memory()
        batch_hsi_lr_tensor = torch.from_numpy(batch_hsi_lr).permute(0, 3, 1, 2).pin_memory()
        batch_msi_hr_tensor = torch.from_numpy(batch_msi_hr).permute(0, 3, 1, 2).pin_memory()

        return batch_idx, (batch_hsi_lr_tensor, batch_hsi_hr_tensor, batch_msi_hr_tensor)

    def _pre_sample_batches(self):
        """预采样所有训练batch - 多线程超高性能版本"""
        print(f"Pre-sampling {self.n_batches} batches with batch_size={self.batch_size} using threads...")
        import time
        start_time = time.time()

        # 预先选择所有位置，避免重复随机选择
        positions_array = np.array(self.valid_positions)
        indices = np.random.randint(0, len(positions_array), self.n_batches * self.batch_size)
        all_positions = positions_array[indices].tolist()

        # 准备batch参数
        batch_args = []
        for batch_idx in range(self.n_batches):
            batch_positions = all_positions[batch_idx*self.batch_size:(batch_idx+1)*self.batch_size]
            batch_args.append((batch_idx, batch_positions))

        self.pre_sampled_batches = [None] * self.n_batches
        failed_batches = []

        # 使用40个线程并行采样
        with ThreadPoolExecutor(max_workers=100) as executor:
            # 提交所有任务
            future_to_batch = {executor.submit(self._sample_single_batch, args): args[0]
                             for args in batch_args}

            # 收集结果
            completed = 0
            for future in as_completed(future_to_batch):
                batch_idx = future_to_batch[future]
                try:
                    idx, batch_data = future.result()
                    self.pre_sampled_batches[idx] = batch_data
                    completed += 1

                    if completed % 100 == 0:
                        elapsed = time.time() - start_time
                        print(f"  Completed {completed}/{self.n_batches} batches ({elapsed:.1f}s)")

                except Exception as exc:
                    print(f'Batch {batch_idx} generated an exception: {exc}')
                    failed_batches.append(batch_idx)

        # 重试失败的batch（最多3次）
        max_retries = 3
        for retry in range(max_retries):
            if not failed_batches:
                break
            
            print(f"\nRetrying {len(failed_batches)} failed batches (attempt {retry+1}/{max_retries})...")
            retry_args = []
            for batch_idx in failed_batches:
                # 使用新的随机位置重试
                new_indices = np.random.randint(0, len(positions_array), self.batch_size)
                new_positions = positions_array[new_indices].tolist()
                retry_args.append((batch_idx, new_positions))
            
            new_failed = []
            with ThreadPoolExecutor(max_workers=100) as executor:
                future_to_batch = {executor.submit(self._sample_single_batch, args): args[0]
                                 for args in retry_args}
                
                for future in as_completed(future_to_batch):
                    batch_idx = future_to_batch[future]
                    try:
                        idx, batch_data = future.result()
                        self.pre_sampled_batches[idx] = batch_data
                        completed += 1
                    except Exception as exc:
                        print(f'Batch {batch_idx} retry failed: {exc}')
                        new_failed.append(batch_idx)
            
            failed_batches = new_failed

        # 如果仍有失败的batch，用成功的batch替换
        if failed_batches:
            print(f"\nWarning: {len(failed_batches)} batches still failed after retries. Replacing with duplicates...")
            # 找到所有成功的batch
            successful_indices = [i for i in range(self.n_batches) if self.pre_sampled_batches[i] is not None]
            if successful_indices:
                for batch_idx in failed_batches:
                    # 随机选择一个成功的batch进行复制
                    replacement_idx = np.random.choice(successful_indices)
                    self.pre_sampled_batches[batch_idx] = self.pre_sampled_batches[replacement_idx]
            else:
                raise RuntimeError("All batches failed to sample! Please check your data and deformation settings.")

        total_time = time.time() - start_time
        print(f"Pre-sampling completed. Total batches: {len(self.pre_sampled_batches)}")
        print(f"Total time: {total_time:.2f}s, Average: {total_time/self.n_batches:.3f}s per batch")
        print(f"Sampling speed: {self.n_batches/total_time:.1f} batches/second")
        print(f"Memory usage: ~{len(self.pre_sampled_batches) * self.batch_size * 3 * 4 * self.size * self.size / 1024 / 1024:.1f} MB")

    def __len__(self):
        """返回预采样batch的数量"""
        return len(self.pre_sampled_batches)

    def __getitem__(self, idx):
        """获取指定索引的batch"""
        return self.pre_sampled_batches[idx]

    def shuffle_batches(self):
        """固定batch顺序 - 每个epoch使用相同的1000次采样"""
        # 不进行shuffle，保持固定的采样顺序
        pass

    def sample_batch(self, idx=None):
        """采样一个batch的数据 - 如果idx为None则随机选择，否则返回指定索引的batch"""
        if idx is None:
            idx = random.randint(0, len(self.pre_sampled_batches) - 1)
        return self.pre_sampled_batches[idx]

    def get_test_data(self):
        """获取测试数据"""
        test_ref = torch.from_numpy(self.test_ref).permute(2, 0, 1).unsqueeze(0)  # [1, C, H, W]
        test_lr = torch.from_numpy(self.test_lr).permute(2, 0, 1).unsqueeze(0)    # [1, C, H_lr, W_lr]
        test_hr = torch.from_numpy(self.test_hr).permute(2, 0, 1).unsqueeze(0)    # [1, C_msi, H, W]

        return test_ref, test_lr, test_hr

    @staticmethod
    def sample_srf(srf, bands):
        """从光谱响应度函数中等间隔采样指定数量的波段"""
        total_bands = srf.shape[0]
        indices = np.linspace(0, total_bands - 1, bands, dtype=int)
        sampled_srf = srf[indices, :]
        return sampled_srf

    @staticmethod
    def normalize(matrix):
        """Normalize each column of the matrix so that the sum of each column is 1."""
        column_sums = matrix.sum(axis=0, keepdims=True)
        column_sums[column_sums == 0] = 1
        normalized_matrix = matrix / column_sums
        return normalized_matrix


def build_datasets(root='data', dataset=None, size=128, n_select_bands=5, scale_ratio=4, gaussian_kernel=5, divation=2, batch_size=8, n_batches=1000,
                   # LR deformation controls
                   with_lr_defor=True,
                   lr_deform_type='',       # 'elastic' | 'homography' | 'turbulence' | 'mixed' | 'none'
                   lr_deform_degree='1',      # 'mild' | 'medium' | 'strong' | 'random' | float (mild for speed)
                   lr_deform_seed=None,
                   save_vis=True,
                   vis_dir='vis'):
    """
    创建数据集采样器，支持batch采样和LR deformation
    返回: sampler (DatasetBatchSampler对象)
    """
    sampler = DatasetBatchSampler(
        root=root, dataset=dataset, size=size, n_select_bands=n_select_bands,
        scale_ratio=scale_ratio, gaussian_kernel=gaussian_kernel, sigma=divation,
        batch_size=batch_size, n_batches=n_batches,
        with_lr_defor=with_lr_defor,
        lr_deform_type=lr_deform_type,
        lr_deform_degree=lr_deform_degree,
        lr_deform_seed=lr_deform_seed,
        save_vis=save_vis,
        vis_dir=vis_dir
    )
    return sampler