#!/usr/bin/env bash
# Merge the primary cache with every node shard, rebuild the corpus, push to the Hub.
#
# Run this after the primary generator exits. Idempotent: it re-merges from the two
# caches every time, so re-running after a resumed generation just produces a larger
# corpus. Neither node's live cache is touched — the merge lands in out_merged/.
#
#   ./finish_merge_push.sh            # build + card, no push  (review, then re-run with push)
#   ./finish_merge_push.sh --push     # build + card + upload
set -euo pipefail
cd /notebooks/50M/reasonsing
P=/notebooks/50M/.venv-lfm2/bin/python
REPO=oddadmix/arabic-math-reasoning-synth
SHARD_DIR=shards_dl/shards

# 1. merge — primary cache first, then every shard pulled from the Hub
mkdir -p out_merged
cat out_synth/generations.jsonl > out_merged/generations.jsonl
for gz in "$SHARD_DIR"/*.jsonl.gz; do
    echo "[*] merging shard $(basename "$gz")"
    gunzip -c "$gz" >> out_merged/generations.jsonl
done
echo "[+] merged cache: $(wc -l < out_merged/generations.jsonl) tasks"

# 2. authoritative re-parse / re-validate / dedup over the whole merged cache
OUT_DIR=out_merged SFT_DIR=data_synth_sft $P build_synth_dataset.py > /tmp/build.log 2>&1
$P - <<'PY'
import json
d = json.load(open("out_merged/build_stats.json")); s = d["stats"]
print(f"[+] kept {s['kept']:,} rows | templates {d['unique_templates']:,} | "
      f"accept {d['accept_rate']:.1%} | dup drops {s.get('dropped_duplicate_template',0):,}")
for k in sorted(s):
    if k.startswith("model_"):
        print(f"      {k[6:]:36s} {s[k]:,}")
PY

# 3. card + upload
if [[ "${1:-}" == "--push" ]]; then
    OUT_DIR=out_merged $P push_synth_dataset.py --repo "$REPO"
else
    OUT_DIR=out_merged $P push_synth_dataset.py --repo "$REPO" --dry-run
fi