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 - single image
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

# New function for analyzing two images
def analyze_two_images(model, tokenizer, image1, image2, prompt):
    try:
        # Check if images are valid
        if image1 is None and image2 is None:
            return "Please upload at least one image."
        
        # Process the images
        processed_images = []
        if image1 is not None:
            processed_images.extend(dynamic_preprocess(image1, image_size=IMAGE_SIZE))
        if image2 is not None:
            processed_images.extend(dynamic_preprocess(image2, image_size=IMAGE_SIZE))
        
        # Prepare the prompt with two image tokens
        text_prompt = f"USER: <image><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 images 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 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 these images in detail.",
        "What can you tell me about these images?",
        "Is there any text in these images? If so, can you read it?",
        "Compare and contrast these two images.",
        "What are the main subjects in these images?", 
        "What emotions or feelings do these images convey?",
        "Describe the composition and visual elements of these images.",
        "Summarize what you see in these images in one paragraph."
    ]
    
    # Create the interface
    demo = gr.Interface(
        fn=lambda img1, img2, prompt: analyze_two_images(model, tokenizer, img1, img2, prompt),
        inputs=[
            gr.Image(type="pil", label="Upload First Image"),
            gr.Image(type="pil", label="Upload Second Image"),
            gr.Dropdown(choices=prompts, value=prompts[0], label="Select a prompt or write your own below", 
                        allow_custom_value=True)
        ],
        outputs=gr.Textbox(label="Analysis Results", lines=15),
        title="InternVL2.5 Dual Image Analyzer",
        description="Upload two images and ask the InternVL2.5 model to analyze them together.",
        examples=[
            ["example_images/example1.jpg", "example_images/example2.jpg", "Compare and contrast these two images."],
            ["example_images/example1.jpg", "example_images/example2.jpg", "What can you tell me about these images?"]
        ],
        theme=gr.themes.Soft(),
        allow_flagging="never"
    )
    
    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()