#!/usr/bin/env bash
# Run the two requested lambda=1.25 ExOPD reference experiments sequentially.
# Each training/evaluation phase owns all six local GPUs.
set -euo pipefail

PROJECT_DIR="$(cd "$(dirname "$0")/.." && pwd)"
cd "${PROJECT_DIR}"

PYTHON_BIN="${PROJECT_DIR}/.venv/bin/python"
ACCELERATE_BIN="${PROJECT_DIR}/.venv/bin/accelerate"
ACCELERATE_CONFIG="${PROJECT_DIR}/default_config.yaml"
CHECKPOINT="${PROJECT_DIR}/outputs/opd-only-7b-chartqa-2gpu-modelonly-smoke10-formatguard-v12-formal/checkpoint-50"
INITIAL_OUTPUT="${PROJECT_DIR}/outputs/exopd-initial-student-lambda125-6gpu-1ep"
TEACHER_OUTPUT="${PROJECT_DIR}/outputs/exopd-teacher-unpriv-lambda125-6gpu-1ep"
PAIR_LOG="${PROJECT_DIR}/outputs/exopd-reference-pair-lambda125-6gpu.log"

unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY all_proxy ALL_PROXY
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export HF_DATASETS_OFFLINE=1
export WANDB_MODE=disabled
export PYTHONPATH="${PROJECT_DIR}"
export PYTHON_BIN ACCELERATE_BIN ACCELERATE_CONFIG

run_training() {
  local base_config="$1"
  local output_dir="$2"
  printf '[%s] starting training: base_config=%s output=%s\n' \
    "$(date -Is)" "${base_config}" "${output_dir}" | tee -a "${PAIR_LOG}"
  NUM_GPUS=6 \
  CANDIDATE_GPUS=0,1,2,3,4,5 \
  GRADIENT_ACCUMULATION_STEPS=8 \
  MIN_FREE_MIB=45056 \
  BASE_CONFIG="${base_config}" \
  MODEL_INIT_CHECKPOINT="${CHECKPOINT}" \
  OUTPUT_DIR="${output_dir}" \
  bash "${PROJECT_DIR}/scripts/train_opd_only_7b_chartqa_3gpu_auto.sh"
  printf '[%s] training completed: output=%s\n' \
    "$(date -Is)" "${output_dir}" | tee -a "${PAIR_LOG}"
}

run_validation() {
  local output_dir="$1"
  local eval_log="${output_dir}/eval-validation.log"
  printf '[%s] starting ChartQA validation: model=%s/final_checkpoint\n' \
    "$(date -Is)" "${output_dir}" | tee -a "${PAIR_LOG}"
  set +o pipefail
  CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 "${ACCELERATE_BIN}" launch \
    --config_file "${ACCELERATE_CONFIG}" \
    --num_processes 6 \
    --main_process_port 0 \
    -m eval.eval_chartqa \
    --model_path "${output_dir}/final_checkpoint" \
    --split validation \
    --batch_size 1 \
    --max_new_tokens 1024 \
    --progress_every 100 \
    2>&1 | tee "${eval_log}"
  local eval_ec="${PIPESTATUS[0]}"
  set -o pipefail
  if [[ "${eval_ec}" -ne 0 ]]; then
    printf '[%s] WARNING: validation exited with code %s; continuing pair. log=%s\n' \
      "$(date -Is)" "${eval_ec}" "${eval_log}" | tee -a "${PAIR_LOG}"
    return 0
  fi
  printf '[%s] validation completed: log=%s\n' \
    "$(date -Is)" "${eval_log}" | tee -a "${PAIR_LOG}"
}

mkdir -p "${INITIAL_OUTPUT}" "${TEACHER_OUTPUT}"
printf '[%s] ExOPD pair started; checkpoint=%s\n' \
  "$(date -Is)" "${CHECKPOINT}" | tee -a "${PAIR_LOG}"

run_training "${INITIAL_OUTPUT}/config.base.yaml" "${INITIAL_OUTPUT}"
run_validation "${INITIAL_OUTPUT}"
run_training "${TEACHER_OUTPUT}/config.base.yaml" "${TEACHER_OUTPUT}"
run_validation "${TEACHER_OUTPUT}"

printf '[%s] ExOPD pair completed successfully.\n' "$(date -Is)" | tee -a "${PAIR_LOG}"