#!/usr/bin/env python
"""
Custom VLM Evaluation Script for Fast_dLLM models
Supports: GQA, POPE, TextVQA, MMMU (basic metrics without GPT judge)
Supports multi-GPU via data parallel (--num_gpus)
Tracks tokens-per-step for MDM models
"""

import argparse
import json
import os
import sys
import types
from pathlib import Path
from typing import Optional, List, Dict, Any, Tuple, Union
import torch
import torch.multiprocessing as mp
from tqdm import tqdm
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoProcessor
from PIL import Image
import re
import numpy as np

# Add models directory to path for MDM generation functions
_script_dir = os.path.dirname(os.path.abspath(__file__))
_dvlm_dir = os.path.dirname(_script_dir)
_models_dir = os.path.join(_dvlm_dir, "models")
if _models_dir not in sys.path:
    sys.path.insert(0, _models_dir)


def normalize_answer(s: str) -> str:
    """Normalize answer for comparison."""
    s = s.lower().strip()
    s = re.sub(r'[^\w\s]', '', s)
    s = re.sub(r'\b(a|an|the)\b', '', s)
    s = ' '.join(s.split())
    return s


def exact_match(pred: str, target: str) -> bool:
    return normalize_answer(pred) == normalize_answer(target)


def contains_match(pred: str, target: str) -> bool:
    return normalize_answer(target) in normalize_answer(pred)


class CustomVLMEvaluator:
    def __init__(
        self,
        model_path: str,
        device: str = "cuda",
        torch_dtype: torch.dtype = torch.bfloat16,
        base_model: str = "Qwen/Qwen2.5-VL-3B-Instruct",
        min_pixels: int = None,
        max_pixels: int = None,
    ):
        print(f"[{device}] Loading model from {model_path}...")
        
        # Build processor kwargs - only set if explicitly provided
        processor_kwargs = {}
        if min_pixels is not None:
            processor_kwargs["min_pixels"] = min_pixels
        if max_pixels is not None:
            processor_kwargs["max_pixels"] = max_pixels
        
        if processor_kwargs:
            print(f"[{device}] Custom image resolution: {processor_kwargs}")
        else:
            print(f"[{device}] Using processor default image resolution")
        
        try:
            from transformers import Qwen2VLProcessor
            self.processor = Qwen2VLProcessor.from_pretrained(base_model, **processor_kwargs)
        except Exception as e:
            self.processor = AutoProcessor.from_pretrained(base_model, **processor_kwargs)
        
        # Try to load model - detect model type and use appropriate class
        # - Original Qwen2.5-VL: use AutoModelForVision2Seq
        # - Fast_dLLM checkpoints: use AutoModelForCausalLM with trust_remote_code
        self.is_mdm = False
        self.block_size = None
        self.mask_id = None
        
        try:
            # First try AutoModelForCausalLM with trust_remote_code (for Fast_dLLM)
            self.model = AutoModelForCausalLM.from_pretrained(
                model_path,
                torch_dtype=torch_dtype,
                device_map=device,
                trust_remote_code=True,
            ).eval()
            
            # Check if this is an MDM model (has bd_size in config)
            if hasattr(self.model.config, 'bd_size') and self.model.config.bd_size is not None:
                self.is_mdm = True
                self.block_size = self.model.config.bd_size
                # Get mask_id from tokenizer
                from transformers import AutoTokenizer
                tokenizer = AutoTokenizer.from_pretrained(model_path)
                self.mask_id = tokenizer.encode("|<MASK>|")[0]
                print(f"[{device}] Detected MDM model with block_size={self.block_size}")
                
                # Bind mdm_sample_cache method
                from generation_functions import Fast_dLLM_QwenForCausalLM
                self.model.mdm_sample = types.MethodType(
                    Fast_dLLM_QwenForCausalLM.mdm_sample_cache, self.model
                )
                
        except ValueError as e:
            # If it fails with "Unrecognized configuration class", use AutoModelForVision2Seq
            if "Qwen2_5_VLConfig" in str(e) or "Qwen2VLConfig" in str(e):
                print(f"[{device}] Detected original Qwen2.5-VL model, using AutoModelForVision2Seq")
                from transformers import AutoModelForVision2Seq
                self.model = AutoModelForVision2Seq.from_pretrained(
                    model_path,
                    torch_dtype=torch_dtype,
                    device_map=device,
                ).eval()
            else:
                raise e
        self.device = device
        print(f"[{device}] Model loaded! (MDM={self.is_mdm})")

    @torch.no_grad()
    def generate(
        self, 
        image: Image.Image, 
        question: str, 
        max_new_tokens: int = 32,
        return_stats: bool = False,
    ) -> Union[str, Tuple[str, Dict[str, Any]]]:
        """
        Generate response for image+question.
        
        Args:
            image: PIL Image
            question: Text question
            max_new_tokens: Max tokens to generate
            return_stats: If True and MDM model, return (response, stats_dict)
        
        Returns:
            response string, or (response, stats) if return_stats=True and MDM model
        """
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text", "text": question},
                ],
            }
        ]
        
        text = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        
        inputs = self.processor(
            text=[text],
            images=[image],
            padding=True,
            return_tensors="pt",
        ).to(self.device)
        
        stats = None
        
        if self.is_mdm:
            # Use MDM sampling with stats tracking
            mdm_kwargs = {
                "input_ids": inputs.input_ids,
                "tokenizer": self.processor.tokenizer,
                "block_size": self.block_size,
                "max_tokens": max_new_tokens,
                "mask_id": self.mask_id,
                "return_stats": return_stats,
            }
            if hasattr(inputs, "pixel_values"):
                mdm_kwargs["pixel_values"] = inputs.pixel_values
            if hasattr(inputs, "image_grid_thw"):
                mdm_kwargs["image_grid_thw"] = inputs.image_grid_thw
            
            result = self.model.mdm_sample(**mdm_kwargs)
            
            if return_stats:
                output_ids, stats = result
            else:
                output_ids = result
            
            generated_ids = output_ids[:, inputs.input_ids.shape[1]:]
        else:
            # Standard AR generation
            output_ids = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
            )
            generated_ids = output_ids[:, inputs.input_ids.shape[1]:]
            
            if return_stats:
                # For AR models, each step generates 1 token
                num_tokens = generated_ids.shape[1]
                stats = {
                    "tokens_per_step": [1] * num_tokens,
                    "total_steps": num_tokens,
                    "gen_length": num_tokens,
                }
        
        response = self.processor.batch_decode(
            generated_ids, skip_special_tokens=True
        )[0]
        
        if return_stats:
            return response.strip(), stats
        return response.strip()


def worker_fn(rank: int, world_size: int, model_path: str, base_model: str,
              task: str, limit: Optional[int], return_dict: dict,
              min_pixels: Optional[int] = None, max_pixels: Optional[int] = None,
              track_tokens_per_step: bool = False):
    """Worker function for each GPU."""
    device = f"cuda:{rank}"
    torch.cuda.set_device(rank)
    
    # Verify GPU assignment
    print(f"[Worker {rank}] Using device: {device}, CUDA current device: {torch.cuda.current_device()}")
    
    evaluator = CustomVLMEvaluator(
        model_path, device=device, base_model=base_model,
        min_pixels=min_pixels, max_pixels=max_pixels
    )
    
    # Verify model is on correct device
    print(f"[Worker {rank}] Model loaded on: {next(evaluator.model.parameters()).device}")
    
    if task == "gqa":
        result = eval_gqa_worker(evaluator, limit, rank, world_size, track_tokens_per_step)
    elif task == "pope":
        result = eval_pope_worker(evaluator, limit, rank, world_size, track_tokens_per_step)
    elif task == "textvqa":
        result = eval_textvqa_worker(evaluator, limit, rank, world_size, track_tokens_per_step)
    elif task == "mmmu":
        result = eval_mmmu_worker(evaluator, limit, rank, world_size, track_tokens_per_step)
    elif task == "mmbench":
        result = eval_mmbench_worker(evaluator, limit, rank, world_size, track_tokens_per_step)
    elif task == "mathvista":
        result = eval_mathvista_worker(evaluator, limit, rank, world_size, track_tokens_per_step)
    else:
        result = {"error": f"Unknown task: {task}"}
    
    return_dict[rank] = result


def eval_gqa_worker(evaluator: CustomVLMEvaluator, limit: Optional[int],
                    rank: int, world_size: int,
                    track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """GQA evaluation worker."""
    images_ds = load_dataset("lmms-lab/GQA", "testdev_balanced_images", split="testdev")
    instr_ds = load_dataset("lmms-lab/GQA", "testdev_balanced_instructions", split="testdev")
    
    image_lookup = {item["id"]: item["image"] for item in images_ds}
    
    # Shard the dataset
    total = len(instr_ds)
    per_rank = total // world_size
    start = rank * per_rank
    end = start + per_rank if rank < world_size - 1 else total
    
    if limit:
        end = min(start + limit, end)
    
    instr_ds = instr_ds.select(range(start, end))
    
    correct = 0
    total_eval = 0
    skipped = 0
    all_tokens_per_step = []  # Collect all tokens_per_step from all samples
    total_steps = 0
    total_gen_tokens = 0
    
    for item in tqdm(instr_ds, desc=f"GQA[GPU{rank}]", position=rank):
        image_id = item.get("imageId")
        image = image_lookup.get(image_id)
        question = item.get("question")
        answer = item.get("answer")
        
        if image is None or question is None or answer is None:
            skipped += 1
            continue
        
        if track_tokens_per_step:
            pred, stats = evaluator.generate(image, question, max_new_tokens=16, return_stats=True)
            all_tokens_per_step.extend(stats.get("tokens_per_step", []))
            total_steps += stats.get("total_steps", 0)
            total_gen_tokens += stats.get("gen_length", 0)
        else:
            pred = evaluator.generate(image, question, max_new_tokens=16)
        
        if exact_match(pred, answer) or contains_match(pred, answer):
            correct += 1
        total_eval += 1
    
    result = {"correct": correct, "total": total_eval, "skipped": skipped}
    if track_tokens_per_step:
        result["all_tokens_per_step"] = all_tokens_per_step
        result["total_steps"] = total_steps
        result["total_gen_tokens"] = total_gen_tokens
    return result


def eval_pope_worker(evaluator: CustomVLMEvaluator, limit: Optional[int],
                     rank: int, world_size: int,
                     track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """POPE evaluation worker."""
    dataset = load_dataset("lmms-lab/POPE", split="test")
    
    total = len(dataset)
    per_rank = total // world_size
    start = rank * per_rank
    end = start + per_rank if rank < world_size - 1 else total
    
    if limit:
        end = min(start + limit, end)
    
    dataset = dataset.select(range(start, end))
    
    correct = 0
    total_eval = 0
    tp, fp, tn, fn = 0, 0, 0, 0
    all_tokens_per_step = []
    total_steps = 0
    total_gen_tokens = 0
    
    for item in tqdm(dataset, desc=f"POPE[GPU{rank}]", position=rank):
        image = item["image"]
        question = item["question"]
        answer = item["answer"].lower()
        
        if track_tokens_per_step:
            pred, stats = evaluator.generate(image, question, return_stats=True)
            all_tokens_per_step.extend(stats.get("tokens_per_step", []))
            total_steps += stats.get("total_steps", 0)
            total_gen_tokens += stats.get("gen_length", 0)
        else:
            pred = evaluator.generate(image, question)
        pred_lower = pred.lower()
        
        if "yes" in pred_lower:
            pred_answer = "yes"
        elif "no" in pred_lower:
            pred_answer = "no"
        else:
            pred_answer = "unknown"
        
        if pred_answer == answer:
            correct += 1
            if answer == "yes":
                tp += 1
            else:
                tn += 1
        else:
            if answer == "yes":
                fn += 1
            else:
                fp += 1
        total_eval += 1
    
    result = {"correct": correct, "total": total_eval, "tp": tp, "fp": fp, "tn": tn, "fn": fn}
    if track_tokens_per_step:
        result["all_tokens_per_step"] = all_tokens_per_step
        result["total_steps"] = total_steps
        result["total_gen_tokens"] = total_gen_tokens
    return result


def eval_textvqa_worker(evaluator: CustomVLMEvaluator, limit: Optional[int],
                        rank: int, world_size: int,
                        track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """TextVQA evaluation worker."""
    dataset = load_dataset("lmms-lab/textvqa", split="validation")
    
    total = len(dataset)
    per_rank = total // world_size
    start = rank * per_rank
    end = start + per_rank if rank < world_size - 1 else total
    
    if limit:
        end = min(start + limit, end)
    
    dataset = dataset.select(range(start, end))
    
    correct = 0
    total_eval = 0
    all_tokens_per_step = []
    total_steps = 0
    total_gen_tokens = 0
    
    for item in tqdm(dataset, desc=f"TextVQA[GPU{rank}]", position=rank):
        image = item["image"]
        question = item["question"]
        answers = item["answers"]
        
        if track_tokens_per_step:
            pred, stats = evaluator.generate(image, question, return_stats=True)
            all_tokens_per_step.extend(stats.get("tokens_per_step", []))
            total_steps += stats.get("total_steps", 0)
            total_gen_tokens += stats.get("gen_length", 0)
        else:
            pred = evaluator.generate(image, question)
        
        pred_norm = normalize_answer(pred)
        matched = any(
            pred_norm == normalize_answer(ans) or 
            normalize_answer(ans) in pred_norm
            for ans in answers
        )
        
        if matched:
            correct += 1
        total_eval += 1
    
    result = {"correct": correct, "total": total_eval}
    if track_tokens_per_step:
        result["all_tokens_per_step"] = all_tokens_per_step
        result["total_steps"] = total_steps
        result["total_gen_tokens"] = total_gen_tokens
    return result


def extract_mcq_answer(pred: str, all_choices: List[str] = None) -> str:
    """
    Extract multiple choice answer (A/B/C/D) from prediction.
    Enhanced matching logic based on VLMEvalKit's parse_multi_choice_response.
    """
    if all_choices is None:
        all_choices = ["A", "B", "C", "D"]
    
    pred_orig = pred.strip()
    pred = pred_orig.upper()
    
    # Direct match: just "A", "B", "C", "D"
    if pred in all_choices:
        return pred
    
    # Step 1: Look for patterns like "A.", "B.", etc at the beginning
    for choice in all_choices:
        if pred.startswith(f"{choice}.") or pred.startswith(f"{choice})"):
            return choice
    
    # Step 2: Look for patterns like "(A)", "(B)", etc
    for choice in all_choices:
        pattern = rf'\({choice}\)'
        if re.search(pattern, pred):
            return choice
    
    # Step 3: Look for "The answer is X" pattern
    patterns = [
        r'answer\s*(?:is|:)\s*([A-D])',
        r'答案\s*(?:是|:)\s*([A-D])',
        r'选择?\s*([A-D])',
        r'\b([A-D])\b\s*(?:is correct|is the answer)',
    ]
    for pattern in patterns:
        match = re.search(pattern, pred, re.IGNORECASE)
        if match:
            return match.group(1).upper()
    
    # Step 4: Look for choice letter followed by content match
    # If the prediction contains something like "A. xxx" anywhere
    for choice in all_choices:
        pattern = rf'\b{choice}[.):]\s'
        if re.search(pattern, pred):
            return choice
    
    # Step 5: Count occurrences and pick most frequent (with preference for earlier)
    candidates = []
    for i, char in enumerate(pred):
        if char in all_choices:
            candidates.append((char, i))
    
    if candidates:
        # Return the first valid choice found
        return candidates[0][0]
    
    return "UNKNOWN"


def parse_options(options_raw):
    """Parse options - handles JSON string or list."""
    if options_raw is None:
        return []
    if isinstance(options_raw, str):
        try:
            return json.loads(options_raw)
        except json.JSONDecodeError:
            return []
    return options_raw if isinstance(options_raw, list) else []


def eval_mmmu_worker(evaluator: CustomVLMEvaluator, limit: Optional[int],
                     rank: int, world_size: int,
                     track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """MMMU evaluation worker."""
    # MMMU validation set
    dataset = load_dataset("lmms-lab/MMMU", split="validation")
    
    total = len(dataset)
    per_rank = total // world_size
    start = rank * per_rank
    end = start + per_rank if rank < world_size - 1 else total
    
    if limit:
        end = min(start + limit, end)
    
    dataset = dataset.select(range(start, end))
    
    correct = 0
    total_eval = 0
    all_tokens_per_step = []
    total_steps = 0
    total_gen_tokens = 0
    
    for item in tqdm(dataset, desc=f"MMMU[GPU{rank}]", position=rank):
        # Get image(s) - MMMU can have multiple images
        images = []
        for i in range(1, 8):  # MMMU has up to 7 images
            img_key = f"image_{i}"
            if img_key in item and item[img_key] is not None:
                images.append(item[img_key])
        
        if not images:
            continue
        
        # Use first image for now (most questions have 1 image)
        image = images[0]
        
        # Build question with options (parse JSON string if needed)
        question = item["question"]
        options = parse_options(item.get("options"))
        
        if options:
            # Format: A. option1  B. option2  C. option3  D. option4
            option_str = "\n".join([f"{chr(65+i)}. {opt}" for i, opt in enumerate(options)])
            question = f"{question}\n\nOptions:\n{option_str}\n\nAnswer with the letter (A, B, C, or D):"
        
        answer = item.get("answer", "")
        
        if track_tokens_per_step:
            pred, stats = evaluator.generate(image, question, return_stats=True)
            all_tokens_per_step.extend(stats.get("tokens_per_step", []))
            total_steps += stats.get("total_steps", 0)
            total_gen_tokens += stats.get("gen_length", 0)
        else:
            pred = evaluator.generate(image, question)
        pred_answer = extract_mcq_answer(pred)
        
        if pred_answer == answer.upper():
            correct += 1
        total_eval += 1
    
    result = {"correct": correct, "total": total_eval}
    if track_tokens_per_step:
        result["all_tokens_per_step"] = all_tokens_per_step
        result["total_steps"] = total_steps
        result["total_gen_tokens"] = total_gen_tokens
    return result


def eval_mmbench_worker(evaluator: CustomVLMEvaluator, limit: Optional[int],
                        rank: int, world_size: int,
                        track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """MMBench evaluation worker."""
    dataset = load_dataset("lmms-lab/MMBench", "en", split="dev")
    
    total = len(dataset)
    per_rank = total // world_size
    start = rank * per_rank
    end = start + per_rank if rank < world_size - 1 else total
    
    if limit:
        end = min(start + limit, end)
    
    dataset = dataset.select(range(start, end))
    
    correct = 0
    total_eval = 0
    all_tokens_per_step = []
    total_steps = 0
    total_gen_tokens = 0
    
    for item in tqdm(dataset, desc=f"MMBench[GPU{rank}]", position=rank):
        image = item.get("image")
        if image is None:
            continue
        
        question = item.get("question", "")
        # MMBench has options A, B, C, D
        options = []
        for opt_key in ["A", "B", "C", "D"]:
            if opt_key in item and item[opt_key]:
                options.append(item[opt_key])
        
        if options:
            option_str = "\n".join([f"{chr(65+i)}. {opt}" for i, opt in enumerate(options)])
            question = f"{question}\n\nOptions:\n{option_str}\n\nAnswer with the letter (A, B, C, or D):"
        
        answer = item.get("answer", "")
        
        if track_tokens_per_step:
            pred, stats = evaluator.generate(image, question, return_stats=True)
            all_tokens_per_step.extend(stats.get("tokens_per_step", []))
            total_steps += stats.get("total_steps", 0)
            total_gen_tokens += stats.get("gen_length", 0)
        else:
            pred = evaluator.generate(image, question)
        pred_answer = extract_mcq_answer(pred)
        
        if pred_answer == answer.upper():
            correct += 1
        total_eval += 1
    
    result = {"correct": correct, "total": total_eval}
    if track_tokens_per_step:
        result["all_tokens_per_step"] = all_tokens_per_step
        result["total_steps"] = total_steps
        result["total_gen_tokens"] = total_gen_tokens
    return result


def eval_mathvista_worker(evaluator: CustomVLMEvaluator, limit: Optional[int],
                          rank: int, world_size: int,
                          track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """MathVista evaluation worker."""
    dataset = load_dataset("AI4Math/MathVista", split="testmini")
    
    total = len(dataset)
    per_rank = total // world_size
    start = rank * per_rank
    end = start + per_rank if rank < world_size - 1 else total
    
    if limit:
        end = min(start + limit, end)
    
    dataset = dataset.select(range(start, end))
    
    correct = 0
    total_eval = 0
    all_tokens_per_step = []
    total_steps = 0
    total_gen_tokens = 0
    
    for item in tqdm(dataset, desc=f"MathVista[GPU{rank}]", position=rank):
        # Try different image field names
        image = item.get("decoded_image") or item.get("image")
        if image is None:
            continue
        
        question = item.get("question", "")
        # Try different field names for choices (choices or options, may be JSON string)
        choices = parse_options(item.get("choices")) or parse_options(item.get("options"))
        question_type = item.get("question_type", "")
        
        # If multiple choice question
        if choices and question_type == "multi_choice":
            option_str = "\n".join([f"{chr(65+i)}. {opt}" for i, opt in enumerate(choices)])
            question = f"{question}\n\nOptions:\n{option_str}\n\nAnswer with the letter:"
            answer = item.get("answer", "")
            
            if track_tokens_per_step:
                pred, stats = evaluator.generate(image, question, max_new_tokens=64, return_stats=True)
                all_tokens_per_step.extend(stats.get("tokens_per_step", []))
                total_steps += stats.get("total_steps", 0)
                total_gen_tokens += stats.get("gen_length", 0)
            else:
                pred = evaluator.generate(image, question, max_new_tokens=64)
            pred_answer = extract_mcq_answer(pred)
            
            is_correct = pred_answer == answer.upper()
        else:
            # Free-form answer
            answer = str(item.get("answer", ""))
            if track_tokens_per_step:
                pred, stats = evaluator.generate(image, question, max_new_tokens=64, return_stats=True)
                all_tokens_per_step.extend(stats.get("tokens_per_step", []))
                total_steps += stats.get("total_steps", 0)
                total_gen_tokens += stats.get("gen_length", 0)
            else:
                pred = evaluator.generate(image, question, max_new_tokens=64)
            
            # Normalize and compare
            pred_norm = normalize_answer(pred)
            ans_norm = normalize_answer(answer)
            is_correct = pred_norm == ans_norm or ans_norm in pred_norm
        
        if is_correct:
            correct += 1
        total_eval += 1
    
    result = {"correct": correct, "total": total_eval}
    if track_tokens_per_step:
        result["all_tokens_per_step"] = all_tokens_per_step
        result["total_steps"] = total_steps
        result["total_gen_tokens"] = total_gen_tokens
    return result


def compute_tokens_per_step_stats(all_tokens_per_step: List[int]) -> Dict[str, Any]:
    """Compute statistics for tokens per step."""
    if not all_tokens_per_step:
        return {}
    
    arr = np.array(all_tokens_per_step)
    total_tokens = int(arr.sum())
    total_steps = len(arr)
    
    # Compute distribution
    unique, counts = np.unique(arr, return_counts=True)
    distribution = {int(k): int(v) for k, v in zip(unique, counts)}
    
    # Compute percentages for common buckets
    one_tok = int((arr == 1).sum())
    two_tok = int((arr == 2).sum())
    three_plus = int((arr >= 3).sum())
    
    return {
        "mean": float(np.mean(arr)),
        "std": float(np.std(arr)),
        "median": float(np.median(arr)),
        "min": int(arr.min()),
        "max": int(arr.max()),
        "total_steps": total_steps,
        "total_tokens": total_tokens,
        "tokens_per_step_ratio": total_tokens / total_steps if total_steps > 0 else 0,
        "distribution": distribution,
        "distribution_pct": {
            "1_token": one_tok / total_steps * 100 if total_steps > 0 else 0,
            "2_tokens": two_tok / total_steps * 100 if total_steps > 0 else 0,
            "3+_tokens": three_plus / total_steps * 100 if total_steps > 0 else 0,
        },
    }


def merge_results(results: List[Dict[str, Any]], task: str) -> Dict[str, Any]:
    """Merge results from all workers."""
    total_correct = sum(r.get("correct", 0) for r in results)
    total_count = sum(r.get("total", 0) for r in results)
    total_skipped = sum(r.get("skipped", 0) for r in results)
    
    if total_count == 0:
        return {"accuracy": 0, "correct": 0, "total": 0, "skipped": total_skipped}
    
    accuracy = total_correct / total_count * 100
    merged = {
        "accuracy": accuracy,
        "correct": total_correct,
        "total": total_count,
        "skipped": total_skipped,
    }
    
    if task == "pope":
        tp = sum(r.get("tp", 0) for r in results)
        fp = sum(r.get("fp", 0) for r in results)
        tn = sum(r.get("tn", 0) for r in results)
        fn = sum(r.get("fn", 0) for r in results)
        
        precision = tp / (tp + fp) * 100 if (tp + fp) > 0 else 0
        recall = tp / (tp + fn) * 100 if (tp + fn) > 0 else 0
        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
        
        merged.update({
            "precision": precision,
            "recall": recall,
            "f1": f1,
        })
    
    # Merge tokens_per_step stats if available
    all_tokens_per_step = []
    for r in results:
        if "all_tokens_per_step" in r:
            all_tokens_per_step.extend(r["all_tokens_per_step"])
    
    if all_tokens_per_step:
        merged["tokens_per_step_stats"] = compute_tokens_per_step_stats(all_tokens_per_step)
        merged["total_gen_tokens"] = sum(r.get("total_gen_tokens", 0) for r in results)
        merged["total_gen_steps"] = sum(r.get("total_steps", 0) for r in results)
    
    return merged


def cache_datasets(tasks: List[str]):
    """Pre-download and cache datasets in main process."""
    print("\n[Main] Pre-caching datasets...")
    for task in tasks:
        if task == "gqa":
            print("  Caching GQA...")
            load_dataset("lmms-lab/GQA", "testdev_balanced_images", split="testdev")
            load_dataset("lmms-lab/GQA", "testdev_balanced_instructions", split="testdev")
        elif task == "pope":
            print("  Caching POPE...")
            load_dataset("lmms-lab/POPE", split="test")
        elif task == "textvqa":
            print("  Caching TextVQA...")
            load_dataset("lmms-lab/textvqa", split="validation")
        elif task == "mmmu":
            print("  Caching MMMU...")
            load_dataset("lmms-lab/MMMU", split="validation")
        elif task == "mmbench":
            print("  Caching MMBench...")
            load_dataset("lmms-lab/MMBench", "en", split="dev")
        elif task == "mathvista":
            print("  Caching MathVista...")
            load_dataset("AI4Math/MathVista", split="testmini")
    print("[Main] Datasets cached!\n")


def run_multi_gpu(model_path: str, base_model: str, task: str, 
                  limit: Optional[int], num_gpus: int,
                  min_pixels: Optional[int] = None, max_pixels: Optional[int] = None,
                  track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """Run evaluation on multiple GPUs."""
    print(f"\n{'='*50}")
    print(f"Running {task.upper()} on {num_gpus} GPUs (track_tokens_per_step={track_tokens_per_step})")
    print("="*50)
    
    mp.set_start_method('spawn', force=True)
    manager = mp.Manager()
    return_dict = manager.dict()
    
    processes = []
    for rank in range(num_gpus):
        p = mp.Process(
            target=worker_fn,
            args=(rank, num_gpus, model_path, base_model, task, limit, return_dict,
                  min_pixels, max_pixels, track_tokens_per_step)
        )
        p.start()
        processes.append(p)
    
    for p in processes:
        p.join()
    
    # Merge results
    results = [return_dict[rank] for rank in range(num_gpus)]
    merged = merge_results(results, task)
    
    return merged


def run_single_gpu(model_path: str, base_model: str, task: str,
                   limit: Optional[int],
                   min_pixels: Optional[int] = None, max_pixels: Optional[int] = None,
                   track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """Run evaluation on single GPU."""
    evaluator = CustomVLMEvaluator(
        model_path, device="cuda", base_model=base_model,
        min_pixels=min_pixels, max_pixels=max_pixels
    )
    
    if task == "gqa":
        return eval_gqa_single(evaluator, limit, track_tokens_per_step)
    elif task == "pope":
        return eval_pope_single(evaluator, limit, track_tokens_per_step)
    elif task == "textvqa":
        return eval_textvqa_single(evaluator, limit, track_tokens_per_step)
    elif task == "mmmu":
        return eval_mmmu_single(evaluator, limit, track_tokens_per_step)
    elif task == "mmbench":
        return eval_mmbench_single(evaluator, limit, track_tokens_per_step)
    elif task == "mathvista":
        return eval_mathvista_single(evaluator, limit, track_tokens_per_step)
    else:
        return {"error": f"Unknown task: {task}"}


def eval_gqa_single(evaluator: CustomVLMEvaluator, limit: Optional[int],
                    track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """Single GPU GQA evaluation."""
    print("\n" + "=" * 50)
    print("Evaluating GQA")
    print("=" * 50)
    
    print("Loading images...")
    images_ds = load_dataset("lmms-lab/GQA", "testdev_balanced_images", split="testdev")
    print("Loading instructions...")
    instr_ds = load_dataset("lmms-lab/GQA", "testdev_balanced_instructions", split="testdev")
    
    image_lookup = {item["id"]: item["image"] for item in images_ds}
    print(f"Loaded {len(image_lookup)} images, {len(instr_ds)} questions")
    
    if limit:
        instr_ds = instr_ds.select(range(min(limit, len(instr_ds))))
    
    correct = 0
    total = 0
    skipped = 0
    all_tokens_per_step = []
    total_steps = 0
    total_gen_tokens = 0
    
    for item in tqdm(instr_ds, desc="GQA"):
        image_id = item.get("imageId")
        image = image_lookup.get(image_id)
        question = item.get("question")
        answer = item.get("answer")
        
        if image is None or question is None or answer is None:
            skipped += 1
            continue
        
        if track_tokens_per_step:
            pred, stats = evaluator.generate(image, question, return_stats=True)
            all_tokens_per_step.extend(stats.get("tokens_per_step", []))
            total_steps += stats.get("total_steps", 0)
            total_gen_tokens += stats.get("gen_length", 0)
        else:
            pred = evaluator.generate(image, question)
        
        if exact_match(pred, answer) or contains_match(pred, answer):
            correct += 1
        total += 1
        
        if total <= 3:
            print(f"\nQ: {question}")
            print(f"A: {answer}")
            print(f"Pred: {pred}")
    
    accuracy = correct / total * 100 if total > 0 else 0
    print(f"\nGQA Results: {correct}/{total} = {accuracy:.2f}%")
    
    result = {"accuracy": accuracy, "correct": correct, "total": total, "skipped": skipped}
    if track_tokens_per_step and all_tokens_per_step:
        result["tokens_per_step_stats"] = compute_tokens_per_step_stats(all_tokens_per_step)
        result["total_gen_tokens"] = total_gen_tokens
        result["total_gen_steps"] = total_steps
    return result


def eval_pope_single(evaluator: CustomVLMEvaluator, limit: Optional[int],
                     track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """Single GPU POPE evaluation."""
    print("\n" + "=" * 50)
    print("Evaluating POPE")
    print("=" * 50)
    
    dataset = load_dataset("lmms-lab/POPE", split="test")
    
    if limit:
        dataset = dataset.select(range(min(limit, len(dataset))))
    
    correct = 0
    total = 0
    tp, fp, tn, fn = 0, 0, 0, 0
    all_tokens_per_step = []
    total_steps = 0
    total_gen_tokens = 0
    
    for item in tqdm(dataset, desc="POPE"):
        image = item["image"]
        question = item["question"]
        answer = item["answer"].lower()
        
        if track_tokens_per_step:
            pred, stats = evaluator.generate(image, question, return_stats=True)
            all_tokens_per_step.extend(stats.get("tokens_per_step", []))
            total_steps += stats.get("total_steps", 0)
            total_gen_tokens += stats.get("gen_length", 0)
        else:
            pred = evaluator.generate(image, question)
        pred_lower = pred.lower()
        
        if "yes" in pred_lower:
            pred_answer = "yes"
        elif "no" in pred_lower:
            pred_answer = "no"
        else:
            pred_answer = "unknown"
        
        if pred_answer == answer:
            correct += 1
            if answer == "yes":
                tp += 1
            else:
                tn += 1
        else:
            if answer == "yes":
                fn += 1
            else:
                fp += 1
        total += 1
        
        if total <= 3:
            print(f"\nQ: {question}")
            print(f"A: {answer}")
            print(f"Pred: {pred} -> {pred_answer}")
    
    accuracy = correct / total * 100 if total > 0 else 0
    precision = tp / (tp + fp) * 100 if (tp + fp) > 0 else 0
    recall = tp / (tp + fn) * 100 if (tp + fn) > 0 else 0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
    
    print(f"\nPOPE Results: Acc={accuracy:.2f}%, F1={f1:.2f}%")
    result = {"accuracy": accuracy, "precision": precision, "recall": recall, "f1": f1,
              "correct": correct, "total": total}
    if track_tokens_per_step and all_tokens_per_step:
        result["tokens_per_step_stats"] = compute_tokens_per_step_stats(all_tokens_per_step)
        result["total_gen_tokens"] = total_gen_tokens
        result["total_gen_steps"] = total_steps
    return result


def eval_textvqa_single(evaluator: CustomVLMEvaluator, limit: Optional[int],
                        track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """Single GPU TextVQA evaluation."""
    print("\n" + "=" * 50)
    print("Evaluating TextVQA")
    print("=" * 50)
    
    dataset = load_dataset("lmms-lab/textvqa", split="validation")
    
    if limit:
        dataset = dataset.select(range(min(limit, len(dataset))))
    
    correct = 0
    total = 0
    all_tokens_per_step = []
    total_steps = 0
    total_gen_tokens = 0
    
    for item in tqdm(dataset, desc="TextVQA"):
        image = item["image"]
        question = item["question"]
        answers = item["answers"]
        
        if track_tokens_per_step:
            pred, stats = evaluator.generate(image, question, return_stats=True)
            all_tokens_per_step.extend(stats.get("tokens_per_step", []))
            total_steps += stats.get("total_steps", 0)
            total_gen_tokens += stats.get("gen_length", 0)
        else:
            pred = evaluator.generate(image, question)
        
        pred_norm = normalize_answer(pred)
        matched = any(
            pred_norm == normalize_answer(ans) or 
            normalize_answer(ans) in pred_norm
            for ans in answers
        )
        
        if matched:
            correct += 1
        total += 1
        
        if total <= 3:
            print(f"\nQ: {question}")
            print(f"A: {answers[:3]}...")
            print(f"Pred: {pred}")
    
    accuracy = correct / total * 100 if total > 0 else 0
    print(f"\nTextVQA Results: {correct}/{total} = {accuracy:.2f}%")
    result = {"accuracy": accuracy, "correct": correct, "total": total}
    if track_tokens_per_step and all_tokens_per_step:
        result["tokens_per_step_stats"] = compute_tokens_per_step_stats(all_tokens_per_step)
        result["total_gen_tokens"] = total_gen_tokens
        result["total_gen_steps"] = total_steps
    return result


def eval_mmmu_single(evaluator: CustomVLMEvaluator, limit: Optional[int],
                     track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """Single GPU MMMU evaluation."""
    print("\n" + "=" * 50)
    print("Evaluating MMMU")
    print("=" * 50)
    
    dataset = load_dataset("lmms-lab/MMMU", split="validation")
    print(f"Loaded {len(dataset)} MMMU questions")
    
    if limit:
        dataset = dataset.select(range(min(limit, len(dataset))))
    
    correct = 0
    total = 0
    all_tokens_per_step = []
    total_steps = 0
    total_gen_tokens = 0
    
    for item in tqdm(dataset, desc="MMMU"):
        # Get image(s)
        images = []
        for i in range(1, 8):
            img_key = f"image_{i}"
            if img_key in item and item[img_key] is not None:
                images.append(item[img_key])
        
        if not images:
            continue
        
        image = images[0]
        
        # Build question with options (parse JSON string if needed)
        question = item["question"]
        options = parse_options(item.get("options"))
        
        if options:
            option_str = "\n".join([f"{chr(65+i)}. {opt}" for i, opt in enumerate(options)])
            question = f"{question}\n\nOptions:\n{option_str}\n\nAnswer with the letter (A, B, C, or D):"
        
        answer = item.get("answer", "")
        
        if track_tokens_per_step:
            pred, stats = evaluator.generate(image, question, return_stats=True)
            all_tokens_per_step.extend(stats.get("tokens_per_step", []))
            total_steps += stats.get("total_steps", 0)
            total_gen_tokens += stats.get("gen_length", 0)
        else:
            pred = evaluator.generate(image, question)
        pred_answer = extract_mcq_answer(pred)
        
        if pred_answer == answer.upper():
            correct += 1
        total += 1
        
        if total <= 3:
            print(f"\nQ: {question[:200]}...")
            print(f"Options: {options}")
            print(f"A: {answer}")
            print(f"Pred: {pred[:100]}... -> {pred_answer}")
    
    accuracy = correct / total * 100 if total > 0 else 0
    print(f"\nMMMU Results: {correct}/{total} = {accuracy:.2f}%")
    result = {"accuracy": accuracy, "correct": correct, "total": total}
    if track_tokens_per_step and all_tokens_per_step:
        result["tokens_per_step_stats"] = compute_tokens_per_step_stats(all_tokens_per_step)
        result["total_gen_tokens"] = total_gen_tokens
        result["total_gen_steps"] = total_steps
    return result


def eval_mmbench_single(evaluator: CustomVLMEvaluator, limit: Optional[int],
                        track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """Single GPU MMBench evaluation."""
    print("\n" + "=" * 50)
    print("Evaluating MMBench")
    print("=" * 50)
    
    dataset = load_dataset("lmms-lab/MMBench", "en", split="dev")
    print(f"Loaded {len(dataset)} MMBench questions")
    
    if limit:
        dataset = dataset.select(range(min(limit, len(dataset))))
    
    correct = 0
    total = 0
    all_tokens_per_step = []
    total_steps = 0
    total_gen_tokens = 0
    
    for item in tqdm(dataset, desc="MMBench"):
        image = item.get("image")
        if image is None:
            continue
        
        question = item.get("question", "")
        options = []
        for opt_key in ["A", "B", "C", "D"]:
            if opt_key in item and item[opt_key]:
                options.append(item[opt_key])
        
        if options:
            option_str = "\n".join([f"{chr(65+i)}. {opt}" for i, opt in enumerate(options)])
            question = f"{question}\n\nOptions:\n{option_str}\n\nAnswer with the letter (A, B, C, or D):"
        
        answer = item.get("answer", "")
        
        if track_tokens_per_step:
            pred, stats = evaluator.generate(image, question, return_stats=True)
            all_tokens_per_step.extend(stats.get("tokens_per_step", []))
            total_steps += stats.get("total_steps", 0)
            total_gen_tokens += stats.get("gen_length", 0)
        else:
            pred = evaluator.generate(image, question)
        pred_answer = extract_mcq_answer(pred)
        
        if pred_answer == answer.upper():
            correct += 1
        total += 1
        
        if total <= 3:
            print(f"\nQ: {question[:200]}...")
            print(f"A: {answer}")
            print(f"Pred: {pred[:100]}... -> {pred_answer}")
    
    accuracy = correct / total * 100 if total > 0 else 0
    print(f"\nMMBench Results: {correct}/{total} = {accuracy:.2f}%")
    result = {"accuracy": accuracy, "correct": correct, "total": total}
    if track_tokens_per_step and all_tokens_per_step:
        result["tokens_per_step_stats"] = compute_tokens_per_step_stats(all_tokens_per_step)
        result["total_gen_tokens"] = total_gen_tokens
        result["total_gen_steps"] = total_steps
    return result


def eval_mathvista_single(evaluator: CustomVLMEvaluator, limit: Optional[int],
                          track_tokens_per_step: bool = False) -> Dict[str, Any]:
    """Single GPU MathVista evaluation."""
    print("\n" + "=" * 50)
    print("Evaluating MathVista")
    print("=" * 50)
    
    dataset = load_dataset("AI4Math/MathVista", split="testmini")
    print(f"Loaded {len(dataset)} MathVista questions")
    
    # Print first item's keys for debugging
    if len(dataset) > 0:
        print(f"Dataset fields: {list(dataset[0].keys())}")
    
    if limit:
        dataset = dataset.select(range(min(limit, len(dataset))))
    
    correct = 0
    total = 0
    all_tokens_per_step = []
    total_steps = 0
    total_gen_tokens = 0
    
    for item in tqdm(dataset, desc="MathVista"):
        # Try different image field names
        image = item.get("decoded_image") or item.get("image")
        if image is None:
            continue
        
        question = item.get("question", "")
        # Try different field names for choices (choices or options, may be JSON string)
        choices = parse_options(item.get("choices")) or parse_options(item.get("options"))
        question_type = item.get("question_type", "")
        
        # If multiple choice question
        if choices and question_type == "multi_choice":
            option_str = "\n".join([f"{chr(65+i)}. {opt}" for i, opt in enumerate(choices)])
            question = f"{question}\n\nOptions:\n{option_str}\n\nAnswer with the letter:"
            answer = item.get("answer", "")
            
            if track_tokens_per_step:
                pred, stats = evaluator.generate(image, question, max_new_tokens=64, return_stats=True)
                all_tokens_per_step.extend(stats.get("tokens_per_step", []))
                total_steps += stats.get("total_steps", 0)
                total_gen_tokens += stats.get("gen_length", 0)
            else:
                pred = evaluator.generate(image, question, max_new_tokens=64)
            pred_answer = extract_mcq_answer(pred)
            
            is_correct = pred_answer == answer.upper()
        else:
            answer = str(item.get("answer", ""))
            if track_tokens_per_step:
                pred, stats = evaluator.generate(image, question, max_new_tokens=64, return_stats=True)
                all_tokens_per_step.extend(stats.get("tokens_per_step", []))
                total_steps += stats.get("total_steps", 0)
                total_gen_tokens += stats.get("gen_length", 0)
            else:
                pred = evaluator.generate(image, question, max_new_tokens=64)
            
            pred_norm = normalize_answer(pred)
            ans_norm = normalize_answer(answer)
            is_correct = pred_norm == ans_norm or ans_norm in pred_norm
        
        if is_correct:
            correct += 1
        total += 1
        
        if total <= 3:
            print(f"\nQ: {question[:200]}...")
            print(f"Type: {question_type}, Choices: {choices}")
            print(f"A: {answer}")
            print(f"Pred: {pred[:100]}...")
            print(f"Correct: {is_correct}")
    
    accuracy = correct / total * 100 if total > 0 else 0
    print(f"\nMathVista Results: {correct}/{total} = {accuracy:.2f}%")
    result = {"accuracy": accuracy, "correct": correct, "total": total}
    if track_tokens_per_step and all_tokens_per_step:
        result["tokens_per_step_stats"] = compute_tokens_per_step_stats(all_tokens_per_step)
        result["total_gen_tokens"] = total_gen_tokens
        result["total_gen_steps"] = total_steps
    return result


def main():
    parser = argparse.ArgumentParser(description="Custom VLM Evaluation")
    parser.add_argument("--model_path", type=str, required=True, help="Path to model checkpoint")
    parser.add_argument("--base_model", type=str, default="Qwen/Qwen2.5-VL-3B-Instruct", help="Base model for processor")
    parser.add_argument("--tasks", type=str, default="gqa", help="Comma-separated tasks: gqa,pope,textvqa,mmmu,mmbench,mathvista")
    parser.add_argument("--limit", type=int, default=None, help="Limit samples (for debugging)")
    parser.add_argument("--output_dir", type=str, default="eval_outputs/custom", help="Output directory")
    parser.add_argument("--num_gpus", type=int, default=1, help="Number of GPUs for data parallel")
    parser.add_argument("--min_pixels", type=int, default=None, help="Min pixels for image (default: 1280*28*28)")
    parser.add_argument("--max_pixels", type=int, default=None, help="Max pixels for image (default: 16384*28*28)")
    parser.add_argument("--track_tokens_per_step", action="store_true", default=True,
                        help="Track tokens per step statistics (enabled by default)")
    parser.add_argument("--no_track_tokens_per_step", action="store_true",
                        help="Disable tokens per step tracking")
    args = parser.parse_args()
    
    # Handle track_tokens_per_step (default True, can be disabled with --no_track_tokens_per_step)
    track_tokens = args.track_tokens_per_step and not args.no_track_tokens_per_step
    
    os.makedirs(args.output_dir, exist_ok=True)
    
    tasks = [t.strip().lower() for t in args.tasks.split(",")]
    results = {}
    
    # Pre-cache datasets in main process (only for multi-GPU)
    if args.num_gpus > 1:
        cache_datasets(tasks)
    
    for task in tasks:
        if args.num_gpus > 1:
            results[task] = run_multi_gpu(
                args.model_path, args.base_model, task, args.limit, args.num_gpus,
                args.min_pixels, args.max_pixels, track_tokens
            )
        else:
            results[task] = run_single_gpu(
                args.model_path, args.base_model, task, args.limit,
                args.min_pixels, args.max_pixels, track_tokens
            )
    
    # Save results
    output_file = os.path.join(args.output_dir, "results.json")
    with open(output_file, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to {output_file}")
    
    # Print summary
    print("\n" + "=" * 50)
    print("Summary")
    print("=" * 50)
    for task, res in results.items():
        acc = res.get('accuracy', res.get('f1', 'N/A'))
        if isinstance(acc, (int, float)):
            print(f"{task}: {acc:.2f}%")
        else:
            print(f"{task}: {acc}")
        
        # Print tokens per step stats if available
        if "tokens_per_step_stats" in res:
            stats = res["tokens_per_step_stats"]
            print(f"  Tokens/Step: mean={stats.get('mean', 0):.2f}, std={stats.get('std', 0):.2f}")
            print(f"  Total: {stats.get('total_tokens', 0)} tokens in {stats.get('total_steps', 0)} steps")
            dist_pct = stats.get('distribution_pct', {})
            print(f"  Distribution: 1-tok={dist_pct.get('1_token', 0):.1f}%, "
                  f"2-tok={dist_pct.get('2_tokens', 0):.1f}%, "
                  f"3+tok={dist_pct.get('3+_tokens', 0):.1f}%")


if __name__ == "__main__":
    main()
