import cv2
import numpy as np
import gradio as gr



def apply_watermark(img, logo, mode):
    
    # change the mode to BGR for processing
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    logo = cv2.cvtColor(logo, cv2.COLOR_RGBA2BGRA)
    
    # Get dimensions of the main image and watermark
    main_height, main_width = img.shape[:2]
    watermark_height, watermark_width = logo.shape[:2]

    # Calculate the top-left corner of the watermark (bottom-right corner placement)
    x = main_width - watermark_width
    y = main_height - watermark_height

    # If you want the watermark to not exactly stick to the corner, you can add some padding
    padding = 10  # 10 pixels padding for example
    x -= padding
    y -= padding

    # Place the watermark
    roi = img[y : y + watermark_height, x : x + watermark_width]
    
    # Separate the color and alpha channels.
    # print(logo.shape)
    logo_bgr = logo[:, :, 0:3]
    logo_alpha = logo[:, :, 3]

    # Make the dimensions of the mask same as the input image.
    # Since the background image is a 3-channel image, we create a 3 channel image for the mask.
    logo_mask = cv2.merge([logo_alpha, logo_alpha, logo_alpha])
    logo_mask_inv = cv2.bitwise_not(logo_mask)

    # Use the mask to create the masked ROI region.
    masked_roi = cv2.bitwise_and(roi, logo_mask_inv)

    # Use the mask to create the masked logo region.
    masked_logo = cv2.bitwise_and(logo_bgr, logo_mask)

    if mode == "Opaque":


        # Combine the masked ROI with the masked logo to get the combined ROI image.
        roi_final = cv2.bitwise_or(masked_roi, masked_logo)

        # make a image copy
        img_1 = img.copy()

        # Insert the ROI patch in the image.
        img_1[y : y + watermark_height, x : x + watermark_width] = roi_final
        # cv2.imshow('output',img_1)
        # cv2.waitKey(0)
        return cv2.cvtColor(img_1, cv2.COLOR_BGR2RGB)
    else:
        # Make a copy of the ROI.
        roi_2 = roi.copy()

        # Blend ROI and the logo.
        watermarked = cv2.addWeighted(roi_2, 1, masked_logo, 0.6, 0)

        # Make a copy of the image.
        img_2 = img.copy()

        # Insert the watermarked ROI patch in the image.
        img_2[y : y + watermark_height, x : x + watermark_width] = watermarked
        # cv2.imshow('output',img_2)
        # cv2.waitKey(0)
        return cv2.cvtColor(img_2, cv2.COLOR_BGR2RGB)


def main(img, logo, mode):
    img_final = apply_watermark(img, logo, mode)
    return img_final


image_input = gr.Image(type="numpy", label="Input Image")
logo_input = gr.Image(type="numpy", label="Logo Image", image_mode="RGBA")
mode = gr.Radio(["Opaque", "Transparent"], label="Watermark Mode", info="Choose the MODE")
final_output = gr.Image(type="numpy", label="Output Image")

interface = gr.Interface(
    fn=main,
    inputs=[image_input, logo_input, mode],
    outputs=[final_output],
    title="OpenCV Watermark",
    description="Upload your Input Image and choose a logo to create your watermark image!",
    examples=[["./images/hp-2.jpg", "./images/opencv-logo-rz.png"], ["./images/hp.jpg", "./images/opencv-university-rz.png"]],
    cache_examples=False,
)

interface.launch()