#!/usr/bin/env python
# coding=utf-8
"""
Analyze ShareGPT4V token statistics for MDM training with bd_size padding.

This script calculates:
1. Total tokens per sample (text + vision tokens)
2. Assistant tokens (tokens that participate in loss calculation)
3. Padding tokens needed for bd_size alignment (bd4, bd8)
4. Padding ratios

Usage:
    # Basic usage with defaults
    python chengyue_script/dvlm/utils/analyze_sharegpt4v_padding.py

    # With custom parameters
    python chengyue_script/dvlm/utils/analyze_sharegpt4v_padding.py \
        --dataset_json /path/to/sharegpt4v.json \
        --image_folder /path/to/images \
        --num_samples 1000 \
        --num_workers 8

    # Analyze with specific resolution (matches training config)
    python chengyue_script/dvlm/utils/analyze_sharegpt4v_padding.py \
        --min_pixels 200704 --max_pixels 200704
"""

import os
import sys
import json
import argparse
from pathlib import Path
from dataclasses import dataclass
from typing import Optional, Dict, List, Tuple
from tqdm import tqdm
import numpy as np

# Add project paths
PROJECT_ROOT = Path(__file__).resolve().parents[4]
sys.path.insert(0, str(PROJECT_ROOT / "src"))
sys.path.insert(0, str(PROJECT_ROOT / "examples"))

import torch
from torch.utils.data import Dataset, DataLoader


def parse_args():
    parser = argparse.ArgumentParser(
        description="Analyze ShareGPT4V token statistics for MDM training"
    )
    parser.add_argument(
        "--dataset_json",
        type=str,
        default="/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/dataset/vila-sft/ShareGPT4V/jason-filter-sharegpt4v_mix665k_cap23k_coco-ap9k_lcs3k_sam9k_div2k.json",
        help="Path to ShareGPT4V JSON file",
    )
    parser.add_argument(
        "--image_folder",
        type=str,
        default="/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/dataset/vila-sft/ShareGPT4V/data",
        help="Path to image folder",
    )
    parser.add_argument(
        "--model_path",
        type=str,
        default="Qwen/Qwen2.5-VL-3B-Instruct",
        help="Model path for tokenizer/processor",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default=None,
        help="Output directory for results (default: PROJECT_ROOT/output)",
    )
    parser.add_argument(
        "--num_samples",
        type=int,
        default=None,
        help="Number of samples to analyze (default: all)",
    )
    parser.add_argument(
        "--min_pixels",
        type=int,
        default=None,
        help="Min pixels for processor (default: processor default)",
    )
    parser.add_argument(
        "--max_pixels",
        type=int,
        default=None,
        help="Max pixels for processor (default: processor default)",
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=4,
        help="Number of DataLoader workers",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=1,
        help="Batch size for processing",
    )
    parser.add_argument(
        "--save_csv",
        action="store_true",
        help="Save per-sample statistics to CSV",
    )
    parser.add_argument(
        "--bd_sizes",
        type=str,
        default="4,8",
        help="Comma-separated bd_size values to analyze (default: 4,8)",
    )
    return parser.parse_args()


class ShareGPT4VDataset(Dataset):
    """Simple dataset for ShareGPT4V that returns messages format."""
    
    def __init__(self, dataset_json: str, image_folder: str, num_samples: Optional[int] = None):
        print(f"Loading dataset from: {dataset_json}", flush=True)
        with open(dataset_json, 'r') as f:
            raw_data = json.load(f)
        
        # Normalize and filter data
        self.data = []
        dropped = 0
        for item in raw_data:
            norm = self._normalize_item(item, image_folder)
            if norm is not None:
                self.data.append(norm)
            else:
                dropped += 1
        
        print(f"Loaded {len(self.data)}/{len(raw_data)} samples (dropped {dropped})", flush=True)
        
        if num_samples is not None and num_samples < len(self.data):
            self.data = self.data[:num_samples]
            print(f"Using first {num_samples} samples for analysis", flush=True)
    
    def _normalize_item(self, item: Dict, image_folder: str) -> Optional[Dict]:
        """Normalize item to messages format with full image path."""
        # Resolve image field
        image_value = item.get('image')
        if not image_value:
            image_path = item.get('image_path')
            image_paths = item.get('image_paths')
            images = item.get('images')
            if isinstance(image_path, str) and image_path:
                image_value = image_path
            elif isinstance(image_paths, list) and len(image_paths) > 0:
                image_value = image_paths[0]
            elif isinstance(images, list) and len(images) > 0:
                image_value = images[0]
        
        if not image_value:
            return None
        
        # Resolve conversations
        conv = item.get('conversations') or item.get('messages')
        if not isinstance(conv, list) or len(conv) == 0:
            return None
        
        # Build full image path
        full_image_path = os.path.join(image_folder, image_value)
        if not os.path.exists(full_image_path):
            return None
        
        # Build messages format
        messages = []
        image_added = False
        
        for m in conv:
            frm = m.get('from')
            if not frm:
                role = m.get('role')
                if role in ('user', 'human'):
                    frm = 'human'
                elif role in ('assistant', 'gpt', 'bot'):
                    frm = 'gpt'
                else:
                    frm = role
            
            value = m.get('value') or m.get('text') or m.get('content') or ''
            
            if frm == 'human':
                if not image_added:
                    messages.append({
                        "role": "user",
                        "content": [
                            {"type": "text", "text": value},
                            {"type": "image", "image": full_image_path},
                        ],
                    })
                    image_added = True
                else:
                    messages.append({
                        "role": "user",
                        "content": [{"type": "text", "text": value}],
                    })
            else:
                messages.append({
                    "role": "assistant",
                    "content": [{"type": "text", "text": value}],
                })
        
        if len(messages) == 0:
            return None
        
        # Count conversation turns (number of assistant responses)
        num_turns = sum(1 for m in messages if m["role"] == "assistant")
        
        return {"messages": messages, "num_turns": num_turns}
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]


@dataclass
class TokenStats:
    """Statistics for a single sample."""
    total_tokens: int
    assistant_tokens: int
    padding: Dict[int, int]  # bd_size -> padding tokens
    num_turns: int = 1  # number of conversation turns (assistant responses)
    per_turn_tokens: List[int] = None  # tokens per turn (assistant response lengths)
    
    def padding_ratio(self, bd_size: int) -> float:
        """Padding ratio relative to total sequence length."""
        pad = self.padding.get(bd_size, 0)
        total_with_pad = self.total_tokens + pad
        return pad / total_with_pad if total_with_pad > 0 else 0.0
    
    def loss_dilution_ratio(self, bd_size: int) -> float:
        """Padding ratio relative to (assistant + padding) - the loss-relevant portion.
        
        This measures how much of the "loss space" is diluted by padding tokens.
        Higher values mean more padding relative to actual trainable tokens.
        """
        pad = self.padding.get(bd_size, 0)
        loss_space = self.assistant_tokens + pad
        return pad / loss_space if loss_space > 0 else 0.0
    
    def padding_overhead(self, bd_size: int) -> float:
        """Padding as a percentage of assistant tokens.
        
        This measures the extra computation overhead relative to actual loss tokens.
        """
        pad = self.padding.get(bd_size, 0)
        return pad / self.assistant_tokens if self.assistant_tokens > 0 else 0.0


def calc_padding(seq_len: int, bd_size: int) -> int:
    """Calculate padding needed to align seq_len to bd_size."""
    return (bd_size - (seq_len % bd_size)) % bd_size


def create_collator_for_stats(processor, tokenizer, bd_sizes: List[int]):
    """Create a collator that computes token statistics."""
    from qwen_vl_utils import process_vision_info
    
    # Get special token IDs
    im_end_token = tokenizer.encode("<|im_end|>", add_special_tokens=False)
    im_end_id = im_end_token[0] if len(im_end_token) == 1 else 151645
    
    # Target sequence for finding assistant responses: <|im_start|>assistant\n
    target_sequence = torch.tensor([151644, 77091, 198], dtype=torch.long)
    
    def collate_fn(instances):
        stats_list = []
        
        for ins in instances:
            messages = ins.get("messages", [])
            if not messages:
                continue
            
            # Get number of conversation turns
            num_turns = ins.get("num_turns", sum(1 for m in messages if m.get("role") == "assistant"))
            
            try:
                # Apply chat template
                text = processor.apply_chat_template(messages, tokenize=False)
                image_inputs, video_inputs = process_vision_info(messages)
                
                # Process
                proc_kwargs = {"text": [text], "padding": False, "return_tensors": "pt"}
                if image_inputs is not None and len(image_inputs) > 0:
                    proc_kwargs["images"] = [image_inputs]
                if video_inputs is not None and len(video_inputs) > 0:
                    proc_kwargs["videos"] = [video_inputs]
                
                inputs = processor(**proc_kwargs)
                input_ids = inputs["input_ids"].squeeze(0)
                
                total_tokens = input_ids.size(0)
                
                # Calculate padding for each bd_size
                padding = {bd: calc_padding(total_tokens, bd) for bd in bd_sizes}
                
                # Count assistant tokens (tokens that would have labels != -100)
                # Also track per-turn token counts
                labels = torch.full_like(input_ids, -100)
                seq = input_ids
                T = target_sequence.numel()
                per_turn_tokens = []
                
                i = 0
                while i <= seq.numel() - T:
                    if torch.equal(seq[i:i+T], target_sequence):
                        response_start = i + T
                        response_end = response_start
                        
                        while response_end < seq.numel():
                            if seq[response_end].item() == im_end_id:
                                response_end += 1
                                break
                            if response_end + T <= seq.numel() and torch.equal(seq[response_end:response_end+T], target_sequence):
                                break
                            response_end += 1
                        
                        if response_start < response_end:
                            labels[response_start:response_end] = seq[response_start:response_end]
                            # Track tokens for this turn
                            per_turn_tokens.append(response_end - response_start)
                        
                        i = response_end + 1
                    else:
                        i += 1
                
                assistant_tokens = (labels != -100).sum().item()
                
                stats_list.append(TokenStats(
                    total_tokens=total_tokens,
                    assistant_tokens=assistant_tokens,
                    padding=padding,
                    num_turns=num_turns,
                    per_turn_tokens=per_turn_tokens,
                ))
                
            except Exception as e:
                print(f"Warning: Error processing sample: {e}", flush=True)
                continue
        
        return stats_list
    
    return collate_fn


def compute_aggregate_stats(all_stats: List[TokenStats], bd_sizes: List[int]) -> Dict:
    """Compute aggregate statistics."""
    if len(all_stats) == 0:
        return {}
    
    total_tokens = np.array([s.total_tokens for s in all_stats])
    assistant_tokens = np.array([s.assistant_tokens for s in all_stats])
    num_turns = np.array([s.num_turns for s in all_stats])
    
    # Count single-turn vs multi-turn
    single_turn_count = int((num_turns == 1).sum())
    multi_turn_count = int((num_turns > 1).sum())
    
    # Collect per-turn token statistics
    # Group tokens by turn index (1st turn, 2nd turn, etc.)
    max_turn_idx = int(num_turns.max())
    per_turn_stats = {}
    for turn_idx in range(1, min(max_turn_idx + 1, 11)):  # Limit to 10 turns
        turn_tokens = []
        for s in all_stats:
            if s.per_turn_tokens and len(s.per_turn_tokens) >= turn_idx:
                turn_tokens.append(s.per_turn_tokens[turn_idx - 1])
        if turn_tokens:
            turn_tokens_arr = np.array(turn_tokens)
            per_turn_stats[f"turn_{turn_idx}"] = {
                "count": len(turn_tokens),
                "min": int(turn_tokens_arr.min()),
                "max": int(turn_tokens_arr.max()),
                "mean": float(turn_tokens_arr.mean()),
                "std": float(turn_tokens_arr.std()),
                "median": float(np.median(turn_tokens_arr)),
                "sum": int(turn_tokens_arr.sum()),
            }
    
    # All turns combined statistics
    all_turn_tokens = []
    for s in all_stats:
        if s.per_turn_tokens:
            all_turn_tokens.extend(s.per_turn_tokens)
    if all_turn_tokens:
        all_turn_tokens_arr = np.array(all_turn_tokens)
        per_turn_stats["all_turns"] = {
            "count": len(all_turn_tokens),
            "min": int(all_turn_tokens_arr.min()),
            "max": int(all_turn_tokens_arr.max()),
            "mean": float(all_turn_tokens_arr.mean()),
            "std": float(all_turn_tokens_arr.std()),
            "median": float(np.median(all_turn_tokens_arr)),
            "sum": int(all_turn_tokens_arr.sum()),
        }
    
    results = {
        "num_samples": len(all_stats),
        "conversation_turns": {
            "single_turn_count": single_turn_count,
            "multi_turn_count": multi_turn_count,
            "single_turn_ratio": float(single_turn_count / len(all_stats)),
            "multi_turn_ratio": float(multi_turn_count / len(all_stats)),
            "min_turns": int(num_turns.min()),
            "max_turns": int(num_turns.max()),
            "mean_turns": float(num_turns.mean()),
            "median_turns": float(np.median(num_turns)),
            "turn_distribution": {int(t): int((num_turns == t).sum()) for t in sorted(set(num_turns)) if t <= 10},
        },
        "per_turn_tokens": per_turn_stats,
        "total_tokens": {
            "sum": int(total_tokens.sum()),
            "min": int(total_tokens.min()),
            "max": int(total_tokens.max()),
            "mean": float(total_tokens.mean()),
            "std": float(total_tokens.std()),
            "median": float(np.median(total_tokens)),
            "p25": float(np.percentile(total_tokens, 25)),
            "p75": float(np.percentile(total_tokens, 75)),
            "p90": float(np.percentile(total_tokens, 90)),
            "p95": float(np.percentile(total_tokens, 95)),
            "p99": float(np.percentile(total_tokens, 99)),
        },
        "assistant_tokens": {
            "sum": int(assistant_tokens.sum()),
            "min": int(assistant_tokens.min()),
            "max": int(assistant_tokens.max()),
            "mean": float(assistant_tokens.mean()),
            "std": float(assistant_tokens.std()),
            "median": float(np.median(assistant_tokens)),
        },
        "assistant_ratio": {
            "mean": float((assistant_tokens / total_tokens).mean()),
            "median": float(np.median(assistant_tokens / total_tokens)),
        },
    }
    
    # Per bd_size stats
    for bd in bd_sizes:
        padding = np.array([s.padding[bd] for s in all_stats])
        padding_ratios = np.array([s.padding_ratio(bd) for s in all_stats])
        
        results[f"padding_bd{bd}"] = {
            "sum": int(padding.sum()),
            "min": int(padding.min()),
            "max": int(padding.max()),
            "mean": float(padding.mean()),
            "std": float(padding.std()),
            "median": float(np.median(padding)),
        }
        results[f"padding_ratio_bd{bd}"] = {
            "mean": float(padding_ratios.mean()),
            "std": float(padding_ratios.std()),
            "median": float(np.median(padding_ratios)),
            "max": float(padding_ratios.max()),
            "p90": float(np.percentile(padding_ratios, 90)),
            "p95": float(np.percentile(padding_ratios, 95)),
        }
        
        # Loss dilution: padding / (assistant + padding)
        # This measures how much of the "loss space" is occupied by padding
        loss_dilution = np.array([s.loss_dilution_ratio(bd) for s in all_stats])
        results[f"loss_dilution_bd{bd}"] = {
            "mean": float(loss_dilution.mean()),
            "std": float(loss_dilution.std()),
            "median": float(np.median(loss_dilution)),
            "max": float(loss_dilution.max()),
            "p90": float(np.percentile(loss_dilution, 90)),
            "p95": float(np.percentile(loss_dilution, 95)),
        }
        
        # Padding overhead: padding / assistant_tokens
        # This measures extra computation relative to actual loss tokens
        padding_overhead = np.array([s.padding_overhead(bd) for s in all_stats])
        results[f"padding_overhead_bd{bd}"] = {
            "mean": float(padding_overhead.mean()),
            "std": float(padding_overhead.std()),
            "median": float(np.median(padding_overhead)),
            "max": float(padding_overhead.max()),
            "p90": float(np.percentile(padding_overhead, 90)),
            "p95": float(np.percentile(padding_overhead, 95)),
        }
    
    return results


def print_summary(stats: Dict, bd_sizes: List[int]):
    """Print a formatted summary of statistics."""
    print("\n" + "=" * 70)
    print("ShareGPT4V MDM Token Statistics Summary")
    print("=" * 70)
    
    print(f"\nTotal Samples Analyzed: {stats['num_samples']:,}")
    
    # Conversation turns statistics
    ct = stats['conversation_turns']
    print("\n" + "-" * 70)
    print("CONVERSATION TURNS")
    print("-" * 70)
    print(f"  Single-turn (1 turn):   {ct['single_turn_count']:>10,} ({ct['single_turn_ratio']:>6.2%})")
    print(f"  Multi-turn (>1 turn):   {ct['multi_turn_count']:>10,} ({ct['multi_turn_ratio']:>6.2%})")
    print()
    print(f"  Min turns:              {ct['min_turns']:>10}")
    print(f"  Max turns:              {ct['max_turns']:>10}")
    print(f"  Mean turns:             {ct['mean_turns']:>10.2f}")
    print(f"  Median turns:           {ct['median_turns']:>10.2f}")
    print()
    print("  Turn distribution:")
    for turns, count in sorted(ct['turn_distribution'].items()):
        pct = count / stats['num_samples'] * 100
        bar = "█" * int(pct / 2)
        label = f"{turns} turn{'s' if turns > 1 else ' '}"
        print(f"    {label}: {count:>8,} ({pct:>5.1f}%) {bar}")
    
    # Per-turn token statistics
    if 'per_turn_tokens' in stats and stats['per_turn_tokens']:
        print("\n" + "-" * 70)
        print("PER-TURN TOKEN STATISTICS (assistant response lengths)")
        print("-" * 70)
        
        pts = stats['per_turn_tokens']
        
        # Print all_turns summary first
        if 'all_turns' in pts:
            at = pts['all_turns']
            print(f"  All turns combined:")
            print(f"    Total responses:  {at['count']:>10,}")
            print(f"    Sum tokens:       {at['sum']:>10,}")
            print(f"    Mean tokens:      {at['mean']:>10.1f}")
            print(f"    Median tokens:    {at['median']:>10.1f}")
            print(f"    Min/Max:          {at['min']:>5} / {at['max']:<5}")
        
        print()
        print("  Per-turn breakdown:")
        print(f"    {'Turn':<8} {'Count':>10} {'Mean':>10} {'Median':>10} {'Min':>8} {'Max':>8}")
        print(f"    {'-'*6:<8} {'-'*8:>10} {'-'*8:>10} {'-'*8:>10} {'-'*6:>8} {'-'*6:>8}")
        
        for key in sorted([k for k in pts.keys() if k.startswith('turn_')], 
                          key=lambda x: int(x.split('_')[1])):
            turn_idx = int(key.split('_')[1])
            t = pts[key]
            print(f"    Turn {turn_idx:<3} {t['count']:>10,} {t['mean']:>10.1f} {t['median']:>10.1f} {t['min']:>8,} {t['max']:>8,}")
    
    print("\n" + "-" * 70)
    print("TOTAL TOKENS (per sample)")
    print("-" * 70)
    ts = stats['total_tokens']
    print(f"  Sum:        {ts['sum']:>15,} tokens")
    print(f"  Min:        {ts['min']:>15,}")
    print(f"  Max:        {ts['max']:>15,}")
    print(f"  Mean:       {ts['mean']:>15,.2f}")
    print(f"  Std:        {ts['std']:>15,.2f}")
    print(f"  Median:     {ts['median']:>15,.2f}")
    print(f"  P25:        {ts['p25']:>15,.2f}")
    print(f"  P75:        {ts['p75']:>15,.2f}")
    print(f"  P90:        {ts['p90']:>15,.2f}")
    print(f"  P95:        {ts['p95']:>15,.2f}")
    print(f"  P99:        {ts['p99']:>15,.2f}")
    
    print("\n" + "-" * 70)
    print("ASSISTANT TOKENS (participate in loss)")
    print("-" * 70)
    at = stats['assistant_tokens']
    ar = stats['assistant_ratio']
    print(f"  Sum:        {at['sum']:>15,} tokens")
    print(f"  Mean:       {at['mean']:>15,.2f}")
    print(f"  Median:     {at['median']:>15,.2f}")
    print(f"  Ratio Mean: {ar['mean']:>15.2%}")
    print(f"  Ratio Med:  {ar['median']:>15.2%}")
    
    for bd in bd_sizes:
        print("\n" + "-" * 70)
        print(f"PADDING FOR BD_SIZE={bd}")
        print("-" * 70)
        pad = stats[f'padding_bd{bd}']
        ratio = stats[f'padding_ratio_bd{bd}']
        loss_dil = stats[f'loss_dilution_bd{bd}']
        pad_over = stats[f'padding_overhead_bd{bd}']
        
        print(f"  Sum Padding:    {pad['sum']:>12,} tokens")
        print(f"  Mean Padding:   {pad['mean']:>12,.2f} tokens/sample")
        print(f"  Max Padding:    {pad['max']:>12,} tokens")
        print(f"  Median Padding: {pad['median']:>12,.2f}")
        print()
        print("  [Padding Ratio: padding / (total + padding)]")
        print(f"    Mean:         {ratio['mean']:>12.4%}")
        print(f"    Median:       {ratio['median']:>12.4%}")
        print(f"    Max:          {ratio['max']:>12.4%}")
        print()
        print("  [Loss Dilution: padding / (assistant + padding)] ⭐")
        print("  (How much of the loss-relevant space is padding)")
        print(f"    Mean:         {loss_dil['mean']:>12.4%}")
        print(f"    Median:       {loss_dil['median']:>12.4%}")
        print(f"    Max:          {loss_dil['max']:>12.4%}")
        print(f"    P90:          {loss_dil['p90']:>12.4%}")
        print(f"    P95:          {loss_dil['p95']:>12.4%}")
        print()
        print("  [Padding Overhead: padding / assistant_tokens]")
        print("  (Extra padding relative to actual loss tokens)")
        print(f"    Mean:         {pad_over['mean']:>12.4%}")
        print(f"    Median:       {pad_over['median']:>12.4%}")
        print(f"    Max:          {pad_over['max']:>12.4%}")
    
    print("\n" + "=" * 70)
    print("IMPACT ANALYSIS (Aggregate)")
    print("=" * 70)
    
    total_tokens = stats['total_tokens']['sum']
    assistant_tokens = stats['assistant_tokens']['sum']
    
    for bd in bd_sizes:
        total_padding = stats[f'padding_bd{bd}']['sum']
        total_with_padding = total_tokens + total_padding
        loss_space = assistant_tokens + total_padding
        
        print(f"\n  BD_SIZE={bd}:")
        print(f"    Total tokens (original):     {total_tokens:>15,}")
        print(f"    Assistant tokens (for loss): {assistant_tokens:>15,}")
        print(f"    Total padding tokens:        {total_padding:>15,}")
        print()
        print(f"    ⭐ Loss Dilution (aggregate): {total_padding/loss_space:>14.4%}")
        print(f"       = padding / (assistant + padding)")
        print(f"       = {total_padding:,} / ({assistant_tokens:,} + {total_padding:,})")
        print()
        print(f"    Padding Overhead (aggregate): {total_padding/assistant_tokens:>13.4%}")
        print(f"       = padding / assistant_tokens")
    
    print("\n" + "=" * 70 + "\n")


def plot_distributions(all_stats: List[TokenStats], bd_sizes: List[int], output_dir: str):
    """Generate distribution plots."""
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    
    total_tokens = np.array([s.total_tokens for s in all_stats])
    assistant_tokens = np.array([s.assistant_tokens for s in all_stats])
    
    # Create figure with subplots
    n_plots = 2 + len(bd_sizes)
    fig, axes = plt.subplots(2, (n_plots + 1) // 2, figsize=(14, 10))
    axes = axes.flatten()
    
    # Plot 1: Total tokens distribution
    ax = axes[0]
    ax.hist(total_tokens, bins=50, edgecolor='black', alpha=0.7)
    ax.axvline(np.mean(total_tokens), color='red', linestyle='--', label=f'Mean: {np.mean(total_tokens):.0f}')
    ax.axvline(np.median(total_tokens), color='green', linestyle='--', label=f'Median: {np.median(total_tokens):.0f}')
    ax.set_xlabel('Total Tokens')
    ax.set_ylabel('Frequency')
    ax.set_title('Total Tokens Distribution')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    # Plot 2: Assistant tokens distribution
    ax = axes[1]
    ax.hist(assistant_tokens, bins=50, edgecolor='black', alpha=0.7, color='orange')
    ax.axvline(np.mean(assistant_tokens), color='red', linestyle='--', label=f'Mean: {np.mean(assistant_tokens):.0f}')
    ax.set_xlabel('Assistant Tokens')
    ax.set_ylabel('Frequency')
    ax.set_title('Assistant Tokens (Loss) Distribution')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    # Plot padding ratios for each bd_size
    for i, bd in enumerate(bd_sizes):
        ax = axes[2 + i]
        ratios = np.array([s.padding_ratio(bd) * 100 for s in all_stats])
        ax.hist(ratios, bins=bd, edgecolor='black', alpha=0.7, color='purple')
        ax.axvline(np.mean(ratios), color='red', linestyle='--', label=f'Mean: {np.mean(ratios):.2f}%')
        ax.set_xlabel(f'Padding Ratio (%)')
        ax.set_ylabel('Frequency')
        ax.set_title(f'Padding Ratio Distribution (bd_size={bd})')
        ax.legend()
        ax.grid(True, alpha=0.3)
    
    # Hide unused subplots
    for i in range(2 + len(bd_sizes), len(axes)):
        axes[i].set_visible(False)
    
    plt.tight_layout()
    output_path = os.path.join(output_dir, 'sharegpt4v_token_distributions.png')
    plt.savefig(output_path, dpi=150)
    plt.close()
    print(f"Distribution plots saved to: {output_path}")


def save_to_csv(all_stats: List[TokenStats], bd_sizes: List[int], output_path: str):
    """Save per-sample statistics to CSV."""
    import csv
    
    headers = ['sample_idx', 'num_turns', 'total_tokens', 'assistant_tokens', 'assistant_ratio']
    for bd in bd_sizes:
        headers.extend([
            f'padding_bd{bd}', 
            f'padding_ratio_bd{bd}',
            f'loss_dilution_bd{bd}',
            f'padding_overhead_bd{bd}'
        ])
    
    with open(output_path, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(headers)
        
        for idx, s in enumerate(all_stats):
            row = [idx, s.num_turns, s.total_tokens, s.assistant_tokens, 
                   s.assistant_tokens / s.total_tokens if s.total_tokens > 0 else 0]
            for bd in bd_sizes:
                row.extend([
                    s.padding[bd], 
                    s.padding_ratio(bd),
                    s.loss_dilution_ratio(bd),
                    s.padding_overhead(bd)
                ])
            writer.writerow(row)
    
    print(f"Per-sample statistics saved to: {output_path}")


def main():
    args = parse_args()
    
    # Parse bd_sizes
    bd_sizes = [int(x.strip()) for x in args.bd_sizes.split(',')]
    
    # Set output directory
    if args.output_dir is None:
        args.output_dir = str(PROJECT_ROOT / "output")
    os.makedirs(args.output_dir, exist_ok=True)
    
    print("=" * 70)
    print("ShareGPT4V MDM Token Padding Analysis")
    print("=" * 70)
    print(f"Dataset JSON:  {args.dataset_json}")
    print(f"Image Folder:  {args.image_folder}")
    print(f"Model:         {args.model_path}")
    print(f"BD Sizes:      {bd_sizes}")
    print(f"Num Samples:   {args.num_samples or 'all'}")
    print(f"Output Dir:    {args.output_dir}")
    if args.min_pixels or args.max_pixels:
        print(f"Pixels:        {args.min_pixels} - {args.max_pixels}")
    print("=" * 70 + "\n")
    
    # Load processor and tokenizer
    from transformers import AutoProcessor, AutoTokenizer
    
    print("Loading tokenizer and processor...", flush=True)
    tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
    
    processor_kwargs = {}
    if args.min_pixels is not None:
        processor_kwargs["min_pixels"] = args.min_pixels
    if args.max_pixels is not None:
        processor_kwargs["max_pixels"] = args.max_pixels
    
    processor = AutoProcessor.from_pretrained(
        args.model_path, trust_remote_code=True, **processor_kwargs
    )
    
    # Log processor config
    if hasattr(processor, 'image_processor'):
        ip = processor.image_processor
        print(f"Processor min_pixels: {getattr(ip, 'min_pixels', 'N/A')}")
        print(f"Processor max_pixels: {getattr(ip, 'max_pixels', 'N/A')}")
    
    # Load dataset
    dataset = ShareGPT4VDataset(
        args.dataset_json, 
        args.image_folder, 
        num_samples=args.num_samples
    )
    
    if len(dataset) == 0:
        print("ERROR: No valid samples found in dataset!")
        return 1
    
    # Create collator
    collate_fn = create_collator_for_stats(processor, tokenizer, bd_sizes)
    
    # Create dataloader
    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=collate_fn,
        num_workers=args.num_workers,
        prefetch_factor=2 if args.num_workers > 0 else None,
    )
    
    # Collect statistics
    all_stats = []
    errors = 0
    
    print(f"\nProcessing {len(dataset)} samples...")
    for batch_stats in tqdm(dataloader, desc="Analyzing tokens"):
        all_stats.extend(batch_stats)
    
    if len(all_stats) == 0:
        print("ERROR: No samples were processed successfully!")
        return 1
    
    print(f"\nSuccessfully processed {len(all_stats)}/{len(dataset)} samples")
    
    # Compute aggregate statistics
    agg_stats = compute_aggregate_stats(all_stats, bd_sizes)
    
    # Print summary
    print_summary(agg_stats, bd_sizes)
    
    # Save JSON summary
    json_path = os.path.join(args.output_dir, 'sharegpt4v_token_stats.json')
    with open(json_path, 'w') as f:
        json.dump(agg_stats, f, indent=2)
    print(f"Statistics saved to: {json_path}")
    
    # Generate plots
    try:
        plot_distributions(all_stats, bd_sizes, args.output_dir)
    except Exception as e:
        print(f"Warning: Could not generate plots: {e}")
    
    # Save CSV if requested
    if args.save_csv:
        csv_path = os.path.join(args.output_dir, 'sharegpt4v_token_stats.csv')
        save_to_csv(all_stats, bd_sizes, csv_path)
    
    print("\nDone!")
    return 0


if __name__ == "__main__":
    sys.exit(main())

