#!/usr/bin/env python3
"""
Design dVLM conversion mix v2 - optimized to meet thresholds.

Changes from v1:
- Pool A weight: 70% (was 60%)
- Pool B weight: 25% (was 30%), filtered to p95 >= 100
- Pool C weight: 5% (was 10%), limited to top 5
- Add more text-only long-output datasets

Usage:
    python3 chengyue_script/dvlm/utils/design_dvlm_mix_v2.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


# Target thresholds
TARGET_P50 = 128
TARGET_P95 = 512

# Optimized pool weights (sum to 1.0)
POOL_A_WEIGHT = 0.70  # Long-output core (increased)
POOL_B_WEIGHT = 0.25  # Math/reasoning (slightly decreased)
POOL_C_WEIGHT = 0.05  # MCQ/knowledge (decreased)


def select_pool_a_datasets(pool_a: List[Dict]) -> List[str]:
    """Select datasets from Pool A (long-output core).
    
    More aggressive selection to maximize p95.
    """
    selected = []
    
    # All datasets with p95 >= 200 (keep all long-output)
    for ds in pool_a:
        if ds["p95"] >= 200:
            selected.append(ds["name"])
    
    return selected


def select_pool_b_datasets(pool_b: List[Dict]) -> List[str]:
    """Select datasets from Pool B (math/reasoning).
    
    Only include datasets with p95 >= 100 to avoid dragging down the average.
    """
    selected = []
    
    for ds in pool_b:
        if ds["p95"] >= 100:
            selected.append(ds["name"])
    
    return selected


def select_pool_c_datasets(pool_c: List[Dict]) -> List[str]:
    """Select datasets from Pool C (MCQ/knowledge).
    
    Very limited selection - top 5 only by p95.
    """
    pool_c_sorted = sorted(pool_c, key=lambda x: x["p95"], reverse=True)
    return [ds["name"] for ds in pool_c_sorted[:5]]


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."""
    
    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
        
        total_samples = sum(ds["n_total"] for ds in stats)
        if total_samples == 0:
            return 0, 0, 0
        
        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
    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],
    pool_a: List[Dict],
    pool_b: List[Dict],
    pool_c: List[Dict],
) -> str:
    """Generate YAML entry for mixtures.yaml with per-dataset stats."""
    
    lines = [
        "# dVLM conversion mix v2 (optimized)",
        "# Target: all_turns.p50 >= 128, p95 >= 512",
        "# Pool weights: A=70%, B=25%, C=5%",
        "#",
        "# Pool A: Long-output core (p95 >= 200)",
        "# Pool B: Math/reasoning (p95 >= 100)",
        "# Pool C: MCQ/knowledge (top 5 by p95)",
        "dvlm_conversion_v2:",
    ]
    
    # Pool A
    lines.append("    # Pool A: Long-output core (70%)")
    for ds_name in selected_a:
        ds_info = next((d for d in pool_a if d["name"] == ds_name), None)
        if ds_info:
            lines.append(f"    - {ds_name}  # p50={ds_info['p50']:.0f}, p95={ds_info['p95']:.0f}")
    
    # Pool B
    lines.append("    # Pool B: Math/reasoning (25%)")
    for ds_name in selected_b:
        ds_info = next((d for d in pool_b if d["name"] == ds_name), None)
        if ds_info:
            lines.append(f"    - {ds_name}  # p50={ds_info['p50']:.0f}, p95={ds_info['p95']:.0f}")
    
    # Pool C
    lines.append("    # Pool C: MCQ/knowledge (5%)")
    for ds_name in selected_c:
        ds_info = next((d for d in pool_c if d["name"] == ds_name), None)
        if ds_info:
            lines.append(f"    - {ds_name}  # p50={ds_info['p50']:.0f}, p95={ds_info['p95']:.0f}")
    
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description="Design dVLM conversion mix v2")
    parser.add_argument("--categories_json", type=str, default="output/vila_stats/dvlm_mix_categories.json")
    parser.add_argument("--output_dir", type=str, default="output/vila_stats")
    args = parser.parse_args()
    
    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 v2 (OPTIMIZED)")
    print("=" * 70)
    print(f"Target: 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"Pool A (long-output, p95>=200): {len(selected_a)} datasets")
    print(f"Pool B (math/reasoning, p95>=100): {len(selected_b)} datasets")
    print(f"Pool C (MCQ/knowledge, top 5): {len(selected_c)} datasets")
    print()
    
    # Compute stats
    stats = compute_weighted_stats(pool_a, pool_b, pool_c, selected_a, selected_b, selected_c)
    
    # Print pool stats
    print("-" * 70)
    print("POOL STATISTICS")
    print("-" * 70)
    
    for pool_name, ps in [("Pool A", stats["pool_a"]), ("Pool B", stats["pool_b"]), ("Pool C", stats["pool_c"])]:
        print(f"{pool_name} ({ps['weight']:.0%} weight): {ps['n_datasets']} datasets, {ps['n_samples']:,} samples")
        print(f"  p50={ps['p50']:.1f}, p95={ps['p95']:.1f}")
    
    # Validation
    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 = "✓ THRESHOLDS MET" if mix["meets_all_thresholds"] else "✗ THRESHOLDS NOT MET"
    
    print(f"p50 >= {TARGET_P50}: {p50_status}")
    print(f"p95 >= {TARGET_P95}: {p95_status}")
    print(f"\nOverall: {overall}")
    
    # If still not meeting, suggest further adjustments
    if not mix["meets_all_thresholds"]:
        print()
        print("-" * 70)
        print("RECOMMENDATION")
        print("-" * 70)
        if not mix["meets_p95_threshold"]:
            deficit = TARGET_P95 - mix["weighted_p95"]
            print(f"p95 deficit: {deficit:.1f} tokens")
            print("Options:")
            print("  1. Increase Pool A weight further (to 80%+)")
            print("  2. Remove Pool B entirely (math/reasoning has low p95)")
            print("  3. Add more long-output text-only datasets (OpenHermes, etc.)")
    
    # Output files
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    mix_design = {
        "version": "v2",
        "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_v2_design.json"
    with open(json_path, "w") as f:
        json.dump(mix_design, f, indent=2)
    print(f"\nWrote: {json_path}")
    
    yaml_content = generate_mixture_yaml_entry(selected_a, selected_b, selected_c, pool_a, pool_b, pool_c)
    yaml_path = output_dir / "dvlm_conversion_v2_mixture.yaml"
    with open(yaml_path, "w") as f:
        f.write(yaml_content)
    print(f"Wrote: {yaml_path}")
    
    # Print the YAML
    print("\n" + "=" * 70)
    print("MIXTURE YAML ENTRY")
    print("=" * 70)
    print()
    print(yaml_content)
    
    return 0 if mix["meets_all_thresholds"] else 1


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


