#!/usr/bin/env bash
# ============================================================================
# ONE-CLICK: build a full zh-TW + English 8 kHz PrimeTTS model in YOUR voice.
# The reference voice is the ONLY input you change; the text pools, ASR gate,
# alignment and training recipe are fixed. Both vocoder and acoustic are
# retrained (both are voice-specific).
#
# RUN FROM THE REPO ROOT:
#   PY=/path/to/venv/bin/python ./scripts/rebuild_voice.sh ref.wav ref.txt mytag
#
#   ref.wav  ~6-12 s clean clip of your target voice (mono; zh-TW that also reads
#            a little English is ideal).  ref.txt  its exact transcript.
#   No recording? synth a reference with Edge-TTS (pip install edge-tts):
#     edge-tts --voice zh-TW-HsiaoYuNeural --text "您好,歡迎來電。Thank you for calling." --write-media ref.mp3
#     ffmpeg -y -i ref.mp3 -ar 24000 -ac 1 ref.wav ; printf '%s' "您好,歡迎來電。Thank you for calling." > ref.txt
#
# PREREQS (one venv, on PYTHONPATH=repo root):
#   pip install torch torchaudio transformers onnxruntime soundfile librosa \
#               g2pw g2p_en cn2an opencc inflect faster-whisper edge-tts voxcpm nltk
#   python -c "import nltk; nltk.download('names')"
#   # base warm-start ckpt: download inflect_nano_v1_acoustic.pt from owensong/Inflect-Nano-v1
# Result corpus: ~12k clips (≈27% zh / 50% mix / 23% en), entity clips ×2 at train time.
# ============================================================================
set -e
REF=${1:?need ref.wav}; REFTXT=${2:?need ref.txt}; TAG=${3:-myvoice}
PY=${PY:-python}                                   # a venv with all deps above
GPUS=${GPUS:-0,1}; G0=${GPUS%,*}; G1=${GPUS#*,}
INIT=${INIT:-inflect_nano_v1_acoustic.pt}          # base acoustic (owensong/Inflect-Nano-v1)
export PYTHONPATH=${PYTHONPATH:-$PWD}
S=scripts; C=corpus_${TAG}; mkdir -p $C/wav

echo "### 1. Text pools (committed in the repo; regenerated only if absent) ###"
[ -f voxcpm_texts.jsonl ] || { echo "need voxcpm_texts.jsonl (base 2500 zh / 2000 en / 2800 mix)"; exit 1; }
[ -f codemix_v2.txt ]     || $PY $S/gen_codemix_v2.py --n 3000 --exclude codemix_corpus.txt --out codemix_v2.txt
[ -f entity_texts.jsonl ] || $PY $S/gen_entity_texts.py --n 2600 --out entity_texts.jsonl
$PY - <<PY
import json
with open("$C/all_texts.jsonl","w",encoding="utf-8") as o:
    for l in open("voxcpm_texts.jsonl"): o.write(l)
    for i,t in enumerate(x.strip() for x in open("codemix_v2.txt") if x.strip()):
        o.write(json.dumps({"id":f"mz{i:05d}","text":t,"lang":"mix"},ensure_ascii=False)+"\n")
    for l in open("entity_texts.jsonl"): o.write(l)
print("assembled", "$C/all_texts.jsonl")
PY

echo "### 2. Clone YOUR voice for every line (VoxCPM2), dual-GPU ###"
$PY -c "ls=[l for l in open('$C/all_texts.jsonl') if l.strip()];open('$C/t0.jsonl','w').writelines(ls[0::2]);open('$C/t1.jsonl','w').writelines(ls[1::2])"
CUDA_VISIBLE_DEVICES=$G0 $PY $S/gen_voxcpm_corpus.py --texts $C/t0.jsonl --ref "$REF" --ref-text "$REFTXT" --out-dir $C/wav --manifest $C/m0.jsonl &
CUDA_VISIBLE_DEVICES=$G1 $PY $S/gen_voxcpm_corpus.py --texts $C/t1.jsonl --ref "$REF" --ref-text "$REFTXT" --out-dir $C/wav --manifest $C/m1.jsonl &
wait; cat $C/m0.jsonl $C/m1.jsonl > $C/manifest.jsonl

echo "### 3. ASR-gate generic clips; trust entity (et*) clips as-is ###"
$PY -c "
import json
g=open('$C/to_gate.jsonl','w'); e=open('$C/entity.jsonl','w')
for l in open('$C/manifest.jsonl'): (e if json.loads(l)['id'].startswith('et') else g).write(l)"
CUDA_VISIBLE_DEVICES=$G0 $PY $S/asr_filter.py --manifest $C/to_gate.jsonl --out $C/to_gate --device cuda
cat $C/to_gate.clean.jsonl $C/entity.jsonl > $C/clean.jsonl

echo "### 4. Normalize entities + phone-level align + entity-upsample 2x ###"
$PY $S/build_corpus_v3.py --out $C/norm.jsonl $C/clean.jsonl
CUDA_VISIBLE_DEVICES=$G0 $PY $S/align_durations_v4.py --manifest $C/norm.jsonl --out $C/align.jsonl --device cuda
$PY -c "import json;r=[l for l in open('$C/align.jsonl') if l.strip()];e=[l for l in r if json.loads(l).get('id','').startswith('et')];open('$C/align.up.jsonl','w').writelines(r+e);print('train rows',len(r)+len(e))"

echo "### 5. Train 8 kHz vocoder on your voice ###"
$PY -c "import json;o=open('$C/voc_rows.jsonl','w');[o.write(json.dumps({'target_audio':json.loads(l)['target_audio'],'target_text':json.loads(l)['text']},ensure_ascii=False)+'\n') for l in open('$C/norm.jsonl')]"
CUDA_VISIBLE_DEVICES=$G1 $PY -m inflect_nano.vocoder --train-jsonl $C/voc_rows.jsonl --out-dir $C/vocoder_8k \
  --variant snake_8k --steps 40000 --batch-size 16 --segment-size 16384 --stft-weight 2.0 --save-interval 5000 --device cuda

echo "### 6. Train acoustic: warm-start v1 + 2D mel-GAN + en-upsample 2 ###"
CUDA_VISIBLE_DEVICES=$G0 $PY -m inflect_nano.acoustic --durations-jsonl $C/align.up.jsonl --out-dir $C/acoustic_8k \
  --vocoder-variant snake_8k --sample-rate 8000 --vocoder-checkpoint $C/vocoder_8k/hifigan-snake_8k-final.pt \
  --vocoder-mel-weight 1.0 --init-checkpoint "$INIT" \
  --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 2 --save-interval 5000 --device cuda

echo "### 7. Export + assess (sweep ckpts; ~35k is usually the held-out sweet spot) ###"
for K in 30000 35000 40000; do
  $PY $S/export_8k.py --acoustic-ckpt $C/acoustic_8k/inflect-micro-fastspeech-$K.pt \
      --vocoder-ckpt $C/vocoder_8k/hifigan-snake_8k-final.pt --out-dir $C/onnx_$K
  $PY $S/synth_from_text.py --onnx-dir $C/onnx_$K --out-dir $C/eval_$K   --texts eval_big.jsonl
  $PY $S/synth_from_text.py --onnx-dir $C/onnx_$K --out-dir $C/entity_$K --texts eval_entity.jsonl
  CUDA_VISIBLE_DEVICES='' $PY $S/assess_quality.py --synth-dir $C/eval_$K --tag ${TAG}_$K || true
done
echo "DONE -> pick best $C/onnx_<K>/ ; ship its *.onnx + meta.json to your model repo."