#!/usr/bin/env python3
# coding=utf-8
"""
Fast-screen VILA registry datasets for dVLM conversion.

Motivation
----------
In dVLM training, "assistant decoding output too short" can cause instability/issues.
This script scans VILA registry datasets and summarizes:
  - assistant output token lengths (labels != IGNORE_INDEX)
  - per-turn assistant response lengths (all turns combined + per-turn breakdown)
  - turn count distribution (single-turn vs multi-turn)
  - failure rate (exceptions during sampling/processing)
  - optional bd_size padding overhead estimates (computed from seq_len only)

Default registry profile
------------------------
Use a clean registry alias:
  export VILA_DATASETS=dataset_analysis
  export DATASET_BACKEND=webdataset_multi_modal

Example
-------
export VILA_DATASETS=dataset_analysis
export DATASET_BACKEND=webdataset_multi_modal
python3 chengyue_script/dvlm/utils/analyze_vila_registry_stats.py \
  --max_per_dataset 2000 --num_workers 8 --output_dir output/vila_stats
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import sys
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
from tqdm import tqdm

try:
    import yaml
except ImportError:
    print("ERROR: PyYAML not installed. Install with: pip install pyyaml", file=sys.stderr)
    raise


# Project paths
PROJECT_ROOT = Path(__file__).resolve().parents[3]  # LMFlow root
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 VILA imports (if present)
try:
    import vila_path_patcher  # noqa: F401
except Exception:
    # Optional; avoid crashing in environments without the patcher
    pass


IGNORE_INDEX = -100


@dataclass
class DatasetSummary:
    dataset: str
    # Full dataset size (len(dataset)) from the VILA dataset object.
    # This is the actual dataset size, not the sampled count (n_ok).
    n_total: int
    n_ok: int
    n_fail: int
    fail_ratio: float
    n_turns_p50: float
    single_turn_ratio: float
    assistant_p50: float
    assistant_p90: float
    assistant_p95: float
    assistant_p99: float
    assistant_max: int
    turn_all_p50: float
    turn_all_p90: float
    turn_all_p95: float
    turn_all_p99: float
    turn_all_max: int
    total_tokens_p50: float
    total_tokens_p95: float
    total_tokens_max: int


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Fast-screen VILA registry datasets for dVLM conversion")
    p.add_argument(
        "--datasets_yaml",
        type=str,
        default=str(PROJECT_ROOT / "VILA/llava/data/registry/datasets/dataset_analysis.yaml"),
        help="Path to VILA datasets registry yaml (default: dataset_analysis.yaml)",
    )
    p.add_argument(
        "--mixtures_yaml",
        type=str,
        default=str(PROJECT_ROOT / "VILA/llava/data/registry/mixtures.yaml"),
        help="Path to VILA mixtures yaml",
    )
    p.add_argument(
        "--model_path",
        type=str,
        default="Qwen/Qwen2.5-VL-3B-Instruct",
        help="HF model path for tokenizer/processor",
    )
    p.add_argument(
        "--min_pixels",
        type=int,
        default=200704,
        help="Processor min_pixels (default: 200704 ~448x448)",
    )
    p.add_argument(
        "--max_pixels",
        type=int,
        default=200704,
        help="Processor max_pixels (default: 200704 ~448x448)",
    )
    p.add_argument(
        "--max_per_dataset",
        type=int,
        default=2000,
        help="Max successful samples to analyze per dataset (fast screen)",
    )
    p.add_argument(
        "--max_fail_per_dataset",
        type=int,
        default=2000,
        help="Stop after this many failures per dataset to avoid hanging on broken datasets",
    )
    p.add_argument(
        "--seed",
        type=int,
        default=1234,
        help="Random seed for index sampling",
    )
    p.add_argument(
        "--num_workers",
        type=int,
        default=0,
        help="DataLoader workers. 0 is usually safest for VILA datasets; increase if stable.",
    )
    p.add_argument(
        "--bd_sizes",
        type=str,
        default="4,8",
        help="Comma-separated bd_size values (padding overhead is estimated from seq_len only)",
    )
    p.add_argument(
        "--output_dir",
        type=str,
        default=str(PROJECT_ROOT / "output/vila_stats"),
        help="Output directory for CSV/JSON/MD reports",
    )
    p.add_argument(
        "--only_datasets",
        type=str,
        default=None,
        help="Optional comma-separated dataset names to analyze (debug)",
    )
    return p.parse_args()


def _load_yaml(path: str) -> dict:
    with open(path, "r") as f:
        return yaml.safe_load(f) or {}


def _parse_mixture_list(mixture: str, mixtures: Dict[str, List[str]]) -> List[str]:
    """
    Expand VILA mixtures.yaml syntax:
      - '+' composition
      - nested mixture names
      - '@ratio' subset notation (we keep base dataset name)
      - '*repeat' repetition (we keep base dataset name)
    """
    names = mixture.split("+") if "+" in mixture else [mixture]
    changed = True
    while changed:
        changed = False
        expanded: List[str] = []
        for name in names:
            base = name
            # strip special suffixes
            if "@" in base:
                base = base.split("@", 1)[0]
            if "*" in base:
                base = base.split("*", 1)[0]
            if base in mixtures:
                expanded.extend(mixtures[base])
                changed = True
            else:
                expanded.append(base)
        names = expanded
    # normalize
    out = []
    for n in names:
        base = n
        if "@" in base:
            base = base.split("@", 1)[0]
        if "*" in base:
            base = base.split("*", 1)[0]
        out.append(base)
    return sorted(set(out))


def _percentile(arr: np.ndarray, q: float) -> float:
    if arr.size == 0:
        return float("nan")
    return float(np.percentile(arr, q))


def _calc_padding(seq_len: int, bd_size: int) -> int:
    return (bd_size - (seq_len % bd_size)) % bd_size


def _compute_basic_stats(values: List[int]) -> Dict[str, Any]:
    if not values:
        return {"count": 0}
    a = np.asarray(values, dtype=np.int64)
    return {
        "count": int(a.size),
        "min": int(a.min()),
        "max": int(a.max()),
        "mean": float(a.mean()),
        "median": float(np.median(a)),
        "p90": _percentile(a, 90),
        "p95": _percentile(a, 95),
        "p99": _percentile(a, 99),
    }


def _ensure_env_for_vila(min_pixels: int, max_pixels: int) -> None:
    # Use the clean alias by default, but don't override if user already set it.
    os.environ.setdefault("VILA_DATASETS", "dataset_analysis")
    os.environ.setdefault("DATASET_BACKEND", "webdataset_multi_modal")
    os.environ["MIN_PIXELS"] = str(min_pixels)
    os.environ["MAX_PIXELS"] = str(max_pixels)


def _load_lmflow_vila_dataset(dataset_name: str, model_path: str):
    """
    Load a single VILA dataset via LMFlow's webdataset_multi_modal backend.
    This goes through VILA build_dataset(mixture=dataset_name) but since dataset_name
    is not a mixture, it becomes a single-dataset ConcatDataset.
    """
    from transformers import AutoProcessor, AutoTokenizer
    from lmflow.datasets.dataset import Dataset
    from lmflow.args import MultiModalDatasetArguments
    from lmflow.datasets.multi_modal_dataset import DataCollatorForQwenVL

    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    processor = AutoProcessor.from_pretrained(
        model_path,
        trust_remote_code=True,
        min_pixels=int(os.environ["MIN_PIXELS"]),
        max_pixels=int(os.environ["MAX_PIXELS"]),
    )

    data_args = MultiModalDatasetArguments(dataset_path=dataset_name)
    data_args.tokenizer = tokenizer
    data_args.image_processor = processor.image_processor if hasattr(processor, "image_processor") else processor

    ds = Dataset(data_args, backend="webdataset_multi_modal")

    collator = DataCollatorForQwenVL(processor=processor, tokenizer=tokenizer, bd_size=None)
    return ds, collator, tokenizer, processor


def _per_turn_lengths_from_input_ids(input_ids: np.ndarray, tokenizer) -> List[int]:
    """
    Reproduce the same per-turn parsing used in analyze_sharegpt4v_padding.py.
    We detect occurrences of <|im_start|>assistant\\n and count tokens until <|im_end|>.
    """
    import torch

    im_end_token = tokenizer.encode("<|im_end|>", add_special_tokens=False)
    im_end_id = im_end_token[0] if len(im_end_token) == 1 else 151645
    target_sequence = torch.tensor([151644, 77091, 198], dtype=torch.long)
    seq = torch.tensor(input_ids, dtype=torch.long)
    T = int(target_sequence.numel())
    per_turn = []
    i = 0
    while i <= seq.numel() - T:
        if torch.equal(seq[i : i + T], target_sequence):
            response_start = i + T
            response_end = response_start
            while response_end < seq.numel():
                if int(seq[response_end].item()) == int(im_end_id):
                    response_end += 1
                    break
                if response_end + T <= seq.numel() and torch.equal(seq[response_end : response_end + T], target_sequence):
                    break
                response_end += 1
            if response_start < response_end:
                per_turn.append(int(response_end - response_start))
            i = response_end + 1
        else:
            i += 1
    return per_turn


def _analyze_dataset(
    dataset_name: str,
    model_path: str,
    max_ok: int,
    max_fail: int,
    seed: int,
    bd_sizes: List[int],
    num_workers: int,
) -> Tuple[Dict[str, Any], Optional[DatasetSummary]]:
    """
    Returns:
      - detailed dict (for JSON)
      - DatasetSummary (for CSV ranking)
    """
    ds, collator, tokenizer, _processor = _load_lmflow_vila_dataset(dataset_name, model_path=model_path)
    n_total = len(ds)

    rng = np.random.default_rng(seed)
    # sample indices without replacement when possible
    if n_total <= max_ok + max_fail:
        indices = np.arange(n_total)
        rng.shuffle(indices)
    else:
        indices = rng.choice(n_total, size=min(n_total, max_ok + max_fail), replace=False)

    assistant_tokens: List[int] = []
    total_tokens: List[int] = []
    num_turns: List[int] = []
    all_turn_tokens: List[int] = []
    padding_by_bd: Dict[int, List[int]] = {bd: [] for bd in bd_sizes}
    n_ok = 0
    n_fail = 0

    # Iterate samples directly (DataLoader often triggers extra complexity for VILA)
    for idx in indices:
        if n_ok >= max_ok or n_fail >= max_fail:
            break
        try:
            item = ds.backend_dataset[int(idx)]
            if not isinstance(item, dict) or "messages" not in item:
                raise ValueError("dataset item missing 'messages'")

            # turns from messages
            turns = sum(1 for m in item.get("messages", []) if m.get("role") == "assistant")
            if turns <= 0:
                turns = 0

            batch = collator([item])
            input_ids = batch["input_ids"][0].detach().cpu().numpy()
            labels = batch["labels"][0].detach().cpu().numpy()

            seq_len = int(input_ids.shape[0])
            total_tokens.append(seq_len)
            num_turns.append(int(turns))

            a_tokens = int(np.sum(labels != IGNORE_INDEX))
            assistant_tokens.append(a_tokens)

            # per-turn lengths (token count until im_end)
            per_turn = _per_turn_lengths_from_input_ids(input_ids, tokenizer)
            all_turn_tokens.extend(per_turn)

            # bd padding estimate (based on total seq_len only)
            for bd in bd_sizes:
                padding_by_bd[bd].append(_calc_padding(seq_len, bd))

            n_ok += 1
        except Exception:
            n_fail += 1
            continue

    details: Dict[str, Any] = {
        "dataset": dataset_name,
        "n_total": int(n_total),
        "n_ok": int(n_ok),
        "n_fail": int(n_fail),
        "fail_ratio": float(n_fail / max(1, (n_ok + n_fail))),
        "assistant_tokens": _compute_basic_stats(assistant_tokens),
        "total_tokens": _compute_basic_stats(total_tokens),
        "turns": {
            **_compute_basic_stats(num_turns),
            "single_turn_ratio": float(np.mean(np.asarray(num_turns) == 1)) if num_turns else float("nan"),
        },
        "per_turn_tokens": {
            "all_turns": _compute_basic_stats(all_turn_tokens),
        },
        "padding": {f"bd{bd}": _compute_basic_stats(padding_by_bd[bd]) for bd in bd_sizes},
    }

    # Even if n_ok == 0 (e.g., due to schema/path issues during sampling),
    # still emit a summary row so downstream consumers can see n_total (size)
    # and fail_ratio in CSV/MD outputs.
    if n_ok == 0:
        empty = float("nan")
        summary = DatasetSummary(
            dataset=dataset_name,
            n_total=int(n_total),
            n_ok=int(n_ok),
            n_fail=int(n_fail),
            fail_ratio=float(n_fail / max(1, (n_ok + n_fail))),
            n_turns_p50=empty,
            single_turn_ratio=empty,
            assistant_p50=empty,
            assistant_p90=empty,
            assistant_p95=empty,
            assistant_p99=empty,
            assistant_max=0,
            turn_all_p50=empty,
            turn_all_p90=empty,
            turn_all_p95=empty,
            turn_all_p99=empty,
            turn_all_max=0,
            total_tokens_p50=empty,
            total_tokens_p95=empty,
            total_tokens_max=0,
        )
        return details, summary

    a = np.asarray(assistant_tokens, dtype=np.int64)
    t = np.asarray(total_tokens, dtype=np.int64)
    turns_arr = np.asarray(num_turns, dtype=np.int64)
    all_turn_arr = np.asarray(all_turn_tokens, dtype=np.int64) if all_turn_tokens else np.asarray([], dtype=np.int64)

    summary = DatasetSummary(
        dataset=dataset_name,
        n_total=int(n_total),
        n_ok=int(n_ok),
        n_fail=int(n_fail),
        fail_ratio=float(n_fail / max(1, (n_ok + n_fail))),
        n_turns_p50=float(np.median(turns_arr)) if turns_arr.size else float("nan"),
        single_turn_ratio=float(np.mean(turns_arr == 1)) if turns_arr.size else float("nan"),
        assistant_p50=float(np.median(a)),
        assistant_p90=_percentile(a, 90),
        assistant_p95=_percentile(a, 95),
        assistant_p99=_percentile(a, 99),
        assistant_max=int(a.max()),
        turn_all_p50=float(np.median(all_turn_arr)) if all_turn_arr.size else float("nan"),
        turn_all_p90=_percentile(all_turn_arr, 90) if all_turn_arr.size else float("nan"),
        turn_all_p95=_percentile(all_turn_arr, 95) if all_turn_arr.size else float("nan"),
        turn_all_p99=_percentile(all_turn_arr, 99) if all_turn_arr.size else float("nan"),
        turn_all_max=int(all_turn_arr.max()) if all_turn_arr.size else 0,
        total_tokens_p50=float(np.median(t)),
        total_tokens_p95=_percentile(t, 95),
        total_tokens_max=int(t.max()),
    )

    return details, summary


def _write_csv(path: Path, rows: List[DatasetSummary]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(asdict(rows[0]).keys()) if rows else ["dataset"])
        w.writeheader()
        for r in rows:
            w.writerow(asdict(r))


def _write_json(path: Path, obj: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w") as f:
        json.dump(obj, f, indent=2, ensure_ascii=False)


def _write_md(path: Path, summaries: List[DatasetSummary], mixture_refs: Dict[str, List[str]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    # Rank by long-tail per-turn length (p95), then by assistant p95
    ranked = sorted(
        summaries,
        key=lambda s: (
            -np.nan_to_num(s.turn_all_p95, nan=-1.0),
            -np.nan_to_num(s.assistant_p95, nan=-1.0),
            s.fail_ratio,
        ),
    )

    def fmt(x):
        if isinstance(x, float) and (np.isnan(x) or np.isinf(x)):
            return "nan"
        if isinstance(x, float):
            return f"{x:.1f}"
        return str(x)

    lines = []
    lines.append("# VILA dataset fast-screen for dVLM conversion\n")
    lines.append("Key signals:\n")
    lines.append("- `turn_all_p95`: 95th percentile of per-turn assistant response length (proxy for decoding length tail)\n")
    lines.append("- `assistant_p95`: 95th percentile of trainable assistant tokens per sample\n")
    lines.append("- `single_turn_ratio`: how often conversations are single-turn (often short-answer)\n")
    lines.append("- `fail_ratio`: sampling failures (bad paths/schema)\n")
    lines.append("\n## Top candidates (ranked by per-turn p95)\n")

    header = "|dataset|n_total|turn_all_p50|turn_all_p95|assistant_p50|assistant_p95|single_turn_ratio|fail_ratio|mixtures|\n"
    sep = "|---|---:|---:|---:|---:|---:|---:|---:|---|\n"
    lines.append(header)
    lines.append(sep)
    for s in ranked[:50]:
        refs = ", ".join(mixture_refs.get(s.dataset, [])[:6])
        lines.append(
            f"|{s.dataset}|{s.n_total:,}|{fmt(s.turn_all_p50)}|{fmt(s.turn_all_p95)}|{fmt(s.assistant_p50)}|{fmt(s.assistant_p95)}|{fmt(100*s.single_turn_ratio)}%|{fmt(100*s.fail_ratio)}%|{refs}|\n"
        )

    lines.append("\n## Likely short-answer datasets (turn_all_p95 <= 32)\n")
    short = [s for s in ranked if not np.isnan(s.turn_all_p95) and s.turn_all_p95 <= 32]
    lines.append(header)
    lines.append(sep)
    for s in short[:50]:
        refs = ", ".join(mixture_refs.get(s.dataset, [])[:6])
        lines.append(
            f"|{s.dataset}|{s.n_total:,}|{fmt(s.turn_all_p50)}|{fmt(s.turn_all_p95)}|{fmt(s.assistant_p50)}|{fmt(s.assistant_p95)}|{fmt(100*s.single_turn_ratio)}%|{fmt(100*s.fail_ratio)}%|{refs}|\n"
        )

    # Decision gate placeholder (user will tune thresholds)
    lines.append("\n## Decision gate: do we need to mix text-only LLM data?\n")
    lines.append("- If even the best multimodal candidates have low `turn_all_p95` / `assistant_p95`, dVLM will still learn short decoding.\n")
    lines.append("- A practical workflow is to set target thresholds (example): `turn_all_p50 >= 128` and `turn_all_p95 >= 256`.\n")
    lines.append("- If your planned VILA mix does not meet the targets, consider mixing text-only long-answer data (e.g., MetaMathQA) and/or CoT-style multimodal SFT.\n")

    with open(path, "w") as f:
        f.writelines(lines)


def main() -> int:
    args = _parse_args()
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    bd_sizes = [int(x) for x in args.bd_sizes.split(",") if x.strip()]

    _ensure_env_for_vila(args.min_pixels, args.max_pixels)

    datasets_yaml = Path(args.datasets_yaml)
    mixtures_yaml = Path(args.mixtures_yaml)
    if not datasets_yaml.exists():
        print(f"ERROR: datasets yaml not found: {datasets_yaml}", file=sys.stderr)
        return 2
    if not mixtures_yaml.exists():
        print(f"ERROR: mixtures yaml not found: {mixtures_yaml}", file=sys.stderr)
        return 2

    datasets_cfg = _load_yaml(str(datasets_yaml))
    mixtures_cfg = _load_yaml(str(mixtures_yaml))

    dataset_names = sorted(list(datasets_cfg.keys()))
    if args.only_datasets:
        only = {x.strip() for x in args.only_datasets.split(",") if x.strip()}
        dataset_names = [d for d in dataset_names if d in only]

    # mixture->dataset refs (for report)
    mixture_refs: Dict[str, List[str]] = {d: [] for d in dataset_names}
    for mix_name in mixtures_cfg.keys():
        expanded = _parse_mixture_list(mix_name, mixtures_cfg)
        for d in expanded:
            if d in mixture_refs:
                mixture_refs[d].append(mix_name)
    # sort refs for deterministic output
    for d in mixture_refs:
        mixture_refs[d] = sorted(set(mixture_refs[d]))

    print("=" * 80)
    print("VILA registry fast-screen for dVLM conversion")
    print(f"datasets_yaml: {datasets_yaml}")
    print(f"mixtures_yaml: {mixtures_yaml}")
    print(f"VILA_DATASETS: {os.environ.get('VILA_DATASETS')}")
    print(f"DATASET_BACKEND: {os.environ.get('DATASET_BACKEND')}")
    print(f"model_path: {args.model_path}")
    print(f"min_pixels/max_pixels: {args.min_pixels}/{args.max_pixels}")
    print(f"max_per_dataset: {args.max_per_dataset}  max_fail_per_dataset: {args.max_fail_per_dataset}")
    print(f"num_workers: {args.num_workers}  bd_sizes: {bd_sizes}")
    print(f"output_dir: {out_dir}")
    print("=" * 80)

    details_out: Dict[str, Any] = {
        "config": {
            "datasets_yaml": str(datasets_yaml),
            "mixtures_yaml": str(mixtures_yaml),
            "model_path": args.model_path,
            "min_pixels": args.min_pixels,
            "max_pixels": args.max_pixels,
            "max_per_dataset": args.max_per_dataset,
            "max_fail_per_dataset": args.max_fail_per_dataset,
            "seed": args.seed,
            "num_workers": args.num_workers,
            "bd_sizes": bd_sizes,
        },
        "datasets": {},
    }

    summaries: List[DatasetSummary] = []

    for d in tqdm(dataset_names, desc="Datasets", unit="ds"):
        try:
            det, summ = _analyze_dataset(
                dataset_name=d,
                model_path=args.model_path,
                max_ok=args.max_per_dataset,
                max_fail=args.max_fail_per_dataset,
                seed=args.seed,
                bd_sizes=bd_sizes,
                num_workers=args.num_workers,
            )
            details_out["datasets"][d] = det
            if summ is not None:
                summaries.append(summ)
        except Exception as e:
            details_out["datasets"][d] = {"dataset": d, "fatal_error": str(e)}
            continue

    # write outputs
    json_path = out_dir / "vila_registry_stats.json"
    _write_json(json_path, details_out)

    if summaries:
        csv_path = out_dir / "vila_registry_stats.csv"
        _write_csv(csv_path, summaries)
        md_path = out_dir / "vila_registry_candidates.md"
        _write_md(md_path, summaries, mixture_refs=mixture_refs)
        print(f"Wrote: {csv_path}")
        print(f"Wrote: {json_path}")
        print(f"Wrote: {md_path}")
    else:
        print(f"Wrote: {json_path}")
        print("No datasets produced usable summaries (n_ok==0).")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())


