#!/usr/bin/env python3
"""
Categorize VILA datasets into pools for dVLM conversion mix design.

Based on the full registry stats, categorizes datasets into:
- Pool A (Long-output core): turn_all_p95 >= 200, fail_ratio < 5%, no video
- Pool B (Math/reasoning): name contains math|gsm|tabmwp|cot|science|reason
- Pool C (MCQ/knowledge): name contains mmmu|mmc|instruct|qa but turn_all_p95 < 100

Usage:
    python3 chengyue_script/dvlm/utils/categorize_vila_datasets.py \
        --stats_json output/vila_stats/vila_registry_stats.json \
        --output_dir output/vila_stats
"""

import json
import re
import argparse
from pathlib import Path
from typing import Dict, List, Any, Tuple


# Video dataset patterns to exclude
VIDEO_PATTERNS = [
    r'video',
    r'activitynet',
    r'charades',
    r'didemo',
    r'youcook',
    r'msrvtt',
    r'msvd',
    r'nextqa',
    r'ssv2',
    r'k710',
    r'egoschema',
    r'shot2story',
]

# Math/reasoning patterns
MATH_PATTERNS = [
    r'math',
    r'gsm',
    r'tabmwp',
    r'_cot',
    r'science',
    r'reason',
    r'geometry',
    r'geo',
]

# MCQ/knowledge patterns
MCQ_PATTERNS = [
    r'mmmu',
    r'mmc',
    r'instruct',
    r'vqa',
    r'_qa',
    r'bench',
]


def is_video_dataset(name: str) -> bool:
    """Check if dataset is video-related."""
    name_lower = name.lower()
    for pattern in VIDEO_PATTERNS:
        if re.search(pattern, name_lower):
            return True
    return False


def is_math_dataset(name: str) -> bool:
    """Check if dataset is math/reasoning related."""
    name_lower = name.lower()
    for pattern in MATH_PATTERNS:
        if re.search(pattern, name_lower):
            return True
    return False


def is_mcq_dataset(name: str) -> bool:
    """Check if dataset is MCQ/knowledge related."""
    name_lower = name.lower()
    for pattern in MCQ_PATTERNS:
        if re.search(pattern, name_lower):
            return True
    return False


def get_dataset_stats(ds_info: Dict) -> Tuple[float, float, float, int]:
    """Extract key stats from dataset info."""
    if "fatal_error" in ds_info or "per_turn_tokens" not in ds_info:
        return None, None, None, None
    
    all_turns = ds_info.get("per_turn_tokens", {}).get("all_turns", {})
    p50 = all_turns.get("median", 0)
    p95 = all_turns.get("p95", 0)
    
    n_ok = ds_info.get("n_ok", 0)
    n_fail = ds_info.get("n_fail", 0)
    fail_ratio = n_fail / (n_ok + n_fail) if (n_ok + n_fail) > 0 else 1.0
    
    n_total = ds_info.get("n_total", n_ok)
    
    return p50, p95, fail_ratio, n_total


def categorize_datasets(stats: Dict) -> Dict[str, List[Dict]]:
    """Categorize datasets into pools."""
    pool_a = []  # Long-output core
    pool_b = []  # Math/reasoning
    pool_c = []  # MCQ/knowledge (short-answer)
    excluded_video = []
    excluded_failed = []
    excluded_short = []
    uncategorized = []
    
    datasets = stats.get("datasets", {})
    
    for name, info in datasets.items():
        p50, p95, fail_ratio, n_total = get_dataset_stats(info)
        
        # Skip datasets with fatal errors or no stats
        if p50 is None:
            excluded_failed.append({
                "name": name,
                "reason": info.get("fatal_error", "No stats available")
            })
            continue
        
        ds_entry = {
            "name": name,
            "p50": p50,
            "p95": p95,
            "fail_ratio": fail_ratio,
            "n_total": n_total,
        }
        
        # Check if video (exclude)
        if is_video_dataset(name):
            excluded_video.append(ds_entry)
            continue
        
        # Check fail ratio
        if fail_ratio >= 0.05:
            excluded_failed.append({
                "name": name,
                "reason": f"fail_ratio={fail_ratio:.1%}"
            })
            continue
        
        # Categorize into pools
        is_math = is_math_dataset(name)
        is_mcq = is_mcq_dataset(name)
        
        # Pool A: Long-output (p95 >= 200)
        if p95 >= 200:
            pool_a.append(ds_entry)
        # Pool B: Math/reasoning (regardless of length)
        elif is_math and p95 >= 50:
            pool_b.append(ds_entry)
        # Pool C: MCQ/knowledge with short output
        elif is_mcq and p95 < 100:
            pool_c.append(ds_entry)
        # Short output, not categorized
        elif p95 < 50:
            excluded_short.append(ds_entry)
        else:
            uncategorized.append(ds_entry)
    
    # Sort pools by p95 descending
    pool_a.sort(key=lambda x: x["p95"], reverse=True)
    pool_b.sort(key=lambda x: x["p95"], reverse=True)
    pool_c.sort(key=lambda x: x["p95"], reverse=True)
    uncategorized.sort(key=lambda x: x["p95"], reverse=True)
    
    return {
        "pool_a_long_output": pool_a,
        "pool_b_math_reasoning": pool_b,
        "pool_c_mcq_knowledge": pool_c,
        "excluded_video": excluded_video,
        "excluded_failed": excluded_failed,
        "excluded_short": excluded_short,
        "uncategorized": uncategorized,
    }


def generate_markdown_report(categories: Dict, output_path: Path):
    """Generate markdown report of categorized datasets."""
    lines = [
        "# dVLM Mix Candidates (Auto-Categorized)",
        "",
        "Based on VILA registry stats. Target: `all_turns.p50 >= 128`, `p95 >= 512`.",
        "",
        "## Pool A: Long-Output Core (p95 >= 200, no video)",
        "**Recommended weight: 60%** - These drive p50/p95 up.",
        "",
        "| Dataset | p50 | p95 | n_total | fail_ratio |",
        "|---------|----:|----:|--------:|-----------:|",
    ]
    
    for ds in categories["pool_a_long_output"]:
        lines.append(f"| {ds['name']} | {ds['p50']:.0f} | {ds['p95']:.0f} | {ds['n_total']:,} | {ds['fail_ratio']:.1%} |")
    
    lines.extend([
        "",
        f"**Total: {len(categories['pool_a_long_output'])} datasets**",
        "",
        "## Pool B: Math/Reasoning (math|gsm|tabmwp|cot|science|reason)",
        "**Recommended weight: 30%** - For MathVista benchmark.",
        "",
        "| Dataset | p50 | p95 | n_total | fail_ratio |",
        "|---------|----:|----:|--------:|-----------:|",
    ])
    
    for ds in categories["pool_b_math_reasoning"]:
        lines.append(f"| {ds['name']} | {ds['p50']:.0f} | {ds['p95']:.0f} | {ds['n_total']:,} | {ds['fail_ratio']:.1%} |")
    
    lines.extend([
        "",
        f"**Total: {len(categories['pool_b_math_reasoning'])} datasets**",
        "",
        "## Pool C: MCQ/Knowledge (short-answer, p95 < 100)",
        "**Recommended weight: 10%** - For MMMU format, but low weight to avoid dragging p50 down.",
        "",
        "| Dataset | p50 | p95 | n_total | fail_ratio |",
        "|---------|----:|----:|--------:|-----------:|",
    ])
    
    for ds in categories["pool_c_mcq_knowledge"]:
        lines.append(f"| {ds['name']} | {ds['p50']:.0f} | {ds['p95']:.0f} | {ds['n_total']:,} | {ds['fail_ratio']:.1%} |")
    
    lines.extend([
        "",
        f"**Total: {len(categories['pool_c_mcq_knowledge'])} datasets**",
        "",
        "## Uncategorized (moderate length, not math/MCQ)",
        "",
        "| Dataset | p50 | p95 | n_total |",
        "|---------|----:|----:|--------:|",
    ])
    
    for ds in categories["uncategorized"][:30]:  # Limit to 30
        lines.append(f"| {ds['name']} | {ds['p50']:.0f} | {ds['p95']:.0f} | {ds['n_total']:,} |")
    
    if len(categories["uncategorized"]) > 30:
        lines.append(f"| ... and {len(categories['uncategorized']) - 30} more | | | |")
    
    lines.extend([
        "",
        f"**Total: {len(categories['uncategorized'])} datasets**",
        "",
        "## Excluded: Video Datasets",
        "",
    ])
    
    for ds in categories["excluded_video"][:10]:
        lines.append(f"- {ds['name']} (p50={ds['p50']:.0f}, p95={ds['p95']:.0f})")
    if len(categories["excluded_video"]) > 10:
        lines.append(f"- ... and {len(categories['excluded_video']) - 10} more")
    
    lines.extend([
        "",
        f"**Total: {len(categories['excluded_video'])} video datasets excluded**",
        "",
        "## Excluded: Short Output (p95 < 50)",
        "",
    ])
    
    for ds in categories["excluded_short"][:15]:
        lines.append(f"- {ds['name']} (p50={ds['p50']:.0f}, p95={ds['p95']:.0f})")
    if len(categories["excluded_short"]) > 15:
        lines.append(f"- ... and {len(categories['excluded_short']) - 15} more")
    
    lines.extend([
        "",
        f"**Total: {len(categories['excluded_short'])} short-output datasets excluded**",
        "",
        "## Excluded: Failed/Unavailable",
        "",
    ])
    
    for ds in categories["excluded_failed"][:10]:
        reason = ds.get("reason", "Unknown")[:60]
        lines.append(f"- {ds['name']}: {reason}")
    if len(categories["excluded_failed"]) > 10:
        lines.append(f"- ... and {len(categories['excluded_failed']) - 10} more")
    
    lines.extend([
        "",
        f"**Total: {len(categories['excluded_failed'])} datasets excluded due to errors/failures**",
    ])
    
    with open(output_path, "w") as f:
        f.write("\n".join(lines))
    
    print(f"Wrote: {output_path}")


def main():
    parser = argparse.ArgumentParser(description="Categorize VILA datasets for dVLM mix")
    parser.add_argument(
        "--stats_json",
        type=str,
        default="output/vila_stats/vila_registry_stats.json",
        help="Path to stats JSON from analyze_vila_registry_stats.py",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="output/vila_stats",
        help="Output directory for categorized results",
    )
    args = parser.parse_args()
    
    # Load stats
    stats_path = Path(args.stats_json)
    if not stats_path.exists():
        print(f"ERROR: Stats file not found: {stats_path}")
        return 1
    
    with open(stats_path, "r") as f:
        stats = json.load(f)
    
    print(f"Loaded stats for {len(stats.get('datasets', {}))} datasets")
    
    # Categorize
    categories = categorize_datasets(stats)
    
    # Print summary
    print("\n" + "=" * 60)
    print("CATEGORIZATION SUMMARY")
    print("=" * 60)
    print(f"Pool A (Long-output, p95>=200):   {len(categories['pool_a_long_output']):>4} datasets")
    print(f"Pool B (Math/reasoning):          {len(categories['pool_b_math_reasoning']):>4} datasets")
    print(f"Pool C (MCQ/knowledge, p95<100):  {len(categories['pool_c_mcq_knowledge']):>4} datasets")
    print(f"Uncategorized:                    {len(categories['uncategorized']):>4} datasets")
    print("-" * 60)
    print(f"Excluded - Video:                 {len(categories['excluded_video']):>4} datasets")
    print(f"Excluded - Short (p95<50):        {len(categories['excluded_short']):>4} datasets")
    print(f"Excluded - Failed/Unavailable:    {len(categories['excluded_failed']):>4} datasets")
    print("=" * 60)
    
    # Output directory
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Save JSON
    json_path = output_dir / "dvlm_mix_categories.json"
    with open(json_path, "w") as f:
        json.dump(categories, f, indent=2)
    print(f"Wrote: {json_path}")
    
    # Generate markdown report
    md_path = output_dir / "dvlm_mix_candidates.md"
    generate_markdown_report(categories, md_path)
    
    # Print top candidates from each pool
    print("\n" + "=" * 60)
    print("TOP CANDIDATES PER POOL")
    print("=" * 60)
    
    print("\nPool A (Long-output) - Top 10:")
    for ds in categories["pool_a_long_output"][:10]:
        print(f"  {ds['name']:<40} p50={ds['p50']:>5.0f}  p95={ds['p95']:>6.0f}")
    
    print("\nPool B (Math/reasoning) - Top 10:")
    for ds in categories["pool_b_math_reasoning"][:10]:
        print(f"  {ds['name']:<40} p50={ds['p50']:>5.0f}  p95={ds['p95']:>6.0f}")
    
    print("\nPool C (MCQ/knowledge) - Top 10:")
    for ds in categories["pool_c_mcq_knowledge"][:10]:
        print(f"  {ds['name']:<40} p50={ds['p50']:>5.0f}  p95={ds['p95']:>6.0f}")
    
    return 0


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

