#!/usr/bin/env bash
# Slow-tau A ablation. Args: <run_name> <ema_end> <ema_warmup_frac> [epochs]
set -euo pipefail
RUN_NAME="${1:?run_name}"
EMA_END="${2:-0.999}"
EMA_WARMUP="${3:-0.60}"
EPOCHS="${4:-25}"

echo "[bootstrap] slow-tau ablation: run=$RUN_NAME ema_end=$EMA_END warmup_frac=$EMA_WARMUP epochs=$EPOCHS"
cd /workspace
REPO_DIR=""
for d in PhysioJEPA physiojepa; do
  if [ -d "$d" ]; then REPO_DIR="$d"; break; fi
done
[ -n "$REPO_DIR" ] || { echo "no repo dir found"; exit 1; }
cd "$REPO_DIR"

PY=/usr/bin/python3
$PY -m pip install --quiet --upgrade pip
$PY -m pip install --quiet \
    'datasets>=4.8.4' 'einops>=0.8.2' 'matplotlib>=3.10.0' \
    'neurokit2>=0.2.13' 'python-dotenv>=1.0' 'pyyaml>=6.0' \
    'scikit-learn>=1.5' 'scipy>=1.13' 'tqdm>=4.66' \
    'wandb>=0.18' 'wfdb>=4.3.1' 'huggingface_hub>=0.25' 'requests'

if [ -f /workspace/.env ]; then cp /workspace/.env .env; fi

if [ ! -f /workspace/cache/mimic_index.json ]; then
  echo "[bootstrap] downloading MIMIC + building index"
  PYTHONPATH=src $PY scripts/prepare_data.py \
    --root /workspace/cache/mimic --index /workspace/cache/mimic_index.json
fi

PYTHONPATH=src $PY -c "
import json, pathlib
roots = sorted([str(p) for p in pathlib.Path('/workspace/cache/mimic').glob('shard_*')
                if (p / 'dataset_info.json').exists()])
pathlib.Path('/workspace/cache/shard_roots.json').write_text(json.dumps(roots))
print('shards:', len(roots))
"

mkdir -p /workspace/runs
echo "[bootstrap] launching A ablation"
PYTHONPATH=src PYTHONUNBUFFERED=1 $PY -u scripts/train.py \
  --config configs/base.yaml \
  --model A \
  --run_name "$RUN_NAME" \
  --epochs "$EPOCHS" \
  --shard_roots_json /workspace/cache/shard_roots.json \
  --index_path /workspace/cache/mimic_index.json \
  --output_dir /workspace/runs \
  --num_workers 8 \
  --subset_frac 0.10 \
  --log_every 25 \
  --ema_end "$EMA_END" \
  --ema_warmup_frac "$EMA_WARMUP" \
  --seed 42 \
  2>&1 | tee "/workspace/runs/${RUN_NAME}.log"

echo "[bootstrap] done"