import torch
from PIL import Image
import requests
from openai import OpenAI
from transformers import (Owlv2Processor, Owlv2ForObjectDetection,
                          AutoProcessor, AutoModelForMaskGeneration)
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import base64
import io
import numpy as np
import gradio as gr
import json
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')


def encode_image_to_base64(image):
    print(f"Encode image type: {type(image)}")  # Debug print
    
    try:
        # If image is a tuple (as sometimes provided by Gradio), take the first element
        if isinstance(image, tuple):
            print(f"Image is tuple with length: {len(image)}")  # Debug print
            if len(image) > 0 and image[0] is not None:
                if isinstance(image[0], np.ndarray):
                    image = Image.fromarray(image[0])
                else:
                    image = image[0]
            else:
                raise ValueError("Invalid image tuple provided")

        # If image is a numpy array, convert to PIL Image
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        # If image is a path string, open it
        elif isinstance(image, str):
            image = Image.open(image)

        print(f"Image type after conversion: {type(image)}")  # Debug print

        # Ensure image is in PIL Image format
        if not isinstance(image, Image.Image):
            raise ValueError(f"Input must be a PIL Image, numpy array, or valid image path. Got {type(image)}")

        # Convert image to RGB if it's in RGBA mode
        if image.mode == 'RGBA':
            image = image.convert('RGB')

        buffered = io.BytesIO()
        image.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode('utf-8')
    except Exception as e:
        print(f"Encode error details: {str(e)}")  # Debug print
        raise
        

def analyze_image(image):
    client = OpenAI(api_key=OPENAI_API_KEY)
    base64_image = encode_image_to_base64(image)

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": """Your task is to determine if the image is surprising or not surprising.    
                    if the image is surprising, determine which element, figure or object in the image is making the image surprising and write it only in one sentence with no more then 6 words, otherwise, write 'NA'.    
                    Also rate how surprising the image is on a scale of 1-5, where 1 is not surprising at all and 5 is highly surprising.
                    Provide the response as a JSON with the following structure:    
                    {
                        "label": "[surprising OR not surprising]",
                        "element": "[element]",
                        "rating": [1-5]
                    }"""
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{base64_image}"
                    }
                }
            ]
        }
    ]

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        max_tokens=100,
        temperature=0.1,
        response_format={
            "type": "json_object"
        }
    )

    return response.choices[0].message.content


def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([1.0, 0.0, 0.0, 0.5])

    if len(mask.shape) == 4:
        mask = mask[0, 0]

    mask_image = np.zeros((*mask.shape, 4), dtype=np.float32)
    mask_image[mask > 0] = color

    ax.imshow(mask_image)


def process_image_detection(image, target_label, surprise_rating):
    # Handle different image input types
    if isinstance(image, tuple):
        if len(image) > 0 and image[0] is not None:
            image = Image.fromarray(image[0])
        else:
            raise ValueError("Invalid image tuple provided")
    elif isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    elif isinstance(image, str):
        image = Image.open(image)
    
    # Ensure image is in PIL Image format
    if not isinstance(image, Image.Image):
        raise ValueError("Input must be a PIL Image, numpy array, or valid image path")

    # Ensure image is in RGB mode
    if image.mode != 'RGB':
        image = image.convert('RGB')

    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Get original image DPI and size
    original_dpi = image.info.get('dpi', (72, 72))
    original_size = image.size

    # Calculate relative font size based on image dimensions
    base_fontsize = min(original_size) / 40

    owlv2_processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16")
    owlv2_model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16").to(device)

    sam_processor = AutoProcessor.from_pretrained("facebook/sam-vit-base")
    sam_model = AutoModelForMaskGeneration.from_pretrained("facebook/sam-vit-base").to(device)

    image_np = np.array(image)

    inputs = owlv2_processor(text=[target_label], images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = owlv2_model(**inputs)

    target_sizes = torch.tensor([image.size[::-1]]).to(device)
    results = owlv2_processor.post_process_object_detection(outputs, target_sizes=target_sizes)[0]

    dpi = 300
    figsize = (original_size[0] / dpi, original_size[1] / dpi)
    fig = plt.figure(figsize=figsize, dpi=dpi)

    ax = plt.Axes(fig, [0., 0., 1., 1.])
    fig.add_axes(ax)

    plt.imshow(image)

    scores = results["scores"]
    if len(scores) > 0:
        max_score_idx = scores.argmax().item()
        max_score = scores[max_score_idx].item()

        if max_score > 0.2:
            box = results["boxes"][max_score_idx].cpu().numpy()

            sam_inputs = sam_processor(
                image,
                input_boxes=[[[box[0], box[1], box[2], box[3]]]],
                return_tensors="pt"
            ).to(device)

            with torch.no_grad():
                sam_outputs = sam_model(**sam_inputs)

            masks = sam_processor.image_processor.post_process_masks(
                sam_outputs.pred_masks.cpu(),
                sam_inputs["original_sizes"].cpu(),
                sam_inputs["reshaped_input_sizes"].cpu()
            )

            mask = masks[0].numpy() if isinstance(masks[0], torch.Tensor) else masks[0]
            show_mask(mask, ax=ax)

            # Draw rectangle with increased line width
            rect = patches.Rectangle(
                (box[0], box[1]),
                box[2] - box[0],
                box[3] - box[1],
                linewidth=max(2, min(original_size) / 500),
                edgecolor='red',
                facecolor='none'
            )
            ax.add_patch(rect)

            # Add confidence score with improved visibility
            plt.text(
                box[0], box[1] - base_fontsize,
                f'{max_score:.2f}',
                color='red',
                fontsize=base_fontsize,
                fontweight='bold',
                bbox=dict(facecolor='white', alpha=0.7, edgecolor='none', pad=2)
            )

            # Add label and rating with improved visibility
            plt.text(
                box[2] + base_fontsize / 2, box[1],
                f'Unexpected (Rating: {surprise_rating}/5)\n{target_label}',
                color='red',
                fontsize=base_fontsize,
                fontweight='bold',
                bbox=dict(facecolor='white', alpha=0.7, edgecolor='none', pad=2),
                verticalalignment='bottom'
            )

    plt.axis('off')

    # Save with high DPI
    buf = io.BytesIO()
    plt.savefig(buf,
                format='png',
                dpi=dpi,
                bbox_inches='tight',
                pad_inches=0,
                metadata={'dpi': original_dpi})
    buf.seek(0)
    plt.close()

    # Process final image
    output_image = Image.open(buf)
    output_image = output_image.resize(original_size, Image.Resampling.LANCZOS)

    final_buf = io.BytesIO()
    output_image.save(final_buf, format='PNG', dpi=original_dpi)
    final_buf.seek(0)

    return final_buf


def process_and_analyze(image):
    if image is None:
        return None, "Please upload an image first."

    print(f"Initial image type: {type(image)}")  # Debug print

    if OPENAI_API_KEY is None:
        return None, "OpenAI API key not found in environment variables."

    try:
        # Convert the image to PIL format if needed
        if isinstance(image, tuple):
            print(f"Image is tuple, length: {len(image)}")  # Debug print
            if len(image) > 0 and image[0] is not None:
                if isinstance(image[0], np.ndarray):
                    image = Image.fromarray(image[0])
                else:
                    print(f"First element type: {type(image[0])}")  # Debug print
                    image = image[0]
            else:
                return None, "Invalid image format provided"
        elif isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        elif isinstance(image, str):
            image = Image.open(image)
        
        print(f"Image type after conversion: {type(image)}")  # Debug print

        if not isinstance(image, Image.Image):
            return None, f"Invalid image format: {type(image)}"

        # Ensure image is in RGB mode
        if image.mode != 'RGB':
            image = image.convert('RGB')

        # Analyze image
        print("Starting GPT analysis...")  # Debug print
        gpt_response = analyze_image(image)
        print(f"GPT response: {gpt_response}")  # Debug print
        
        try:
            response_data = json.loads(gpt_response)
        except json.JSONDecodeError:
            return None, "Error: Invalid response format from GPT"

        if not all(key in response_data for key in ["label", "element", "rating"]):
            return None, "Error: Missing required fields in analysis response"

        print(f"Response data: {response_data}")  # Debug print

        if response_data["label"].lower() == "surprising" and response_data["element"].lower() != "na":
            try:
                print("Starting image detection...")  # Debug print
                result_buf = process_image_detection(image, response_data["element"], response_data["rating"])
                result_image = Image.open(result_buf)
                analysis_text = (
                    f"Label: {response_data['label']}\n"
                    f"Element: {response_data['element']}\n"
                    f"Rating: {response_data['rating']}/5"
                )
                return result_image, analysis_text
            except Exception as detection_error:
                print(f"Detection error details: {str(detection_error)}")  # Debug print
                return None, f"Error in image detection processing: {str(detection_error)}"
        else:
            return image, "Not Surprising"

    except Exception as e:
        error_type = type(e).__name__
        error_msg = str(e)
        detailed_error = f"Error ({error_type}): {error_msg}"
        print(detailed_error)  # Debug print
        return None, f"Error processing image: {error_msg}"


# Create Gradio interface
def create_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# Image Surprise Analysis")

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Upload Image")
                analyze_btn = gr.Button("Analyze Image")

            with gr.Column():
                output_image = gr.Image(label="Processed Image")
                output_text = gr.Textbox(label="Analysis Results")

        analyze_btn.click(
            fn=process_and_analyze,
            inputs=[input_image],
            outputs=[output_image, output_text]
        )

    return demo


if __name__ == "__main__":
    demo = create_interface()
    demo.launch()