import numpy as  np
import copy
import cv2 as cv
from tqdm import tqdm
import os
import torch
import torch.nn.functional as F
from PIL import Image
import uuid
from pathlib import Path

class segmentation():
    
    def __init__(self, model_name="xx"):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.base_dir = os.path.dirname(os.path.abspath(__file__))
        self.model_name = model_name
        self.palette = [0, 0, 0, 255, 0, 0, 0, 255, 0]
        self.fusin = True
    
    def create_tmp_path(self, path):
        folder_name = str(uuid.uuid4())
        folder_path = os.path.join(path,'tmp', folder_name)
        os.makedirs(folder_path)
        return folder_path

    

    # 图片分割识别结果
    def start_segmentation(self,path,net):
        assert os.path.isfile(path)
        image = Image.open(path)
        orininal_h = image.size[1]
        orininal_w = image.size[0]
        item = {"orininal_h": orininal_h, "orininal_w": orininal_w, "image": image}
        image = item["image"]
        orininal_w = item["orininal_w"]
        orininal_h = item["orininal_h"]
        old_img = copy.deepcopy(image)
        imaged = cv.resize(np.array(image), dsize=(512, 512), interpolation=cv.INTER_LINEAR)
        image_data = np.expand_dims(
            np.transpose(self.preprocess_input(np.array(imaged, np.float32), md=False), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            image = images.to(device="cuda", dtype=torch.float32)
            model = net.to(device="cuda")
            out = model(image)                                   # batch_size, 2, 512, 512
            if isinstance(out, list) or isinstance(out, tuple):  # 可能有多个输出（这里把辅助解码头的也输出的所以是多个）
                out = out[0]  # 就取第一个
            out = out[0]  # 去掉batch
            pr = F.softmax(out.permute(1, 2, 0), dim=-1).cpu().numpy()
            result = pr.argmax(axis=-1)
        # 二值图像
        output_im = Image.fromarray(np.uint8(result)).convert('P')
        output_im.putpalette(self.palette)
        file_name = os.path.basename(path)
        file_directory = os.path.dirname(path)
        file_name_ = "binary_"+ file_name
        file_path = os.path.join(file_directory,file_name_)
        output_im.save(file_path)
        # 融合图像
        PALETTE = [(0, 0, 0), (255, 0, 0), (0, 255, 0)]  # 这里主动改分类了，将识别图像与底图融合

        seg_img0 = np.reshape(np.array(PALETTE, np.uint8)[np.reshape(result, [-1])],
                                [orininal_h, orininal_w, -1])

        image0 = Image.fromarray(np.uint8(seg_img0))

        fusion = Image.blend(old_img, image0, 0.4)  # 图像融合
        file_name_ = "segment_"+ file_name
        file_path = os.path.join(file_directory,file_name_)
        fusion.save(file_path)


    # 获取大型图片，需要进行裁剪
    def cut_big_image(self,path):
        file_name = os.path.basename(path)
        file_directory = os.path.dirname(path)
        output_floder_normal = os.path.join(file_directory,'ori')
        # 创建输出文件夹（如果不存在）
        os.makedirs(output_floder_normal, exist_ok=True)
        # 打开图像
        image = Image.open(path)
        width, height = image.size
        # 定义块的大小
        block_size = 512

        # 计算需要的块数
        num_blocks_x = (width + block_size - 1) // block_size  # 向上取整
        num_blocks_y = (height + block_size - 1) // block_size  # 向上取整
        # 裁剪并保存图像块
        for i in range(num_blocks_x):
            for j in range(num_blocks_y):
                # 计算裁剪区域
                left = i * block_size
                upper = j * block_size
                right = min(left + block_size, width)
                lower = min(upper + block_size, height)

                # 裁剪图像
                block = image.crop((left, upper, right, lower))

                # 创建一个新的512x512白色图像
                new_block = Image.new('RGB', (block_size, block_size), (255, 255, 255))
                # 将裁剪的图像粘贴到白色图像上
                new_block.paste(block, (0, 0))

                # 保存图像块
                block_filename = f'block_{i}_{j}.png'
                new_block.save(os.path.join(output_floder_normal, block_filename))

        print("裁剪完成，图像块已保存。")


    def get_images(self, path):
        imglist = []
        lines = os.listdir(path)
        # print(lines)
        for ii, line in enumerate(lines):
            name = line
            _imagepath = os.path.join(path, name)
            assert os.path.isfile(_imagepath)
            image = Image.open(_imagepath)
            orininal_h = image.size[1]
            orininal_w = image.size[0]
            item = {"name": name, "orininal_h": orininal_h, "orininal_w": orininal_w, "image": image}
            imglist.append(item)
        # print("共监测到{}张原始图像和标签".format(len(imglist)))
        return imglist

    def folder_segment(self,path,net):
        # file_name = os.path.basename(path)
        # file_directory = os.path.dirname(path)
        # ori_path = os.path.join(file_directory,'ori')
        imglist = self.get_images(path)
        assert len(imglist) != 0
        for i in tqdm(imglist):
            image = i["image"]
            name = i["name"]
            orininal_w = i["orininal_w"]
            orininal_h = i["orininal_h"]
            old_img = copy.deepcopy(image)
            imaged = cv.resize(np.array(image), dsize=(512, 512), interpolation=cv.INTER_LINEAR)
            image_data = np.expand_dims(
                np.transpose(self.preprocess_input(np.array(imaged, np.float32), md=False), (2, 0, 1)), 0)
            with torch.no_grad():
                images = torch.from_numpy(image_data)
                image = images.to(device="cuda", dtype=torch.float32)
                model = net.to(device="cuda")
                out = model(image)                                   # batch_size, 2, 512, 512
                if isinstance(out, list) or isinstance(out, tuple):  # 可能有多个输出（这里把辅助解码头的也输出的所以是多个）
                    out = out[0]  # 就取第一个
                out = out[0]  # 去掉batch

                pr = F.softmax(out.permute(1, 2, 0), dim=-1).cpu().numpy()
                result = pr.argmax(axis=-1)
            # 结果图
            path_ = Path(path)
            parent_directory = path_.parent
            directory = os.path.join(parent_directory,'binary')
            if not os.path.exists(directory):
                os.makedirs(directory)
            output_im = Image.fromarray(np.uint8(result)).convert('P')

            output_im.putpalette(self.palette)
            tmp_path_binary= os.path.join(directory,name)
            output_im.save(tmp_path_binary)

            directory = os.path.join(parent_directory,'fusion')
            if not os.path.exists(directory):
                os.makedirs(directory)
            PALETTE = [(0, 0, 0), (255, 0, 0), (0, 255, 0)]  # 这里主动改分类了，将识别图像与底图融合
            seg_img0 = np.reshape(np.array(PALETTE, np.uint8)[np.reshape(result, [-1])],
                                    [orininal_h, orininal_w, -1])

            image0 = Image.fromarray(np.uint8(seg_img0))

            fusion = Image.blend(old_img, image0, 0.4)  # 图像融合
            tmp_path_fusion= os.path.join(directory,name)
            fusion.save(tmp_path_fusion)


    def folder_segment_all(self,path,net,mode):
        # file_name = os.path.basename(path)
        # file_directory = os.path.dirname(path)
        # ori_path = os.path.join(file_directory,'ori')
        imglist = self.get_images(path)
        assert len(imglist) != 0
        for i in tqdm(imglist):
            image = i["image"]
            name = i["name"]
            orininal_w = i["orininal_w"]
            orininal_h = i["orininal_h"]
            old_img = copy.deepcopy(image)
            imaged = cv.resize(np.array(image), dsize=(512, 512), interpolation=cv.INTER_LINEAR)
            image_data = np.expand_dims(
                np.transpose(self.preprocess_input(np.array(imaged, np.float32), md=False), (2, 0, 1)), 0)
            with torch.no_grad():
                images = torch.from_numpy(image_data)
                image = images.to(device="cuda", dtype=torch.float32)
                model = net.to(device="cuda")
                out = model(image)                                   # batch_size, 2, 512, 512
                if isinstance(out, list) or isinstance(out, tuple):  # 可能有多个输出（这里把辅助解码头的也输出的所以是多个）
                    out = out[0]  # 就取第一个
                out = out[0]  # 去掉batch

                pr = F.softmax(out.permute(1, 2, 0), dim=-1).cpu().numpy()
                result = pr.argmax(axis=-1)
            # 结果图
            path_ = Path(path)
            parent_directory = path_.parent
            directory = os.path.join(parent_directory, mode+'_binary')
            if not os.path.exists(directory):
                os.makedirs(directory)
            output_im = Image.fromarray(np.uint8(result)).convert('P')

            output_im.putpalette(self.palette)
            tmp_path_binary= os.path.join(directory,name)
            output_im.save(tmp_path_binary)

            # directory = os.path.join(parent_directory,mode + '_fusion')
            # if not os.path.exists(directory):
            #     os.makedirs(directory)
            # PALETTE = [(0, 0, 0), (255, 0, 0), (0, 255, 0)]  # 这里主动改分类了，将识别图像与底图融合
            # seg_img0 = np.reshape(np.array(PALETTE, np.uint8)[np.reshape(result, [-1])],
            #                         [orininal_h, orininal_w, -1])

            # image0 = Image.fromarray(np.uint8(seg_img0))

            # fusion = Image.blend(old_img, image0, 0.4)  # 图像融合
            # tmp_path_fusion= os.path.join(directory,name)
            # fusion.save(tmp_path_fusion)
            

    def merge_pic(self,path):
        image = Image.open(path)
        width, height = image.size
        file_name = os.path.basename(path)
        file_directory = os.path.dirname(path)
        output_image_path = os.path.join(file_directory,file_name)
        for name in ['fusion','binary']:
            file_name_ = "merge_" +name+ "_" + file_name
            input_folder = os.path.join(file_directory,name)
            # 获取所有小图片文件名
            image_files = [f for f in os.listdir(input_folder) if f.endswith('.png')]
            image_files.sort()  # 按文件名排序，确保按顺序合并
            # 假设小图的尺寸是512x512
            block_size = 512
            # 计算需要的块数
            num_blocks_x = (width + block_size - 1) // block_size  # 向上取整
            num_blocks_y = (height + block_size - 1) // block_size  # 向上取整

            # 创建一个新的空白图像，用于合并块
            merged_image = Image.new('RGB', (width, height))

            # 遍历所有块并将它们粘贴到合并图像中
            for i in range(num_blocks_x):
                for j in range(num_blocks_y):
                    # 计算块的文件名
                    block_filename = f'block_{i}_{j}.png'
                    block_path = os.path.join(input_folder, block_filename)

                    # 打开块图像
                    if os.path.exists(block_path):
                        block = Image.open(block_path)
                        # 计算粘贴位置
                        left = i * block_size
                        upper = j * block_size
                        merged_image.paste(block, (left, upper))
                # 保存合并后的图像
            target = os.path.join(file_directory,file_name_)
            merged_image.save(target)  # 替换为你想保存的路径
            print("合并完成，保存为:", target)


    def merge_pic_all(self,path):
        image = Image.open(path)
        width, height = image.size
        file_name = os.path.basename(path)
        file_directory = os.path.dirname(path)
        path_list = []
        for name in ['pv_binary','roof_binary']:
            file_name_ = "merge_" +name+ "_" + file_name
            input_folder = os.path.join(file_directory,name)
            # 获取所有小图片文件名
            image_files = [f for f in os.listdir(input_folder) if f.endswith('.png')]
            image_files.sort()  # 按文件名排序，确保按顺序合并
            # 假设小图的尺寸是512x512
            block_size = 512
            # 计算需要的块数
            num_blocks_x = (width + block_size - 1) // block_size  # 向上取整
            num_blocks_y = (height + block_size - 1) // block_size  # 向上取整

            # 创建一个新的空白图像，用于合并块
            merged_image = Image.new('RGB', (width, height))

            # 遍历所有块并将它们粘贴到合并图像中
            for i in range(num_blocks_x):
                for j in range(num_blocks_y):
                    # 计算块的文件名
                    block_filename = f'block_{i}_{j}.png'
                    block_path = os.path.join(input_folder, block_filename)

                    # 打开块图像
                    if os.path.exists(block_path):
                        block = Image.open(block_path)
                        # 计算粘贴位置
                        left = i * block_size
                        upper = j * block_size
                        merged_image.paste(block, (left, upper))
                # 保存合并后的图像
            target = os.path.join(file_directory,file_name_)
            merged_image.save(target)  # 替换为你想保存的路径
            print("合并完成，保存为:", target)
            path_list.append(target)
        return path_list



    def preprocess_input(self, image, md=False):
        mean = (0.231, 0.217, 0.22)  # 针对北京的数据集
        std = (0.104, 0.086, 0.085)
        if md:
            image /= 255.0
            image -= mean
            image /= std
            return image
        else:
            image /= 255.0
            return image


    # 将屋顶与光伏的二值图进行合并
    def merge_binary(self, path_list):
        
        # 打开二值图像
        image1_path  = path_list[1]  # 替换为你的屋顶图像路径
        image2_path  = path_list[0]  # 替换为你的光伏图像路径
        file_directory = os.path.dirname(image1_path)
        # 加载图像并确保它们是RGB模式
        image1 = Image.open(image1_path).convert('RGB')
        image2 = Image.open(image2_path).convert('RGB')

        # 创建一个新的图像，RGB模式，初始为黑色
        output_image = Image.new('RGB', image1.size, (0, 0, 0))

        # 获取像素数据
        pixels1 = image1.load()
        pixels2 = image2.load()
        output_pixels = output_image.load()

        # 遍历所有像素点
        for x in range(image1.width):
            for y in range(image1.height):
                # 检查两个图像在当前像素点上的颜色
                if pixels1[x, y] == (255, 0, 0) and pixels2[x, y] == (255, 0, 0):  # 红色
                    output_pixels[x, y] = (255, 0, 0)  # 设置为红色
                else:
                    output_pixels[x, y] = (0, 0, 0)  # 设置为黑色

        # 保存输出图像
        final_path = os.path.join(file_directory, 'merge_pvroof_bianry.png')
        output_image.save(final_path)  # 替换为你想保存的路径
        return final_path
    
    # 将二值图与原始图合并
    def merge_final(self,path_list):
        ori_path  = path_list[1]  # 原始图像
        binary_path  = path_list[0]  # 二值图像
        # 加载图像并确保它们是RGBA模式（如果需要透明度）
        ori_image = Image.open(ori_path).convert('RGBA')
        binary_image = Image.open(binary_path).convert('RGBA')

        # 确保两张图像的大小相同
        if ori_image.size != binary_image.size:
            raise ValueError("两张图像的尺寸必须相同！")

        # 设置混合比例，0.5表示两张图像各占50%
        alpha = 0.4

        # 融合图像
        blended_image = Image.blend(ori_image, binary_image, alpha)
        file_directory = os.path.dirname(ori_path)
        file_name = os.path.basename(ori_path)
        final_path = os.path.join(file_directory,'fusion_'+file_name)
        
        # 保存融合后的图像
        blended_image.save(final_path)  # 替换为你想保存的路径