import gradio as gr
import torch
import boto3
import requests
import re
from bs4 import BeautifulSoup
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
import os
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import hashlib
import spaces # Make sure to import spaces



#First output the thinking process in <think> </think> tags and then output the final answer in <answer> </answer> tags.
#Answer the following question based on the information above and the given image, and provide citations for your response.


# Cache for storing recognition results
recognition_cache = {}

# Function to calculate image hash for caching
def calculate_image_hash(image):
    """Calculate a hash for the image to use as a cache key"""
    if image is None:
        return None
    img_array = np.array(image)
    # Use a simple hash of the image data
    return hashlib.md5(img_array.tobytes()).hexdigest()

# Helper function to remove numbers from brackets
def remove_bracketed_nums(input_string):
    """
    Removes all occurrences of [num] from the input string.
    """
    return re.sub(r'\[\d+\]', '', input_string)

# Function to get Wikipedia links
def fetch_wikipedia_links(entity_id):
    url = f"https://www.wikidata.org/wiki/Special:EntityData/{entity_id}.json"
    links = {}
    try:
        response = requests.get(url)
        response.raise_for_status()
        data = response.json()
        entity_data = data.get("entities", {}).get(entity_id, {})
        if not entity_data:
            print("No data found for the given entity ID.")
            return links

        sitelinks = entity_data.get("sitelinks", {})
        if not sitelinks:
            print("No sitelinks found for the given entity ID.")
            return links

        wikipedia_links = {}
        for site, site_info in sitelinks.items():
            if site.endswith("wiki"):
                language = site.replace("wiki", "")
                url = site_info.get("url")
                wikipedia_links[language] = url

        for lang, link in wikipedia_links.items():
            links[lang] = link

    except requests.RequestException as e:
        print(f"Error fetching data: {e}")
    return links

# Function to get Wikipedia information
def fetch_wikipedia_info(url):
    try:
        # Send HTTP request
        response = requests.get(url)
        response.raise_for_status()

        # Parse HTML content using BeautifulSoup
        soup = BeautifulSoup(response.content, 'html.parser')
        
        # Try multiple methods to find the title
        title = None
        
        # Method 1: Find title in the <title> tag of the page
        if not title:
            title_tag = soup.find('title')
            if title_tag:
                title = title_tag.text.replace(' - Wikipedia', '').strip()
        
        # Method 2: Find title in h1 tag
        if not title:
            title_tag = soup.find('h1', {'id': 'firstHeading'})
            if title_tag:
                title = title_tag.text.strip()
        
        # Method 3: If previous methods fail, use a more flexible method
        if not title:
            title_tag = soup.find('h1')
            if title_tag:
                title = title_tag.text.strip()
        
        # Extract sections and their content
        sections = {}
        current_section = None
        paragraphs = soup.find_all('p')
        
        # Traverse all headings and paragraphs, mapping sections to content
        for element in soup.find_all(['h2', 'p']):
            if element.name == 'h2':
                # New section starts, get section title
                section_title = element.text.strip()
                current_section = section_title
                sections[current_section] = []
            elif element.name == 'p' and current_section:
                # Add paragraph to current section
                paragraph_text = element.text.strip()
                if paragraph_text:
                    sections[current_section].append(paragraph_text)
        
        # Convert section content to format where each section is connected into a single string
        sections = {section: '\n'.join(content) for section, content in sections.items()}
        
        return {
            'title': title or 'No title found',
            'sections': sections
        }
    
    except requests.RequestException as e:
        return {'error': f'Request failed: {e}'}
    except Exception as e:
        return {'error': f'An error occurred: {e}'}

# Use AWS Rekognition to recognize celebrities in the image
def recognize_celebrities(image_path, confidence_threshold=90):
    client = boto3.client(
        "rekognition",
        aws_access_key_id=os.getenv('aws_access_key_id'),
        aws_secret_access_key=os.getenv('aws_secret_access_key'),
        region_name='us-east-1'
    )
    
    with open(image_path, 'rb') as image:
        response = client.recognize_celebrities(Image={'Bytes': image.read()})
    
    # Filter out celebrities with None or low confidence
    filtered_celebs = [cele for cele in response.get('CelebrityFaces', []) 
                      if cele.get('MatchConfidence') is not None and cele.get('MatchConfidence') > confidence_threshold]
    
    names = [cele['Name'] for cele in filtered_celebs]
    bounding_boxes = [cele['Face']['BoundingBox'] for cele in filtered_celebs]
    wikidatas = [cele['Urls'][0] for cele in filtered_celebs]
    
    return names, bounding_boxes, wikidatas

# Draw bounding boxes on the image
def draw_bounding_boxes(image_path, bounding_boxes, names):
    img = Image.open(image_path)
    draw = ImageDraw.Draw(img)
    
    width, height = img.size
    
    for i, bbox in enumerate(bounding_boxes):
        left = int(bbox['Left'] * width)
        top = int(bbox['Top'] * height)
        right = int((bbox['Left'] + bbox['Width']) * width)
        bottom = int((bbox['Top'] + bbox['Height']) * height)
        
        # Draw rectangle
        draw.rectangle([(left, top), (right, bottom)], outline="red", width=3)
        
        # Add name label
        text = f"[{i+1}]: {names[i]}"
        font = ImageFont.truetype("./Arial.ttf", 20)  # Adjust font and size as needed
        text_bbox = draw.textbbox((left, top - 20), text, font=font)
        draw.rectangle(text_bbox, fill="white")
        draw.text((left, top - 20), text, fill="red", font=font)
    return np.array(img)

# Resize the image to a smaller size for display
def resize_image_for_display(img_array, max_width=500):
    img = Image.fromarray(img_array)
    width, height = img.size
    if width > max_width:
        ratio = max_width / width
        new_height = int(height * ratio)
        img = img.resize((max_width, new_height), Image.LANCZOS)
    return np.array(img)

model = Qwen2VLForConditionalGeneration.from_pretrained(
    "weihongliang/RC-Qwen2VL-7b", 
    torch_dtype=torch.float16, 
    device_map="auto"
)
processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")

# Use Qwen model for Q&A
def qwen_qa(image_path, question, names, bounding_boxes, en_wiki_pedia_links, en_intros):
    # Prepare prompt text
    prompt = '\n'.join([
        f"[{i+1}]: The information of the person located at <|box_start|>({int(bbox['Left']*1000)},{int(bbox['Top']*1000)}),"
        f"({int((bbox['Left']+bbox['Width'])*1000)},{int((bbox['Top']+bbox['Height'])*1000)})<|box_end|> in the image: {intro}"
        for i, (name, intro, bbox) in enumerate(zip(
            names, 
            [remove_bracketed_nums(intro).replace('\n','') for intro in en_intros], 
            bounding_boxes
        ))
    ])
    print(prompt)
    # Prepare messages
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": image_path,
                },
                {"type": "text", "text": f"{prompt}\nAnswer the following question based on the information above and the given image, and provide citations for your response.\n{question}"},
            ],
        }
    ]
    
    # Prepare inference
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)

    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=None,
        padding=True,
        return_tensors="pt",
    )
    
    inputs = inputs.to("cuda")
    
    # Inference: generate output
    generated_ids = model.generate(**inputs, max_new_tokens=4096)
    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=False, clean_up_tokenization_spaces=False
    )[0]
    print(output_text)
    return output_text, en_wiki_pedia_links, prompt

# Compare image with examples by filename or content when available
def is_example_image(image, examples):
    if isinstance(image, str) and os.path.exists(image):
        # If we have a path, try to match by filename
        image_name = os.path.basename(image)
        for example in examples:
            example_name = os.path.basename(example[0])
            if image_name == example_name:
                print(f"Matched example by filename: {image_name}")
                return True, example
    
    # If we couldn't match by filename or if image is an object, return False
    return False, None

# Main processing function
# This is the function that will run on the GPU
@spaces.GPU
def process_image(image, question, confidence_threshold, examples=None):
    if image is None:
        return None, "Please upload an image", []
    
    # Calculate image hash for caching
    image_hash = calculate_image_hash(image)
    
    # Check if we already have cached recognition results for this image
    cache_key = f"{image_hash}_{confidence_threshold}"
    if cache_key in recognition_cache:
        print("Using cached recognition results")
        names, bounding_boxes, en_wiki_pedia_links, en_intros, result_image = recognition_cache[cache_key]
    else:
        print("Performing new recognition")
        # Save uploaded image
        temp_image_path = "temp_image.jpg"
        image.save(temp_image_path)
        
        # Check if this image matches one of our example images
        # First, get original filename from gradio image object if available
        original_filename = None
        if hasattr(image, 'name'):
            original_filename = image.name
        
        # For each example, check if the file exists and compare
        example_match = None
        if examples:
            for example in examples:
                example_path = example[0]
                if os.path.exists(example_path):
                    example_img = Image.open(example_path)
                    # Compare images by size first (quick check)
                    if example_img.size == image.size:
                        # More thorough comparison - hash the pixels
                        example_hash = calculate_image_hash(example_img)
                        if example_hash == image_hash:
                            example_match = example
                            break
                    example_img.close()
        
        if example_match:
            print("Using example data instead of recognition")
            names = example_match[2]
            bounding_boxes = example_match[3]
            wikidatas = example_match[4]
            print(f"Example data: {names}, {bounding_boxes}, {wikidatas}")
        else:
            # Recognize celebrities with adjusted confidence threshold
            names, bounding_boxes, wikidatas = recognize_celebrities(temp_image_path, confidence_threshold)
        
        print(f'识别_{names}_{bounding_boxes}_{wikidatas}')
        en_wiki_pedia_links = []
        en_intros = []
        
        for wikidata_url in wikidatas:
            entity_id = wikidata_url.split('/')[-1]
            links = fetch_wikipedia_links(entity_id)
            
            if 'en' in links:
                en_link = links['en']
                en_wiki_pedia_links.append(en_link)
                
                wiki_info = fetch_wikipedia_info(en_link)
                # Try to get 'Contents' section, if it doesn't exist, try to get the first section
                intro = wiki_info['sections'].get('Contents', '')
                if not intro and wiki_info['sections']:
                    intro = next(iter(wiki_info['sections'].values()), '')
                
                en_intros.append(intro)
            else:
                print(links)
                en_wiki_pedia_links.append("No English Wikipedia link found")
                en_intros.append(f"This person is {name}.")
        if not names:
            if os.path.exists(temp_image_path):
                os.remove(temp_image_path)
            return None, f"No celebrities could be recognized with confidence threshold {confidence_threshold}%", []
        
        indices = list(range(len(en_intros)))
        
        # Sort indices by the length of en_intros elements
        indices.sort(key=lambda i: len(en_intros[i]))
        
        # Use sorted indices to rearrange the three lists
        names = [names[i] for i in indices]
        bounding_boxes = [bounding_boxes[i] for i in indices]
        en_intros = [en_intros[i] for i in indices]
        en_wiki_pedia_links = [en_wiki_pedia_links[i] for i in indices]
        
        # Draw bounding boxes on the image
        result_image = draw_bounding_boxes(temp_image_path, bounding_boxes, names)
        
        # Cache the recognition results
        recognition_cache[cache_key] = (names, bounding_boxes, en_wiki_pedia_links, en_intros, result_image)
        
        # Clean up temporary files
        if os.path.exists(temp_image_path):
            os.remove(temp_image_path)
    
    # Save image for qwen_qa (it requires a file path)
    temp_image_path = "temp_image_for_qa.jpg"
    Image.fromarray(result_image if isinstance(image, np.ndarray) else np.array(image)).save(temp_image_path)
    
    # Resize the result image for display
    resized_image = resize_image_for_display(result_image, max_width=500)
    
    # Use Qwen model for Q&A
    answer, wiki_links, prompt = qwen_qa(temp_image_path, question, names, bounding_boxes, en_wiki_pedia_links, en_intros)
    
    # Prepare recognized person information for frontend
    people_info = [[prompt]]
    [people_info.append([f"{name}: {link}"]) for name, link in zip(names, en_wiki_pedia_links)]
   
    # Clean up temporary files
    if os.path.exists(temp_image_path):
        os.remove(temp_image_path)
    
    return resized_image, answer, people_info

# Create Gradio interface
def create_interface():
    # Define examples - with proper wikidata URLs
    examples = [
        ["./musk_and_huang.jpg", "What is Jensen doing?", 
         ['Jen-Hsun Huang', 'Elon Musk'], 
         [{'Width': 0.14037726819515228, 'Height': 0.20855841040611267, 'Left': 0.20549052953720093, 'Top': 0.05287889018654823}, 
          {'Width': 0.07806053012609482, 'Height': 0.12407417595386505, 'Left': 0.7048892974853516, 'Top': 0.12597641348838806}], 
         ['https://www.wikidata.org/wiki/Q305177', 'https://www.wikidata.org/wiki/Q317521']],
        ["./dtms.jpg", "Describe the image.", 
         ['Claude Shannon', 'John McCarthy', 'Marvin Minsky'], 
         [{'Width': 0.06745696812868118, 'Height': 0.1393456608057022, 'Left': 0.8298415541648865, 'Top': 0.2935926616191864}, 
          {'Width': 0.058158036321401596, 'Height': 0.10835719853639603, 'Left': 0.7029165029525757, 'Top': 0.16815270483493805}, 
          {'Width': 0.05310939624905586, 'Height': 0.1121278703212738, 'Left': 0.4884793162345886, 'Top': 0.19025743007659912}], 
         ['https://www.wikidata.org/wiki/Q92760', 'https://www.wikidata.org/wiki/Q92739', 'https://www.wikidata.org/wiki/Q204815']],
        ["./hinton.jpeg", "Describe the image in detail.", 
         ['Yoshua Bengio', 'Geoffrey Hinton', 'Yann LeCun'], 
         [{'Width': 0.077728271484375, 'Height': 0.16116079688072205, 'Left': 0.8799420595169067, 'Top': 0.4856656789779663}, 
          {'Width': 0.07422236353158951, 'Height': 0.15943190455436707, 'Left': 0.4633428454399109, 'Top': 0.07901764661073685}, 
          {'Width': 0.07562466710805893, 'Height': 0.13936467468738556, 'Left': 0.025178398936986923, 'Top': 0.4953641891479492}], 
         ['https://www.wikidata.org/wiki/Q3572699', 'https://www.wikidata.org/wiki/Q92894', 'https://www.wikidata.org/wiki/Q3571662']],
    ]
    
    # Filter examples to only include files that exist
    existing_examples = []
    for example in examples:
        if os.path.exists(example[0]):
            existing_examples.append(example)
    
    with gr.Blocks(title="Celebrity Recognition and VQA System") as app:
        gr.Markdown("<div style='text-align: center;'><h1 style=' font-size: 28px; '>Celebrity Recognition and VQA</h1></div>")
        gr.Markdown("**RC-MLLM** model is developed based on the Qwen2-VL model through a novel method called **RCVIT (Region-level Context-aware Visual Instruction Tuning)**, using the specially constructed **RCMU dataset** for training. Its core feature is the capability for **Region-level Context-aware Multimodal Understanding (RCMU)**. This means it can simultaneously understand both the visual content of specific regions/objects within an image and their associated textual information (utilizing bounding boxes coordinates), allowing it to respond to user instructions in a more context-aware manner. Simply put, RC-MLLM not only understands images but can also integrate the textual information linked to specific objects within the image for understanding. It achieves outstanding performance on RCMU tasks and is suitable for applications like personalized conversation.")

        markdown_content = """
            📑 [Region-Level Context-Aware Multimodal Understanding](https://arxiv.org/abs/2508.12263) | 
            🤗 Models:[RC-Qwen2VL-2b](https://huggingface.co/weihongliang/RC-Qwen2VL-2b/blob/main/README.md) [RC-Qwen2VL-7b](https://huggingface.co/weihongliang/RC-Qwen2VL-7b/blob/main/README.md)| 
            📁 [Dataset](https://huggingface.co/your-model-name) |
            [Github](https://github.com/hongliang-wei/RC-MLLM) | 
            🚀 [Personalized Conversation Demo](https://1684c5f6e1c5a19b2c.gradio.live/)
            """
        gr.Markdown(markdown_content)
        gr.Markdown("📌 Upload an image containing celebrities, the system will recognize them and provide Wikipedia-based VQA using the RC-Qwen2-VL model.")
        
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(type="pil", label="Upload Image")
                question_input = gr.Textbox(label="Question", placeholder="Enter your question...")
                confidence_slider = gr.Slider(
                    minimum=50, 
                    maximum=100, 
                    value=90, 
                    step=1, 
                    label="Confidence Threshold (%)", 
                    info="Adjust the minimum confidence level for celebrity recognition"
                )
                submit_button = gr.Button("Ask RC-Qwen2VL-7B")
                
                # Add examples section
                if existing_examples:
                    gr.Examples(
                        examples=[[example[0], example[1]] for example in existing_examples],
                        inputs=[image_input, question_input],
                        label="Example Images with Questions",
                        fn=None
                    )
            with gr.Column(scale=1):
                image_output = gr.Image(label="Recognition Result")
                answer_output = gr.Textbox(label="RC-Qwen2-VL Answer")

                people_info = gr.Dataframe(
                    headers=["Recognized Person Information and Wikipedia Links"],
                    datatype=["str"],
                    label="Person Information",
                    wrap=True,  # Enable text wrapping
                )
        
        # Function to clear the cache
        def clear_cache():
            global recognition_cache
            old_count = len(recognition_cache)
            recognition_cache = {}
            return f"Cache cleared. {old_count} entries were removed."
        
        # Handle image upload - clear cache when new image is uploaded
        def on_image_upload(image):
            if image is not None:
                clear_cache()
                return f"New image detected. Cache has been cleared."
            return f"No image uploaded. Cache status unchanged."
        
        # Modified submit function to handle examples without the current_example parameter
        def submit_fn(image, question, confidence_threshold):
            # Process image using the recognition system
            return process_image(image, question, confidence_threshold, examples=examples)
        
        # Connect buttons to functions
        submit_button.click(
            fn=submit_fn,
            inputs=[image_input, question_input, confidence_slider],
            outputs=[image_output, answer_output, people_info]
        )
        
        # Clear cache automatically when new image is uploaded
        image_input.change(
            fn=on_image_upload,
            inputs=[image_input],
        )
        
        gr.Markdown("## Instructions")
        gr.Markdown("""
        1. Upload an image containing celebrities
        2. Enter your question, for example:
           - "Who are the people in the image?"
           - "What achievements does the person on the left have?"
           - "What is the relationship between these people?"
        3. Adjust the confidence threshold slider if needed (lower values will recognize more faces but might be less accurate)
        4. Click the submit button to get the answer
        """)
    
    return app

# Launch the application
if __name__ == "__main__":
    app = create_interface()
    app.launch(share=True)