#!/usr/bin/env bash
# BILINGUAL NATIVE experiment: apply the proven recipe (v4 phone-align + diverse data) symmetrically.
# Merge zh-expanded corpus (15423, mark voice, incl code-mix) + diverse English-expanded (mark voice)
# -> one balanced single-voice bilingual corpus -> train ONE 4.63M model good at zh AND en (+code-mix).
# Run when English synth done + a GPU free. GPU via $1 (default 0).
set -e
ZT=/home/luigi/jetson-tts/mossnano/zhtw8k; LOG=$ZT/bilingual.log
TRAIN=/home/luigi/moss-train-venv/bin/python
GPU=${1:-0}
VOC=$ZT/m2_vocoder_8k/hifigan-snake_8k-latest.pt
exec >>"$LOG" 2>&1
echo "===== BILINGUAL START $(date) GPU=$GPU ====="

# 1) canonical-fix English-expand manifest wav paths from id
$TRAIN - <<PY
import json, os
p="$ZT/teacher_corpus_en_expand/manifest.jsonl"; D="$ZT/teacher_corpus_en_expand"
rows=[json.loads(l) for l in open(p) if l.strip()]
ok=0
for r in rows: r["wav"]=f"{D}/{r['id']}.wav"; ok+=os.path.exists(r["wav"])
open(p,"w").write("\n".join(json.dumps(r,ensure_ascii=False) for r in rows)+"\n")
print(f"[en-expand manifest] {len(rows)} rows, {ok} files exist")
PY

# 2) v4 phone-align the English-expand clips (espeak handles English natively)
echo "[align en-expand] $(date)"
CUDA_VISIBLE_DEVICES=$GPU $TRAIN $ZT/align_durations_v4.py \
  --manifest $ZT/teacher_corpus_en_expand/manifest.jsonl \
  --out $ZT/en_expand_v4_align.jsonl --device cuda

# 3) merge zh-expanded + en-expanded -> bilingual corpus (all mark voice, speaker 0)
$TRAIN - <<PY
import json
def fix(p):
    rows=[json.loads(l) for l in open(p) if l.strip()]
    for r in rows:
        if not r["target_audio"].startswith("/"): r["target_audio"]="/home/luigi/jetson-tts/"+r["target_audio"]
    return rows
a=fix("$ZT/m_v4_v2_align.jsonl"); b=fix("$ZT/en_expand_v4_align.jsonl")
out="$ZT/m_bilingual_align.jsonl"
open(out,"w").write("\n".join(json.dumps(r,ensure_ascii=False) for r in a+b)+"\n")
import re
en=sum(1 for r in a+b if all(ord(c)<128 for c in r["text"]))
print(f"[bilingual] {len(a)} + {len(b)} = {len(a)+len(b)} rows -> {out} | ~{en} pure-ASCII(en) rows")
PY

# 4) train the bilingual model (en-upsample 1: English now abundant via expansion)
echo "[bilingual acoustic] $(date)"
cd /tmp/inflect-nano
CUDA_VISIBLE_DEVICES=$GPU $TRAIN -m inflect_nano.acoustic --durations-jsonl $ZT/m_bilingual_align.jsonl \
  --out-dir $ZT/bili_acoustic_8k --vocoder-variant snake_8k --sample-rate 8000 \
  --steps 60000 --batch-size 16 --lr 2e-4 --max-frames 1000 --en-upsample 1 \
  --vocoder-checkpoint $VOC --vocoder-mel-weight 1.0 \
  --save-interval 5000 --log-interval 200 --device cuda
echo "===== BILINGUAL TRAIN DONE $(date) ====="