#!/usr/bin/env bash
set -euo pipefail

cd "$(dirname "$0")/.."

mkdir -p logs
PIPELINE_LOG="${PIPELINE_LOG:-logs/presentation_ablation_pipeline.log}"
PUBLIC_VAL="${PUBLIC_VAL:-val.bin}"
EVAL_BATCH_SIZE="${EVAL_BATCH_SIZE:-8}"

exec > >(tee -a "$PIPELINE_LOG") 2>&1

log() {
  printf '[%s] %s\n' "$(date '+%F %T')" "$*"
}

run_train() {
  local name="$1"
  local model_file="$2"
  local config="$3"
  local train_log="$4"

  log "training ${name} with ${config}"
  if [[ -n "$model_file" ]]; then
    NANOGPT_MODEL_FILE="$model_file" \
    PYTHONUNBUFFERED=1 \
    PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
      python -u train.py "$config" 2>&1 | tee -a "$train_log"
  else
    env -u NANOGPT_MODEL_FILE \
    PYTHONUNBUFFERED=1 \
    PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
      python -u train.py "$config" 2>&1 | tee -a "$train_log"
  fi
}

export_and_eval() {
  local name="$1"
  local exporter="$2"
  local model_file="$3"
  local out_dir="$4"
  local submission_dir="$5"
  local eval_log="$6"

  if [[ ! -f "${out_dir}/ckpt.pt" ]]; then
    log "missing checkpoint for ${name}: ${out_dir}/ckpt.pt"
    return 1
  fi

  log "exporting ${name} to ${submission_dir}"
  python "$exporter" \
    --checkpoint "${out_dir}/ckpt.pt" \
    --model_file "$model_file" \
    --out_dir "$submission_dir"

  log "evaluating ${name} on course public val"
  PYTHONUNBUFFERED=1 python -u evaluate.py \
    --model_dir "$submission_dir" \
    --data "$PUBLIC_VAL" \
    --batch_size="$EVAL_BATCH_SIZE" \
    2>&1 | tee "$eval_log"
}

if [[ ! -f "$PUBLIC_VAL" ]]; then
  log "public val not found: ${PUBLIC_VAL}"
  exit 1
fi

if [[ "$(wc -c < "$PUBLIC_VAL")" -lt 1000000 ]]; then
  log "public val file ${PUBLIC_VAL} is too small; expected real binary, not an LFS pointer"
  exit 1
fi

SPECS=(
  "a0_nanogpt_fineweb_adamw|Naive NanoGPT baseline|config/pres_iso_a0_nanogpt_fineweb_adamw.py||scripts/export_gpt2_cse251b.py|model.py|out-pres-iso-a0-nanogpt-fineweb-adamw"
  "a1_nanogpt_mixed_adamw|Baseline model + mixed data|config/pres_iso_a1_nanogpt_mixed_adamw.py||scripts/export_gpt2_cse251b.py|model.py|out-pres-iso-a1-nanogpt-mixed-adamw"
  "a2_nanogpt_fineweb_muon|Baseline model + Muon optimizer|config/pres_iso_a2_nanogpt_fineweb_muon.py|model_nanogpt_muon.py|scripts/export_gpt2_cse251b.py|model.py|out-pres-iso-a2-nanogpt-fineweb-muon"
  "a3_lyra_fineweb_adamw|Lyra architecture + baseline data/optimizer|config/pres_iso_a3_lyra_fineweb_adamw.py|model_lyra.py|scripts/export_lyra_cse251b.py|model_lyra.py|out-pres-iso-a3-lyra-fineweb-adamw"
)

log "presentation ablation pipeline started"
for spec in "${SPECS[@]}"; do
  IFS='|' read -r name label config model_env exporter export_model out_dir <<< "$spec"
  train_log="logs/pres_ablation_${name}_train.log"
  submission_dir="submission-pres-ablation-${name}"
  eval_log="logs/pres_ablation_${name}_course_val_eval.log"

  log "=== ${label} ==="
  run_train "$name" "$model_env" "$config" "$train_log"
  export_and_eval "$name" "$exporter" "$export_model" "$out_dir" "$submission_dir" "$eval_log"
done

log "generating ablation figure"
python scripts/plot_presentation_ablations.py
log "presentation ablation pipeline finished"