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


# Enhanced debug printing
import logging
import traceback

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger("InternVL2.5-Debug")

# Print environment info
logger.info("Python version: %s", sys.version)
logger.info("PyTorch version: %s", torch.__version__)
logger.info("Transformers version: %s", __import__("transformers").__version__)
try:
    logger.info("Einops version: %s", __import__("einops").__version__)
except ImportError:
    logger.error("Einops is not installed!")
# 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

# Load and preprocess image for the model - following the official documentation pattern
def load_image(image_pil, max_num=12):
    # Process the image using dynamic_preprocess
    processed_images = dynamic_preprocess(image_pil, image_size=IMAGE_SIZE, max_num=max_num)
    
    # Convert PIL images to tensor format expected by the model
    transform = build_transform(IMAGE_SIZE)
    pixel_values = [transform(img) for img in processed_images]
    pixel_values = torch.stack(pixel_values)
    
    # Convert to appropriate data type
    if torch.cuda.is_available():
        pixel_values = pixel_values.cuda().to(torch.bfloat16)
    else:
        pixel_values = pixel_values.to(torch.float32)
        
    return pixel_values

# 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

# Get model dtype
def get_model_dtype():
    return torch.bfloat16 if torch.cuda.is_available() else torch.float32

# Model loading function
def load_model():
    print(f"\n=== Loading {MODEL_NAME} ===")
    print(f"CUDA available: {torch.cuda.is_available()}")
    
    model_dtype = get_model_dtype()
    print(f"Using model dtype: {model_dtype}")
    
    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=model_dtype,
            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:
        logger.error(f"❌ Error loading model: {e}")
        logger.error("Detailed traceback:")
        import traceback
        traceback.print_exc()
        
        # Check if einops is available
        try:
            import einops
            logger.info(f"einops is available, version: {einops.__version__}")
        except ImportError:
            logger.error("ImportError: einops is not installed! This is required for InternVL2.5.")
            
        # Check for CUDA availability
        if torch.cuda.is_available():
            logger.info(f"CUDA is available. Device count: {torch.cuda.device_count()}")
            for i in range(torch.cuda.device_count()):
                logger.info(f"Device {i}: {torch.cuda.get_device_name(i)}")
                logger.info(f"Memory allocated: {torch.cuda.memory_allocated(i) / 1e9:.2f} GB")
                logger.info(f"Memory reserved: {torch.cuda.memory_reserved(i) / 1e9:.2f} GB")
        else:
            logger.warning("CUDA is not available. Running on CPU.")        return None, None

# Image analysis function using the chat method from documentation
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 following official pattern
        pixel_values = load_image(image)
        
        # Debug info
        print(f"Image processed: tensor shape {pixel_values.shape}, dtype {pixel_values.dtype}")
        
        # Define generation config
        generation_config = {
            "max_new_tokens": 512,
            "do_sample": False
        }
        
        # Use the model.chat method as shown in the official documentation
        question = f"<image>\n{prompt}"
        response, _ = model.chat(
            tokenizer=tokenizer,
            pixel_values=pixel_values,
            question=question,
            generation_config=generation_config,
            history=None,
            return_history=True
        )
        
        return response
    except Exception as e:
        import traceback
        error_msg = f"Error analyzing image: {str(e)}\n{traceback.format_exc()}"
        return error_msg

# Main function
def main():
    # Add debug info at the start of main
    logger.info("Starting main() function...")
    logger.info(f"MODEL_NAME: {MODEL_NAME}")
    
    # 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 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."
    ]
    
    # Create the interface
    demo = gr.Interface(
        fn=lambda img, prompt: analyze_image(model, tokenizer, img, prompt),
        inputs=[
            gr.Image(type="pil", label="Upload 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 Image Analyzer",
        description="Upload an image and ask the InternVL2.5 model to analyze it.",
        examples=[
            ["example_images/example1.jpg", "Describe this image in detail."],
            ["example_images/example2.jpg", "What can you tell me about this image?"]
        ],
        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()