import cv2
import numpy as np
import os
import torch
import scipy.io as sio
import scipy.interpolate as spi
import tensorly as tl
import xlrd
import logging
import glob

# ===== 预处理 =====
def norm01(arr):
    """
    把一个 ndarray 归一化到 [0,1] 并返回新的 ndarray（不修改输入）。
    保持原始归一化策略（直接除最大值），确保数值分布一致性。
    """
    x = np.asarray(arr, dtype=np.float32)
    x_min, x_max = x.min(), x.max()

    if x_max > 0 and x_min >= 0 and x_max != 1:
        # 直接除最大值（保持原始策略）
        print(f'直接除最大值归一化: [0, {x_max:.4f}] -> [0, 1]')
        return (x / x_max).astype(np.float32)
    else:
        if x_max > x_min:
            # min-max 归一化
            print(f'Min-Max归一化: [{x_min:.4f}, {x_max:.4f}] -> [0, 1]')
            return ((x - x_min) / (x_max - x_min)).astype(np.float32)
        else:
            # 所有值相同的情况
            print(f'警告: 数据所有值相同 ({x_min:.4f})，返回全零数组')
            return np.zeros_like(x, dtype=np.float32)


def gauss_kernel(row_size, col_size, sigma):
    kernel = cv2.getGaussianKernel(row_size, sigma)
    kernel = kernel * cv2.getGaussianKernel(col_size, sigma).T
    return kernel

def intersect(list1, list2):
    list1 = list(list1)
    elem = list(set(list1).intersection(set(list2)))
    elem.sort()
    res = np.zeros(len(elem))
    for i in range(0, len(elem)):
        res[i] = list1.index(elem[i])
    res = res.astype("int32")
    return res

# ===== 加载光谱响应函数 =====
def create_spec_resp(dataset, genPath):
    if dataset == 'CAVE':  # CAVE  31 X 3
        band = 31
        file = os.path.join(genPath, 'srf/D700.mat')  # 377-948
        mat = sio.loadmat(file)
        spec_rng = np.arange(400, 700 + 1, 10)
        spec_resp = mat['spec_resp']
        #print("spec_rng:", spec_rng)
        R = spec_resp[spec_rng - 377, 1:4]
        
    elif dataset == 'Harvard':  # harvard  31 X 3
        file = os.path.join(genPath, 'srf/D700.mat')  # 377-948
        mat = sio.loadmat(file)
        spec_rng = np.arange(420, 720 + 1, 10)
        spec_resp = mat['spec_resp']
        R = spec_resp[spec_rng - 377, 1:4].T

    elif dataset == 'PaviaC':
        band = 102  # paviaC
        file = os.path.join(genPath,'srf/ikonos.mat')  # 350 : 5 : 1035
        mat = sio.loadmat(file)
        spec_rng = np.arange(430, 861)
        spec_resp = mat['spec_resp']
        ms_bands = range(1, 5)
        valid_ik_bands = intersect(spec_resp[:, 0], spec_rng)
        no_wa = len(valid_ik_bands)
        # Spline interpolation
        xx = np.linspace(1, no_wa, band)
        x = range(1, no_wa + 1)
        R = np.zeros([5, band])
        for i in range(0, 5):
            ipo3 = spi.splrep(x, spec_resp[valid_ik_bands, i + 1], k=3)
            R[i, :] = spi.splev(xx, ipo3)
        R = R[ms_bands, :].T

    elif dataset == 'PaviaU':
        # PaviaU 103 X 4
        band = 103 
        file = os.path.join(genPath, 'srf/ikonos_SRF.mat')
        mat = sio.loadmat(file)
        spec_resp = mat['R']
        R  = spec_resp.T
        print(R.shape)

    elif dataset == 'Chikusei':
        # Chikusei  128 X 4
        band = 128
        file = os.path.join(genPath, 'srf/ikonos.mat')
        print(file)
        mat = sio.loadmat(file)
        spec_rng = np.arange(375, 1015, 5)
        spec_resp = mat['spec_resp']
        R = spec_resp[(spec_rng - 350) // 5, 2:6]

    elif dataset == 'TianGong-1':
        # TianGong-1 54 X 8
        band = 54
        file = os.path.join(genPath,'srf/TG.xls')
        #print(file)
        data = xlrd.open_workbook(file)
        table = data.sheets()[0]
    
        num_cols = table.ncols
        cols_list = [np.array(table.col_values(i)).reshape(-1,1) for i in range(0,num_cols)]
    
        sp_data = np.concatenate(cols_list, axis=1)
        R = sp_data
        print(f'srf:{sp_data.shape}')

    elif dataset == 'WDC':
        band = 191
        file = os.path.join(genPath,'WDC/WDC_srf.mat')
        print(file)
        data = sio.loadmat(file)['R']
        R = data

    else:
        raise ValueError("没有该数据集的光谱响应函数")


    c = 1 / np.sum(R, axis=0) # 按列求和，并进行归一化
    assert not np.isnan(c).any(), f"Invalid values in inverse column sums (NaN found): {c}"
    R = np.multiply(R, c.reshape([1,c.size])) # R:54x8  C:1X8 每一行进行缩放
    print(f'R[hsi_c,msi_c]:{R.shape}')
    return R

# ===== 光谱退化 =====
def generate_MSI(img, srf_gt):
    h, w, c = img.shape
    #print(f'Generating MSI from HSI with shape: {img.shape} using SRF shape: {srf_gt.shape}')
    if srf_gt.shape[0] == c:
        img_msi = np.dot(img.reshape(h*w, c), srf_gt).reshape(h, w, srf_gt.shape[1])
    else:
        raise Exception("The shape of sp matrix does not match the image")
    return img_msi

# ===== 高斯模糊核 =====
def matlab_style_gauss2D(shape=(3,3),sigma=2): 
    m,n = [(ss-1.)/2. for ss in shape]
    y,x = np.ogrid[-m:m+1,-n:n+1]
    h = np.exp( -(x*x + y*y) / (2.*sigma*sigma) )
    h[ h < np.finfo(h.dtype).eps*h.max() ] = 0
    sumh = h.sum()
    if sumh != 0:
        h /= sumh
    return h 

# ===== 空间退化 =====
from scipy import signal
def downsamplePSF(img, sigma,stride):
    # generate filter same with fspecial('gaussian') function
    h = matlab_style_gauss2D((stride,stride),sigma)
    # 正确的形状解释：img.shape -> (height, width, channels)
    if img.ndim == 3:
        img_h, img_w, img_c = img.shape
    elif img.ndim == 2:
        img_c = 1
        img_h, img_w = img.shape
        img = img.reshape((img_h, img_w, 1))
    
    out_img = np.zeros((img_h//(stride), img_w//(stride), img_c))
    for i in range(img_c):
        # 以 (rows, cols) 顺序对每个通道做二维卷积
        out = signal.convolve2d(img[:,:,i], h, 'valid')
        # 下采样：行/列步长均为 stride
        out_img[:,:,i] = out[::stride, ::stride]
    return out_img

def generate_low_HSI(img,sigma,scale_factor):
    (h, w, c) = img.shape
    img_lr = downsamplePSF(img, sigma=sigma, stride=scale_factor)
    #print(f'lrhsi:{img_lr.shape}')
    return img_lr 

# 已弃用
def create_hrms_lrhs(hs, B, R, ratio):
    hrms = tl.tenalg.mode_dot(hs, R, mode=2) # 张量的2-model乘法
    # blur
    lrhs = cv2.filter2D(hs, -1, B, borderType=cv2.BORDER_REFLECT)
    # downsample
    lrhs = lrhs[0::ratio, 0::ratio, :]
    return hrms, lrhs

def time2file_name(time):
    year = time[0:4]
    month = time[5:7]
    day = time[8:10]
    hour = time[11:13]
    minute = time[14:16]
    second = time[17:19]
    time_filename = year + '_' + month + '_' + day + '_' + hour + '_' + minute + '_' + second
    return time_filename

class AverageMeter(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n # 记录当前的样本数
        self.avg = self.sum / self.count # 计算当前的指标平均值
 
def initialize_logger(file_dir):
    logger = logging.getLogger()
    fhandler = logging.FileHandler(filename=file_dir, mode='a')
    formatter = logging.Formatter('%(asctime)s - %(message)s', "%Y-%m-%d %H:%M:%S")
    fhandler.setFormatter(formatter)
    logger.addHandler(fhandler)
    logger.setLevel(logging.INFO)
    return logger

def save_checkpoint(model_path, epoch, model, optimizer):
    state = {
        'epoch': epoch,
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict(),
    }
    torch.save(state, os.path.join(model_path, 'net_%depoch.pth' % epoch))


import re
def get_best_model_name_from_log(log_path):
    """
    从train.log中自动获取Eval PSNR最高的epoch，并返回对应的模型文件名


    :param log_path: 日志文件路径
    :return: 最佳模型文件名（如 net_149epoch.pth）
    """
    best_psnr = -float('inf')
    best_epoch = None
    pattern = re.compile(r'Epoch\[(\d+)\].*Eval PSNR: ([\d\.]+)')

    with open(log_path, 'r') as f:
        for line in f:
            match = pattern.search(line)
            if match:
                epoch = int(match.group(1))
                psnr = float(match.group(2))
                if psnr > best_psnr:
                    best_psnr = psnr
                    best_epoch = epoch

    if best_epoch is not None:
        return f'net_{best_epoch}epoch.pth'
    else:
        raise ValueError("未找到有效的PSNR记录！")

def keep_only_best_model(model_dir, best_model_name):
    """
    只保留指定的最佳模型文件，删除其他 .pth 文件
    :param model_dir: 模型文件夹路径
    :param best_model_name: 最佳模型文件名（如 'net_149epoch.pth'）
    """
    best_model_path = os.path.join(model_dir, best_model_name)
    print(best_model_path)
    all_pth_files = glob.glob(os.path.join(model_dir, "*.pth"))
    for file_path in all_pth_files:
        if os.path.abspath(file_path) != os.path.abspath(best_model_path):
            os.remove(file_path)
            print(f"Deleted: {file_path}")
    print(f"Kept best model: {best_model_path}")

def save_best_pth(args):
    model_dir = os.path.dirname(args.model_path)                # 获得存放model的目录
    log_path = os.path.join(args.log_path, 'train.log')              # 添加后缀
    best_model_name = get_best_model_name_from_log(log_path)    # 从log中找到最佳的Epoch
    keep_only_best_model(model_dir, best_model_name)            # 保存最佳的pth
    print("success")