import os
import torch
import numpy as np
import cv2
import math
import torch.nn as nn
import sys
import datetime
import logging
from collections import defaultdict
import re

def calculate_rmse(img1, img2):
    """
    Root Mean Squared Error
    Calculated individually for all bands, then averaged
    """
    img1 = img1.astype(np.float32)
    img2 = img2.astype(np.float32)
    mse = np.mean((img1 - img2) ** 2)
    if mse == 0:
        return float('inf')

    rmse = np.sqrt(mse)

    return np.mean(rmse)


def calculate_mae(img1, img2):

    img1 = img1.astype(np.float32)
    img2 = img2.astype(np.float32)
    apd = np.mean(np.abs(img1 - img2))
    if apd == 0:
        return float('inf')

    return np.mean(apd)


def calculate_psnr(img1, img2):
    # img1 and img2 have range [0, 255]
    img1 = img1.astype(np.float32)
    img2 = img2.astype(np.float32)
    mse = np.mean((img1 - img2)**2)
    if mse == 0:
        return float('inf')
    return 20 * math.log10(255.0 / math.sqrt(mse))


def ssim(img1, img2):
    C1 = (0.01 * 255)**2
    C2 = (0.03 * 255)**2

    img1 = img1.astype(np.float32)
    img2 = img2.astype(np.float32)
    kernel = cv2.getGaussianKernel(11, 1.5)
    window = np.outer(kernel, kernel.transpose())

    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
    mu1_sq = mu1**2
    mu2_sq = mu2**2
    mu1_mu2 = mu1 * mu2
    sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
    sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2

    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
                                                            (sigma1_sq + sigma2_sq + C2))
    return ssim_map.mean()


def calculate_ssim(img1, img2):
    '''calculate SSIM
    the same outputs as MATLAB's
    img1, img2: [0, 255]
    '''
    # img1 = img1.transpose((1, 2, 0))
    # img2 = img2.transpose((1, 2, 0))

    # print(img1.shape)
    # print(img2.shape)
    # bk
    if not img1.shape == img2.shape:
        raise ValueError('Input images must have the same dimensions.')
    if img1.ndim == 2:
        return ssim(img1, img2)
    elif img1.ndim == 3:
        if img1.shape[2] == 3:
            ssims = []
            for i in range(3):
                ssims.append(ssim(img1, img2))
            return np.array(ssims).mean()
        elif img1.shape[2] == 1:
            return ssim(np.squeeze(img1), np.squeeze(img2))
    else:
        raise ValueError('Wrong input image dimensions.')
    


def quantization(tensor):
    return torch.round(torch.clamp(tensor*255, min=0., max=255.))/255

def mkdir(path):
    if not os.path.exists(path):
        os.makedirs(path)


def mkdirs(paths):
    if isinstance(paths, str):
        mkdir(paths)
    else:
        for path in paths:
            mkdir(path)




'''
# --------------------------------------------
# logger
# --------------------------------------------
'''
def log(*args, **kwargs):
    print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:"), *args, **kwargs)



def beijing(sec, what):
    beijing_time = datetime.datetime.now() + datetime.timedelta(hours=8)
    return beijing_time.timetuple()

logging.Formatter.converter = beijing


def logger_info(logger_name, log_path='default_logger.log'):
    ''' set up logger
    modified by Kai Zhang (github: https://github.com/cszn)
    '''
    log = logging.getLogger(logger_name)
    if log.hasHandlers():
        print('LogHandlers exist!')
    else:
        print('LogHandlers setup!')
        level = logging.INFO
        
        formatter = logging.Formatter('%(asctime)s.%(msecs)03d : %(message)s', datefmt='%y-%m-%d %H:%M:%S')
        fh = logging.FileHandler(log_path, mode='a')
        fh.setFormatter(formatter)
        log.setLevel(level)
        log.addHandler(fh)
        # print(len(log.handlers))

        sh = logging.StreamHandler()
        sh.setFormatter(formatter)
        log.addHandler(sh)


'''
# --------------------------------------------
# print to file and std_out simultaneously
# --------------------------------------------
'''

class logger_print(object):
    def __init__(self, log_path="default.log"):
        self.terminal = sys.stdout
        self.log = open(log_path, 'a')

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)  # write the message

    def flush(self):
        pass



class MetricMonitor:
    '''
    MetricMonitor helps to track metrics such as accuracy or loss during training and validation and shows them on terminal.
     '''
    def __init__(self, float_precision=3):
        self.float_precision = float_precision
        self.reset()

    def reset(self):
        self.metrics = defaultdict(lambda: {"val": 0, "count": 0, "avg": 0})

    def update(self, metric_name, val):
        metric = self.metrics[metric_name]

        metric["val"] += val
        metric["count"] += 1
        metric["avg"] = metric["val"] / metric["count"]

    def __str__(self):
        return " | ".join(
            [
                "{metric_name}: {avg:.{float_precision}f}".format(
                    metric_name=metric_name, avg=metric["avg"], float_precision=self.float_precision
                )
                for (metric_name, metric) in self.metrics.items()
            ]
        )


def init_weights(model, random_seed=None):
    if random_seed != None:
        np.random.seed(random_seed)
    for m in model.modules():
        if isinstance(m, nn.Conv2d):
            m.weight.data = nn.Parameter(torch.tensor(np.random.normal(0, 1, m.weight.shape)).float())  
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.BatchNorm2d):
            nn.init.normal_(m.weight)
            nn.init.constant_(m.bias, 0)

        # elif isinstance(m, nn.Linear):
        #     m.weight.data = nn.Parameter(torch.tensor(np.random.normal(0, 1, m.weight.shape)).float())  
    np.random.seed() # Restoring randomness   



def get_first_n_digits_as_integer(value, n=3):
    # 将 float 值转换为 torch tensor
    abs_value = abs(value)  # 直接使用 Python 的 abs 函数处理 float
    exponent = torch.log10(torch.tensor(abs_value))
    exponent = torch.floor(torch.log10(torch.tensor(abs_value)))
    scaled_value = torch.tensor(abs_value / (10 ** exponent.item())) # .item() 用于提取标量值
    result = torch.floor(scaled_value * (10 ** (n - 1)))  # 保留前 n 位并转换为整数
    return int(result.item())  # 转换为 Python int 类型


def check_loss(loss_history_list):

    # We think loss is converged if its first 3 significant digits do not change over 5 iterations. 
    
    loss_history_list = loss_history_list[len(loss_history_list)-5:]
    loss_history_list = [get_first_n_digits_as_integer(x) for x in loss_history_list]
    beachmark_value = loss_history_list[0]

    loss_is_converged = True
    for x in loss_history_list:
        if x != beachmark_value:
            loss_is_converged = False
            break
    return loss_is_converged



def find_index(cur_wm_list, cur_jpeg_30_acc_list):
    # 找到 cur_jpeg_30_acc_list 中的最小值
    min_acc = min(cur_jpeg_30_acc_list)
    
    # 找到所有最小值的 index
    min_indices = [i for i, x in enumerate(cur_jpeg_30_acc_list) if x == min_acc]
    
    # 如果只有一个最小值，直接返回该 index
    if len(min_indices) == 1:
        return min_indices[0]
    
    # 如果有多个最小值，找到 cur_wm_list 中最大的值对应的 index
    max_wm_index = min_indices[0]
    for i in min_indices:
        if cur_wm_list[i] > cur_wm_list[max_wm_index]:
            max_wm_index = i
    
    return max_wm_index