import os
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from PIL import Image
import gradio as gr
from transformers import AutoProcessor, AutoModelForCausalLM
import torch
import numpy as np
import spaces
import subprocess
from io import BytesIO

# Ensure flash-attn is installed correctly
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

# Initialize Florence-2-large model and processor
model_id = 'microsoft/Florence-2-large'
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to("cuda").eval()
processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)

# Function to resize and preprocess image
def preprocess_image(image_path, max_size=(800, 800)):
    image = Image.open(image_path).convert('RGB')
    if image.size[0] > max_size[0] or image.size[1] > max_size[1]:
        image.thumbnail(max_size, Image.LANCZOS)

    # Convert image to numpy array
    image_np = np.array(image, dtype=np.float32)  # Ensure the array is float32

    # Ensure the image is in the format [height, width, channels]
    if image_np.ndim == 2:  # Grayscale image
        image_np = np.expand_dims(image_np, axis=-1)
    elif image_np.shape[0] == 3:  # Image in [channels, height, width] format
        image_np = np.transpose(image_np, (1, 2, 0))

    return image_np, image.size

# Function to run Florence-2-large model
@spaces.GPU
def run_florence_model(image_np, image_size, task_prompt, text_input=None):
    if text_input is None:
        prompt = task_prompt
    else:
        prompt = task_prompt + text_input

    inputs = processor(text=prompt, images=image_np, return_tensors="pt")

    with torch.no_grad():
        outputs = model.generate(
            input_ids=inputs["input_ids"].cuda(),
            pixel_values=inputs["pixel_values"].cuda(),
            max_new_tokens=1024,
            early_stopping=False,
            do_sample=False,
            num_beams=3,
        )

    generated_text = processor.batch_decode(outputs, skip_special_tokens=False)[0]
    parsed_answer = processor.post_process_generation(
        generated_text,
        task=task_prompt,
        image_size=image_size
    )

    return parsed_answer, generated_text

# Function to plot image with bounding boxes
def plot_image_with_bboxes(image_np, bboxes, labels=None):
    fig, ax = plt.subplots(1)
    ax.imshow(image_np / 255.0)  # Normalize the image for plotting
    colors = ['red', 'blue', 'green', 'yellow', 'purple', 'cyan']
    for i, bbox in enumerate(bboxes):
        color = colors[i % len(colors)]
        x, y, width, height = bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]
        rect = patches.Rectangle((x, y), width, height, linewidth=2, edgecolor=color, facecolor='none')
        ax.add_patch(rect)
        if labels and i < len(labels):
            ax.text(x, y, labels[i], color=color, fontsize=8, bbox=dict(facecolor='white', alpha=0.7))
    plt.axis('off')
    
    # Save the plot to a BytesIO object
    buf = BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    buf.seek(0)
    
    # Convert the BytesIO buffer to PIL Image
    pil_image = Image.open(buf)
    return pil_image

# Gradio function to process uploaded images
@spaces.GPU
def process_image(image_path):
    image_np, image_size = preprocess_image(image_path)
    
    # Convert image_np to float32
    image_np = image_np.astype(np.float32)

    # Image Captioning
    caption_result, _ = run_florence_model(image_np, image_size, '<CAPTION>')
    detailed_caption_result, _ = run_florence_model(image_np, image_size, '<DETAILED_CAPTION>')

    # Object Detection
    od_result, _ = run_florence_model(image_np, image_size, '<OD>')
    od_bboxes = od_result['<OD>'].get('bboxes', [])
    od_labels = od_result['<OD>'].get('labels', [])

    # OCR
    ocr_result, _ = run_florence_model(image_np, image_size, '<OCR>')

    # Phrase Grounding
    pg_result, _ = run_florence_model(image_np, image_size, '<CAPTION_TO_PHRASE_GROUNDING>', text_input=caption_result['<CAPTION>'])
    pg_bboxes = pg_result['<CAPTION_TO_PHRASE_GROUNDING>'].get('bboxes', [])
    pg_labels = pg_result['<CAPTION_TO_PHRASE_GROUNDING>'].get('labels', [])

    # Cascaded Tasks (Detailed Caption + Phrase Grounding)
    cascaded_result, _ = run_florence_model(image_np, image_size, '<CAPTION_TO_PHRASE_GROUNDING>', text_input=detailed_caption_result['<DETAILED_CAPTION>'])
    cascaded_bboxes = cascaded_result['<CAPTION_TO_PHRASE_GROUNDING>'].get('bboxes', [])
    cascaded_labels = cascaded_result['<CAPTION_TO_PHRASE_GROUNDING>'].get('labels', [])

    # Create plots
    od_fig = plot_image_with_bboxes(image_np, od_bboxes, od_labels)
    pg_fig = plot_image_with_bboxes(image_np, pg_bboxes, pg_labels)
    cascaded_fig = plot_image_with_bboxes(image_np, cascaded_bboxes, cascaded_labels)

    # Prepare response
    response = f"""
    Image Captioning:
    - Simple Caption: {caption_result['<CAPTION>']}
    - Detailed Caption: {detailed_caption_result['<DETAILED_CAPTION>']}
    Object Detection:
    - Detected {len(od_bboxes)} objects
    OCR:
    {ocr_result['<OCR>']}
    Phrase Grounding:
    - Grounded {len(pg_bboxes)} phrases from the simple caption
    Cascaded Tasks:
    - Grounded {len(cascaded_bboxes)} phrases from the detailed caption
    """

    return response, od_fig, pg_fig, cascaded_fig

# Gradio interface
with gr.Blocks(theme='NoCrypt/miku') as demo:
    gr.Markdown("""
    # Image Processing with Florence-2-large
    Upload an image to perform image captioning, object detection, OCR, phrase grounding, and cascaded tasks.
    """)

    image_input = gr.Image(type="filepath")
    text_output = gr.Textbox()
    plot_output_1 = gr.Image()
    plot_output_2 = gr.Image()
    plot_output_3 = gr.Image()

    image_input.upload(process_image, inputs=[image_input], outputs=[text_output, plot_output_1, plot_output_2, plot_output_3])

    footer = """
    <div style="text-align: center; margin-top: 20px;">
        <a href="https://www.linkedin.com/in/pejman-ebrahimi-4a60151a7/" target="_blank">LinkedIn</a> |
        <a href="https://github.com/arad1367" target="_blank">GitHub</a> |
        <a href="https://arad1367.pythonanywhere.com/" target="_blank">Live demo of my PhD defense</a>
        <br>
        Made with 💖 by Pejman Ebrahimi
    </div>
    """
    gr.HTML(footer)

demo.launch()