import os
import sys
import math
import numpy as np
import tempfile
import torch
import torchvision.transforms as T
from torchvision.transforms.functional import InterpolationMode
from PIL import Image
import gradio as gr
from transformers import AutoModel, AutoTokenizer
import pdf2image

# Constants
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)

# Configuration
MODEL_NAME = "OpenGVLab/InternVL2_5-8B"
IMAGE_SIZE = 448

# Set up environment variables
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

# Utility functions for image processing
def build_transform(input_size):
    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
    transform = T.Compose([
        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(mean=MEAN, std=STD)
    ])
    return transform

def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
    best_ratio_diff = float('inf')
    best_ratio = (1, 1)
    area = width * height
    for ratio in target_ratios:
        target_aspect_ratio = ratio[0] / ratio[1]
        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
        if ratio_diff < best_ratio_diff:
            best_ratio_diff = ratio_diff
            best_ratio = ratio
        elif ratio_diff == best_ratio_diff:
            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
                best_ratio = ratio
    return best_ratio

def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
    orig_width, orig_height = image.size
    aspect_ratio = orig_width / orig_height

    # calculate the existing image aspect ratio
    target_ratios = set(
        (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
        i * j <= max_num and i * j >= min_num)
    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

    # find the closest aspect ratio to the target
    target_aspect_ratio = find_closest_aspect_ratio(
        aspect_ratio, target_ratios, orig_width, orig_height, image_size)

    # calculate the target width and height
    target_width = image_size * target_aspect_ratio[0]
    target_height = image_size * target_aspect_ratio[1]
    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]

    # resize the image
    resized_img = image.resize((target_width, target_height))
    processed_images = []
    for i in range(blocks):
        box = (
            (i % (target_width // image_size)) * image_size,
            (i // (target_width // image_size)) * image_size,
            ((i % (target_width // image_size)) + 1) * image_size,
            ((i // (target_width // image_size)) + 1) * image_size
        )
        # split the image
        split_img = resized_img.crop(box)
        processed_images.append(split_img)
    assert len(processed_images) == blocks
    if use_thumbnail and len(processed_images) != 1:
        thumbnail_img = image.resize((image_size, image_size))
        processed_images.append(thumbnail_img)
    return processed_images

# Function to split model across GPUs
def split_model(model_name):
    device_map = {}
    world_size = torch.cuda.device_count()
    if world_size <= 1:
        return "auto"
    
    num_layers = {
        'InternVL2_5-1B': 24, 
        'InternVL2_5-2B': 24, 
        'InternVL2_5-4B': 36, 
        'InternVL2_5-8B': 32,
        'InternVL2_5-26B': 48, 
        'InternVL2_5-38B': 64, 
        'InternVL2_5-78B': 80
    }[model_name]
    
    # Since the first GPU will be used for ViT, treat it as half a GPU.
    num_layers_per_gpu = math.ceil(num_layers / (world_size - 0.5))
    num_layers_per_gpu = [num_layers_per_gpu] * world_size
    num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)
    layer_cnt = 0
    for i, num_layer in enumerate(num_layers_per_gpu):
        for j in range(num_layer):
            device_map[f'language_model.model.layers.{layer_cnt}'] = i
            layer_cnt += 1
    device_map['vision_model'] = 0
    device_map['mlp1'] = 0
    device_map['language_model.model.tok_embeddings'] = 0
    device_map['language_model.model.embed_tokens'] = 0
    device_map['language_model.model.rotary_emb'] = 0
    device_map['language_model.output'] = 0
    device_map['language_model.model.norm'] = 0
    device_map['language_model.lm_head'] = 0
    device_map[f'language_model.model.layers.{num_layers - 1}'] = 0

    return device_map

# Model loading function
def load_model():
    print(f"\n=== Loading {MODEL_NAME} ===")
    print(f"CUDA available: {torch.cuda.is_available()}")
    
    if torch.cuda.is_available():
        print(f"GPU count: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
            
        # Memory info
        print(f"Total GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
        print(f"Allocated GPU memory: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
        print(f"Reserved GPU memory: {torch.cuda.memory_reserved() / 1e9:.2f} GB")
    
    # Determine device map
    device_map = "auto"
    if torch.cuda.is_available() and torch.cuda.device_count() > 1:
        model_short_name = MODEL_NAME.split('/')[-1]
        device_map = split_model(model_short_name)
    
    # Load model and tokenizer
    try:
        model = AutoModel.from_pretrained(
            MODEL_NAME,
            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
            low_cpu_mem_usage=True,
            trust_remote_code=True,
            device_map=device_map
        )
        
        tokenizer = AutoTokenizer.from_pretrained(
            MODEL_NAME,
            use_fast=False,
            trust_remote_code=True
        )
        
        print(f"✓ Model and tokenizer loaded successfully!")
        return model, tokenizer
    except Exception as e:
        print(f"❌ Error loading model: {e}")
        import traceback
        traceback.print_exc()
        return None, None

# Extract slides from uploaded PDF file
def extract_slides_from_pdf(file_obj):
    try:
        file_bytes = file_obj.read()
        file_extension = os.path.splitext(file_obj.name)[1].lower()
        
        # Check if it's a PDF
        if file_extension != '.pdf':
            return []
        
        # Create temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as temp_file:
            temp_file.write(file_bytes)
            temp_path = temp_file.name
        
        # Extract images from PDF using pdf2image
        slides = []
        try:
            images = pdf2image.convert_from_path(temp_path, dpi=300)
            slides = [(f"Slide {i+1}", img) for i, img in enumerate(images)]
        except Exception as e:
            print(f"Error converting PDF: {e}")
        
        # Clean up temporary file
        os.unlink(temp_path)
        
        return slides
    
    except Exception as e:
        import traceback
        error_msg = f"Error extracting slides: {str(e)}\n{traceback.format_exc()}"
        print(error_msg)
        return []

# Image analysis function
def analyze_image(model, tokenizer, image, prompt):
    try:
        # Check if image is valid
        if image is None:
            return "Please upload an image first."
        
        # Process the image
        processed_images = dynamic_preprocess(image, image_size=IMAGE_SIZE)
        
        # Prepare the prompt
        text_prompt = f"USER: <image>\n{prompt}\nASSISTANT:"
        
        # Convert inputs for the model
        inputs = tokenizer([text_prompt], return_tensors="pt")
        
        # Move inputs to the right device
        if torch.cuda.is_available():
            inputs = {k: v.cuda() for k, v in inputs.items()}
        
        # Add image to the inputs
        inputs["images"] = processed_images
        
        # Generate a response
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=512,
            )
        
        # Decode the outputs
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract only the assistant's response
        assistant_response = generated_text.split("ASSISTANT:")[-1].strip()
        
        return assistant_response
    except Exception as e:
        import traceback
        error_msg = f"Error analyzing image: {str(e)}\n{traceback.format_exc()}"
        return error_msg

# Analyze multiple slides from a PDF
def analyze_pdf_slides(model, tokenizer, file_obj, prompt, num_slides=2):
    try:
        if file_obj is None:
            return "Please upload a PDF file."
        
        # Extract slides from PDF
        slides = extract_slides_from_pdf(file_obj)
        
        if not slides:
            return "No slides were extracted from the file. Please check that it's a valid PDF."
        
        # Limit to the requested number of slides
        slides = slides[:num_slides]
        
        # Analyze each slide
        analyses = []
        for slide_title, slide_image in slides:
            analysis = analyze_image(model, tokenizer, slide_image, prompt)
            analyses.append((slide_title, analysis))
        
        # Format the results
        result = ""
        for slide_title, analysis in analyses:
            result += f"## {slide_title}\n\n{analysis}\n\n---\n\n"
        
        return result
    
    except Exception as e:
        import traceback
        error_msg = f"Error analyzing slides: {str(e)}\n{traceback.format_exc()}"
        return error_msg

# Main function
def main():
    # Load the model
    model, tokenizer = load_model()
    
    if model is None:
        # Create an error interface if model loading failed
        demo = gr.Interface(
            fn=lambda x: "Model loading failed. Please check the logs for details.",
            inputs=gr.Textbox(),
            outputs=gr.Textbox(),
            title="InternVL2.5 Analyzer - Error",
            description="The model failed to load. Please check the logs for more information."
        )
        return demo
    
    # Create an interface with tabs
    with gr.Blocks(title="InternVL2.5 Analyzer") as demo:
        gr.Markdown("# InternVL2.5 Image and Slide Analyzer")
        
        with gr.Tabs():
            # Single Image Analysis Tab
            with gr.TabItem("Single Image Analysis"):
                # Predefined prompts for analysis
                image_prompts = [
                    "Describe this image in detail.",
                    "What can you tell me about this image?",
                    "Is there any text in this image? If so, can you read it?",
                    "What is the main subject of this image?", 
                    "What emotions or feelings does this image convey?",
                    "Describe the composition and visual elements of this image.",
                    "Summarize what you see in this image in one paragraph."
                ]
                
                with gr.Row():
                    image_input = gr.Image(type="pil", label="Upload Image")
                    image_prompt = gr.Dropdown(
                        choices=image_prompts, 
                        value=image_prompts[0], 
                        label="Select a prompt",
                        allow_custom_value=True
                    )
                
                image_analyze_btn = gr.Button("Analyze Image")
                image_output = gr.Textbox(label="Analysis Results", lines=15)
                
                # Handle the image analysis action
                image_analyze_btn.click(
                    fn=lambda img, prompt: analyze_image(model, tokenizer, img, prompt),
                    inputs=[image_input, image_prompt],
                    outputs=image_output
                )
            
            # PDF Slides Analysis Tab
            with gr.TabItem("PDF Slides Analysis"):
                slide_prompts = [
                    "Analyze this slide and describe its contents.",
                    "What is the main message of this slide?",
                    "Extract all the text visible in this slide.",
                    "What are the key points presented in this slide?",
                    "Describe the visual elements and layout of this slide."
                ]
                
                with gr.Row():
                    file_input = gr.File(label="Upload PDF")
                    slide_prompt = gr.Dropdown(
                        choices=slide_prompts, 
                        value=slide_prompts[0], 
                        label="Select a prompt",
                        allow_custom_value=True
                    )
                
                num_slides = gr.Slider(
                    minimum=1, 
                    maximum=5, 
                    value=2, 
                    step=1, 
                    label="Number of Slides to Analyze"
                )
                
                slides_analyze_btn = gr.Button("Analyze Slides")
                slides_output = gr.Markdown(label="Analysis Results")
                
                # Handle the slides analysis action
                slides_analyze_btn.click(
                    fn=lambda file, prompt, num: analyze_pdf_slides(model, tokenizer, file, prompt, num),
                    inputs=[file_input, slide_prompt, num_slides],
                    outputs=slides_output
                )
                
                # Add example if available
                if os.path.exists("example_slides/test_slides.pdf"):
                    gr.Examples(
                        examples=[
                            ["example_slides/test_slides.pdf", "Extract all the text visible in this slide.", 2]
                        ],
                        inputs=[file_input, slide_prompt, num_slides]
                    )
    
    return demo

# Run the application
if __name__ == "__main__":
    try:
        # Create and launch the interface
        demo = main()
        demo.launch(server_name="0.0.0.0")
    except Exception as e:
        print(f"Error starting the application: {e}")
        import traceback
        traceback.print_exc()