#!/usr/bin/env bash
# ============================================================================
# Reproduce the SHIPPED PrimeTTS — young-female bilingual zh-TW + English, 8 kHz, 4.63M FROZEN.
# Teacher = VoxCPM2 voice-cloning a single Edge-TTS zh-TW reference. Acoustic warm-started from
# Inflect-Nano-v1. End-to-end: reference -> texts -> teacher corpus -> ASR gate -> align ->
# vocoder -> warm-started GAN acoustic -> export -> assess. Shipped checkpoint = step 35000.
#
# Held-out (eval_big.jsonl, 36 sentences): zh-CER 0.090 | mix 0.178 | en-WER 0.083 | PESQ 3.31 | MOS 4.24
#
# Prereqs: Python 3.12; pip install torch torchaudio transformers onnxruntime soundfile librosa \
#          g2pw g2p_en cn2an opencc faster-whisper edge-tts voxcpm ; plus the inflect_nano trainer
#          on PYTHONPATH and the base ckpt inflect_nano_v1_acoustic.pt (from owensong/Inflect-Nano-v1).
# Edit the paths below for your environment, then run stage-by-stage (each is resumable).
# ============================================================================
set -e
PY=python                       # a venv with the deps above
INIT_CKPT=inflect_nano_v1_acoustic.pt   # base acoustic from owensong/Inflect-Nano-v1 (warm-start)

# --- 0. Reference voice: one young Taiwan-female clip (also reads English) ---
REF_TEXT="您好,歡迎來電。請問今天需要什麼服務呢?好的,我馬上幫您轉接給 Jason 經理,請稍等一下,謝謝。"
edge-tts --voice zh-TW-HsiaoChenNeural --rate=-5% --text "$REF_TEXT" --write-media ref.mp3
ffmpeg -y -i ref.mp3 -ar 24000 -ac 1 ref.wav
echo "$REF_TEXT" > ref.txt

# --- 1. Text pool: diverse zh + diverse en + frame-bank code-mix (one {id,text,lang} jsonl) ---
#   zh / en: select_diverse_text.py (Tatoeba -> OpenCC s2twp -> char/word-coverage greedy).
#   mix:     gen_codemix.py (Taiwan office / phone-attendant frame bank, English in varied positions).
$PY select_diverse_text.py --lang zh --n 2500 --out train_zh.tsv
$PY select_diverse_text.py --lang en --n 2000 --out train_en.tsv
$PY gen_codemix.py --n 2800 --out codemix_corpus.txt
#   -> assemble texts.jsonl with lang tags {zh|en|mix}; ids zh*/en*/mix* (see your own merge step).

# --- 2. Teacher corpus: VoxCPM2 clones the ONE reference for every line (48 kHz, one voice) ---
$PY gen_voxcpm_corpus.py --texts texts.jsonl --ref ref.wav --ref-text ref.txt \
    --out-dir corpus --manifest corpus/manifest.jsonl

# --- 3. ASR quality gate (Taiwan-tuned so it never penalizes the target accent) ---
#   zh/mix -> Breeze-ASR-25 Han-CER (<=0.12 / 0.15) ; en -> Whisper-medium WER (<=0.20).
CUDA_VISIBLE_DEVICES=0 $PY asr_filter.py --manifest corpus/manifest.jsonl --out corpus/manifest --device cuda
#   -> corpus/manifest.clean.jsonl   (shipped: 6623 kept / 677 dropped = 9.3%)

# --- 4. Phone-level alignment (bopomofo + arpabet per-phone durations) -- THE key step ---
CUDA_VISIBLE_DEVICES=0 $PY align_durations_v4.py --manifest corpus/manifest.clean.jsonl \
    --out align.jsonl --device cuda

# --- 5. 8 kHz vocoder (snake_8k), retrained on the teacher audio ---
#   build voc_rows.jsonl = {target_audio,target_text} from manifest.clean.jsonl, then:
CUDA_VISIBLE_DEVICES=1 $PY -m inflect_nano.vocoder --train-jsonl voc_rows.jsonl \
    --out-dir vocoder_8k --variant snake_8k --steps 40000 --batch-size 16 \
    --segment-size 16384 --min-seconds 0.8 --max-seconds 20 --num-workers 4 \
    --stft-weight 2.0 --save-interval 5000 --device cuda

# --- 6. Acoustic: WARM-START from Inflect-Nano-v1 + 2D mel-GAN recipe (25k recon warmup) ---
CUDA_VISIBLE_DEVICES=0 $PY -m inflect_nano.acoustic --durations-jsonl align.jsonl \
    --out-dir acoustic_8k --vocoder-variant snake_8k --sample-rate 8000 \
    --vocoder-checkpoint vocoder_8k/hifigan-snake_8k-final.pt --vocoder-mel-weight 1.0 \
    --init-checkpoint "$INIT_CKPT" \
    --mel-gan-weight 0.1 --gan-2d --gan-fm-auto --gan-r1-gamma 1.0 --gan-crop 128 --gan-warmup-steps 25000 \
    --steps 60000 --batch-size 16 --lr 2e-4 --max-frames 1000 --en-upsample 1 \
    --save-interval 5000 --device cuda
#   (expect "199 tensors copied, 0 skipped" -> full v1 transfer incl. English)

# --- 7. Export ONNX (auto-detects snake_8k) + evaluate on held-out eval_big.jsonl ---
#   Sweep ckpts; pick the HELD-OUT sweet spot (GAN keeps sharpening train past the held-out optimum).
#   Shipped = step 35000.
$PY export_8k.py --acoustic-ckpt acoustic_8k/inflect-micro-fastspeech-35000.pt \
    --vocoder-ckpt vocoder_8k/hifigan-snake_8k-final.pt --out-dir onnx
$PY synth_from_text.py --onnx-dir onnx --out-dir eval --texts eval_big.jsonl
CUDA_VISIBLE_DEVICES='' $PY assess_quality.py --synth-dir eval --tag primetts_tw_8k
echo "DONE — onnx/ holds acoustic_encoder.onnx acoustic_decoder.onnx vocoder.onnx meta.json"