#!/usr/bin/env bash
# v4zh-v2: retrain on DIVERSITY-EXPANDED corpus (existing 9042 + ~6500 broad-char Tatoeba/zh-TW).
# Tests the data-diversity hypothesis: held-out zh should drop from 0.40 toward the 0.21 seen-text floor.
# Keeps the proven v4 phone-level alignment. Run when teacher synth + GPU are free. GPU via $1 (default 1).
set -e
ZT=/home/luigi/jetson-tts/mossnano/zhtw8k; LOG=$ZT/v4zh_v2.log
TRAIN=/home/luigi/moss-train-venv/bin/python
GPU=${1:-1}
VOC=$ZT/m2_vocoder_8k/hifigan-snake_8k-latest.pt
exec >>"$LOG" 2>&1
echo "===== V4ZH_V2 START $(date) GPU=$GPU ====="

# 1) canonical-fix expansion manifest wav paths from id (idempotent, no doubling)
$TRAIN - <<PY
import json, os
p="$ZT/teacher_corpus_expand/manifest.jsonl"; D="$ZT/teacher_corpus_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"[expand manifest] {len(rows)} rows, {ok} files exist")
PY

# 2) v4-align the expansion clips (phone-level)
echo "[align expand] $(date)"
CUDA_VISIBLE_DEVICES=$GPU $TRAIN $ZT/align_durations_v4.py \
  --manifest $ZT/teacher_corpus_expand/manifest.jsonl \
  --out $ZT/m_v4_expand_align.jsonl --device cuda

# 3) abspath-fix + merge with existing v4 alignment
$TRAIN - <<PY
import json
R="/home/luigi/jetson-tts/"
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"]=R+r["target_audio"]
    return rows
a=fix("$ZT/m_v4_align.jsonl"); b=fix("$ZT/m_v4_expand_align.jsonl")
out="$ZT/m_v4_v2_align.jsonl"
open(out,"w").write("\n".join(json.dumps(r,ensure_ascii=False) for r in a+b)+"\n")
print(f"[merged] {len(a)} + {len(b)} = {len(a)+len(b)} -> {out}")
PY

# 4) retrain on merged corpus (same recipe, frozen arch)
echo "[v4zh-v2 acoustic] $(date)"
cd /tmp/inflect-nano
CUDA_VISIBLE_DEVICES=$GPU $TRAIN -m inflect_nano.acoustic --durations-jsonl $ZT/m_v4_v2_align.jsonl \
  --out-dir $ZT/v4zh_v2_acoustic_8k --vocoder-variant snake_8k --sample-rate 8000 \
  --steps 60000 --batch-size 16 --lr 2e-4 --max-frames 1000 --en-upsample 2 \
  --vocoder-checkpoint $VOC --vocoder-mel-weight 1.0 \
  --save-interval 5000 --log-interval 200 --device cuda
echo "===== V4ZH_V2 TRAIN DONE $(date) ====="