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

# Run this script inside an allocated/interactive GPU session.
# It does not request SLURM resources.

PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
HF_HOME="${HF_HOME:-/public/home/scnb9biwet/.cache/huggingface}"
CONDA_ENV="${CONDA_ENV:-struct-evo}"

DATASET_JSON="${DATASET_JSON:-${PROJECT_DIR}/scripts/data/processed/finetune_data.json}"
CHECKPOINT_DIR="${CHECKPOINT_DIR:-${PROJECT_DIR}/weight/checkpoints/finetune}"
CHECKPOINT_FILENAME="${CHECKPOINT_FILENAME:-finetune.ckpt}"

BATCH_SIZE="${BATCH_SIZE:-6}"
MAX_EPOCHS="${MAX_EPOCHS:-15}"
NUM_WORKERS="${NUM_WORKERS:-5}"
ACCUMULATE_GRAD_BATCHES="${ACCUMULATE_GRAD_BATCHES:-1}"
NUM_GPUS="${NUM_GPUS:-4}"
LEARNING_RATE="${LEARNING_RATE:-0.00005}"
WARMUP_FRACTION="${WARMUP_FRACTION:-0.1}"
SAVE_EVERY_N_STEPS="${SAVE_EVERY_N_STEPS:-512}"
SEED="${SEED:-123}"
DEBUG="${DEBUG:-0}"
OFFLINE="${OFFLINE:-1}"

cd "${PROJECT_DIR}"
mkdir -p "${CHECKPOINT_DIR}"

export HF_HOME
export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
export CUBLAS_WORKSPACE_CONFIG=:4096:2
export NCCL_DEBUG="${NCCL_DEBUG:-INFO}"
export PYTHONFAULTHANDLER=1

if [[ "${OFFLINE}" == "1" ]]; then
    export HF_HUB_OFFLINE=1
    export TRANSFORMERS_OFFLINE=1
fi

if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
    # shellcheck disable=SC1091
    source "$(conda info --base)/etc/profile.d/conda.sh"
    conda activate "${CONDA_ENV}"
fi

if [[ ! -f "${DATASET_JSON}" ]]; then
    echo "Missing DATASET_JSON: ${DATASET_JSON}" >&2
    echo "Create it first with CodonTransformer.CodonData.prepare_training_data." >&2
    exit 1
fi

args=(
    "${PROJECT_DIR}/scripts/finetune.py"
    --dataset_dir "${DATASET_JSON}"
    --checkpoint_dir "${CHECKPOINT_DIR}"
    --checkpoint_filename "${CHECKPOINT_FILENAME}"
    --batch_size "${BATCH_SIZE}"
    --max_epochs "${MAX_EPOCHS}"
    --num_workers "${NUM_WORKERS}"
    --accumulate_grad_batches "${ACCUMULATE_GRAD_BATCHES}"
    --num_gpus "${NUM_GPUS}"
    --learning_rate "${LEARNING_RATE}"
    --warmup_fraction "${WARMUP_FRACTION}"
    --save_every_n_steps "${SAVE_EVERY_N_STEPS}"
    --seed "${SEED}"
)

if [[ "${DEBUG}" == "1" ]]; then
    args+=(--debug)
fi

python "${args[@]}"