from typing import Tuple
import gradio as gr
import numpy as np
import cv2
import SaRa.saraRC1 as sara
import warnings
warnings.filterwarnings("ignore")


ALPHA = 0.3
GENERATORS = ['itti', 'deepgaze']

MARKDOWN = """
<h1 style='text-align: center'>Saliency Ranking: Itti vs. Deepgaze</h1>
"""

IMAGE_EXAMPLES = [
    ['https://media.roboflow.com/supervision/image-examples/people-walking.png', 9],
    ['https://media.roboflow.com/supervision/image-examples/vehicles.png', 9],
    ['https://media.roboflow.com/supervision/image-examples/basketball-1.png', 9],
]

def detect_and_annotate(image: np.ndarray, 
                        GRID_SIZE: int, 
                        generator: str, 
                        ALPHA:float =ALPHA)-> np.ndarray:
    # Convert image from BGR to RGB
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Copy and convert the image for sara processing
    sara_image = image.copy()
    sara_image = cv2.cvtColor(sara_image, cv2.COLOR_RGB2BGR)

    # Resetting sara
    sara.reset()

    # Running sara (Original implementation on itti)
    sara_info = sara.return_sara(sara_image, GRID_SIZE, generator, mode=1)

    # Generate saliency map
    saliency_map = sara.return_saliency(image, generator=generator)
    # Resize saliency map to match the image size
    saliency_map = cv2.resize(saliency_map, (image.shape[1], image.shape[0]))

    # Apply color map and convert to RGB
    saliency_map = cv2.applyColorMap(saliency_map, cv2.COLORMAP_JET)
    saliency_map = cv2.cvtColor(saliency_map, cv2.COLOR_BGR2RGB)

    # Overlay the saliency map on the original image
    saliency_map = cv2.addWeighted(saliency_map, ALPHA, image, 1-ALPHA, 0)

    # Extract and convert heatmap to RGB
    heatmap = sara_info[0]
    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)

    return saliency_map, heatmap

def process_image(
    input_image: np.ndarray,
    GRIDSIZE: int,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    # Validate GRID_SIZE
    if GRIDSIZE is None and GRIDSIZE < 4:
        GRIDSIZE = 9

    itti_saliency_map, itti_heatmap = detect_and_annotate(
        input_image, sara, GRIDSIZE, 'itti')
    deepgaze_saliency_map, deepgaze_heatmap = detect_and_annotate(
        input_image, sara, GRIDSIZE, 'deepgaze')

    return (
        itti_saliency_map,
        itti_heatmap,
        deepgaze_saliency_map,
        deepgaze_heatmap,
    )

grid_size_Component = gr.Slider(
    minimum=4,
    maximum=70,
    value=9,
    step=1,
    label="Grid Size",
    info=(
        "The grid size for the Saliency Ranking (SaRa) model. The grid size determines "
        "the number of regions the image is divided into. A higher grid size results in "
        "more regions and a lower grid size results in fewer regions. The default grid "
        "size is 9."
    ))


with gr.Blocks() as demo:
    gr.Markdown(MARKDOWN)
    with gr.Accordion("Configuration", open=False):
        with gr.Row():
            grid_size_Component.render()
    with gr.Row():
        input_image_component = gr.Image(
            type='pil',
            label='Input'
        )
    with gr.Row():
        itti_saliency_map = gr.Image(
            type='pil',
            label='Itti Saliency Map'
        )
        itti_heatmap = gr.Image(
            type='pil',
            label='Itti Saliency Ranking Heatmap'
        )
    with gr.Row():
        deepgaze_saliency_map = gr.Image(
            type='pil',
            label='DeepGaze Saliency Map'
        )
        deepgaze_heatmap = gr.Image(
            type='pil',
            label='DeepGaze Saliency Ranking Heatmap'
        )
    submit_button_component = gr.Button(
        value='Submit',
        scale=1,
        variant='primary'
    )
    gr.Examples(
        fn=process_image,
        examples=IMAGE_EXAMPLES,
        inputs=[
            input_image_component,
            grid_size_Component,
        ],
        outputs=[
            itti_saliency_map,
            itti_heatmap,
            deepgaze_saliency_map,
            deepgaze_heatmap,
        ]
    )

    submit_button_component.click(
        fn=process_image,
        inputs=[
            input_image_component,
            grid_size_Component,
        ],
        outputs=[
            itti_saliency_map,
            itti_heatmap,
            deepgaze_saliency_map,
            deepgaze_heatmap,
        ]
    )

demo.launch(debug=False, show_error=True, max_threads=1)