from PIL import Image, ImageDraw
import numpy as np

settings = {
    "256narrow": {
        "p_irr": 1,
        "min_n_irr": 4,
        "max_n_irr": 50,
        "max_l_irr": 40,
        "max_w_irr": 10,
        "min_n_box": None,
        "max_n_box": None,
        "min_s_box": None,
        "max_s_box": None,
        "marg": None,
    },
    "256train": {
        "p_irr": 0.5,
        "min_n_irr": 1,
        "max_n_irr": 5,
        "max_l_irr": 200,
        "max_w_irr": 100,
        "min_n_box": 1,
        "max_n_box": 4,
        "min_s_box": 30,
        "max_s_box": 150,
        "marg": 10,
    },
    "512train": {    # TODO: experimental
            "p_irr": 0.5,
            "min_n_irr": 1,
            "max_n_irr": 5,
            "max_l_irr": 450,
            "max_w_irr": 250,
            "min_n_box": 1,
            "max_n_box": 4,
            "min_s_box": 30,
            "max_s_box": 300,
            "marg": 10,
        },
    "512train-large": {    # TODO: experimental
            "p_irr": 0.5,
            "min_n_irr": 1,
            "max_n_irr": 5,
            "max_l_irr": 450,
            "max_w_irr": 400,
            "min_n_box": 1,
            "max_n_box": 4,
            "min_s_box": 75,
            "max_s_box": 450,
            "marg": 10,
        },
}


def gen_segment_mask(mask, start, end, brush_width):
    mask = mask > 0
    mask = (255 * mask).astype(np.uint8)
    mask = Image.fromarray(mask)
    draw = ImageDraw.Draw(mask)
    draw.line([start, end], fill=255, width=brush_width, joint="curve")
    mask = np.array(mask) / 255
    return mask


def gen_box_mask(mask, masked):
    x_0, y_0, w, h = masked
    mask[y_0:y_0 + h, x_0:x_0 + w] = 1
    return mask


def gen_round_mask(mask, masked, radius):
    x_0, y_0, w, h = masked
    xy = [(x_0, y_0), (x_0 + w, y_0 + w)]

    mask = mask > 0
    mask = (255 * mask).astype(np.uint8)
    mask = Image.fromarray(mask)
    draw = ImageDraw.Draw(mask)
    draw.rounded_rectangle(xy, radius=radius, fill=255)
    mask = np.array(mask) / 255
    return mask


def gen_large_mask(prng, img_h, img_w,
                   marg, p_irr, min_n_irr, max_n_irr, max_l_irr, max_w_irr,
                   min_n_box, max_n_box, min_s_box, max_s_box):
    """
    img_h: int, an image height
    img_w: int, an image width
    marg: int, a margin for a box starting coordinate
    p_irr: float, 0 <= p_irr <= 1, a probability of a polygonal chain mask

    min_n_irr: int, min number of segments
    max_n_irr: int, max number of segments
    max_l_irr: max length of a segment in polygonal chain
    max_w_irr: max width of a segment in polygonal chain

    min_n_box: int, min bound for the number of box primitives
    max_n_box: int, max bound for the number of box primitives
    min_s_box: int, min length of a box side
    max_s_box: int, max length of a box side
    """

    mask = np.zeros((img_h, img_w))
    uniform = prng.randint

    if np.random.uniform(0, 1) < p_irr:  # generate polygonal chain
        n = uniform(min_n_irr, max_n_irr)  # sample number of segments

        for _ in range(n):
            y = uniform(0, img_h)  # sample a starting point
            x = uniform(0, img_w)

            a = uniform(0, 360)  # sample angle
            l = uniform(10, max_l_irr)  # sample segment length
            w = uniform(5, max_w_irr)  # sample a segment width

            # draw segment starting from (x,y) to (x_,y_) using brush of width w
            x_ = x + l * np.sin(a)
            y_ = y + l * np.cos(a)

            mask = gen_segment_mask(mask, start=(x, y), end=(x_, y_), brush_width=w)
            x, y = x_, y_
    else:  # generate Box masks
        n = uniform(min_n_box, max_n_box)  # sample number of rectangles

        for _ in range(n):
            h = uniform(min_s_box, max_s_box)  # sample box shape
            w = uniform(min_s_box, max_s_box)

            x_0 = uniform(marg, img_w - marg - w)  # sample upper-left coordinates of box
            y_0 = uniform(marg, img_h - marg - h)

            if np.random.uniform(0, 1) < 0.5:
                mask = gen_box_mask(mask, masked=(x_0, y_0, w, h))
            else:
                r = uniform(0, 60)  # sample radius
                mask = gen_round_mask(mask, masked=(x_0, y_0, w, h), radius=r)
    return mask


make_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["256train"])
make_narrow_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["256narrow"])
make_512_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["512train"])
make_512_lama_mask_large = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["512train-large"])


MASK_MODES = {
    "256train": make_lama_mask,
    "256narrow": make_narrow_lama_mask,
    "512train": make_512_lama_mask,
    "512train-large": make_512_lama_mask_large
}

if __name__ == "__main__":
    import sys

    out = sys.argv[1]

    prng = np.random.RandomState(1)
    kwargs = settings["256train"]
    mask = gen_large_mask(prng, 256, 256, **kwargs)
    mask = (255 * mask).astype(np.uint8)
    mask = Image.fromarray(mask)
    mask.save(out)