set -uo pipefail
cd /home/ruizzhao/diffusion_baseline
export DLB_ROOT="$PWD"
export DLB_PYTHON=/opt/anaconda3/envs/dlb-bootstrap/bin/python

FAILURES="results/logs/eval_bench_failures_finished_distills_gpu7.tsv"
printf 'stage\tmodel\tdataset\tsteps\tgpu\n' > "$FAILURES"

run_eval() {
  local model="$1"
  local dataset="$2"
  local steps="$3"
  local gpu=7

  local sample="results/samples/${dataset}/${model}/steps_${steps}/samples.jsonl"
  local metrics="results/metrics/${dataset}/${model}/steps_${steps}/metrics.json"
  mkdir -p "$(dirname "$metrics")"

  if [[ ! -f "$sample" ]] || [[ "$(wc -l < "$sample")" -ne 1024 ]]; then
    echo "SKIP eval ${model}/${dataset} steps=${steps}: missing/incomplete samples"
    printf 'eval_missing\t%s\t%s\t%s\t%s\n' "$model" "$dataset" "$steps" "$gpu" >> "$FAILURES"
    return 0
  fi

  echo "[$(date)] GPU${gpu} EVAL ${model}/${dataset} steps=${steps}"
  if ! CUDA_VISIBLE_DEVICES="$gpu" /opt/anaconda3/envs/dlb-eval/bin/python -m evaluation.evaluate \
    --root "$PWD" \
    --samples "$sample" \
    --metrics gen_ppl,entropy,self_bleu \
    --dataset "$dataset" \
    --output "$metrics" \
    --device cuda \
    > "results/logs/eval_${model}_${dataset}_steps_${steps}_gpu${gpu}.log" 2>&1; then
    printf 'eval\t%s\t%s\t%s\t%s\n' "$model" "$dataset" "$steps" "$gpu" >> "$FAILURES"
    echo "FAILED eval ${model}/${dataset} steps=${steps}"
  fi
}

run_bench() {
  local model="$1"
  local dataset="$2"
  local steps="$3"
  local gpu=7

  echo "[$(date)] GPU${gpu} BENCH ${model}/${dataset} steps=${steps}"
  if ! CUDA_VISIBLE_DEVICES="$gpu" bash scripts/benchmark_one.sh \
    --model "$model" --dataset "$dataset" --steps "$steps" \
    --seed 42 --precision author \
    > "results/logs/bench_${model}_${dataset}_steps_${steps}_gpu${gpu}.log" 2>&1; then
    printf 'bench\t%s\t%s\t%s\t%s\n' "$model" "$dataset" "$steps" "$gpu" >> "$FAILURES"
    echo "FAILED bench ${model}/${dataset} steps=${steps}"
  fi
}

for model_dataset in "duo_di4c owt" "mdlm_sdtt lm1b" "duo_dcd lm1b"; do
  read -r model dataset <<< "$model_dataset"
  for steps in 1 2 4 8 16 32; do
    run_eval "$model" "$dataset" "$steps"
  done
done

for model_dataset in "duo_di4c owt" "mdlm_sdtt lm1b" "duo_dcd lm1b"; do
  read -r model dataset <<< "$model_dataset"
  for steps in 1 2 4 8 16 32; do
    run_bench "$model" "$dataset" "$steps"
  done
done

echo "===== failures ====="
cat "$FAILURES"
echo "ALL EVAL/BENCH ATTEMPTED"