#!/usr/bin/env python3
"""FLEURS lingala (ln_cd) + shona (sn_zw) : de la parole EXTERNE, avec des
LOCUTEURS NOUVEAUX.

Hypothese : on a mesure TROIS FOIS que reentrainer sur WAXAL degrade le test
(pseudo-labels lin -0.0062, tour 2 sna -0.00018, SWA -0.0010) alors que la
validation s'ameliorait a chaque fois. Mecanisme identifie : l'ajout de donnees
des MEMES locuteurs colle le modele a eux, or le test est fait de locuteurs
INEDITS. Corollaire jamais teste : des donnees EXTERNES apportent des locuteurs
nouveaux, donc l'effet inverse -- plus de diversite locuteur, meilleure
generalisation.

Conformite : FLEURS (google/fleurs) est public, non restreint, CC-BY-4.0, et
l'organisateur autorise explicitement les corpus externes publics et licencies
s'ils sont DECLARES (fil Zindi 33616). A ajouter a COMPLIANCE.md s'il finit dans
le systeme soumis.

Reserve connue : FLEURS est de la parole LUE (phrases Wikipedia), alors que WAXAL
est de la description d'images spontanee. Decalage de domaine reel -- c'est
precisement ce que la soumission tranchera.
"""
import glob, os, tarfile, csv as _csv, json, sys

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

tok = open(os.path.expanduser("~/.cache/huggingface/token")).read().strip()
DEST = "/scratch/fleurs"
LANGS = {"ln_cd": "lin", "sn_zw": "sna"}

pats = []
for lg in LANGS:
    pats += ["data/%s/audio/train.tar.gz" % lg, "data/%s/train.tsv" % lg,
             "data/%s/audio/dev.tar.gz" % lg, "data/%s/dev.tsv" % lg]
print("telechargement FLEURS %s ..." % ", ".join(LANGS), flush=True)
snapshot_download("google/fleurs", repo_type="dataset", allow_patterns=pats,
                  local_dir=DEST, max_workers=8, token=tok)
got = glob.glob(DEST + "/data/*/audio/*.tar.gz")
print("archives : %d (%.2f Go)" % (len(got), sum(os.path.getsize(f) for f in got) / 1e9), flush=True)

import soundfile as sf
import numpy as np

OUTD = "/scratch/prep/audio/fleurs"
os.makedirs(OUTD, exist_ok=True)
M = "/scratch/prep/manifests"
os.makedirs(M, exist_ok=True)

for lg, ourlang in LANGS.items():
    rows = []
    for split in ("train", "dev"):
        tsv = "%s/data/%s/%s.tsv" % (DEST, lg, split)
        tgz = "%s/data/%s/audio/%s.tar.gz" % (DEST, lg, split)
        if not (os.path.exists(tsv) and os.path.exists(tgz)):
            print("  %s/%s absent, saute" % (lg, split), flush=True); continue
        meta = {}
        with open(tsv, encoding="utf-8") as f:
            for line in f:
                p = line.rstrip("\n").split("\t")
                if len(p) >= 4:
                    meta[p[1]] = p[3]          # nom de fichier -> transcription normalisee
        d = os.path.join(OUTD, lg, split)
        os.makedirs(d, exist_ok=True)
        with tarfile.open(tgz) as t:
            t.extractall(d)
        wavs = glob.glob(d + "/**/*.wav", recursive=True)
        n_ok = 0
        for w in wavs:
            b = os.path.basename(w)
            txt = meta.get(b)
            if not txt or not txt.strip():
                continue
            try:
                info = sf.info(w)
            except Exception:
                continue
            dur = info.frames / float(info.samplerate)
            if not (1.5 <= dur <= 30.0):
                continue
            if info.samplerate != 16000:
                x, sr = sf.read(w, dtype="float32")
                if x.ndim > 1:
                    x = x.mean(1)
                idx = np.linspace(0, len(x) - 1, int(len(x) * 16000 / sr))
                x = np.interp(idx, np.arange(len(x)), x).astype("float32")
                sf.write(w, x, 16000)
                dur = len(x) / 16000.0
            rows.append({"id": "fleurs_%s_%s" % (lg, b.replace(".wav", "")),
                         "audio": w, "duration": dur, "text": txt.strip(),
                         "speaker": "fleurs_%s_unk" % lg, "source": "fleurs_%s" % lg,
                         "lang": ourlang})
            n_ok += 1
        print("  %s/%s : %d clips retenus sur %d" % (lg, split, n_ok, len(wavs)), flush=True)
    out = "%s/fleurs_%s.jsonl" % (M, ourlang)
    with open(out, "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print("  -> %s : %d clips / %.1f h" % (out, len(rows), sum(r["duration"] for r in rows) / 3600), flush=True)

print("FLEURS_DONE", flush=True)