#!/usr/bin/env python3
"""
Design and validate dVLM conversion mix based on categorized VILA datasets.

This script:
1. Selects datasets from each pool based on criteria
2. Computes weighted p50/p95 for the mix
3. Validates against thresholds (p50 >= 128, p95 >= 512)
4. Outputs mixture config and validation report

Usage:
    python3 chengyue_script/dvlm/utils/design_dvlm_mix.py \
        --categories_json output/vila_stats/dvlm_mix_categories.json \
        --output_dir output/vila_stats
"""

import json
import argparse
from pathlib import Path
from typing import Dict, List, Tuple
import numpy as np


# Target thresholds
TARGET_P50 = 128
TARGET_P95 = 512

# Pool weights (sum to 1.0)
POOL_A_WEIGHT = 0.60  # Long-output core
POOL_B_WEIGHT = 0.30  # Math/reasoning
POOL_C_WEIGHT = 0.10  # MCQ/knowledge


def select_pool_a_datasets(pool_a: List[Dict]) -> List[str]:
    """Select datasets from Pool A (long-output core).
    
    Criteria:
    - Prioritize high p50 (>= 100) to drive median up
    - Include high p95 datasets for tail
    - Limit total to avoid over-representation
    """
    selected = []
    
    # First pass: datasets with both high p50 (>=100) AND high p95 (>=300)
    for ds in pool_a:
        if ds["p50"] >= 100 and ds["p95"] >= 300:
            selected.append(ds["name"])
    
    # Second pass: add remaining high-p95 datasets (for tail coverage)
    for ds in pool_a:
        if ds["name"] not in selected and ds["p95"] >= 400:
            selected.append(ds["name"])
    
    # Third pass: add more if we need volume (p95 >= 200)
    for ds in pool_a:
        if ds["name"] not in selected and ds["p95"] >= 200 and ds["p50"] >= 50:
            selected.append(ds["name"])
    
    return selected


def select_pool_b_datasets(pool_b: List[Dict]) -> List[str]:
    """Select datasets from Pool B (math/reasoning).
    
    Include all available math/reasoning datasets.
    """
    return [ds["name"] for ds in pool_b]


def select_pool_c_datasets(pool_c: List[Dict]) -> List[str]:
    """Select datasets from Pool C (MCQ/knowledge).
    
    Criteria:
    - Only include those with relatively higher p95 (>= 20)
    - Limit to avoid dragging p50 down
    """
    selected = []
    
    for ds in pool_c:
        # Only include if p95 is at least 20 (not extremely short)
        if ds["p95"] >= 20:
            selected.append(ds["name"])
    
    # Limit to top 10 by p95
    pool_c_sorted = sorted([ds for ds in pool_c if ds["name"] in selected], 
                           key=lambda x: x["p95"], reverse=True)
    return [ds["name"] for ds in pool_c_sorted[:10]]


def compute_weighted_stats(
    pool_a: List[Dict], 
    pool_b: List[Dict], 
    pool_c: List[Dict],
    selected_a: List[str],
    selected_b: List[str],
    selected_c: List[str],
) -> Dict:
    """Compute weighted p50/p95 for the mix."""
    
    # Get stats for selected datasets
    def get_pool_stats(pool: List[Dict], selected: List[str]) -> Tuple[float, float, int]:
        stats = [ds for ds in pool if ds["name"] in selected]
        if not stats:
            return 0, 0, 0
        
        # Weight by n_total within pool
        total_samples = sum(ds["n_total"] for ds in stats)
        if total_samples == 0:
            return 0, 0, 0
        
        # Weighted average p50 and p95
        weighted_p50 = sum(ds["p50"] * ds["n_total"] for ds in stats) / total_samples
        weighted_p95 = sum(ds["p95"] * ds["n_total"] for ds in stats) / total_samples
        
        return weighted_p50, weighted_p95, total_samples
    
    a_p50, a_p95, a_samples = get_pool_stats(pool_a, selected_a)
    b_p50, b_p95, b_samples = get_pool_stats(pool_b, selected_b)
    c_p50, c_p95, c_samples = get_pool_stats(pool_c, selected_c)
    
    # Cross-pool weighted average based on pool weights
    # We weight by pool weight, not by sample count (to control the mix)
    mix_p50 = POOL_A_WEIGHT * a_p50 + POOL_B_WEIGHT * b_p50 + POOL_C_WEIGHT * c_p50
    mix_p95 = POOL_A_WEIGHT * a_p95 + POOL_B_WEIGHT * b_p95 + POOL_C_WEIGHT * c_p95
    
    return {
        "pool_a": {
            "p50": a_p50,
            "p95": a_p95,
            "n_datasets": len(selected_a),
            "n_samples": a_samples,
            "weight": POOL_A_WEIGHT,
        },
        "pool_b": {
            "p50": b_p50,
            "p95": b_p95,
            "n_datasets": len(selected_b),
            "n_samples": b_samples,
            "weight": POOL_B_WEIGHT,
        },
        "pool_c": {
            "p50": c_p50,
            "p95": c_p95,
            "n_datasets": len(selected_c),
            "n_samples": c_samples,
            "weight": POOL_C_WEIGHT,
        },
        "mix": {
            "weighted_p50": mix_p50,
            "weighted_p95": mix_p95,
            "meets_p50_threshold": mix_p50 >= TARGET_P50,
            "meets_p95_threshold": mix_p95 >= TARGET_P95,
            "meets_all_thresholds": mix_p50 >= TARGET_P50 and mix_p95 >= TARGET_P95,
        },
    }


def generate_mixture_yaml_entry(
    selected_a: List[str],
    selected_b: List[str],
    selected_c: List[str],
) -> str:
    """Generate YAML entry for mixtures.yaml."""
    all_datasets = selected_a + selected_b + selected_c
    
    lines = [
        "# dVLM conversion mix v1",
        "# Designed to achieve: all_turns.p50 >= 128, p95 >= 512",
        "# Pool A (60%): Long-output core",
        "# Pool B (30%): Math/reasoning",
        "# Pool C (10%): MCQ/knowledge",
        "dvlm_conversion_v1:",
    ]
    
    # Add Pool A with comments
    lines.append("    # Pool A: Long-output core (60%)")
    for ds in selected_a:
        lines.append(f"    - {ds}")
    
    # Add Pool B
    lines.append("    # Pool B: Math/reasoning (30%)")
    for ds in selected_b:
        lines.append(f"    - {ds}")
    
    # Add Pool C
    lines.append("    # Pool C: MCQ/knowledge (10%)")
    for ds in selected_c:
        lines.append(f"    - {ds}")
    
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description="Design dVLM conversion mix")
    parser.add_argument(
        "--categories_json",
        type=str,
        default="output/vila_stats/dvlm_mix_categories.json",
        help="Path to categorized datasets JSON",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="output/vila_stats",
        help="Output directory",
    )
    args = parser.parse_args()
    
    # Load categories
    with open(args.categories_json, "r") as f:
        categories = json.load(f)
    
    pool_a = categories["pool_a_long_output"]
    pool_b = categories["pool_b_math_reasoning"]
    pool_c = categories["pool_c_mcq_knowledge"]
    
    print("=" * 70)
    print("dVLM CONVERSION MIX DESIGN")
    print("=" * 70)
    print(f"Target thresholds: p50 >= {TARGET_P50}, p95 >= {TARGET_P95}")
    print(f"Pool weights: A={POOL_A_WEIGHT:.0%}, B={POOL_B_WEIGHT:.0%}, C={POOL_C_WEIGHT:.0%}")
    print()
    
    # Select datasets
    selected_a = select_pool_a_datasets(pool_a)
    selected_b = select_pool_b_datasets(pool_b)
    selected_c = select_pool_c_datasets(pool_c)
    
    print(f"Selected from Pool A (long-output): {len(selected_a)} datasets")
    print(f"Selected from Pool B (math/reasoning): {len(selected_b)} datasets")
    print(f"Selected from Pool C (MCQ/knowledge): {len(selected_c)} datasets")
    print()
    
    # Compute weighted stats
    stats = compute_weighted_stats(pool_a, pool_b, pool_c, selected_a, selected_b, selected_c)
    
    # Print results
    print("-" * 70)
    print("POOL STATISTICS")
    print("-" * 70)
    
    for pool_name, pool_stats in [("Pool A", stats["pool_a"]), 
                                   ("Pool B", stats["pool_b"]), 
                                   ("Pool C", stats["pool_c"])]:
        print(f"{pool_name} ({pool_stats['weight']:.0%} weight):")
        print(f"  Datasets: {pool_stats['n_datasets']}")
        print(f"  Samples: {pool_stats['n_samples']:,}")
        print(f"  p50: {pool_stats['p50']:.1f}")
        print(f"  p95: {pool_stats['p95']:.1f}")
        print()
    
    print("-" * 70)
    print("MIX VALIDATION")
    print("-" * 70)
    mix = stats["mix"]
    print(f"Weighted p50: {mix['weighted_p50']:.1f} (target: {TARGET_P50})")
    print(f"Weighted p95: {mix['weighted_p95']:.1f} (target: {TARGET_P95})")
    print()
    
    p50_status = "✓ PASS" if mix["meets_p50_threshold"] else "✗ FAIL"
    p95_status = "✓ PASS" if mix["meets_p95_threshold"] else "✗ FAIL"
    overall_status = "✓ VILA-ONLY SUFFICIENT" if mix["meets_all_thresholds"] else "✗ NEED TEXT-ONLY LLM DATA"
    
    print(f"p50 >= {TARGET_P50}: {p50_status}")
    print(f"p95 >= {TARGET_P95}: {p95_status}")
    print(f"\nOverall: {overall_status}")
    
    # Output
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Save mix design
    mix_design = {
        "pool_a_datasets": selected_a,
        "pool_b_datasets": selected_b,
        "pool_c_datasets": selected_c,
        "all_datasets": selected_a + selected_b + selected_c,
        "stats": stats,
        "thresholds": {
            "target_p50": TARGET_P50,
            "target_p95": TARGET_P95,
        },
        "pool_weights": {
            "pool_a": POOL_A_WEIGHT,
            "pool_b": POOL_B_WEIGHT,
            "pool_c": POOL_C_WEIGHT,
        },
    }
    
    json_path = output_dir / "dvlm_mix_v1_design.json"
    with open(json_path, "w") as f:
        json.dump(mix_design, f, indent=2)
    print(f"\nWrote: {json_path}")
    
    # Generate mixture YAML
    yaml_content = generate_mixture_yaml_entry(selected_a, selected_b, selected_c)
    yaml_path = output_dir / "dvlm_conversion_v1_mixture.yaml"
    with open(yaml_path, "w") as f:
        f.write(yaml_content)
    print(f"Wrote: {yaml_path}")
    
    # Print selected datasets
    print("\n" + "=" * 70)
    print("SELECTED DATASETS")
    print("=" * 70)
    
    print("\nPool A (Long-output core):")
    for ds in selected_a:
        ds_info = next((d for d in pool_a if d["name"] == ds), None)
        if ds_info:
            print(f"  {ds:<40} p50={ds_info['p50']:>5.0f}  p95={ds_info['p95']:>6.0f}")
    
    print("\nPool B (Math/reasoning):")
    for ds in selected_b:
        ds_info = next((d for d in pool_b if d["name"] == ds), None)
        if ds_info:
            print(f"  {ds:<40} p50={ds_info['p50']:>5.0f}  p95={ds_info['p95']:>6.0f}")
    
    print("\nPool C (MCQ/knowledge):")
    for ds in selected_c:
        ds_info = next((d for d in pool_c if d["name"] == ds), None)
        if ds_info:
            print(f"  {ds:<40} p50={ds_info['p50']:>5.0f}  p95={ds_info['p95']:>6.0f}")
    
    print("\n" + yaml_content)
    
    return 0 if mix["meets_all_thresholds"] else 1


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


