#!/bin/bash
# RECONSTRUCTION COMPLETE depuis w2v-bert-2.0, avec TOUTES les donnees de la page Data.
#
# Justification : ALLDATA (continuation) porte la penalite de continuation mesuree
# (evals ~0.299 vs 0.2848 pour le champion). Repartir de zero est le SEUL moyen d y
# echapper. Les deux paris sont complementaires et tournent en parallele.
#
# Donnees : lin + sna + lug, splits train + validation + test (Phase 1 autorisee pour
# la Phase 2, fil 34195). Le champion n avait ni les test ni les validation.
#
# Pas de selection de checkpoint (LOAD_BEST=0) : devhard est a 100 % dans le train,
# choisir le "meilleur" revient a choisir celui qui memorise le mieux.
set -e
export HF_HUB_ENABLE_HF_TRANSFER=0

echo "########## 1. parquets manquants (test + luganda) ##########"
/root/venv/bin/python - <<'PY'
from huggingface_hub import snapshot_download
snapshot_download(repo_id="google/WaxalNLP", repo_type="dataset",
    local_dir="/scratch/data/waxal", max_workers=8,
    allow_patterns=["data/ASR/lin/lin-test-*.parquet",
                    "data/ASR/sna/sna-test-*.parquet",
                    "data/ASR/lug/lug-train-*.parquet",
                    "data/ASR/lug/lug-validation-*.parquet",
                    "data/ASR/lug/lug-test-*.parquet"])
print("PARQUETS_OK")
PY

echo "########## 2. extraction FLAC (CPU, sans gener le GPU) ##########"
/root/venv/bin/python /root/prep_data.py 2>&1 | tail -6

echo "########## 3. inventaire des donnees etiquetees ##########"
/root/venv/bin/python - <<'PY'
import json, glob, os
M="/scratch/prep/manifests"
tot=0; keep=[]
for f in sorted(glob.glob(M+"/waxal_*.jsonl")):
    n=os.path.basename(f)[:-6]
    if not any(k in n for k in ["lin","sna","lug"]): continue
    if "unlabeled" in n: continue
    rows=[json.loads(l) for l in open(f,encoding="utf-8")]
    lab=[r for r in rows if r.get("text","").strip()]
    if not lab: continue
    if not os.path.exists(lab[0]["audio"]): print("  SANS AUDIO:",n); continue
    h=sum(r.get("duration",0) for r in lab)/3600
    tot+=h; keep.append(f)
    print("  %-28s %6d clips | %6.1f h"%(n,len(lab),h))
print("  TOTAL : %.1f h sur %d manifests"%(tot,len(keep)))
open("/root/rebuild_manifests.txt","w").write(" ".join(keep))
PY

echo "########## 4. attente du GPU ##########"
while pgrep -f "train_wbert.py|train_ahb.py" > /dev/null; do sleep 60; done
sleep 20
echo "GPU libre $(date -u '+%H:%M UTC')"

echo "########## 5. entrainement depuis zero ##########"
TRAINSET=$(cat /root/rebuild_manifests.txt)
echo "manifests : $TRAINSET"
SAVE_LIMIT=16 LOAD_BEST=0 /root/venv/bin/python /root/train_wbert.py \
  --lang lin --out /scratch/runs/rebuild \
  --init facebook/w2v-bert-2.0 \
  --train $TRAINSET \
  --eval  /scratch/prep/manifests/waxal_lin_validation.jsonl \
  --lr 1e-4 --epochs 8 --bs 4 --grad_accum 15 --warmup 800 --eval_steps 300

echo REBUILD_TRAIN_DONE