import cv2
import numpy as np
from PIL import Image
Image.MAX_IMAGE_PIXELS = None
import gradio as gr

sb_logo = Image.open("./SB_logo_horizontal.png")

def freemium_watermark(img, sb_logo, alpha):
    width_px, height_px = img.size
    # Convert the logo to "RGBA" mode and resize it to 80% of the image width
    sb_logo = sb_logo.convert("RGBA")
    new_logo_width = int(width_px * 0.7)  # 80% of the image width
    logo_ratio = new_logo_width / sb_logo.width
    new_logo_height = int(sb_logo.height * logo_ratio)
    sb_logo = sb_logo.resize((new_logo_width, new_logo_height))
    # Make the logo semi-translucent
    transparent_img = Image.new('RGBA', (sb_logo.width, sb_logo.height), (0, 0, 0, 0))  # Create a transparent image of the same size
    sb_logo = Image.blend(transparent_img, sb_logo, alpha=alpha)  # Blend the logo with the transparent image
    # Calculate the position to paste the logo at the center of the image
    paste_x_position = (width_px - new_logo_width) // 2
    paste_y_position = int(height_px * 0.4) - (new_logo_height // 2)
    img.paste(sb_logo, (paste_x_position, paste_y_position), sb_logo)
    # Save the image
    return img

def lighter(input_path):
    print('lighter')
    if input_path is None:
        return None
    img_array = np.array(input_path)
    # Line processing
    blurred = cv2.GaussianBlur(img_array, (7, 7), 0)
    #lighter
    edges = cv2.Canny(blurred, 150, 180)
    structuring_element = cv2.getStructuringElement(cv2.MORPH_RECT, (2,2))
    dilated_edges = cv2.dilate(edges, structuring_element, iterations=2)
    line_drawing = 255 - dilated_edges
    line_img = Image.fromarray(line_drawing)
    alpha = 0.3
    return freemium_watermark(line_img, sb_logo, alpha)
    
def darker(input_path):
    print('darker')
    if input_path is None:
        return None
    img_array = np.array(input_path)
    # Line processing
    blurred = cv2.GaussianBlur(img_array, (5, 5), 0)
    #darker
    edges = cv2.Canny(blurred, 70, 120) #darker
    #lighter
    # edges = cv2.Canny(blurred, 100, 180)
    structuring_element = cv2.getStructuringElement(cv2.MORPH_RECT, (2,2))
    dilated_edges = cv2.dilate(edges, structuring_element, iterations=2)
    line_drawing = 255 - dilated_edges
    line_img = Image.fromarray(line_drawing)
    alpha = 0.5
    return freemium_watermark(line_img, sb_logo, alpha)

with gr.Blocks(analytics_enabled=False, theme=gr.themes.Soft()) as demo:
    with gr.Row():
        with gr.Column():
            input = gr.Image(type="pil", label="Upload Image", sources=["upload"])
            gr.ClearButton(components=input)
            # gr.Examples(examples=["hotel.jpg", "intersection.jpg", "road_small.jpg", "crash_small.jpg", "trailer.jpg"], inputs=input)
            # gr.Examples(examples=["hotel.jpg", "intersection.jpg", "road_small.jpg"], inputs=input)
        with gr.Column():
            output = gr.Image(type = "filepath", label="Sketch Drawing", show_share_button=False)
            with gr.Row():
                lighter_btn=gr.Button("Less Detail")
                darker_btn=gr.Button("More Detail")
                
            gr.Markdown("<p style='text-align: center; font-size: 20px;'>Want to remove the watermark, add a scale bar, and compass heading?</p>\n")
            gr.Markdown("<p style='text-align: center; font-size: 20px;'><a href='https://www.skyebrowse.com/pricing#sectionaboutpurchase'>Subscribe Here</a><p>")
            lighter_btn.click(lighter, inputs=[input], outputs=[output], show_progress="minimal")
            darker_btn.click(darker, inputs=[input], outputs=[output], show_progress="minimal")
    gr.Examples(examples=["hotel.jpg", "intersection.jpg", "road_small.jpg"], inputs=input)
    
    @gr.on(inputs=[input], outputs=[output], show_progress="minimal")
    def sketch(input_path):
        if input_path is None:
            return None
        img_array = np.array(input_path)
        # Line processing
        blurred = cv2.GaussianBlur(img_array, (7, 7), 0)
        edges = cv2.Canny(blurred, 90, 150)
        structuring_element = cv2.getStructuringElement(cv2.MORPH_RECT, (2,2))
        dilated_edges = cv2.dilate(edges, structuring_element, iterations=2)
        line_drawing = 255 - dilated_edges
        line_img = Image.fromarray(line_drawing)
        alpha = 0.3
        return freemium_watermark(line_img, sb_logo, alpha)

demo.queue(default_concurrency_limit=10, api_open=False).launch(show_api=False)