from random import sample, choice
from PIL import Image


class create_pic:
    def __init__(
        self,
        broad_w: int = 3000,
        broad_h: int = 5000,
        part_w: int = 30,
        part_h: int = 50,
        surfaces: list = ["A", "B", "C"],
        rotations: list = [x * 15 for x in range(24)],
    ) -> None:
        """
        param
        broad_w : 柔性振动板的宽
        broad_h : 柔性振动板的高
        part_w : 工件的宽
        part_h : 工件的高
        surfaces : 工件有多少种面,例如['A','B','C','D']四种面
        rotations : 代表工件有多少种旋转姿态,默认以15°为单位旋转360°,一共24种姿态
        设定画布中能够让摄像头有效识别的工件密度是1/10的总面积,称为有效容量parts_no / 3.
        计算出单个工件的有效容量onepart_no
        """
        self.broad_w = broad_w
        self.broad_h = broad_h
        self.surfaces = surfaces
        self.rotations = rotations
        self.dict_max_matrix = dict_max_matrix = self.count_max_matrix(
            (broad_w, broad_h), (part_w, part_h)
        )
        # 总容量
        self.parts_no = parts_no = (
            dict_max_matrix["matrix_r"] * dict_max_matrix["matrix_c"]
        )
        self.effective_parts_no = int(parts_no / 10)

    def count_max_matrix(self, broad_area, part_area):
        """
        本方法的作用是计算柔性振动盘内能够容纳的最大工件数量
        """
        plan_1_row = int(broad_area[0] / part_area[0])
        plan_1_col = int(broad_area[1] / part_area[1])
        plan_1 = plan_1_row * plan_1_col

        plan_2_row = int(broad_area[0] / part_area[1])
        plan_2_col = int(broad_area[1] / part_area[0])
        plan_2 = plan_2_row * plan_2_col
        if plan_1 >= plan_2:
            return {
                "matrix_r": plan_1_row,
                "matrix_c": plan_1_col,
                "rotation": 0,
            }
        else:
            return {
                "matrix_r": plan_2_row,
                "matrix_c": plan_2_col,
                "rotation": 90,
            }

    def random_2D_point(self):
        """
        parts_no表示柔性振动板能够摆放最大的工件数量
        如果按照依次铺满柔性振动板,lst_point就可以转化为二维数列,即坐标系
        同时,还得将这一坐标系乘以倍数常量c_x/c_y,转化成柔性振动板真实的尺寸
        本方法的作用是,从parts_no中随机选中有效容量(即effective_parts_no)个不重复的点形成一维list(即lst_1D)
        并转化为二维坐标系(即lst_2D),再乘以c_x/c_y后,就可以转化为对应柔性振动板尺寸的坐标系
        """
        effective_parts_no = self.effective_parts_no
        parts_no = self.parts_no
        broad_w = self.broad_w
        broad_h = self.broad_h
        lst_point = range(parts_no)
        dict_max_matrix = self.dict_max_matrix
        c_x = int(broad_w / dict_max_matrix["matrix_r"])
        c_y = int(broad_h / dict_max_matrix["matrix_c"])
        lst_1D = sample(lst_point, effective_parts_no)

        def transf_1D_to_2D(lst_1D, dict_max_matrix):
            lst_2D = []
            append = lst_2D.append
            for point_1D in lst_1D:
                append(
                    (
                        int(point_1D / dict_max_matrix["matrix_r"]) * c_x,  # 取整
                        (point_1D % dict_max_matrix["matrix_r"]) * c_y,  # 取余
                    )
                )
            return lst_2D

        return transf_1D_to_2D(lst_1D, dict_max_matrix)

    def get_coordinate(self):
        """
        本方法的作用是获取随机放置工件姿态的坐标序列,输出为dict:{(x,y):工件姿态名称}
        """
        surfaces = self.surfaces
        rotations = self.rotations
        effective_parts_no = self.effective_parts_no
        # 从surfaces中随机抽取可重复的元素,随机抽取的数量为effective_parts_no的数量
        lst_parts = [choice(surfaces) for _ in range(effective_parts_no)]
        # 从rotations中随机抽取可重复的元素,随机抽取的数量为effective_parts_no的数量
        lst_rotations = [choice(rotations) for _ in range(effective_parts_no)]
        # 随机选取用于放置工件的坐标点
        lst_coordinates = self.random_2D_point()
        # 将随机坐标点与随机的工件姿态图片合并
        lstdic_coordinate = list(
            map(
                lambda x, y, z: {"coordinate": x, "part": y, "rotation": z},
                lst_coordinates,
                lst_parts,
                lst_rotations,
            )
        )
        return lstdic_coordinate

    def create_base_canvas(self, lstdict_coordinate, pic_path):
        broad_w = self.broad_w
        broad_h = self.broad_h
        surfaces = self.surfaces
        rotation = self.dict_max_matrix["rotation"]
        base_canvas = Image.new("RGB", (broad_w, broad_h), (255, 255, 255))
        pic_paste = base_canvas.paste
        dict_pic_temp = {}

        def rotate_pic(img, rotation, filled_color=(255,) * 4):
            """
            使用Pillow进行旋转后空白处填充
            """
            # 转换为有alpha层
            img_alpha = img.convert("RGBA")
            # 旋转
            img_rot = img_alpha.rotate(rotation)
            # 创建一个与旋转图像大小相同的白色图像
            img_base = Image.new("RGBA", img_rot.size, filled_color)
            # 使用alpha层的rot作为掩码创建一个复合图像
            img_filled = Image.composite(img_rot, img_base, img_rot)
            # 保存
            return img_filled.convert(img.mode)

        for pic in surfaces:
            img = Image.open("./Pictures/template/temp_%s.png" % pic)
            if rotation == 90:
                img = img.transpose(Image.ROTATE_90)
            dict_pic_temp[pic] = img
        for dic in lstdict_coordinate:
            rotation = dic["rotation"]
            img_tmp = dict_pic_temp[dic["part"]]
            if rotation > 0:
                img_tmp = rotate_pic(img_tmp, rotation)
            pic_paste(img_tmp, dic["coordinate"])
        base_canvas.save(pic_path)
        base_canvas.show()


if __name__ == "__main__":
    broad_w = 3000
    broad_h = 5000
    part_w = 30
    part_h = 50
    obj = create_pic(broad_w, broad_h, part_w, part_h, ["A", "B", "C"])
    lstdic_coordinate = obj.get_coordinate()
    obj.create_base_canvas(lstdic_coordinate, "./picture_test.png")