# -*- coding: utf-8 -*-
'''
@Author: Yehui
@██╗   ██╗███████╗██╗  ██╗██╗   ██╗██╗
@╚██╗ ██╔╝██╔════╝██║  ██║██║   ██║██║
@ ╚████╔╝ █████╗  ███████║██║   ██║██║
@  ╚██╔╝  ██╔══╝  ██╔══██║██║   ██║██║
@   ██║   ███████╗██║  ██║╚██████╔╝██║
@   ╚═╝   ╚══════╝╚═╝  ╚═╝ ╚═════╝ ╚═╝
@Date: 2025-06-05 15:08:08
@LastEditTime: 2025-09-26 15:26:19
@FilePath: /Global_Module/mlgb/datasets.py
@Copyright (c) 2025 by , All Rights Reserved.
'''
from math import pi
from turtle import forward
from unittest import TestLoader
import numpy as np
import os
import scipy.io as sio
from torch.utils.data import Dataset
import torch
from mlgb.utils import create_hrms_lrhs, create_spec_resp, gauss_kernel
from mlgb.config import Args
from torch.utils.data import DataLoader
import glob
import random
import matplotlib
import matplotlib.pyplot as plt
import time


matplotlib.use('Agg')   # 无头后端
args = Args()

class CaveDataset(Dataset): 
    def __init__(self, data_path, gen_path, patch_size=128, stride=6, ratio=8, sigma = 2, type='train'):
        super(CaveDataset, self).__init__()
        self.data_path = data_path
        self.stride = stride
        self.rows = 512 
        self.cols = 512
        self.patch_size = patch_size            # HR patch size
        self.lr_patch = patch_size // ratio
        self.kerSize = ratio        
        self.sigma = sigma
        self.stride = stride                    # stride for HR-MSI patch extraction
        self.ratio = ratio
        self.type = type   
        self.B = gauss_kernel(self.kerSize, self.kerSize, self.sigma)
        self.R = create_spec_resp('CAVE', gen_path)

        # Generate samples and labels
        if self.type=='train':
            self.hsi_data, self.msi_data, self.label = self.generateTrain(self.patch_size, self.ratio, num_star=1, num_end=28, s=1,auto_patch = True)
        if self.type=='eval':
            self.hsi_data, self.msi_data, self.label = self.generateEval(self.patch_size, self.ratio, num_star=23, num_end=23, s=1,auto_patch = True)
        if self.type=='test':
            self.hsi_data, self.msi_data, self.label = self.generateTest(self.patch_size, self.ratio, num_star=28, num_end=33, s=1, auto_patch = True)

    def generateTrain(self, patch_size,ratio, num_star, num_end, s ,auto_patch=True):
        lr_patch = self.lr_patch
        num_patch = ((self.rows - patch_size) // self.stride + 1) * ((self.cols - patch_size) // self.stride + 1)
        print('-------------------------------------------------')
        print('One Image patch number:', num_patch)
        if auto_patch:
            s = int(np.sqrt(num_patch))
        label_patch = np.zeros((s * s * (num_end - num_star), patch_size , patch_size , 31), dtype=np.float32)
        hrmsi_patch = np.zeros((s * s * (num_end - num_star), patch_size , patch_size , 3), dtype=np.float32)
        lrhsi_patch = np.zeros((s * s * (num_end - num_star), patch_size//self.ratio, patch_size//self.ratio, 31), dtype=np.float32)
        
        # label_patch = []
        # hrmsi_patch = []
        # lrhsi_patch = []
        count = 0
        for i in range(num_star, num_end):
            hrhsi = sio.loadmat(self.data_path + '%d.mat' % i)['HS']
            lrhsi = generate_low_HSI(hrhsi, self.B, self.ratio)
            hrmsi = generate_MSI(hrhsi, self.R)

            # Data type conversion
            if hrhsi.dtype != np.float32: hrhsi = hrhsi.astype(np.float32)
            if lrhsi.dtype != np.float32: lrhsi = lrhsi.astype(np.float32)
            if hrmsi.dtype != np.float32: hrmsi = hrmsi.astype(np.float32)

            for x in range(0, self.rows - patch_size + 1, self.stride):
                for y in range(0, self.cols - patch_size + 1, self.stride):
                    lx,ly = x // self.ratio, y // self.ratio
                    label_patch[count] = hrhsi[x :(x + patch_size) , y :(y + patch_size) , :]
                    hrmsi_patch[count] = hrmsi[x :(x + patch_size) , y :(y + patch_size) , :]
                    lrhsi_patch[count] = lrhsi[lx :(lx + lr_patch), ly:(ly + lr_patch), :]
                    count += 1

        return lrhsi_patch, hrmsi_patch, label_patch

    def generateEval(self, patch_size, ratio, num_star, num_end, s,auto_patch=True):
        lr_patch = self.lr_patch
        num_patch = ((self.rows - patch_size) // self.stride + 1) * ((self.cols - patch_size) // self.stride + 1)
        if auto_patch:
            s = int(np.sqrt(num_patch))
        label_patch = np.zeros((s * s * (num_end - num_star), patch_size, patch_size, 31), dtype=np.float32)
        hrmsi_patch = np.zeros((s * s * (num_end - num_star), patch_size, patch_size, 3), dtype=np.float32)
        lrhsi_patch = np.zeros((s * s * (num_end - num_star), lr_patch, lr_patch, 31), dtype=np.float32)
        count = 0
        for i in range(num_star, num_end):
            hrhsi = sio.loadmat(self.data_path + '%d.mat' % i)['HS']
            lrhsi = generate_low_HSI(hrhsi, self.B, self.ratio)
            hrmsi = generate_MSI(hrhsi, self.R)

            # Data type conversion
            if hrhsi.dtype != np.float32: hrhsi = hrhsi.astype(np.float32)
            if lrhsi.dtype != np.float32: lrhsi = lrhsi.astype(np.float32)
            if hrmsi.dtype != np.float32: hrmsi = hrmsi.astype(np.float32)

            for x in range(0, self.rows - patch_size + 1, self.stride):
                for y in range(0, self.cols - patch_size + 1, self.stride):
                    lx,ly = x // self.ratio, y // self.ratio
                    label_patch[count] = hrhsi[x :(x + patch_size) , y :(y + patch_size) , :]
                    hrmsi_patch[count] = hrmsi[x :(x + patch_size) , y :(y + patch_size) , :]
                    lrhsi_patch[count] = lrhsi[lx :(lx + lr_patch), ly:(ly + lr_patch), :]
                    count += 1

        return lrhsi_patch, hrmsi_patch, label_patch

    def generateTest(self, patch_size, ratio, num_star, num_end, s,auto_patch=False):
        lr_patch = self.lr_patch
        num_patch = ((self.rows - patch_size) // self.stride + 1) * ((self.cols - patch_size) // self.stride + 1)
        if auto_patch:
            s = int(np.sqrt(num_patch))
        label_patch = np.zeros((s * s * (num_end - num_star), patch_size, patch_size, 31), dtype=np.float32)
        hrmsi_patch = np.zeros((s * s * (num_end - num_star), patch_size, patch_size, 3), dtype=np.float32)
        lrhsi_patch = np.zeros((s * s * (num_end - num_star), lr_patch, lr_patch, 31), dtype=np.float32)
        count = 0
        for i in range(num_star, num_end):
            hrhsi = sio.loadmat(self.data_path + '%d.mat' % i)['HS']
            lrhsi = generate_low_HSI(hrhsi, self.B, self.ratio)
            hrmsi = generate_MSI(hrhsi, self.R)

            # Data type conversion
            if hrhsi.dtype != np.float32: hrhsi = hrhsi.astype(np.float32)
            if lrhsi.dtype != np.float32: lrhsi = lrhsi.astype(np.float32)
            if hrmsi.dtype != np.float32: hrmsi = hrmsi.astype(np.float32)

            for x in range(0, self.rows - patch_size + 1, self.stride):
                for y in range(0, self.cols - patch_size + 1, self.stride):
                    # HR -> LR 对齐坐标
                    lx, ly = x // self.ratio, y // self.ratio

                    # 边界保护（避免 LR 切片越界变空）
                    if (lx + lr_patch) > lrhsi.shape[0] or (ly + lr_patch) > lrhsi.shape[1]:
                        continue # 跳过该 patch
                    # rotTimes = random.randint(0, 3)
                    # vFlip = random.randint(0, 1)
                    # hFlip = random.randint(0, 1)
                    # label_patch[count] = self.arguement(hrhsi[x * ratio:(x + patch_size) * ratio, y * ratio:(y + patch_size) * ratio, :], rotTimes, vFlip, hFlip)
                    # hrmsi_patch[count] = self.arguement(hrmsi[x * ratio:(x + patch_size) * ratio, y * ratio:(y + patch_size) * ratio, :], rotTimes, vFlip, hFlip)
                    # lrhsi_patch[count] = self.arguement(lrhsi[x:x + patch_size, y:y + patch_size, :], rotTimes, vFlip, hFlip)
                    label_patch[count] = hrhsi[x :(x + patch_size) , y :(y + patch_size) , :]
                    hrmsi_patch[count] = hrmsi[x :(x + patch_size) , y :(y + patch_size) , :]
                    lrhsi_patch[count] = lrhsi[lx :(lx + lr_patch), ly:(ly + lr_patch), :]
                    count += 1
        return lrhsi_patch, hrmsi_patch, label_patch

    def arguement(self, img, rotTimes, vFlip, hFlip):
        # Random rotation
        for j in range(rotTimes):
            img = np.rot90(img.copy(), axes=(0, 1))
        # Random vertical Flip
        for j in range(vFlip):
            img = img[::-1, :, :].copy()
        # Random horizontal Flip
        for j in range(hFlip):
            img = img[:, ::-1, :].copy()
        return img

    def __getitem__(self, index):
        hrhsi = np.transpose(self.label[index], (2,0,1))
        hrmsi = np.transpose(self.msi_data[index], (2,0,1))
        lrhsi = np.transpose(self.hsi_data[index], (2,0,1))

        hrhsi = torch.tensor(hrhsi, dtype=torch.float32)
        hrmsi = torch.tensor(hrmsi, dtype=torch.float32)
        lrhsi = torch.tensor(lrhsi, dtype=torch.float32)

        return lrhsi, hrmsi, hrhsi

    def __len__(self):
        return self.label.shape[0]


# 由于PaviaU数据集是单图,数据集划分较为多样
# 1.我们将其打成patch之后，再进行训练集和测试集的划分
# 2.直接随机切指定尺寸的patch作为测试集，剩余作为训练集
# 补充说明，可以加1个bool变量来控制划分方式
class PaviaUDataset(Dataset):
    def __init__(self, data_path, gen_path, patch_size=64, stride=24, ratio=8, sigma=2, split_ratio=(0.8, 0.1, 0.1), type='train'):
        super().__init__()
        self.data_path = data_path
        self.stride = stride
        self.ratio = ratio
        self.sigma = sigma
        self.patch_size = patch_size
        self.lr_patch = patch_size // ratio
        self.kerSize = ratio
        self.type = type
        self.B = gauss_kernel(self.kerSize, self.kerSize, sigma)
        self.R = create_spec_resp('PaviaU', gen_path)

        # 读取单幅高光谱图像
        hrhsi = sio.loadmat(data_path)['HS'].astype(np.float32)
        hrhsi = norm01(hrhsi)
        lrhsi = generate_low_HSI(hrhsi, self.B, self.ratio)
        hrmsi = generate_MSI(hrhsi, self.R)

        # 生成patch（以 HR 为基准）
        hr_size = patch_size
        lr_size = self.lr_patch
        stride = self.stride
        H, W, _ = hrhsi.shape
        hrhs_list, hrms_list, lrhs_list = [], [], []
        n_rows = (H - hr_size) // stride + 1
        n_cols = (W - hr_size) // stride + 1
        for i in range(n_rows):
            for j in range(n_cols):
                hr_row = i * stride
                hr_col = j * stride
                lr_row = hr_row // ratio
                lr_col = hr_col // ratio
                if (lr_row + lr_size) > lrhsi.shape[0] or (lr_col + lr_size) > lrhsi.shape[1]:
                    continue
                hrhs_patch = hrhsi[hr_row: hr_row + hr_size, hr_col: hr_col + hr_size, :]
                hrms_patch = hrmsi[hr_row: hr_row + hr_size, hr_col: hr_col + hr_size, :]
                lrhs_patch = lrhsi[lr_row: lr_row + lr_size, lr_col: lr_col + lr_size, :]
                hrhs_list.append(hrhs_patch)
                hrms_list.append(hrms_patch)
                lrhs_list.append(lrhs_patch)

        # 划分train/val/test
        (train_hrhs, train_hrms, train_lrhs), (val_hrhs, val_hrms, val_lrhs), (test_hrhs, test_hrms, test_lrhs) = self.split_patches(
            hrhs_list, hrms_list, lrhs_list, split_ratio=split_ratio
        )

        if self.type == 'train':
            self.hrhs_list, self.hrms_list, self.lrhs_list = train_hrhs, train_hrms, train_lrhs
        elif self.type == 'eval' or self.type == 'val':
            self.hrhs_list, self.hrms_list, self.lrhs_list = val_hrhs, val_hrms, val_lrhs
        elif self.type == 'test':
            self.hrhs_list, self.hrms_list, self.lrhs_list = test_hrhs, test_hrms, test_lrhs
        else:
            raise ValueError(f"Unknown type: {self.type}")

    def split_patches(self, hrhs_list, hrms_list, lrhs_list, split_ratio=(0.8, 0.1, 0.1)):
        total = len(hrhs_list)
        indices = np.arange(total)
        np.random.seed(args.seed)
        np.random.shuffle(indices)
        n_train = int(total * split_ratio[0])
        n_val = int(total * split_ratio[1])
        train_idx = indices[:n_train]
        val_idx = indices[n_train:n_train + n_val]
        test_idx = indices[n_train + n_val:]
        train = ([hrhs_list[i] for i in train_idx],
                 [hrms_list[i] for i in train_idx],
                 [lrhs_list[i] for i in train_idx])
        val = ([hrhs_list[i] for i in val_idx],
               [hrms_list[i] for i in val_idx],
               [lrhs_list[i] for i in val_idx])
        test = ([hrhs_list[i] for i in test_idx],
                [hrms_list[i] for i in test_idx],
                [lrhs_list[i] for i in test_idx])
        return train, val, test

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

    def __getitem__(self, index):
        hrhs = torch.from_numpy(np.ascontiguousarray(self.hrhs_list[index].astype(np.float32).transpose(2, 0, 1)))
        hrms = torch.from_numpy(np.ascontiguousarray(self.hrms_list[index].astype(np.float32).transpose(2, 0, 1)))
        lrhs = torch.from_numpy(np.ascontiguousarray(self.lrhs_list[index].astype(np.float32).transpose(2, 0, 1)))
        return  lrhs, hrms, hrhs



class PaviaCDataset(Dataset):
    def __init__(self, data_path, gen_path, patch_size=64, stride=32, ratio=8, sigma=2, split_ratio=(0.8, 0.1, 0.1), type='train', seed=None):
        super().__init__()
        self.data_path = data_path
        self.stride = stride
        self.ratio = ratio
        self.sigma = sigma
        self.patch_size = patch_size
        self.lr_patch = patch_size // ratio
        self.kerSize = ratio
        self.type = type
        self.seed = args.seed if seed is None else int(seed)
        self.B = gauss_kernel(self.kerSize, self.kerSize, sigma)
        self.R = create_spec_resp('PaviaC', gen_path)

        hrhsi = sio.loadmat(data_path)['pavia'].astype(np.float32)
        hrhsi = norm01(hrhsi)
        #hrmsi, lrhsi = create_hrms_lrhs(hrhsi, self.B, self.R, self.ratio)
        lrhsi = generate_low_HSI(hrhsi, self.B, self.ratio)
        hrmsi = generate_MSI(hrhsi, self.R)
        hr_size = patch_size
        lr_size = self.lr_patch
        stride = self.stride
        H, W, _ = hrhsi.shape
        hrhs_list, hrms_list, lrhs_list = [], [], []
        n_rows = (H - hr_size) // stride + 1
        n_cols = (W - hr_size) // stride + 1
        for i in range(n_rows):
            for j in range(n_cols):
                hr_row = i * stride
                hr_col = j * stride
                lr_row = hr_row // ratio
                lr_col = hr_col // ratio
                if (lr_row + lr_size) > lrhsi.shape[0] or (lr_col + lr_size) > lrhsi.shape[1]:
                    continue
                hrhs_patch = hrhsi[hr_row: hr_row + hr_size, hr_col: hr_col + hr_size, :]
                hrms_patch = hrmsi[hr_row: hr_row + hr_size, hr_col: hr_col + hr_size, :]
                lrhs_patch = lrhsi[lr_row: lr_row + lr_size, lr_col: lr_col + lr_size, :]
                hrhs_list.append(hrhs_patch)
                hrms_list.append(hrms_patch)
                lrhs_list.append(lrhs_patch)

        (train_hrhs, train_hrms, train_lrhs), (val_hrhs, val_hrms, val_lrhs), (test_hrhs, test_hrms, test_lrhs) = self.split_patches(
            hrhs_list, hrms_list, lrhs_list, split_ratio=split_ratio
        )

        if self.type == 'train':
            self.hrhs_list, self.hrms_list, self.lrhs_list = train_hrhs, train_hrms, train_lrhs
        elif self.type == 'eval' or self.type == 'val':
            self.hrhs_list, self.hrms_list, self.lrhs_list = val_hrhs, val_hrms, val_lrhs
        elif self.type == 'test':
            self.hrhs_list, self.hrms_list, self.lrhs_list = test_hrhs, test_hrms, test_lrhs
        else:
            raise ValueError(f"Unknown type: {self.type}")

    def split_patches(self, hrhs_list, hrms_list, lrhs_list, split_ratio=(0.8, 0.1, 0.1)):
        total = len(hrhs_list)
        indices = np.arange(total)
        np.random.seed(self.seed)
        np.random.shuffle(indices)
        n_train = int(total * split_ratio[0])
        n_val = int(total * split_ratio[1])
        train_idx = indices[:n_train]
        val_idx = indices[n_train:n_train + n_val]
        test_idx = indices[n_train + n_val:]
        train = ([hrhs_list[i] for i in train_idx],
                 [hrms_list[i] for i in train_idx],
                 [lrhs_list[i] for i in train_idx])
        val = ([hrhs_list[i] for i in val_idx],
               [hrms_list[i] for i in val_idx],
               [lrhs_list[i] for i in val_idx])
        test = ([hrhs_list[i] for i in test_idx],
                [hrms_list[i] for i in test_idx],
                [lrhs_list[i] for i in test_idx])
        return train, val, test

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

    def __getitem__(self, index):
        hrhs = torch.from_numpy(np.ascontiguousarray(self.hrhs_list[index].astype(np.float32).transpose(2, 0, 1)))
        hrms = torch.from_numpy(np.ascontiguousarray(self.hrms_list[index].astype(np.float32).transpose(2, 0, 1)))
        lrhs = torch.from_numpy(np.ascontiguousarray(self.lrhs_list[index].astype(np.float32).transpose(2, 0, 1)))
        return lrhs, hrms, hrhs

class PaviaUDataset_Region(Dataset):
    """
    无监督区域随机采样数据集（单幅 PaviaU 图像）
    - 输入：data_path, gen_path, patch_size, ratio, kerSize, sigma
    - 不划分 train/val/test；通过随机坐标采样 patch，并处理边界
    - 返回与 PaviaUDataset 相同：(lrhs[C,h,w], hrms[C,H,W], hrhs[C,H,W])
    """
    def __init__(self, data_path, gen_path, patch_size=8, ratio=8, kerSize=8, sigma=2,
                 num_patches=1, seed=None):
        super().__init__()
        self.data_path = data_path
        self.gen_path = gen_path
        self.patch_size = int(patch_size)           # LR patch 边长
        self.ratio = int(ratio)
        self.kerSize = int(kerSize)
        self.sigma = float(sigma)
        self.num_patches = int(num_patches)
        self.seed = args.seed if seed is None else int(seed)
        self.rng = np.random.RandomState(self.seed)

        # 生成退化核与光谱响应
        self.B = gauss_kernel(self.kerSize, self.kerSize, self.sigma)
        self.R = create_spec_resp(3, gen_path)

        # 读取单幅 HR-HSI，并生成 HR-MSI 与 LR-HSI
        hrhsi = sio.loadmat(self.data_path)['HS'].astype(np.float32)  # (H,W,C_hsi)
        hrmsi, lrhsi = create_hrms_lrhs(hrhsi, self.B, self.R, self.ratio)

        self.hrhsi = hrhsi     # (H,W,C_hsi)
        self.hrmsi = hrmsi     # (H,W,C_msi)
        self.lrhsi = lrhsi     # (h,w,C_hsi)，其中 h=H/ratio, w=W/ratio

        H, W, _ = self.hrhsi.shape
        h, w, _ = self.lrhsi.shape
        self.hr_patch = self.patch_size * self.ratio  # HR 边长

        # 边界检查
        if self.hr_patch > H or self.hr_patch > W:
            raise ValueError(f"HR patch({self.hr_patch}) 大于图像尺寸 {(H,W)}，请减小 patch_size 或 ratio。")
        if self.patch_size > h or self.patch_size > w:
            raise ValueError(f"LR patch({self.patch_size}) 大于图像尺寸 {(h,w)}，请减小 patch_size。")

        # 预生成随机坐标（复现友好）：在 LR 网格上采样，再放大到 HR 对齐
        lr_max_row = max(0, h - self.patch_size)
        lr_max_col = max(0, w - self.patch_size)
        self.coords = []
        for _ in range(self.num_patches):
            lr_r = int(self.rng.randint(0, lr_max_row + 1)) if lr_max_row > 0 else 0
            lr_c = int(self.rng.randint(0, lr_max_col + 1)) if lr_max_col > 0 else 0
            hr_r = lr_r * self.ratio
            hr_c = lr_c * self.ratio
            self.coords.append((lr_r, lr_c, hr_r, hr_c))

    def __len__(self):
        return self.num_patches

    def __getitem__(self, index):
        lr_r, lr_c, hr_r, hr_c = self.coords[index]
        # 切片
        lr_patch = self.lrhsi[lr_r: lr_r + self.patch_size, lr_c: lr_c + self.patch_size, :]          # (lr, lr, C_hsi)
        hrhs_patch = self.hrhsi[hr_r: hr_r + self.hr_patch, hr_c: hr_c + self.hr_patch, :]            # (HR, HR, C_hsi)
        hrms_patch = self.hrmsi[hr_r: hr_r + self.hr_patch, hr_c: hr_c + self.hr_patch, :]            # (HR, HR, C_msi)

        # (H,W,C) -> (C,H,W)
        hrhs = torch.from_numpy(np.ascontiguousarray(hrhs_patch.transpose(2, 0, 1))).float()
        hrms = torch.from_numpy(np.ascontiguousarray(hrms_patch.transpose(2, 0, 1))).float()
        lrhs = torch.from_numpy(np.ascontiguousarray(lr_patch.transpose(2, 0, 1))).float()
        return lrhs, hrms, hrhs

class HoustonDataset:
    pass

class ChikuseiDataset:
    pass

class TianGong_1Dataset:
    def __init__(self,data_path,gen_path,patch_size=8, ratio=8, kerSize=8, sigma=2,
                  seed=None):
        super().__init__()
        self.R = create_spec_resp(5,gen_path)
        self.B = gauss_kernel(kerSize,kerSize,sigma)
        self.ratio = ratio
        self.patch_size = patch_size
        self.num_patches = 1 # 默认单图
        # self.rows = hrhsi.shape
        # self.cols = 
        # 读取单幅高光谱图像：

        # 240 x 240 x 54
        hrhsi = sio.loadmat(data_path)['REF'].astype(np.float32)

        # 归一化到 [0, 1]
        gt_min, gt_max = hrhsi.min(), hrhsi.max()
        print(f'Max Spectral Value:{gt_max}')


        hrmsi,lrhsi = create_hrms_lrhs(hrhsi,self.B,self.R,self.ratio)
        self.hrhsi = hrhsi     # (H,W,C_hsi)
        self.hrmsi = hrmsi     # (H,W,C_msi)
        self.lrhsi = lrhsi     # (h,w,C_hsi)，其中 h=H/ratio, w=W/ratio

        H, W, _ = self.hrhsi.shape
        h, w, _ = self.lrhsi.shape
        self.hr_patch = self.patch_size * self.ratio  # HR 边长

        # 边界检查
        if self.hr_patch > H or self.hr_patch > W:
            raise ValueError(f"HR patch({self.hr_patch}) 大于图像尺寸 {(H,W)}，请减小 patch_size 或 ratio。")
        if self.patch_size > h or self.patch_size > w:
            raise ValueError(f"LR patch({self.patch_size}) 大于图像尺寸 {(h,w)}，请减小 patch_size。")

        # 预生成随机坐标（复现友好）：在 LR 网格上采样，再放大到 HR 对齐
        lr_max_row = max(0, h - self.patch_size)
        lr_max_col = max(0, w - self.patch_size)

        self.coords = [0,0]
        self.israndom = True  # 是否随机选取坐标
        if self.israndom:
            self.coords = []
            lr_r = int(self.rng.randint(0, lr_max_row + 1)) if lr_max_row > 0 else 0
            lr_c = int(self.rng.randint(0, lr_max_col + 1)) if lr_max_col > 0 else 0
            hr_r = lr_r * self.ratio
            hr_c = lr_c * self.ratio
            self.coords.append((lr_r, lr_c, hr_r, hr_c))
        else:
            # 单图情况默认取中心patch（确保不越界）
            lr_r = min(max(0, lr_max_row // 2), lr_max_row)
            lr_c = min(max(0, lr_max_col // 2), lr_max_col)
            hr_r = lr_r * self.ratio
            hr_c = lr_c * self.ratio
            self.coords.append((lr_r, lr_c, hr_r, hr_c))
 
    def __len__(self):
        return len(self.coords)

    def __getitem__(self, index):
        lr_r, lr_c, hr_r, hr_c = self.coords[index]
        # 切片
        lr_patch = self.lrhsi[lr_r: lr_r + self.patch_size, lr_c: lr_c + self.patch_size, :]          # (lr, lr, C_hsi)
        hrhs_patch = self.hrhsi[hr_r: hr_r + self.hr_patch, hr_c: hr_c + self.hr_patch, :]            # (HR, HR, C_hsi)
        hrms_patch = self.hrmsi[hr_r: hr_r + self.hr_patch, hr_c: hr_c + self.hr_patch, :]            # (HR, HR, C_msi)

        # (H,W,C) -> (C,H,W)
        hrhs = torch.from_numpy(np.ascontiguousarray(hrhs_patch.transpose(2, 0, 1))).float()
        hrms = torch.from_numpy(np.ascontiguousarray(hrms_patch.transpose(2, 0, 1))).float()
        lrhs = torch.from_numpy(np.ascontiguousarray(lr_patch.transpose(2, 0, 1))).float()
        return lrhs, hrms, hrhs


class BotswanaDataset:
    pass

class XiongAnDataset:
    pass

class WDCMDataset:
    pass



def _minmax(img):
    img = img.astype(np.float32)
    mn, mx = img.min(), img.max()
    return (img - mn) / (mx - mn + 1e-8)


def visualize_paviau_sample(lrhs_bchw, hrms_bchw, hrhs_bchw, band_idx=30, rgb_idx=(60, 30, 10), save_dir='/home/yehui/Documents/code/supervised/Global_Module_For_MHIF/mlgb  ', show=False):
    """
    lrhs_bchw: (B, C_hsi, h, w)
    hrms_bchw: (B, 3, H, W)
    hrhs_bchw: (B, C_hsi, H, W)
    band_idx:  展示的单通道索引（HSI）
    rgb_idx:   HR-HSI 伪彩色合成所用的三个通道索引 (R,G,B)
    save_dir:  保存路径（None 则不保存）
    show:      是否调用 plt.show()（无头环境通常为 False）
    """
    # 取 batch 中第一个样本
    lrhs = lrhs_bchw[0].detach().cpu().numpy()
    hrms = hrms_bchw[0].detach().cpu().numpy()
    hrhs = hrhs_bchw[0].detach().cpu().numpy()

    C_hsi = lrhs.shape[0]
    band_idx = int(np.clip(band_idx, 0, C_hsi - 1))
    rgb_idx = tuple(int(np.clip(i, 0, C_hsi - 1)) for i in rgb_idx)

    # 单通道显示
    lr_band = _minmax(lrhs[band_idx])                   # (h, w)
    hr_band = _minmax(hrhs[band_idx])                   # (H, W)

    # HR-MSI RGB
    hrms_rgb = np.stack([_minmax(hrms[0]), _minmax(hrms[1]), _minmax(hrms[2])], axis=-1)  # (H, W, 3)

    # HR-HSI 伪彩色
    hsi_rgb = np.stack([_minmax(hrhs[rgb_idx[0]]),
                        _minmax(hrhs[rgb_idx[1]]),
                        _minmax(hrhs[rgb_idx[2]])], axis=-1)

    # 可视化 第一张：4 子图
    fig, axes = plt.subplots(1, 4, figsize=(14, 4))
    axes[0].imshow(lr_band, cmap='gray')
    axes[0].set_title(f'LR-HSI band {band_idx}')
    axes[0].axis('off')

    axes[1].imshow(hr_band, cmap='gray')
    axes[1].set_title(f'HR-HSI band {band_idx}')
    axes[1].axis('off')

    axes[2].imshow(hrms_rgb)
    axes[2].set_title('HR-MSI RGB')
    axes[2].axis('off')

    axes[3].imshow(hsi_rgb)
    axes[3].set_title(f'HR-HSI pseudo RGB {rgb_idx}')
    axes[3].axis('off')

    plt.tight_layout()

    # 保存或显示/关闭
    timestamp = int(time.time())
    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
        grid_path = os.path.join(save_dir, f'paviau_sample_grid_{timestamp}.png')
        plt.savefig(grid_path, dpi=150, bbox_inches='tight')
    if show:
        plt.show()
    plt.close(fig)

    # 第二张：光谱曲线（中心像素）
    H, W = hrhs.shape[1], hrhs.shape[2]
    cy, cx = H // 2, W // 2
    spec = hrhs[:, cy, cx]
    spec = (spec - spec.min()) / (spec.max() - spec.min() + 1e-8)

    fig2 = plt.figure(figsize=(6, 3))
    plt.plot(np.arange(C_hsi), spec, '-o', ms=2)
    plt.title(f'Spectrum at center pixel ({cy},{cx})')
    plt.xlabel('HSI Band')
    plt.ylabel('Normalized Intensity')
    plt.grid(True, alpha=0.3)
    plt.tight_layout()

    if save_dir:
        spec_path = os.path.join(save_dir, f'paviau_sample_spectrum_{timestamp}.png')
        plt.savefig(spec_path, dpi=150, bbox_inches='tight')
    if show:
        plt.show()
    plt.close(fig2)


if __name__ == '__main__':
    #! Launch dataset.py independently: 
    #! -------------terminal---------------
    #$ python -m model.datasets               
    #! ------------------------------------

    #data_path = '/home/yehui/Documents/code/datasets/CAVE/'
    data_path = '/home/yehui/Documents/code/datasets/PaviaU/PaviaU.mat'
    genPath = '/home/yehui/Documents/code/datasets/'


    args = Args()
    data_train = PaviaUDataset(data_path, genPath, type='train')
    trainLoader = DataLoader(data_train,batch_size=args.batch_size, num_workers=args.num_workers, shuffle=True)
    # data_eval = PaviaUDataset(data_path,genPath, type='eval')
    # evalLoader = DataLoader(data_eval, batch_size=1, num_workers=args.num_workers)
    # data_test = PaviaUDataset(data_path, genPath, type='test')
    # TestLoader = DataLoader(data_test, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=False)
    # dataloader = {'train': trainLoader, 'eval': evalLoader,'test':TestLoader}

    #data_train = PaviaUDataset_Region(data_path,genPath,60,4,4,2,1,args.seed)
    #trainLoader = DataLoader(data_train,batch_size=args.batch_size,num_workers=args.num_workers, shuffle=True)



    #data_train = CaveDataset(data_path,genPath,64,24,8,2,'train')
    #trainLoader = DataLoader(data_train,batch_size=args.batch_size,num_workers=args.num_workers, shuffle=True)

    # 取一个 batch 并可视化
    lrhs_b, hrms_b, hrhs_b = next(iter(trainLoader))
    print('Batch shapes:', lrhs_b.shape, hrms_b.shape, hrhs_b.shape)
    visualize_paviau_sample(lrhs_b, hrms_b, hrhs_b, band_idx=30, rgb_idx=(60, 30, 10))



    print('-------------------------------------------------')
    print("训练集 batch_size：", trainLoader.batch_size)
    print("训练集样本数：", len(data_train))
    print("训练集 batch 数量：", len(trainLoader))
    # print('-------------------------------------------------')
    # print("训练集 batch_size：", dataloader['train'].batch_size)
    # print("训练集样本数：", len(data_train))
    # print("训练集 batch 数量：", len(dataloader['train']))
    # print("验证集 batch_size：", dataloader['eval'].batch_size)
    # print("验证集样本数：", len(data_eval))
    # print("验证集 batch 数量：", len(dataloader['eval']))
    # print("测试集 batch_size：", TestLoader.batch_size)
    # print("测试集样本数：", len(data_test))
    # print("测试集 batch 数量：", len(TestLoader))
    # print('-------------------------------------------------')




