#!/usr/bin/env python
# coding=utf-8
"""
Analyze input_ids length distribution for j63_recipe (or other VILA mixtures).

Usage:
    # Single thread
    python chengyue_script/dvlm/utils/analyze_dataset_length.py --mixture j63_recipe
    
    # With DataLoader num_workers for faster I/O
    python chengyue_script/dvlm/utils/analyze_dataset_length.py --mixture j63_recipe --num_workers 4 --batch_size 8

This script:
1. Loads the specified VILA mixture dataset
2. Uses DataCollatorForQwenVL to tokenize each sample
3. Computes and prints length statistics
4. Generates a histogram visualization
"""

import os
import sys
import argparse
from pathlib import Path
from tqdm import tqdm
import numpy as np
import torch
from torch.utils.data import DataLoader

# 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 / "VILA"))
sys.path.insert(0, str(PROJECT_ROOT / "examples"))

# Apply VILA path patches before any VILA imports
try:
    import vila_path_patcher
except ImportError:
    print("Warning: vila_path_patcher not found, VILA paths may not be remapped", flush=True)


def parse_args():
    parser = argparse.ArgumentParser(description="Analyze dataset input_ids length distribution")
    parser.add_argument(
        "--mixture",
        type=str,
        default="j63_recipe",
        help="VILA mixture name (default: j63_recipe)",
    )
    parser.add_argument(
        "--model_path",
        type=str,
        default="Qwen/Qwen2.5-VL-3B-Instruct",
        help="Model path for tokenizer (default: Qwen/Qwen2.5-VL-3B-Instruct)",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default=None,
        help="Output directory for histogram (default: PROJECT_ROOT/output)",
    )
    parser.add_argument(
        "--min_pixels",
        type=int,
        default=None,
        help="Min pixels for processor (default: 100352, ~316x316)",
    )
    parser.add_argument(
        "--max_pixels",
        type=int,
        default=None,
        help="Max pixels for processor (default: 100352, ~316x316)",
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=64,
        help="Number of DataLoader workers for parallel I/O (default: 4)",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=1,
        help="Batch size for DataLoader (default: 8)",
    )
    parser.add_argument(
        "--prefetch_factor",
        type=int,
        default=32,
        help="Prefetch factor for DataLoader (default: 2)",
    )
    parser.add_argument(
        "--save_lengths",
        action="store_true",
        help="Save raw lengths to a .npy file for later analysis",
    )
    return parser.parse_args()


def load_dataset_and_collator(args):
    """Load dataset and create collator."""
    from transformers import AutoProcessor, AutoTokenizer
    from lmflow.datasets.dataset import Dataset
    from lmflow.datasets.multi_modal_dataset import DataCollatorForQwenVL
    from lmflow.args import MultiModalDatasetArguments
    
    # Set environment variables for VILA
    os.environ["VILA_DATASETS"] = "cs-oci-ord"
    os.environ["DATASET_BACKEND"] = "webdataset_multi_modal"
    os.environ["MIN_PIXELS"] = str(args.min_pixels)
    os.environ["MAX_PIXELS"] = str(args.max_pixels)
    
    print(f"Loading tokenizer from: {args.model_path}", flush=True)
    tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
    
    print(f"Loading processor from: {args.model_path}", flush=True)
    processor = AutoProcessor.from_pretrained(
        args.model_path,
        trust_remote_code=True,
        min_pixels=args.min_pixels,
        max_pixels=args.max_pixels,
    )
    
    # Create data args
    data_args = MultiModalDatasetArguments(
        dataset_path=args.mixture,
    )
    data_args.tokenizer = tokenizer
    data_args.image_processor = processor.image_processor if hasattr(processor, 'image_processor') else processor
    
    print(f"Loading dataset mixture: {args.mixture}", flush=True)
    dataset = Dataset(data_args, backend="webdataset_multi_modal")
    
    print(f"Dataset loaded with {len(dataset)} samples", flush=True)
    
    # Create collator
    collator = DataCollatorForQwenVL(
        processor=processor,
        tokenizer=tokenizer,
        bd_size=None,  # No MDM padding
    )
    
    return dataset, collator


def compute_length_stats(lengths):
    """Compute and print length statistics."""
    lengths = np.array(lengths)
    
    stats = {
        "count": len(lengths),
        "min": int(np.min(lengths)),
        "max": int(np.max(lengths)),
        "mean": float(np.mean(lengths)),
        "std": float(np.std(lengths)),
        "median": float(np.median(lengths)),
        "p25": float(np.percentile(lengths, 25)),
        "p75": float(np.percentile(lengths, 75)),
        "p90": float(np.percentile(lengths, 90)),
        "p95": float(np.percentile(lengths, 95)),
        "p99": float(np.percentile(lengths, 99)),
    }
    
    print("\n" + "=" * 60)
    print("Input IDs Length Statistics")
    print("=" * 60)
    print(f"  Total samples:  {stats['count']:,}")
    print(f"  Min length:     {stats['min']:,}")
    print(f"  Max length:     {stats['max']:,}")
    print(f"  Mean length:    {stats['mean']:,.2f}")
    print(f"  Std deviation:  {stats['std']:,.2f}")
    print(f"  Median:         {stats['median']:,.2f}")
    print("-" * 60)
    print("Percentiles:")
    print(f"  25th:           {stats['p25']:,.2f}")
    print(f"  75th:           {stats['p75']:,.2f}")
    print(f"  90th:           {stats['p90']:,.2f}")
    print(f"  95th:           {stats['p95']:,.2f}")
    print(f"  99th:           {stats['p99']:,.2f}")
    print("=" * 60 + "\n")
    
    return stats


def plot_histogram(lengths, output_path, stats, mixture_name):
    """Generate and save histogram."""
    import matplotlib
    matplotlib.use('Agg')  # Non-interactive backend
    import matplotlib.pyplot as plt
    
    plt.figure(figsize=(12, 6))
    
    # Plot histogram
    n, bins, patches = plt.hist(lengths, bins=100, edgecolor='black', alpha=0.7)
    
    # Add vertical lines for key statistics
    plt.axvline(stats['mean'], color='red', linestyle='--', linewidth=2, label=f"Mean: {stats['mean']:.0f}")
    plt.axvline(stats['median'], color='green', linestyle='--', linewidth=2, label=f"Median: {stats['median']:.0f}")
    plt.axvline(stats['p95'], color='orange', linestyle='--', linewidth=2, label=f"95th: {stats['p95']:.0f}")
    
    plt.xlabel('Input IDs Length', fontsize=12)
    plt.ylabel('Frequency', fontsize=12)
    plt.title(f'{mixture_name} Input IDs Length Distribution\n(N={stats["count"]:,}, Min={stats["min"]}, Max={stats["max"]})', fontsize=14)
    plt.legend(loc='upper right', fontsize=10)
    plt.grid(True, alpha=0.3)
    
    # Add text box with stats
    textstr = f'Mean: {stats["mean"]:.1f}\nStd: {stats["std"]:.1f}\nMedian: {stats["median"]:.1f}'
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
    plt.gca().text(0.98, 0.75, textstr, transform=plt.gca().transAxes, fontsize=10,
                   verticalalignment='top', horizontalalignment='right', bbox=props)
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=150)
    plt.close()
    
    print(f"Histogram saved to: {output_path}")


def main():
    args = parse_args()
    
    # 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("=" * 60)
    print(f"Analyzing dataset: {args.mixture}")
    print(f"Model path: {args.model_path}")
    print(f"Output directory: {args.output_dir}")
    print(f"Processor pixels: {args.min_pixels} - {args.max_pixels}")
    print(f"DataLoader: batch_size={args.batch_size}, num_workers={args.num_workers}, prefetch={args.prefetch_factor}")
    print("=" * 60 + "\n")
    
    # Load dataset and collator
    dataset, collator = load_dataset_and_collator(args)
    
    # Create DataLoader for parallel I/O
    # num_workers > 0 enables background data loading
    dataloader_kwargs = {
        "batch_size": args.batch_size,
        "shuffle": False,
        "collate_fn": collator,
        "num_workers": args.num_workers,
        "pin_memory": False,
        "drop_last": False,
    }
    if args.num_workers > 0:
        dataloader_kwargs["prefetch_factor"] = args.prefetch_factor
        dataloader_kwargs["persistent_workers"] = True
    
    dataloader = DataLoader(dataset.backend_dataset, **dataloader_kwargs)
    
    # Collect lengths
    lengths = []
    errors = 0
    total_samples = len(dataset)
    
    print(f"\nProcessing {total_samples} samples with DataLoader...")
    
    try:
        for batch in tqdm(dataloader, desc="Computing lengths", total=len(dataloader)):
            try:
                # Get input_ids lengths for this batch
                input_ids = batch["input_ids"]  # [batch_size, seq_len]
                attention_mask = batch.get("attention_mask", None)
                
                if attention_mask is not None:
                    # Count non-padding tokens per sample
                    batch_lengths = attention_mask.sum(dim=1).tolist()
                else:
                    batch_lengths = [len(ids) for ids in input_ids]
                
                lengths.extend(batch_lengths)
                
            except Exception as e:
                errors += args.batch_size  # Approximate error count
                if errors <= 10 * args.batch_size:
                    print(f"\nWarning: Error processing batch: {e}", flush=True)
    except Exception as e:
        print(f"\nFatal error during iteration: {e}", flush=True)
        # Try to get partial results
        if len(lengths) == 0:
            print("No samples processed, falling back to single-sample mode...")
            # Fallback to single-sample processing
            for idx in tqdm(range(min(1000, total_samples)), desc="Fallback processing"):
                try:
                    sample = dataset.backend_dataset[idx]
                    batch = collator([sample])
                    if "attention_mask" in batch:
                        length = batch["attention_mask"][0].sum().item()
                    else:
                        length = len(batch["input_ids"][0])
                    lengths.append(length)
                except:
                    errors += 1
    
    if errors > 0:
        print(f"\nTotal errors: ~{errors} / {total_samples} samples")
    
    if len(lengths) == 0:
        print("ERROR: No samples were processed successfully!")
        return 1
    
    # Compute statistics
    stats = compute_length_stats(lengths)
    
    # Save raw lengths if requested
    if args.save_lengths:
        lengths_path = os.path.join(args.output_dir, f"{args.mixture}_lengths.npy")
        np.save(lengths_path, np.array(lengths))
        print(f"Raw lengths saved to: {lengths_path}")
    
    # Generate histogram
    hist_path = os.path.join(args.output_dir, f"{args.mixture}_length_distribution.png")
    plot_histogram(lengths, hist_path, stats, args.mixture)
    
    print("\nDone!")
    return 0


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

