#!/usr/bin/env python3
"""
Design dVLM conversion mix v3 - final optimized version.

Key insight: Pool A already includes math datasets (metamathqa, scienceqa, math)
with good p95, so we can minimize Pool B/C weight.

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

# Final pool weights
# NOTE: after applying the explicit denylist, Pool A remains strong but Pool B/C
# (short-answer heavy) can drag p95 below the target if over-weighted.
# Use a more aggressive weighting toward Pool A while keeping small coverage for
# math/MCQ formats.
POOL_A_WEIGHT = 0.97  # Long-output core
POOL_B_WEIGHT = 0.02  # Math/reasoning (minimal)
POOL_C_WEIGHT = 0.01  # MCQ/knowledge (minimal)

# Explicit denylist for dVLM conversion.
# Rationale: these datasets are domain-mismatched / not relevant for our target
# benchmarks (QA/MathVista/MMMU) in this conversion setting, even if they have
# good length stats. Keep this list as a hard constraint when generating mixes.
EXCLUDE_DATASETS = {
    # per user feedback: “明显不相关，去掉”
    "vqa_viquae_train",
    "llave_onevision_images_sft",
    "olmo_pointing_high_freq",
    "pdfa-1m",
    "squad_rendering",
}


def _filter_excluded(names: List[str]) -> List[str]:
    """Filter out explicitly excluded datasets, preserving order."""
    return [n for n in names if n not in EXCLUDE_DATASETS]


def select_pool_a_datasets(pool_a: List[Dict]) -> List[str]:
    """
    Select datasets from Pool A (p95 >= 200) with two constraints:
    - apply explicit denylist
    - avoid a few massive-but-low-p95 datasets that drag down p95 (tail length)
      when we compute weighted averages by n_total.
    """
    selected = []
    for ds in pool_a:
        name = ds["name"]
        if name in EXCLUDE_DATASETS:
            continue
        # Heuristic: if a dataset is extremely large and has relatively low p95,
        # it can dominate the weighted tail and cause the mix to miss p95>=512.
        n_total = int(ds.get("n_total", 0) or 0)
        p95 = float(ds.get("p95", 0) or 0)
        if n_total >= 500_000 and p95 < 300:
            continue
        selected.append(name)
    return selected


def select_pool_b_datasets(pool_b: List[Dict]) -> List[str]:
    """Select only the highest p95 datasets from Pool B (p95 >= 140)."""
    return _filter_excluded([ds["name"] for ds in pool_b if ds["p95"] >= 140])


def select_pool_c_datasets(pool_c: List[Dict]) -> List[str]:
    """Select top 3 datasets from Pool C."""
    pool_c_sorted = sorted(pool_c, key=lambda x: x["p95"], reverse=True)
    pool_c_sorted = [ds for ds in pool_c_sorted if ds["name"] not in EXCLUDE_DATASETS]
    return [ds["name"] for ds in pool_c_sorted[:3]]


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)
    
    # Handle case where pool might be empty
    active_weight = 0
    mix_p50 = 0
    mix_p95 = 0
    
    if selected_a:
        mix_p50 += POOL_A_WEIGHT * a_p50
        mix_p95 += POOL_A_WEIGHT * a_p95
        active_weight += POOL_A_WEIGHT
    
    if selected_b:
        mix_p50 += POOL_B_WEIGHT * b_p50
        mix_p95 += POOL_B_WEIGHT * b_p95
        active_weight += POOL_B_WEIGHT
    
    if selected_c:
        mix_p50 += POOL_C_WEIGHT * c_p50
        mix_p95 += POOL_C_WEIGHT * c_p95
        active_weight += POOL_C_WEIGHT
    
    # Normalize if not all pools have data
    if active_weight > 0 and active_weight < 1.0:
        mix_p50 /= active_weight
        mix_p95 /= active_weight
    
    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."""
    
    lines = [
        "# dVLM conversion mix v3 (final)",
        "# Target: all_turns.p50 >= 128, p95 >= 512",
        f"# Pool weights: A={POOL_A_WEIGHT:.0%}, B={POOL_B_WEIGHT:.0%}, C={POOL_C_WEIGHT:.0%}",
        "#",
        "# Pool A already includes math datasets: metamathqa, scienceqa, math, etc.",
        "dvlm_conversion_v3:",
    ]
    
    # Pool A
    lines.append(f"    # Pool A: Long-output core ({POOL_A_WEIGHT:.0%})")
    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
    if selected_b:
        lines.append(f"    # Pool B: Math/reasoning ({POOL_B_WEIGHT:.0%})")
        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
    if selected_c:
        lines.append(f"    # Pool C: MCQ/knowledge ({POOL_C_WEIGHT:.0%})")
        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 v3")
    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 v3 (FINAL)")
    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()
    print("Key insight: Pool A already includes math datasets (metamathqa, scienceqa, math)")
    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): {len(selected_a)} datasets")
    print(f"Pool B (math/reasoning, p95>=140): {len(selected_b)} datasets")
    print(f"Pool C (MCQ/knowledge, top 3): {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"])]:
        if ps['n_datasets'] > 0:
            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}")
        else:
            print(f"{pool_name} ({ps['weight']:.0%} weight): 0 datasets (empty)")
    
    # 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 = "✓ ALL THRESHOLDS MET - VILA-ONLY SUFFICIENT" 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"\n{overall}")
    
    # Output files
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    mix_design = {
        "version": "v3",
        "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_v3_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_v3_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("FINAL MIXTURE YAML ENTRY")
    print("=" * 70)
    print()
    print(yaml_content)
    
    # Total datasets count
    total = len(selected_a) + len(selected_b) + len(selected_c)
    print(f"\nTotal datasets in mix: {total}")
    
    return 0 if mix["meets_all_thresholds"] else 1


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


