#!/usr/bin/env bash
# Fine-tuning experiment queue (the instructor's dials: loss function, epochs, layers, plus the
# pretraining and synthetic-data ablations). Waits for the pretraining run to finish, then runs
# each experiment sequentially on the GPU, evaluates its best/ on the TREC-06 binary holdout and
# the three-way holdout, and exports the main run. Launch in tmux:
#
#   tmux new-session -d -s phish-finetune 'scripts/run_experiments.sh 2>&1 | tee logs/experiments_$(date +%Y%m%d-%H%M%S).out'
#
# Resumable: a run whose best/ already exists is skipped; re-run the script to continue after a crash.
set -uo pipefail
cd "$(dirname "$0")/.."
PY=".venv/bin/python"
PRETRAIN="checkpoints/pretrain/${PRETRAIN_RUN:-large-24L}/final"
EVAL_EVERY="${EVAL_EVERY:-1250}"          # ~quarter epoch at batch 32 over 160k rows (val eval costs ~1 min)
MAIN="${MAIN_RUN:-ft-ce}"

log() { printf '%s %s\n' "$(date +%H:%M:%S)" "$*"; }

if [[ ! -f "$PRETRAIN/model.safetensors" ]]; then
  log "waiting for $PRETRAIN/model.safetensors (pretraining still running)"
  while [[ ! -f "$PRETRAIN/model.safetensors" ]]; do sleep 300; done
  log "pretraining final/ found; waiting 60 s for the GPU to free"; sleep 60
fi
$PY src/pretrain.py plot --run-name "${PRETRAIN_RUN:-large-24L}" || true

# name | extra flags            (all: --init pretrained:$PRETRAIN unless --init scratch is given)
EXPERIMENTS=(
  "ft-ce        --loss cross_entropy --epochs 3"
  "ft-weighted  --loss weighted_cross_entropy --epochs 2"
  "ft-focal     --loss focal --epochs 2"
  "ft-synthetic --loss cross_entropy --epochs 2 --synthetic on"
  "ft-12L       --loss cross_entropy --epochs 2 --n-layers 12"
  "ft-scratch   --loss cross_entropy --epochs 2 --init scratch"
)
STATUS=()
for spec in "${EXPERIMENTS[@]}"; do
  name="${spec%% *}"; flags="${spec#* }"
  if [[ -f "checkpoints/finetune/$name/best/model.safetensors" ]]; then
    log "$name: best/ exists, skipping training"
  else
    init="--init pretrained:$PRETRAIN"; [[ "$flags" == *"--init scratch"* ]] && init=""
    log "$name: $PY src/finetune.py train --run-name $name $init $flags --eval-every $EVAL_EVERY"
    $PY src/finetune.py train --run-name "$name" $init $flags --eval-every "$EVAL_EVERY"
    rc=$?; log "$name: finetune exit $rc"; STATUS+=("$name:train=$rc")
    [[ $rc -ne 0 ]] && continue
  fi
  $PY src/finetune.py plot --run-name "$name" --out "reports/finetune_curves_$name.png" || true
  if [[ ! -f "reports/eval_done_$name" ]]; then
    # one invocation = one experiments.jsonl line + one RUN_LOG entry covering both holdouts
    $PY src/evaluate.py run --checkpoint "checkpoints/finetune/$name/best" --split test,test3 --run-name "$name" --out-dir "reports/eval_$name"
    rc=$?; log "$name: evaluate exit $rc"; STATUS+=("$name:eval=$rc")
    [[ $rc -eq 0 ]] && touch "reports/eval_done_$name"
  fi
done
# main run: top-level reports + export + app files
if [[ -f "checkpoints/finetune/$MAIN/best/model.safetensors" ]]; then
  $PY src/finetune.py plot --run-name "$MAIN" || true
  $PY src/evaluate.py run --checkpoint "checkpoints/finetune/$MAIN/best" --split test,test3 --run-name "$MAIN" \
      --finetune-run-dir "checkpoints/finetune/$MAIN" --pretrain-run-dir "checkpoints/pretrain/${PRETRAIN_RUN:-large-24L}"
  $PY src/export.py run --checkpoint "checkpoints/finetune/$MAIN/best" --out-dir "model_export/$MAIN" --app-dir app
  STATUS+=("$MAIN:export=$?")
fi
$PY src/finetune.py list
log "queue finished: ${STATUS[*]:-nothing ran}"