# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import numpy as np

class CompositionAnnotator:
    def __init__(self, cfg):
        self.process_types = ["repaint", "extension", "control"]
        self.process_map = {
            "repaint": "repaint",
            "extension": "extension",
            "control": "control",
            "inpainting": "repaint",
            "outpainting": "repaint",
            "frameref": "extension",
            "clipref": "extension",
            "depth": "control",
            "flow": "control",
            "gray": "control",
            "pose": "control",
            "scribble": "control",
            "layout": "control"
        }

    def forward(self, process_type_1, process_type_2, frames_1, frames_2, masks_1, masks_2):
        total_frames = min(len(frames_1), len(frames_2), len(masks_1), len(masks_2))
        combine_type = (self.process_map[process_type_1], self.process_map[process_type_2])
        if combine_type in [("extension", "repaint"), ("extension", "control"), ("extension", "extension")]:
            output_video = [frames_2[i] * masks_1[i] + frames_1[i] * (1 - masks_1[i]) for i in range(total_frames)]
            output_mask = [masks_1[i] * masks_2[i] * 255 for i in range(total_frames)]
        elif combine_type in [("repaint", "extension"), ("control", "extension"), ("repaint", "repaint")]:
            output_video = [frames_1[i] * (1 - masks_2[i]) + frames_2[i] * masks_2[i] for i in range(total_frames)]
            output_mask = [(masks_1[i] * (1 - masks_2[i]) + masks_2[i] * masks_2[i]) * 255 for i in range(total_frames)]
        elif combine_type in [("repaint", "control"), ("control", "repaint")]:
            if combine_type in [("control", "repaint")]:
                frames_1, frames_2, masks_1, masks_2 = frames_2, frames_1, masks_2, masks_1
            output_video = [frames_1[i] * (1 - masks_1[i]) + frames_2[i] * masks_1[i] for i in range(total_frames)]
            output_mask = [masks_1[i] * 255 for i in range(total_frames)]
        elif combine_type in [("control", "control")]:  # apply masks_2
            output_video = [frames_1[i] * (1 - masks_2[i]) + frames_2[i] * masks_2[i] for i in range(total_frames)]
            output_mask = [(masks_1[i] * (1 - masks_2[i]) + masks_2[i] * masks_2[i]) * 255 for i in range(total_frames)]
        else:
            raise Exception("Unknown combine type")
        return output_video, output_mask


class ReferenceAnythingAnnotator:
    def __init__(self, cfg):
        from .subject import SubjectAnnotator
        self.sbjref_ins = SubjectAnnotator(cfg['SUBJECT'] if 'SUBJECT' in cfg else cfg)
        self.key_map = {
            "image": "images",
            "mask": "masks"
        }
    def forward(self, images, mode=None, return_mask=None, mask_cfg=None):
        ret_data = {}
        for image in images:
            ret_one_data = self.sbjref_ins.forward(image=image, mode=mode, return_mask=return_mask, mask_cfg=mask_cfg)
            if isinstance(ret_one_data, dict):
                for key, val in ret_one_data.items():
                    if key in self.key_map:
                        new_key = self.key_map[key]
                    else:
                        continue
                    if new_key in ret_data:
                        ret_data[new_key].append(val)
                    else:
                        ret_data[new_key] = [val]
            else:
                if 'images' in ret_data:
                    ret_data['images'].append(ret_data)
                else:
                    ret_data['images'] = [ret_data]
        return ret_data


class AnimateAnythingAnnotator:
    def __init__(self, cfg):
        from .pose import PoseBodyFaceVideoAnnotator
        self.pose_ins = PoseBodyFaceVideoAnnotator(cfg['POSE'])
        self.ref_ins = ReferenceAnythingAnnotator(cfg['REFERENCE'])

    def forward(self, frames=None, images=None, mode=None, return_mask=None, mask_cfg=None):
        ret_data = {}
        ret_pose_data = self.pose_ins.forward(frames=frames)
        ret_data.update({"frames": ret_pose_data})

        ret_ref_data = self.ref_ins.forward(images=images, mode=mode, return_mask=return_mask, mask_cfg=mask_cfg)
        ret_data.update({"images": ret_ref_data['images']})

        return ret_data


class SwapAnythingAnnotator:
    def __init__(self, cfg):
        from .inpainting import InpaintingVideoAnnotator
        self.inp_ins = InpaintingVideoAnnotator(cfg['INPAINTING'])
        self.ref_ins = ReferenceAnythingAnnotator(cfg['REFERENCE'])

    def forward(self, video=None, frames=None, images=None, mode=None, mask=None, bbox=None, label=None, caption=None, return_mask=None, mask_cfg=None):
        ret_data = {}
        mode = mode.split(',') if ',' in mode else [mode, mode]

        ret_inp_data = self.inp_ins.forward(video=video, frames=frames, mode=mode[0], mask=mask, bbox=bbox, label=label, caption=caption, mask_cfg=mask_cfg)
        ret_data.update(ret_inp_data)

        ret_ref_data = self.ref_ins.forward(images=images, mode=mode[1], return_mask=return_mask, mask_cfg=mask_cfg)
        ret_data.update({"images": ret_ref_data['images']})

        return ret_data


class ExpandAnythingAnnotator:
    def __init__(self, cfg):
        from .outpainting import OutpaintingAnnotator
        from .frameref import FrameRefExpandAnnotator
        self.ref_ins = ReferenceAnythingAnnotator(cfg['REFERENCE'])
        self.frameref_ins = FrameRefExpandAnnotator(cfg['FRAMEREF'])
        self.outpainting_ins = OutpaintingAnnotator(cfg['OUTPAINTING'])

    def forward(self, images=None, mode=None, return_mask=None, mask_cfg=None, direction=None, expand_ratio=None, expand_num=None):
        ret_data = {}
        expand_image, reference_image= images[0], images[1:]
        mode = mode.split(',') if ',' in mode else ['firstframe', mode]

        outpainting_data = self.outpainting_ins.forward(expand_image,expand_ratio=expand_ratio, direction=direction)
        outpainting_image, outpainting_mask = outpainting_data['image'], outpainting_data['mask']

        frameref_data = self.frameref_ins.forward(outpainting_image,  mode=mode[0], expand_num=expand_num)
        frames, masks = frameref_data['frames'], frameref_data['masks']
        masks[0] = outpainting_mask
        ret_data.update({"frames": frames, "masks": masks})

        ret_ref_data = self.ref_ins.forward(images=reference_image, mode=mode[1], return_mask=return_mask, mask_cfg=mask_cfg)
        ret_data.update({"images": ret_ref_data['images']})

        return ret_data


class MoveAnythingAnnotator:
    def __init__(self, cfg):
        from .layout import LayoutBboxAnnotator
        self.layout_bbox_ins = LayoutBboxAnnotator(cfg['LAYOUTBBOX'])

    def forward(self, image=None, bbox=None, label=None, expand_num=None):
        frame_size = image.shape[:2]   # [H, W]
        ret_layout_data = self.layout_bbox_ins.forward(bbox, frame_size=frame_size, num_frames=expand_num, label=label)

        out_frames = [image] + ret_layout_data
        out_mask = [np.zeros(frame_size, dtype=np.uint8)] + [np.ones(frame_size, dtype=np.uint8) * 255] * len(ret_layout_data)

        ret_data = {
            "frames": out_frames,
            "masks": out_mask
        }
        return ret_data