import argparse
import os

import numpy as np
import torch
from PIL import Image

from self.utils import get_dirs_files
from union import eval_seg
from utils import args


path = args.path
thres = args.threshold
device = torch.device('cuda', 1)

def get_test_mult_results_single(path, device, thres, mode='arteriole', ):
    _, labels = get_dirs_files(os.path.join('DualModal2019', mode, 'Test'))
    threshold = (thres,)
    temp_sum = (0, 0, 0, 0, 0, 0)

    for label in labels:
        # 获取正确分割结果
        gt_mask = Image.open(label).convert('L')
        # 转化图像尺寸
        np_gt_mask = np.array(gt_mask)/255
        gt_mask = torch.tensor(np_gt_mask, dtype=torch.float32, device=device)
        label = label.split('/')[-1].split(".png")[0]

        label = os.path.basename(label)
        # dual_img_path = os.path.join('DualModal2019/RGB/Test', label + '.png')
        dual_img_path = os.path.join('DualModal2019/RGB/Test', label + '.png')
        dual_img = Image.open(dual_img_path).convert('RGB')
        dual_img = np.array(dual_img) / 255

        dual_img = torch.tensor(dual_img, dtype=torch.float32, device=device)

        masks = []

        for i in range(0, 49):
            label_new = label + '-' + str(i) + ".png"
            label_path = os.path.join(path, label_new)
            img = Image.open(label_path).convert('RGB')
            if mode == 'vessel':
                np_label = np.array(img)[:, :, 0]
            elif mode == 'arteriole':
                np_label = np.array(img)[:, :, 1]
            elif mode == 'venule':
                np_label = np.array(img)[:, :, 2]
            # 转化图像尺寸
            np_label = np_label/255

            masks.append(np_label)

        combined_mask = torch.zeros((1024, 1024), dtype=torch.float32, device=device)
        mask_count = torch.zeros((1024, 1024), dtype=torch.uint8, device=device)

        # 设置每个子图像的宽度和高度
        # 循环遍历每个子图像
        for i in range(7):
            for j in range(7):
                # 获取当前子图像的索引
                index = i * 7 + j
                # 计算当前子图像的起始坐标
                start_x = j * 128
                start_y = i * 128

                # 获取当前子图像的mask，并将其移到GPU上
                mask = torch.tensor(masks[index], dtype=torch.float32, device=device)

                # 将mask叠加到combined_mask上，并更新mask_count
                combined_mask[start_y:start_y + 256, start_x:start_x + 256] += mask
                mask_count[start_y:start_y + 256, start_x:start_x + 256] += 1

        # 计算重叠区域的平均值
        combined_mask /= mask_count.float().clamp(min=1)  # 防止除零错误
        combined_mask.to(device)

        print(combined_mask.shape)

        temp = eval_seg(combined_mask, gt_mask, threshold)

        # print(label, ["{:.2%}".format(x) for x in temp])

        temp_sum = tuple([sum(a) for a in zip(temp_sum, temp)])

        # 将结果移到CPU上，并转换为numpy数组
        combined_mask = combined_mask.cpu().numpy()*255

        # 创建一个新的Image对象来保存合并后的mask
        combined_mask_image = Image.fromarray((combined_mask).astype(np.uint8))

    print(["{:.2%}".format(x / 5) for x in temp_sum])

get_test_mult_results_single(path, device, thres, 'arteriole')
get_test_mult_results_single(path, device, thres, 'venule')