#!/usr/bin/env python3
"""Build the vendored WARMTH data bundle read by ``mteval/warmth.py``.

WARMTH (https://github.com/alvations/warmth) ships ~544 MB of parquet that
includes per-system MT hypotheses, human scores and annotations we never use.
This script distills each collection to just what the leaderboard needs —
``source`` + ``reference`` per segment — producing ``warmth_data/<collection>.parquet``
(~106 MB total, all releases and language pairs), so the Space needs no network
access to load test sets.

Usage:
    git clone --depth 1 https://github.com/alvations/warmth.git /tmp/warmth
    python scripts/build_warmth_bundle.py /tmp/warmth
"""

from __future__ import annotations

import glob
import os
import sys

import pyarrow as pa
import pyarrow.dataset as pads
import pyarrow.parquet as pq

# Must match mteval.warmth.COLLECTIONS. `wmt-metrics` is excluded on purpose:
# its `source` column is empty (metrics/QE task), so it can't drive translation.
COLLECTIONS = [
    "wmt24pp", "wmt-general", "ntrex", "flores-plus", "wmt-mqm", "wmt-metrics-hi",
    "wmt-biomed", "wmt-terminology", "bio-mqm", "iwslt", "mtedx", "mtnt",
    "multi30k", "tatoeba", "diabla",
]
KEEP = ["release", "langpair", "segment_id", "source", "reference", "tgt_lang"]


def main(src_root: str, out_dir: str) -> int:
    os.makedirs(out_dir, exist_ok=True)
    total = 0
    for coll in COLLECTIONS:
        files = sorted(glob.glob(os.path.join(src_root, "data", coll, "*.parquet")))
        if not files:
            print(f"  {coll}: no source parquet found, skipped")
            continue
        dset = pads.dataset(files, format="parquet")
        not_canary = pads.field("domain").is_null() | (pads.field("domain") != "canary")
        filt = (
            pads.field("source").is_valid()
            & pads.field("reference").is_valid()
            & not_canary
        )
        df = dset.to_table(columns=KEEP + ["domain"], filter=filt).select(KEEP).to_pandas()
        df = df.sort_values("segment_id", kind="stable").drop_duplicates(
            subset=["release", "langpair", "segment_id"], keep="first"
        )
        out = os.path.join(out_dir, f"{coll}.parquet")
        pq.write_table(pa.Table.from_pandas(df, preserve_index=False), out, compression="zstd")
        total += len(df)
        print(f"  {coll:16} {len(df):>8} rows  {os.path.getsize(out) / 1048576:6.1f} MB")
    print(f"Total {total} rows -> {out_dir}")
    return 0


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__)
        raise SystemExit(2)
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    raise SystemExit(main(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else os.path.join(root, "warmth_data")))