#!/usr/bin/env python3
"""LES AUTRES LANGUES DE WAXAL : 16 langues etiquetees jamais utilisees.

On n'a exploite que lin, sna, lug -- or le corpus en contient 19. Les 13 autres
splits train pesent 78,7 Go.

Pourquoi c'est meilleur que FLEURS pour notre probleme : memes conditions
d'enregistrement, meme domaine (description d'images), meme convention
d'annotation, et surtout des LOCUTEURS ENTIEREMENT NOUVEAUX -- exactement ce qui
manque, puisque nos echecs viennent tous du sur-apprentissage locuteur.
FLEURS apporte des locuteurs neufs mais de la parole LUE, domaine different.

Preuve directe que le multilingue transfere : joint_cont (multilingue) BAT le
monolingue lin_s4 de 0.0072 au leaderboard. C'est le seul ajout de donnees qui
ait jamais gagne sur ce projet. On l'a fait avec 3 langues ; il y en a 19.

100 % donnee du concours => aucune declaration de conformite requise.

Premier lot : bantou et voisines, celles dont la phonotactique transfere le mieux
vers lin/sna. Le bloc ethiopien (amh, orm, sid, tir, wal, 52 Go) reste en reserve.
"""
import glob, os, sys

os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
os.environ.setdefault("HF_HOME", "/scratch/hf_home")
from huggingface_hub import snapshot_download

LANGS = os.environ.get("LANGS", "nyn,ach,sog,mas,aka,ewe").split(",")
dest = "/scratch/data/waxal"
tok = open(os.path.expanduser("~/.cache/huggingface/token")).read().strip()

pats = ["data/ASR/%s/%s-train-*.parquet" % (l, l) for l in LANGS]
print("telechargement : %s" % ", ".join(LANGS), flush=True)
snapshot_download(repo_id="google/WaxalNLP", repo_type="dataset",
                  allow_patterns=pats, local_dir=dest, max_workers=16, token=tok)
fs = []
for l in LANGS:
    fs += glob.glob("%s/data/ASR/%s/%s-train-*.parquet" % (dest, l, l))
print("DL_OK : %d parquets, %.1f Go" % (len(fs), sum(os.path.getsize(f) for f in fs) / 1e9), flush=True)

sys.path.insert(0, "/root")
import prep_data
prep_data.JOBS = [("waxal_%s_train" % l,
                   "%s/data/ASR/%s/%s-train-*.parquet" % (dest, l, l),
                   "transcription") for l in LANGS]
prep_data.main()

import json
M = "/scratch/prep/manifests/"
tot_clips = tot_h = 0
allrows = []
for l in LANGS:
    p = M + "waxal_%s_train.jsonl" % l
    if not os.path.exists(p):
        print("  %-5s MANIFESTE ABSENT" % l, flush=True); continue
    rows = [json.loads(x) for x in open(p, encoding="utf-8")]
    ok = [r for r in rows if os.path.exists(r["audio"])]
    h = sum(r["duration"] for r in ok) / 3600
    spk = len(set(r.get("speaker") for r in ok))
    print("  %-5s %6d clips | %6.1f h | %4d locuteurs" % (l, len(ok), h, spk), flush=True)
    tot_clips += len(ok); tot_h += h
    allrows += ok
with open(M + "waxal_multi_train.jsonl", "w", encoding="utf-8") as f:
    for r in allrows:
        f.write(json.dumps(r, ensure_ascii=False) + "\n")
print("TOTAL NOUVEAU : %d clips / %.1f h -> %swaxal_multi_train.jsonl" % (tot_clips, tot_h, M), flush=True)
print("MULTI_DONE", flush=True)