from cProfile import label
import gradio as gr
import random
import numpy as np
import os
import requests
import torch
import torchvision.transforms as T
from PIL import Image
from transformers import AutoProcessor, AutoModelForVision2Seq
import cv2
import ast
import re
from datetime import datetime
from datetime import datetime
import time
import tqdm
import argparse

parser = argparse.ArgumentParser()

parser.add_argument("--share", action="store_true", help="Enable Gradio app sharing")
parser.add_argument("--bit4", action="store_true", help="Start 4 Bit")
parser.add_argument("--bit8", action="store_true", help="Start 8 Bit")
parser.add_argument("--bit16", action="store_true", help="Start 16 Bit - Half Precision")
args = parser.parse_args()

bit_4=False
bit_8=False
bit_16=False

if(args.bit4):
	bit_4=True
	
if(args.bit8):
	bit_8=True

if(args.bit16):
	bit_16=True
	
status_message = ""
status_message_v2 = gr.State("")

colors = [
    (0, 255, 0),
    (0, 0, 255),
    (255, 255, 0),
    (255, 0, 255),
    (0, 255, 255),
    (114, 128, 250),
    (0, 165, 255),
    (0, 128, 0),
    (144, 238, 144),
    (238, 238, 175),
    (255, 191, 0),
    (0, 128, 0),
    (226, 43, 138),
    (255, 0, 255),
    (0, 215, 255),
    (255, 0, 0),
]

color_map = {
    f"{color_id}": f"#{hex(color[2])[2:].zfill(2)}{hex(color[1])[2:].zfill(2)}{hex(color[0])[2:].zfill(2)}" for color_id, color in enumerate(colors)
}


def is_overlapping(rect1, rect2):
    x1, y1, x2, y2 = rect1
    x3, y3, x4, y4 = rect2
    return not (x2 < x3 or x1 > x4 or y2 < y3 or y1 > y4)


def draw_entity_boxes_on_image(image, entities, show=False, save_path=None, entity_index=-1):
    """_summary_
    Args:
        image (_type_): image or image path
        collect_entity_location (_type_): _description_
    """
    if isinstance(image, Image.Image):
        image_h = image.height
        image_w = image.width
        image = np.array(image)[:, :, [2, 1, 0]]
    elif isinstance(image, str):
        if os.path.exists(image):
            pil_img = Image.open(image).convert("RGB")
            image = np.array(pil_img)[:, :, [2, 1, 0]]
            image_h = pil_img.height
            image_w = pil_img.width
        else:
            raise ValueError(f"invaild image path, {image}")
    elif isinstance(image, torch.Tensor):
        # pdb.set_trace()
        image_tensor = image.cpu()
        reverse_norm_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073])[:, None, None]
        reverse_norm_std = torch.tensor([0.26862954, 0.26130258, 0.27577711])[:, None, None]
        image_tensor = image_tensor * reverse_norm_std + reverse_norm_mean
        pil_img = T.ToPILImage()(image_tensor)
        image_h = pil_img.height
        image_w = pil_img.width
        image = np.array(pil_img)[:, :, [2, 1, 0]]
    else:
        raise ValueError(f"invaild image format, {type(image)} for {image}")

    if len(entities) == 0:
        return image

    indices = list(range(len(entities)))
    if entity_index >= 0:
        indices = [entity_index]

    # Not to show too many bboxes
    entities = entities[:len(color_map)]

    new_image = image.copy()
    previous_bboxes = []
    # size of text
    text_size = 1
    # thickness of text
    text_line = 1  # int(max(1 * min(image_h, image_w) / 512, 1))
    box_line = 3
    (c_width, text_height), _ = cv2.getTextSize("F", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)
    base_height = int(text_height * 0.675)
    text_offset_original = text_height - base_height
    text_spaces = 3

    # num_bboxes = sum(len(x[-1]) for x in entities)
    used_colors = colors  # random.sample(colors, k=num_bboxes)

    color_id = -1
    for entity_idx, (entity_name, (start, end), bboxes) in enumerate(entities):
        color_id += 1
        if entity_idx not in indices:
            continue
        for bbox_id, (x1_norm, y1_norm, x2_norm, y2_norm) in enumerate(bboxes):
            # if start is None and bbox_id > 0:
            #     color_id += 1
            orig_x1, orig_y1, orig_x2, orig_y2 = int(x1_norm * image_w), int(y1_norm * image_h), int(x2_norm * image_w), int(y2_norm * image_h)

            # draw bbox
            # random color
            color = used_colors[color_id]  # tuple(np.random.randint(0, 255, size=3).tolist())
            new_image = cv2.rectangle(new_image, (orig_x1, orig_y1), (orig_x2, orig_y2), color, box_line)

            l_o, r_o = box_line // 2 + box_line % 2, box_line // 2 + box_line % 2 + 1

            x1 = orig_x1 - l_o
            y1 = orig_y1 - l_o

            if y1 < text_height + text_offset_original + 2 * text_spaces:
                y1 = orig_y1 + r_o + text_height + text_offset_original + 2 * text_spaces
                x1 = orig_x1 + r_o

            # add text background
            (text_width, text_height), _ = cv2.getTextSize(f"  {entity_name}", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)
            text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2 = x1, y1 - (text_height + text_offset_original + 2 * text_spaces), x1 + text_width, y1

            for prev_bbox in previous_bboxes:
                while is_overlapping((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2), prev_bbox):
                    text_bg_y1 += (text_height + text_offset_original + 2 * text_spaces)
                    text_bg_y2 += (text_height + text_offset_original + 2 * text_spaces)
                    y1 += (text_height + text_offset_original + 2 * text_spaces)

                    if text_bg_y2 >= image_h:
                        text_bg_y1 = max(0, image_h - (text_height + text_offset_original + 2 * text_spaces))
                        text_bg_y2 = image_h
                        y1 = image_h
                        break

            alpha = 0.5
            for i in range(text_bg_y1, text_bg_y2):
                for j in range(text_bg_x1, text_bg_x2):
                    if i < image_h and j < image_w:
                        if j < text_bg_x1 + 1.35 * c_width:
                            # original color
                            bg_color = color
                        else:
                            # white
                            bg_color = [255, 255, 255]
                        new_image[i, j] = (alpha * new_image[i, j] + (1 - alpha) * np.array(bg_color)).astype(np.uint8)

            cv2.putText(
                new_image, f"  {entity_name}", (x1, y1 - text_offset_original - 1 * text_spaces), cv2.FONT_HERSHEY_COMPLEX, text_size, (0, 0, 0), text_line, cv2.LINE_AA
            )
            # previous_locations.append((x1, y1))
            previous_bboxes.append((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2))

    pil_image = Image.fromarray(new_image[:, :, [2, 1, 0]])
    if save_path:
        pil_image = pil_image.convert("RGB")
        pil_image.save(save_path)
    if show:
        pil_image.show()

    return pil_image

def update_status_label(status):
    return status

torch_type=torch.float32

if bit_16:
    torch_type=torch.float16

if bit_8:
    torch_type=torch.float16

def main():

    model_source = "microsoft/kosmos-2-patch14-224"

    model = AutoModelForVision2Seq.from_pretrained(model_source, 
	load_in_4bit=bit_4,
        load_in_8bit=bit_8,torch_dtype=torch_type)
    processor = AutoProcessor.from_pretrained(model_source,
	        load_in_4bit=bit_4,
        load_in_8bit=bit_8,torch_dtype=torch_type)

    if not bit_4 and not bit_8:
        model = model.to("cuda")

    def batch_caption_images(folder_path, text_input, status_label, progress=gr.Progress(), skip_existing=False):
        if not os.path.isdir(folder_path):
            return "Invalid folder path"
        valid_extensions = ('.png', '.jpg', '.jpeg')
        global status_message
        images = [f for f in os.listdir(folder_path) if f.lower().endswith(valid_extensions)]
        total_images = len(images)
        captions = {}
        start_time = time.time()
        counter = 1.0
        for idx, img_name in enumerate(images):
            caption_file_name = os.path.splitext(img_name)[0] + '.txt'
            caption_file_path = os.path.join(folder_path, caption_file_name)

            # Skip captioning if the .txt file already exists and skip_existing is True
            if skip_existing and os.path.exists(caption_file_path):
                continue

            img_path = os.path.join(folder_path, img_name)
            img = Image.open(img_path)
            caption = generate_predictions(img, text_input, is_batch=True)
            caption = caption.replace(text_input, "").strip()
            captions[img_name] = caption

            # Save individual caption to a text file
            with open(caption_file_path, 'w', encoding='utf-8') as caption_file:
                caption_file.write(caption)

            # Update Gradio label
            processed_images = idx + 1
            remaining_images = total_images - processed_images
            average_time_per_image = (time.time() - start_time) / processed_images
            estimated_time_remaining = average_time_per_image * remaining_images
            status_message = f"Processed {processed_images}/{total_images} images, " \
                             f"{remaining_images} images left, " \
                             f"AVG processing time: {average_time_per_image:.2f} seconds, " \
                             f"ETA: {estimated_time_remaining:.2f} seconds"
            print(status_message)
            
            progress((idx + 1) / total_images, desc=status_message)

        timestamp = datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
        batch_file_name = f"{timestamp}_batch_list.txt"
        with open(os.path.join(folder_path, batch_file_name), 'w', encoding='utf-8') as f:
            for img_name, caption in captions.items():
                f.write(f"{img_name}: {caption}\n")

        return "Batch Captioning Completed"


    def strip_tags(text, init_text):
        init_text=init_text.replace("<grounding>","")
        # Remove all content within <image> and </image> tags
        text = re.sub(r'<image>.*?</image>', '', text, flags=re.DOTALL)

        # Check if init_text is in text
        init_text_index = text.find(init_text)
        if init_text_index != -1:
            # Split from init_text and discard the initial part
            text = text[init_text_index + len(init_text):]
        else:
            # Split from first <phrase> and discard the previous part
            phrase_index = text.find('<phrase>')
            if phrase_index != -1:
                text = text[phrase_index:]

        # Remove all parts starting with < and ending with >
        text = re.sub(r'<.*?>', '', text)

        return text.strip()

    def generate_predictions(image_input, text_input, is_batch=False, progress=gr.Progress()):
        start_time = time.time()  # Start timing

        user_image_path = "user_input_test_image.png"
        if image_input.mode != 'RGBA':
            image_input = image_input.convert('RGBA')  # Convert to RGBA if it's not already

        image_input.save(user_image_path)
        image_input = Image.open(user_image_path)

        text_input = f"<grounding>{text_input}"

        inputs = processor(text=text_input, images=image_input, return_tensors="pt").to("cuda")

        generated_ids = model.generate(
            pixel_values=inputs["pixel_values"],
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            image_embeds=None,
            image_embeds_position_mask=inputs["image_embeds_position_mask"],
            use_cache=True,
            max_new_tokens=512,
        )

        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

        if is_batch and isinstance(generated_text, str) and isinstance(text_input, str):
            return (strip_tags(generated_text, text_input))

        processed_text, entities = processor.post_process_generation(generated_text)

        annotated_image = draw_entity_boxes_on_image(image_input, entities, show=False)

        color_id = -1
        entity_info = []
        filtered_entities = []
        for entity in entities:
            entity_name, (start, end), bboxes = entity
            if start == end:
                continue
            color_id += 1
            entity_info.append(((start, end), color_id))
            filtered_entities.append(entity)

        colored_text = []
        prev_start = 0
        end = 0
        for idx, ((start, end), color_id) in enumerate(entity_info):
            if start > prev_start:
                colored_text.append((processed_text[prev_start:start], None))
            colored_text.append((processed_text[start:end], f"{color_id}"))
            prev_start = end

        if end < len(processed_text):
            colored_text.append((processed_text[end:len(processed_text)], None))

        duration = time.time() - start_time  # Calculate duration
        formatted_duration = round(duration, 2)  # Format to 2 decimal places
        image_size = image_input.size[0] * image_input.size[1]  # Calculate image size (width * height)
        return annotated_image, colored_text, str(filtered_entities), f"{formatted_duration} second/image"  # Return duration and speed


    term_of_use = """
    """

    with gr.Blocks(title="Kosmos-2", theme=gr.themes.Base()).queue() as demo:

        gr.Markdown(("""
            # Kosmos-2: Grounding Multimodal Large Language Models to the World
            [[Paper]](https://arxiv.org/abs/2306.14824) [[Code]](https://github.com/microsoft/unilm/blob/master/kosmos-2)
            """))
        with gr.Row():
            with gr.Row():
                with gr.Column():
                    batch_folder_path = gr.Textbox(label='Batch Captioning Folder Path')
                    skip_existing_checkbox = gr.Checkbox(label="Skip existing captions", value=False)
            with gr.Column():
                    batch_run_button = gr.Button("Batch Caption Images Uses Selected Brief or Detailed")
            

            status_label = gr.Label() 

        with gr.Row():
            with gr.Column():
                image_input = gr.Image(type="pil", label="Test Image", height=500)
                text_input = gr.Textbox(label="Enter your prompt (Default prompt works good)", value="Describe the image in details and tell what is in the image:")

                run_button = gr.Button("Caption Image")

            with gr.Column():
                image_output = gr.Image(type="pil", height=500)
                text_output1 = gr.HighlightedText(
                                    label="Generated Description",
                                    combine_adjacent=False,
                                    show_legend=True,
                                )

        gr.Markdown(term_of_use)

        # record which text span (label) is selected
        selected = gr.Number(-1, show_label=False, visible=False)

        # record the current `entities`
        entity_output = gr.Textbox(visible=False)

        # get the current selected span label
        def get_text_span_label(evt: gr.SelectData):
            if evt.value[-1] is None:
                return -1
            return int(evt.value[-1])
        # and set this information to `selected`
        text_output1.select(get_text_span_label, None, selected)

        # update output image when we change the span (enity) selection
        def update_output_image(img_input, image_output, entities, idx):
            entities = ast.literal_eval(entities)
            updated_image = draw_entity_boxes_on_image(img_input, entities, entity_index=idx)
            return updated_image
        selected.change(update_output_image, [image_input, image_output, entity_output, selected], [image_output])

        run_button.click(fn=generate_predictions,
                         inputs=[image_input, text_input],
                         outputs=[image_output, text_output1, entity_output, gr.Label()],  # Add labels for duration and speed
                         show_progress=True, queue=True)
        batch_run_button.click(fn=batch_caption_images,
                               inputs=[batch_folder_path, text_input, status_label, skip_existing_checkbox],
                               outputs=status_label,
                               show_progress=True, queue=True)

    demo.queue().launch(share=args.share,inbrowser=True)


if __name__ == "__main__":
    main()
    # trigger