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

from .utils import convert_to_numpy


class SubjectAnnotator:
    def __init__(self, cfg, device=None):
        self.mode = cfg.get('MODE', "salientmasktrack")
        self.use_aug = cfg.get('USE_AUG', False)
        self.use_crop = cfg.get('USE_CROP', False)
        self.roi_only = cfg.get('ROI_ONLY', False)
        self.return_mask = cfg.get('RETURN_MASK', True)

        from .inpainting import InpaintingAnnotator
        self.inp_anno = InpaintingAnnotator(cfg['INPAINTING'], device=device)
        if self.use_aug:
            from .maskaug import MaskAugAnnotator
            self.maskaug_anno = MaskAugAnnotator(cfg={})
        assert self.mode in ["plain", "salient", "mask", "bbox", "salientmasktrack", "salientbboxtrack", "masktrack",
                             "bboxtrack", "label", "caption", "all"]

    def forward(self, image=None, mode=None, return_mask=None, mask_cfg=None, mask=None, bbox=None, label=None, caption=None):
        return_mask = return_mask if return_mask is not None else self.return_mask

        if mode == "plain":
            return {"image": image, "mask": None} if return_mask else image

        inp_res = self.inp_anno.forward(image,  mask=mask, bbox=bbox, label=label, caption=caption, mode=mode, return_mask=True, return_source=True)
        src_image = inp_res['src_image']
        mask = inp_res['mask']

        if self.use_aug and mask_cfg is not None:
            mask = self.maskaug_anno.forward(mask, mask_cfg)

        _, binary_mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)
        if (binary_mask is None or binary_mask.size == 0 or cv2.countNonZero(binary_mask) == 0):
            x, y, w, h = 0, 0, binary_mask.shape[1], binary_mask.shape[0]
        else:
            x, y, w, h = cv2.boundingRect(binary_mask)

        ret_mask = mask.copy()
        ret_image = src_image.copy()

        if self.roi_only:
            ret_image[mask == 0] = 255

        if self.use_crop:
            ret_image = ret_image[y:y + h, x:x + w]
            ret_mask = ret_mask[y:y + h, x:x + w]

        if return_mask:
            return {"image": ret_image, "mask": ret_mask}
        else:
            return ret_image