import numpy as np
import cv2
from config import RESIZE_TO, MASKIMAGE

# image mask

# free form mask
# bbox mask


def create_ff_mask():
    config = {
        "img_shape": list(RESIZE_TO),
        "mv": 15,
        "ma": 4.0,
        "ml": 40,
        "mbw": 5,
    }

    h, w = config["img_shape"]
    mask = np.zeros((h, w))
    num_v = np.random.randint(config["mv"])

    for i in range(num_v):
        start_x = np.random.randint(w)
        start_y = np.random.randint(h)
        for j in range(1 + np.random.randint(5)):
            angle = 0.01 + np.random.randint(config["ma"])
            if i % 2 == 0:
                angle = 2 * 3.1415926 - angle
            length = 10 + np.random.randint(config["ml"])
            brush_w = 5 + np.random.randint(config["mbw"])
            end_x = (start_x + length * np.sin(angle)).astype(np.int32)
            end_y = (start_y + length * np.cos(angle)).astype(np.int32)

            cv2.line(mask, (start_y, start_x), (end_y, end_x), 255.0, brush_w)
            start_x, start_y = end_x, end_y

    mask = mask.astype(np.uint8)
    cv2.imwrite(MASKIMAGE, mask)


def create_bbox_mask():
    shape = list(RESIZE_TO)
    margin = [10, 10]
    bbox_shape = [30, 30]

    def random_bbox(shape, margin, bbox_shape):
        """Generate a random tlhw with configuration.
        Args:
            config: Config should have configuration including IMG_SHAPES,
                    VERTICAL_MARGIN, HEIGHT, HORIZONTAL_MARGIN, WIDTH.
        Returns:
            tuple: (top, left, height, width)
        """
        img_height, img_width = shape
        height, width = bbox_shape
        ver_margin, hor_margin = margin
        maxt = img_height - ver_margin - height
        maxl = img_width - hor_margin - width
        t = np.random.randint(low=ver_margin, high=maxt)
        l = np.random.randint(low=hor_margin, high=maxl)
        h = height
        w = width
        return (t, l, h, w)

    bboxs = []
    for i in range(20):
        bbox = random_bbox(shape, margin, bbox_shape)
        bboxs.append(bbox)

    height, width = shape
    mask = np.zeros((height, width), np.float32)
    # print(mask.shape)
    for bbox in bboxs:
        h = int(bbox[2] * 0.1) + np.random.randint(int(bbox[2] * 0.2 + 1))
        w = int(bbox[3] * 0.1) + np.random.randint(int(bbox[3] * 0.2) + 1)
        mask[
            (bbox[0] + h) : (bbox[0] + bbox[2] - h),
            (bbox[1] + w) : (bbox[1] + bbox[3] - w),
        ] = 255.0

    mask = mask.astype(np.uint8)
    cv2.imwrite(MASKIMAGE, mask)