import cv2
import numpy as np
import math


# Stolen from https://stackoverflow.com/questions/40895785/using-opencv-to-overlay-transparent-image-onto-another-image
def apply_overlay(background, foreground):
    # normalize alpha channels from 0-255 to 0-1
    alpha_background = background[:,:,3] / 255.0
    alpha_foreground = foreground[:,:,3] / 255.0

    # set adjusted colors
    for color in range(0, 3):
        background[:,:,color] = alpha_foreground * foreground[:,:,color] + \
            alpha_background * background[:,:,color] * (1 - alpha_foreground)

    # set adjusted alpha and denormalize back to 0-255
    background[:,:,3] = (1 - (1 - alpha_foreground) * (1 - alpha_background)) * 255
    return background


# Generates a camo overlay for an image. Changing this function will change how the overlay is generated.
# The algorithm below "blocks" the image into "super pixels" of equal size. It then selects a random subset
# of the super pixels to be replaced with "replace_block". If "replace_block" is not set, it defaults to black
# with 25% opacity, as per recommendation in (https://wearethene.ws/notable/139468)
def generate_overlay(dims, blocks, percent_overlaid=0.23, replace_block=None):
    if replace_block is None:
        replace_block = [0, 0, 0, 64]

    def blockify_dimension( dim_val, blocks ):
        if dim_val < blocks:
            blocks = dim_val
        split = np.array_split(range(dim_val), blocks)
        return split

    def generate_mask( rows, cols ):
        num_blocks = rows * cols
        overlaid_blocks = math.ceil( num_blocks * percent_overlaid )
        mask = np.arange(rows*cols)
        mask = np.where(mask<overlaid_blocks, 1, 0)
        np.random.shuffle(mask)
        mask = mask.reshape(rows,cols)
        return mask

    d_row = blockify_dimension( dims[0], blocks )
    d_col = blockify_dimension( dims[1], blocks )

    masked_image = np.zeros(dims, np.uint8)

    mask = generate_mask(len(d_row), len(d_col))
    for row_idx, row in enumerate(d_row):
        for col_idx, col in enumerate(d_col):
            if mask[row_idx][col_idx] == 1:
                for r in row:
                    for c in col:
                        masked_image[r][c] = replace_block
    return masked_image



# Driver function. This first reads the input image, ensures it's in RGBA mode,
# generates an overlay, applies the overlay, then writes the new image.
def camo_image(input_fn, output_fn, blocks=20, replace_block=None):
    img = cv2.imread(input_fn)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2RGBA)
    overlay = generate_overlay( img.shape, blocks)
    overlaid = apply_overlay(img, overlay)
    cv2.imwrite(output_fn, overlaid)