from PIL import Image, ImageDraw, ImageFont

def draw_boxes_on_2d_image(img, boxes, **kwargs):
    pilimg = Image.fromarray(np.copy(img))  # drawing modifies array in-place
    draw = ImageDraw.Draw(pilimg)
    font_size = kwargs.get('font_size', 18)
    linewidth = kwargs.get('linewidth', 4)

    draw.font = ImageFont.truetype(font="arial.ttf", size=font_size)

    for box in boxes:
        y0 = box['info'].y0
        y1 = box['info'].y1
        x0 = box['info'].x0
        x1 = box['info'].x1
        xm = round((x0 + x1) / 2)

        la = box['info'].label
        zi = box['info'].zi

        draw.rectangle([(x0, y0), (x1, y1)], outline='white', width=linewidth)

        if kwargs.get('add_label') is True:
            draw.text((xm, y0), f'Z{zi:04d}-L{la:04d}', fill='white', anchor='mb')

    return pilimg

def generate_patches(
        desc, stack, boxes, rescale_clip=0.0,
        pad_to=256,
        proj=lambda x: x.max(axis=0),
        prefix='patch',
        **kwargs
):
    patch_dir = root / 'output' / 'patches' / desc
    patch_dir.mkdir(parents=True, exist_ok=True)

    for box in boxes:
        obj = box['info']
        sl = box['slice']
        rbb = box['relative_bounding_box']

        patch = proj(stack[sl])
        patch_fname = f'{prefix}-la{obj.label:04d}-zi{obj.zi:04d}'

        if rescale_clip is not None:
            patch = rescale(patch, rescale_clip)

        if kwargs.get('draw_bounding_box') is True:
            x0 = rbb['x0']
            y0 = rbb['y0']
            x1 = rbb['x1']
            y1 = rbb['y1']

            pilimg = Image.fromarray(patch)  # drawing modifies array in-place
            draw = ImageDraw.Draw(pilimg)
            draw.rectangle([(x0, y0), (x1, y1)], outline='white', width=kwargs.get('linewidth', 1))
            patch = np.array(pilimg)

        if pad_to:
            patch = pad(patch, pad_to)

        imsave(
            patch_dir / (patch_fname + '.png'),
            resample(patch),
            check_contrast=False,
        )
    print(f'Successfully wrote {len(boxes)} patches to:\n{patch_dir}')


def generate_3d_patches(  # in extensions.chaeo.products
        desc, stack, boxes, rescale_clip=0.0,
        pad_to=256,
        prefix='patch',
        proj=lambda x: x,
):
    patch_dir = root / 'output' / '3d_patches' / desc
    patch_dir.mkdir(parents=True, exist_ok=True)

    for box in boxes:
        obj = box['info']
        sl = box['slice']
        patch = proj(stack[sl])
        patch_fname = f'{prefix}-la{obj.label:04d}-zi{obj.zi:04d}'

        if rescale_clip is not None:
            patch = rescale(patch, rescale_clip)

        if pad_to:
            patch = pad_3d(patch, pad_to)

        imwrite(
            patch_dir / (patch_fname + '.tif'),
            patch,
            imagej=True
        )
    print(f'Successfully wrote {len(boxes)} patches to:\n{patch_dir}')