#!/bin/bash
# Build isolated icefall training venv (torch 2.4.1+cu124 + matching k2 + lhotse + icefall deps)
set -e
VENV=/workspace/venvs/icefall
echo "=== $(date) setup_env START ==="
mkdir -p /workspace/venvs
python3 -m venv $VENV
source $VENV/bin/activate
python -m pip install -U pip wheel setuptools 2>&1 | tail -1

echo "=== torch 2.4.1 + torchaudio (cu124) ==="
pip install torch==2.4.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu124 2>&1 | tail -2

echo "=== k2 (exact matching wheel) ==="
pip install k2==1.24.4.dev20250715+cuda12.4.torch2.4.1 -f https://k2-fsa.github.io/k2/cuda.html 2>&1 | tail -3

echo "=== lhotse + icefall deps ==="
pip install lhotse kaldialign sentencepiece tensorboard librosa soundfile \
    kaldi-native-fbank kaldilm numpy'<2' 2>&1 | tail -3
pip install -r /root/icefall/requirements.txt 2>&1 | tail -3 || echo "icefall reqs partial (ok)"

echo "=== VERIFY ==="
python - <<'PY'
import torch, k2, lhotse
print("torch", torch.__version__, "cuda_avail", torch.cuda.is_available(), "dev", torch.cuda.get_device_name(0) if torch.cuda.is_available() else None)
print("k2", k2.__version__)
print("lhotse", lhotse.__version__)
# quick k2 cuda smoke test
a = k2.Fsa.from_str("0 1 1 0.1\n1 2 -1 0.2\n2", num_aux_labels=0)
print("k2 fsa ok, ragged ok:", k2.RaggedTensor([[1,2],[3]]).sum())
PY
echo "=== $(date) setup_env DONE ==="
touch $VENV/.SETUP_COMPLETE