import os
import sys
import math
import numpy as np
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

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

# Configuration
MODEL_NAME = "OpenGVLab/InternVL2_5-8B"  # Smaller model for faster loading
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
        )
        
        # Fix for image context token ID - needed to make the model work with images
        print("Setting image context token ID...")
        if hasattr(tokenizer, 'encode'):
            # Get special token ID from tokenizer
            img_context_token_id = tokenizer.encode("<image>", add_special_tokens=False)[0]
            model.img_context_token_id = img_context_token_id
            print(f"Set img_context_token_id to {img_context_token_id}")
        
        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

# 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

# Function to handle two images
def analyze_two_images(model, tokenizer, image1, image2, prompt):
    try:
        # Check if at least one image is provided
        if image1 is None and image2 is None:
            return "Please upload at least one image."
        
        results = []
        
        # Process first image if provided
        if image1 is not None:
            result1 = analyze_image(model, tokenizer, image1, prompt)
            results.append(f"# Image 1 Analysis\n\n{result1}")
        else:
            results.append("# Image 1\n\nNo image uploaded.")
        
        # Process second image if provided
        if image2 is not None:
            result2 = analyze_image(model, tokenizer, image2, prompt)
            results.append(f"# Image 2 Analysis\n\n{result2}")
        else:
            results.append("# Image 2\n\nNo image uploaded.")
        
        # Combine results
        combined_result = f"{results[0]}\n\n---\n\n{results[1]}"
        
        return combined_result
    except Exception as e:
        import traceback
        error_msg = f"Error analyzing images: {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 Dual Image Analyzer - Error",
            description="The model failed to load. Please check the logs for more information."
        )
        return demo
    
    # Predefined prompts for analysis
    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.",
        "Compare these images and describe the differences."
    ]
    
    # Create the interface with two images
    with gr.Blocks(title="InternVL2.5 Dual Image Analyzer") as demo:
        gr.Markdown("# 🖼️ InternVL2.5 Dual Image Analyzer")
        gr.Markdown("Upload one or two images and ask the InternVL2.5 model to analyze them.")
        
        with gr.Row():
            with gr.Column(scale=1):
                image1 = gr.Image(type="pil", label="Upload Image 1")
                image2 = gr.Image(type="pil", label="Upload Image 2")
                prompt = gr.Dropdown(
                    choices=prompts, 
                    value=prompts[0], 
                    label="Select a prompt or write your own below",
                    allow_custom_value=True
                )
                analyze_button = gr.Button("Analyze Images", variant="primary")
                
            with gr.Column(scale=1):
                output = gr.Markdown(label="Analysis Results")
        
        analyze_button.click(
            fn=lambda img1, img2, p: analyze_two_images(model, tokenizer, img1, img2, p),
            inputs=[image1, image2, prompt],
            outputs=output
        )
        
        # Example images
        if os.path.exists("example_images"):
            example_files = [f for f in os.listdir("example_images") if f.endswith((".jpg", ".jpeg", ".png"))]
            if len(example_files) >= 2:
                example1 = os.path.join("example_images", example_files[0])
                example2 = os.path.join("example_images", example_files[1])
                
                examples = [
                    [example1, None, "Describe this image in detail."],
                    [None, example2, "Describe this image in detail."],
                    [example1, example2, "Compare these images and describe the differences."]
                ]
                
                gr.Examples(
                    examples=examples,
                    inputs=[image1, image2, prompt]
                )
    
    return demo

# Run the application
if __name__ == "__main__":
    try:
        # Check for GPU
        if not torch.cuda.is_available():
            print("WARNING: CUDA is not available. The model requires a GPU to function properly.")
            
        # 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()