"""
Stage 5 — Filter FineWeb by classifier score and write release datasets.

Reads per-dump score parquets produced by Stage 4 and streams through FineWeb
to produce two filtered datasets:

  fineweb-marketing         top ~8%  by Head A percentile (~1.3T tokens)
  fineweb-marketing-score-2 top ~35% by Head A percentile (~5.4T tokens)

Filtering is percentile-based, not absolute int_score (spec §5.2).  The three
score columns (score, int_score, percentile) are kept in every output row
alongside all original FineWeb document columns.

Crash-safe: both output files must exist for a dump to be considered complete;
if either is missing the dump is reprocessed from scratch.

Usage — single dump:
    python -m filtering.filter \\
        --scores-dir data/scores \\
        --output-dir data/release \\
        --dump CC-MAIN-2024-10

Usage — all dumps (discovers via HF Hub):
    python -m filtering.filter \\
        --scores-dir data/scores \\
        --output-dir data/release \\
        --all-dumps

Usage — stats only (threshold statistics without streaming FineWeb):
    python -m filtering.filter \\
        --scores-dir data/scores \\
        --output-dir data/release \\
        --all-dumps --stats-only
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
from datasets import get_dataset_config_names, load_dataset
from tqdm import tqdm

FINEWEB_REPO = "HuggingFaceFW/fineweb"

PRIMARY_PCT = 0.08    # top  8% → FineWeb-Marketing
SECONDARY_PCT = 0.35  # top 35% → FineWeb-Marketing-score-2

PRIMARY_SUBDIR = "fineweb-marketing"
SECONDARY_SUBDIR = "fineweb-marketing-score-2"

WRITE_CHUNK = 10_000


# ── Pure functions (testable without network) ──────────────────────────────────

def compute_percentile_threshold(target_fraction: float) -> float:
    """
    Convert a retention fraction to the minimum percentile for kept documents.

    Top 8% → 0.92.  Raises ValueError for out-of-range input.
    """
    if not 0.0 < target_fraction < 1.0:
        raise ValueError(
            f"target_fraction must be in (0, 1), got {target_fraction!r}"
        )
    return 1.0 - target_fraction


def build_score_index(
    scores_path: Path,
    min_percentile: float,
) -> dict[str, tuple[float, int, float]]:
    """
    Load a Stage 4 score parquet and return {id: (score, int_score, percentile)}
    for all rows with percentile >= min_percentile.

    Uses PyArrow predicate pushdown to avoid loading below-threshold rows into RAM.
    The threshold is converted to float32 space before filtering to stay consistent
    with the stored float32 column and avoid boundary mismatches.
    """
    # Convert to float32 space so the filter matches the stored column type
    threshold_f32 = float(np.float32(min_percentile))
    table = pq.read_table(
        scores_path,
        columns=["id", "score", "int_score", "percentile"],
        filters=[("percentile", ">=", threshold_f32)],
    )
    ids = table.column("id").to_pylist()
    scores = table.column("score").to_pylist()
    int_scores = table.column("int_score").to_pylist()
    percentiles = table.column("percentile").to_pylist()
    return {
        id_: (float(score), int(int_score), float(pct))
        for id_, score, int_score, pct in zip(ids, scores, int_scores, percentiles)
    }


def route_doc(
    doc_id: str,
    score_index: dict[str, tuple[float, int, float]],
    primary_threshold: float,
) -> tuple[str, tuple[float, int, float]] | None:
    """
    Classify a document as "primary" or "secondary" and return its score entry.

    "primary"   → in index with percentile >= primary_threshold
    "secondary" → in index with percentile < primary_threshold
    None        → not in index (below secondary threshold)

    The score_index is already filtered to >= secondary_threshold, so any ID
    absent from it is automatically excluded. Returning the entry alongside
    the tier avoids a second dict lookup at the call site.
    """
    entry = score_index.get(doc_id)
    if entry is None:
        return None
    _, _, pct = entry
    tier = "primary" if pct >= primary_threshold else "secondary"
    return tier, entry


def compute_threshold_stats(
    scores_path: Path,
    primary_threshold: float,
    secondary_threshold: float,
) -> dict[str, Any]:
    """
    Compute per-dump doc counts and minimum score at each percentile cutoff.

    The returned int_score distributions reveal the equivalent int_score
    thresholds for FineWeb-EDU users (written to threshold_summary.json).
    """
    table = pq.read_table(scores_path, columns=["score", "int_score", "percentile"])
    df = table.to_pandas()

    p_mask = df["percentile"] >= primary_threshold
    s_mask = df["percentile"] >= secondary_threshold

    def _int_counts(mask: pd.Series) -> dict[str, int]:
        if not mask.any():
            return {}
        return {
            str(int(k)): int(v)
            for k, v in df.loc[mask, "int_score"].value_counts().sort_index().items()
        }

    return {
        "total_docs": int(len(df)),
        "primary_docs": int(p_mask.sum()),
        "secondary_docs": int(s_mask.sum()),
        "primary_min_score": (
            float(df.loc[p_mask, "score"].min()) if p_mask.any() else None
        ),
        "secondary_min_score": (
            float(df.loc[s_mask, "score"].min()) if s_mask.any() else None
        ),
        "primary_int_score_counts": _int_counts(p_mask),
        "secondary_int_score_counts": _int_counts(s_mask),
    }


# ── Buffered parquet writer ────────────────────────────────────────────────────

class _TierWriter:
    """Accumulates rows in memory and flushes to a single parquet file in chunks."""

    def __init__(self, path: Path, chunk_size: int = WRITE_CHUNK) -> None:
        self._path = path
        self._chunk_size = chunk_size
        self._buf: list[dict[str, Any]] = []
        self._writer: pq.ParquetWriter | None = None
        self._schema: pa.Schema | None = None
        self.count: int = 0

    def add(self, row: dict[str, Any]) -> None:
        self._buf.append(row)
        self.count += 1
        if len(self._buf) >= self._chunk_size:
            self._flush()

    def _flush(self) -> None:
        if not self._buf:
            return
        columns = list(self._buf[0].keys())
        table = pa.Table.from_arrays(
            [pa.array([r[c] for r in self._buf]) for c in columns],
            names=columns,
        )
        if self._writer is None:
            self._schema = table.schema
            self._writer = pq.ParquetWriter(
                self._path, self._schema, compression="snappy"
            )
        else:
            table = table.cast(self._schema)
        self._writer.write_table(table)
        self._buf.clear()

    def close(self) -> None:
        """Flush remaining rows and close. Writes a zero-row file if no data was added."""
        self._flush()
        if self._writer is not None:
            self._writer.close()
        else:
            _write_empty(self._path)


def _write_empty(path: Path) -> None:
    """Write a zero-row parquet to mark a dump as complete when no docs pass the threshold."""
    schema = pa.schema([
        ("id", pa.string()),
        ("dump", pa.string()),
        ("score", pa.float32()),
        ("int_score", pa.int8()),
        ("percentile", pa.float32()),
    ])
    pq.write_table(
        pa.table(
            {f: pa.array([], type=t) for f, t in zip(schema.names, schema.types)},
            schema=schema,
        ),
        path,
    )


# ── Main filter function ───────────────────────────────────────────────────────

def filter_dump(
    dump: str,
    scores_dir: Path,
    output_dir: Path,
    primary_pct: float = PRIMARY_PCT,
    secondary_pct: float = SECONDARY_PCT,
    hf_dataset: str = FINEWEB_REPO,
) -> dict[str, Any]:
    """
    Filter one CC-MAIN dump into the primary and secondary release datasets.

    Phase 1: load the Stage 4 score parquet, build an in-memory index for all
             documents above the secondary percentile threshold.
    Phase 2: stream through FineWeb once, routing each matching document to the
             primary and/or secondary output writer.

    Both output files must exist for the dump to be considered complete.
    Temp files are used for atomic rename; stale temps from a prior crash are
    removed at the start of each run.
    """
    primary_out = output_dir / PRIMARY_SUBDIR
    secondary_out = output_dir / SECONDARY_SUBDIR
    primary_file = primary_out / f"{dump}.parquet"
    secondary_file = secondary_out / f"{dump}.parquet"

    if primary_file.exists() and secondary_file.exists():
        print(f"  {dump}: already done, skipping.")
        return {"dump": dump, "status": "skipped"}

    scores_path = scores_dir / f"{dump}.parquet"
    if not scores_path.exists():
        print(f"  {dump}: no score file found at {scores_path}, skipping.")
        return {"dump": dump, "status": "no_scores"}

    primary_out.mkdir(parents=True, exist_ok=True)
    secondary_out.mkdir(parents=True, exist_ok=True)

    primary_threshold = compute_percentile_threshold(primary_pct)
    secondary_threshold = compute_percentile_threshold(secondary_pct)

    score_index = build_score_index(scores_path, secondary_threshold)
    print(
        f"  {dump}: {len(score_index):,} docs pass secondary threshold "
        f"({secondary_pct:.0%})"
    )

    primary_tmp = primary_file.with_suffix(".parquet.tmp")
    secondary_tmp = secondary_file.with_suffix(".parquet.tmp")
    for tmp in (primary_tmp, secondary_tmp):
        tmp.unlink(missing_ok=True)

    primary_writer = _TierWriter(primary_tmp, chunk_size=WRITE_CHUNK)
    secondary_writer = _TierWriter(secondary_tmp, chunk_size=WRITE_CHUNK)
    streamed = 0

    ds = load_dataset(hf_dataset, name=dump, split="train", streaming=True)
    try:
        for doc in tqdm(ds, desc=f"Filtering {dump}"):
            streamed += 1
            routed = route_doc(doc["id"], score_index, primary_threshold)
            if routed is None:
                continue

            tier, (score, int_score, percentile) = routed
            row = dict(doc)
            row["score"] = np.float32(score)
            row["int_score"] = np.int8(int_score)
            row["percentile"] = np.float32(percentile)

            if tier == "primary":
                primary_writer.add(row)
            secondary_writer.add(row)

    finally:
        primary_writer.close()
        secondary_writer.close()

    # Only reached if streaming completed without exception
    primary_tmp.rename(primary_file)
    secondary_tmp.rename(secondary_file)

    print(
        f"  {dump}: {primary_writer.count:,} primary, "
        f"{secondary_writer.count:,} secondary "
        f"(of {streamed:,} streamed)"
    )
    return {
        "dump": dump,
        "status": "done",
        "primary": primary_writer.count,
        "secondary": secondary_writer.count,
        "streamed": streamed,
    }


# ── Threshold summary ──────────────────────────────────────────────────────────

def write_threshold_summary(
    per_dump_stats: list[dict[str, Any]],
    primary_pct: float,
    secondary_pct: float,
    output_path: Path,
) -> None:
    """
    Write a JSON file documenting per-dump score and int_score statistics at
    both percentile cutoffs.

    This documents the equivalent int_score thresholds so FineWeb-EDU users
    can compare datasets without re-running inference.
    """
    output_path.write_text(
        json.dumps(
            {
                "primary_fraction": primary_pct,
                "secondary_fraction": secondary_pct,
                "primary_percentile_threshold": compute_percentile_threshold(primary_pct),
                "secondary_percentile_threshold": compute_percentile_threshold(
                    secondary_pct
                ),
                "per_dump": per_dump_stats,
            },
            indent=2,
        )
    )
    print(f"Threshold summary → {output_path}")


# ── Dump discovery ─────────────────────────────────────────────────────────────

def get_cc_dumps(hf_dataset: str = FINEWEB_REPO) -> list[str]:
    """Return sorted list of CC-MAIN-* config names for the given HF dataset."""
    configs = get_dataset_config_names(hf_dataset)
    return sorted(c for c in configs if c.startswith("CC-MAIN-"))


# ── CLI ────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    ap = argparse.ArgumentParser(
        description="Stage 5: filter FineWeb by Head A score and write release datasets."
    )
    ap.add_argument(
        "--scores-dir",
        type=Path,
        required=True,
        help="Directory with per-dump score parquets from Stage 4",
    )
    ap.add_argument(
        "--output-dir",
        type=Path,
        required=True,
        help=(
            "Root output directory.  Two subdirectories are created: "
            f"{PRIMARY_SUBDIR!r} and {SECONDARY_SUBDIR!r}."
        ),
    )

    dump_group = ap.add_mutually_exclusive_group(required=True)
    dump_group.add_argument("--dump", help="Single CC-MAIN dump name")
    dump_group.add_argument(
        "--all-dumps",
        action="store_true",
        help="Discover and process all CC-MAIN configs from HF Hub",
    )
    dump_group.add_argument(
        "--dump-list",
        type=Path,
        help="Plain-text file with one dump name per line",
    )

    ap.add_argument(
        "--primary-pct",
        type=float,
        default=PRIMARY_PCT,
        help=f"Primary retention fraction (default {PRIMARY_PCT})",
    )
    ap.add_argument(
        "--secondary-pct",
        type=float,
        default=SECONDARY_PCT,
        help=f"Secondary retention fraction (default {SECONDARY_PCT})",
    )
    ap.add_argument("--hf-dataset", default=FINEWEB_REPO)
    ap.add_argument(
        "--stats-only",
        action="store_true",
        help="Print threshold statistics from score parquets only (no FineWeb streaming)",
    )
    args = ap.parse_args()

    if args.dump:
        dumps = [args.dump]
    elif args.dump_list:
        dumps = [
            d.strip()
            for d in args.dump_list.read_text().splitlines()
            if d.strip()
        ]
    else:
        print(f"Discovering CC-MAIN dumps from {args.hf_dataset} …")
        dumps = get_cc_dumps(args.hf_dataset)
        print(f"Found {len(dumps)} CC-MAIN dumps.")

    primary_threshold = compute_percentile_threshold(args.primary_pct)
    secondary_threshold = compute_percentile_threshold(args.secondary_pct)

    args.output_dir.mkdir(parents=True, exist_ok=True)
    summary_path = args.output_dir / "threshold_summary.json"

    if args.stats_only:
        per_dump_stats = []
        for d in dumps:
            p = args.scores_dir / f"{d}.parquet"
            if not p.exists():
                print(f"  {d}: no score file, skipping.")
                continue
            stats = compute_threshold_stats(p, primary_threshold, secondary_threshold)
            stats["dump"] = d
            per_dump_stats.append(stats)
            print(
                f"  {d}: {stats['primary_docs']:,} primary / "
                f"{stats['secondary_docs']:,} secondary of "
                f"{stats['total_docs']:,} total"
            )
        write_threshold_summary(
            per_dump_stats, args.primary_pct, args.secondary_pct, summary_path
        )
    else:
        all_stats = []
        for d in dumps:
            result = filter_dump(
                dump=d,
                scores_dir=args.scores_dir,
                output_dir=args.output_dir,
                primary_pct=args.primary_pct,
                secondary_pct=args.secondary_pct,
                hf_dataset=args.hf_dataset,
            )
            # Include stats for both newly processed and already-skipped dumps
            # so threshold_summary.json is complete even on resumed runs.
            if result.get("status") in ("done", "skipped"):
                score_p = args.scores_dir / f"{d}.parquet"
                if score_p.exists():
                    dump_stats = compute_threshold_stats(
                        score_p, primary_threshold, secondary_threshold
                    )
                    dump_stats["dump"] = d
                    all_stats.append(dump_stats)

        if all_stats:
            write_threshold_summary(
                all_stats, args.primary_pct, args.secondary_pct, summary_path
            )

    print("Stage 5 complete.")