import cv2
import glob
import os
import random
import numpy as np


def create_composite_image(_fg_image , _bg_image):
    """
    Creates composite images from a given set of foreground and background images
    :param _fg_image: path to foreground images folder
    :param _bg_image: path to background images folder
    :return:
    """
    flag = 0
    for plate in glob.glob(_bg_image):
        composite_folder_name = os.path.splitext(plate)[0] + "_composite"
        mask_folder_name = os.path.splitext(plate)[0] + "_mask"

        if not os.path.exists(composite_folder_name):
            os.makedirs(composite_folder_name)

        if not os.path.exists(mask_folder_name):
            os.makedirs(mask_folder_name)

        for food in glob.glob(_fg_image):
            plate_base_image = cv2.imread(plate, cv2.IMREAD_UNCHANGED)
            plate_base_image = cv2.resize(plate_base_image, (640, 480))
            h_1, w_1 = plate_base_image.shape[:2]

            random_size = random.randint(1, 100)
            food_image = cv2.imread(food, cv2.IMREAD_UNCHANGED)
            food_image = cv2.resize(food_image, (200 + random_size, 200 + random_size))

            h1, w1 = food_image.shape[:2]
            b, g, r, a = cv2.split(food_image)

            x = int((w_1 / 2) - (w1 / 2))
            y = int((h_1 / 2) - (h1 / 2))

            print(food_image.shape, plate_base_image.shape)

            cropped_bg_image = plate_base_image[y:y + h1, x:x + w1]
            inv_a = cv2.bitwise_not(a)

            image1_out = cv2.bitwise_and(cropped_bg_image, cropped_bg_image, mask=inv_a)
            image2_out = cv2.bitwise_and(food_image, food_image, mask=a)
            image_cropped = cv2.bitwise_or(image1_out, image2_out)

            segmented_mask = get_mask(a)

            plate_base_image[y:y + h1, x:x + w1] = image_cropped

            cv2.imwrite(composite_folder_name + "/" + str(flag) + "_" + os.path.basename(food), plate_base_image)
            cv2.imwrite(mask_folder_name + "/" + str(flag) + "_" + os.path.basename(food), segmented_mask)

            print(composite_folder_name + "/" + str(flag) + "_" + os.path.basename(food))

        flag += 1

def get_mask(_in):

    bg = np.zeros((480,640))
    h_1, w_1 =bg.shape[:2]

    #b, g, r, a = cv2.split(_in)
    h1 , w1 = _in.shape[:2]

    x = int((w_1 / 2) - (w1 / 2))
    y = int((h_1 / 2) - (h1 / 2))

    bg[y:y + h1, x:x + w1] = _in

    return bg


if __name__ == "__main__":

    BG_FOLDER_PATH = "/home/anirudh/HBRS/HomeLab/Dataset/dataset_green_segmentation/bg_png/*.png"
    FG_FOLDERS_PATH = "/home/anirudh/HBRS/HomeLab/Dataset/dataset_green_segmentation/fg/*.png"

    create_composite_image(FG_FOLDERS_PATH,BG_FOLDER_PATH)