#!/bin/bash
# Campagne finale WAXAL : gains legitimes public + robustesse Phase 2
set -uo pipefail
export HF_HOME=/scratch/hf PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True TMPDIR=/scratch/tmp
mkdir -p /scratch/tmp /scratch/runs
step(){ echo "=== $1 ($(date +%H:%M:%S)) ==="; }

# attendre le pool de bruit urbain
while ! grep -q "NOISE_URBAN_DONE" /root/extract_noise.log 2>/dev/null; do sleep 30; done
# pool combine : babble in-domain + bruit urbain africain
mkdir -p /scratch/noise_all
cp -sn /scratch/noise/*.flac /scratch/noise_all/ 2>/dev/null || true
cp -sn /scratch/noise_urban/*.flac /scratch/noise_all/ 2>/dev/null || true
echo "pool bruit total: $(ls /scratch/noise_all/*.flac 2>/dev/null | wc -l) clips"

# ---------- VOLET 1 : pseudo-labeling luganda IN-DOMAIN (le seul type qui marche) ----------
step "VOLET 1 : pseudo luganda in-domain (WAXAL unlabeled)"
if [ ! -f /scratch/prep/manifests/pseudo_lug_filtered.jsonl ]; then
  if [ ! -f /scratch/prep/manifests/waxal_lug_unlabeled.jsonl ]; then
    /root/venv/bin/python -c "
import sys; sys.path.insert(0,'/root'); import prep_data
prep_data.JOBS=[('waxal_lug_unlabeled','/scratch/data/waxal/data/ASR/lug/lug-unlabeled-*.parquet','transcription')]
prep_data.main()" > /root/prep_lug_unlab.log 2>&1 || echo "prep lug KO"
  fi
  /root/venv/bin/python /root/pseudo_label.py --model /root/models/lug_r_best \
    --manifest /scratch/prep/manifests/waxal_lug_unlabeled.jsonl --out /root/pseudo_lug_raw.jsonl > /root/pseudo_lug.log 2>&1
  if grep -q PSEUDO_DONE /root/pseudo_lug.log; then
    /root/venv/bin/python -c "
import json
rows=[json.loads(l) for l in open('/root/pseudo_lug_raw.jsonl',encoding='utf-8')]
rows=[r for r in rows if r['text'] and 1.5<=r['duration']<=40 and len(r['text'].split())/max(r['duration'],.1)<=4 and len(r['text'].split())>=3]
c=sorted((r['conf'] for r in rows),reverse=True); s=c[int(len(c)*0.5)] if c else 0
k=[r for r in rows if r['conf']>=s]
open('/scratch/prep/manifests/pseudo_lug_filtered.jsonl','w',encoding='utf-8').write('\n'.join(json.dumps(r,ensure_ascii=False) for r in k)+'\n')
print(f'filtre lug: {len(k)} clips, {sum(r[\"duration\"] for r in k)/3600:.1f}h')"
  fi
fi
if [ -f /scratch/prep/manifests/pseudo_lug_filtered.jsonl ] && [ ! -d /root/models/lug_ps_best ]; then
  /root/venv/bin/python /root/train_wbert.py --lang lug --out /scratch/runs/lug_ps --init /root/models/lug_r_best \
    --train /scratch/prep/manifests/waxal_lug_train.jsonl /scratch/prep/manifests/pseudo_lug_filtered.jsonl /scratch/prep/manifests/cv_lug_train.jsonl \
    --eval /scratch/prep/manifests/waxal_lug_validation.jsonl \
    --exclude_texts /scratch/prep/manifests/waxal_lug_validation.jsonl,/scratch/prep/manifests/waxal_lug_test.jsonl \
    --lr 2e-5 --epochs 3 --bs 4 --grad_accum 16 --num_workers 7 --eval_steps 300 --eval_subset 664 > /root/lug_ps.log 2>&1
  grep -q TRAIN_DONE /root/lug_ps.log && { mkdir -p /root/models/lug_ps_best; cp -r /scratch/runs/lug_ps/best/* /root/models/lug_ps_best/; cp /scratch/runs/lug_ps/vocab.json /root/models/lug_ps_best/ 2>/dev/null; rm -rf /scratch/runs/lug_ps/checkpoint-*; bash /root/hf_backup.sh; }
  grep -aA4 "EVAL FINALE" /root/lug_ps.log | head -5
fi
echo "VOLET1_DONE"

# ---------- VOLET 2 : variantes robustes v2 (bruit urbain + babble + 8kHz) -> Phase 2 ----------
for CFG in "lin lin_s4_best" "lug lug_r_best" "sna sna_ps_best"; do
  set -- $CFG; LG=$1; INIT=$2
  [ -d /root/models/${LG}_r2_best ] && { echo "${LG}_r2 deja fait"; continue; }
  step "VOLET 2 : robuste v2 $LG (bruit urbain+babble+8k)"
  EXTRA=""; [ "$LG" = "lug" ] && EXTRA="/scratch/prep/manifests/cv_lug_train.jsonl"
  /root/venv/bin/python /root/train_wbert.py --lang $LG --out /scratch/runs/${LG}_r2 --init /root/models/$INIT \
    --train /scratch/prep/manifests/waxal_${LG}_train.jsonl $EXTRA \
    --eval /scratch/prep/manifests/waxal_${LG}_validation.jsonl \
    --augment --noise_dir /scratch/noise_all \
    --lr 2e-5 --epochs 2 --bs 4 --grad_accum 16 --num_workers 7 --eval_steps 300 --eval_subset 800 > /root/${LG}_r2.log 2>&1
  grep -q TRAIN_DONE /root/${LG}_r2.log && { mkdir -p /root/models/${LG}_r2_best; cp -r /scratch/runs/${LG}_r2/best/* /root/models/${LG}_r2_best/; cp /scratch/runs/${LG}_r2/vocab.json /root/models/${LG}_r2_best/ 2>/dev/null; rm -rf /scratch/runs/${LG}_r2/checkpoint-*; bash /root/hf_backup.sh; }
  grep -aA4 "EVAL FINALE" /root/${LG}_r2.log | head -5
done
echo "VOLET2_DONE"

# ---------- VOLET 3 : selection finale sur VALIDATION COMPLETE + DEV DIFFICILE ----------
step "VOLET 3 : selection champions (public + Phase 2)"
/root/venv/bin/python /root/select_final.py > /root/select_final.log 2>&1
grep -aE "PUBLIC|PHASE2|GAGNANT|SELECT_FINAL_DONE" /root/select_final.log | tail -20
bash /root/hf_backup.sh
echo "CAMPAGNE_DONE"