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

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

PYTHON_BIN="${PYTHON_BIN:-/home/suiqk/anaconda3/envs/scalerag-ts-v4/bin/python}"
ACCELERATE_BIN="$(dirname "$PYTHON_BIN")/accelerate"
ACCELERATE_CONFIG="${ACCELERATE_CONFIG:-${PROJECT_DIR}/accelerate_config_2gpu.yaml}"
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1}"
NUM_PROCESSES="${NUM_PROCESSES:-2}"
RUN_ID="${RUN_ID:-$(date +%Y%m%d_%H%M%S)}"
MAIN_PROCESS_PORT="${MAIN_PROCESS_PORT:-30621}"

CHECKPOINT="${CHECKPOINT:-${PROJECT_DIR}/save/sft_rats40k_20260608_184144/checkpoint-2838}"
DATA_PATH="${DATA_PATH:-/mnt/share01/sqk/datasets/RATs40K/RATs-Uni-TSImage_Reason.json}"
LLM_MODEL_PATH="${LLM_MODEL_PATH:-/mnt/share01/sqk/models/qwen2.5-7b-instruct}"
LLM_ATTN_IMPLEMENTATION="${LLM_ATTN_IMPLEMENTATION:-flash_attention_2}"

EVAL_BATCH="${EVAL_BATCH:-1}"
EVAL_MAX_NEW_TOKENS="${EVAL_MAX_NEW_TOKENS:-256}"
EVAL_NUM_WORKERS="${EVAL_NUM_WORKERS:-4}"
MAX_EVAL_SAMPLES="${MAX_EVAL_SAMPLES:-}"
EVAL_OUTPUT_DIR="${EVAL_OUTPUT_DIR:-${PROJECT_DIR}/inference_results_rats40k_ckpt2838_${RUN_ID}}"
LOG_DIR="${LOG_DIR:-${PROJECT_DIR}/logs}"
LOG_FILE="${LOG_FILE:-${LOG_DIR}/rats40k_ckpt2838_eval_2gpu_${RUN_ID}.log}"

mkdir -p "$LOG_DIR" "$EVAL_OUTPUT_DIR"
touch "$LOG_FILE"

on_error() {
  local exit_code="$1"
  local line_no="$2"
  local failed_command="$3"
  echo "ITFormer RATs40K eval failed: exit=${exit_code} line=${line_no} command=${failed_command}" >&2
  echo "Log file: ${LOG_FILE}" >&2
  exit "$exit_code"
}
trap 'on_error "$?" "$LINENO" "$BASH_COMMAND"' ERR
exec > >(tee -a "$LOG_FILE") 2>&1

export CUDA_VISIBLE_DEVICES
export PYTHONWARNINGS="ignore::FutureWarning:transformers.utils.hub"
export TOKENIZERS_PARALLELISM=false
export PYTHONUNBUFFERED=1
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"

[ -x "$PYTHON_BIN" ] || { echo "Python not found: $PYTHON_BIN" >&2; exit 1; }
[ -x "$ACCELERATE_BIN" ] || { echo "Accelerate not found: $ACCELERATE_BIN" >&2; exit 1; }
[ -f "$ACCELERATE_CONFIG" ] || { echo "Accelerate config not found: $ACCELERATE_CONFIG" >&2; exit 1; }
[ -f "$CHECKPOINT/model.safetensors" ] || { echo "Checkpoint not found: $CHECKPOINT" >&2; exit 1; }
[ -f "$DATA_PATH" ] || { echo "Dataset not found: $DATA_PATH" >&2; exit 1; }

extra_args=()
if [ -n "$LLM_ATTN_IMPLEMENTATION" ]; then
  extra_args+=(--llm_attn_implementation "$LLM_ATTN_IMPLEMENTATION")
fi
if [ -n "$MAX_EVAL_SAMPLES" ]; then
  extra_args+=(--max_eval_samples "$MAX_EVAL_SAMPLES")
fi

run_eval() {
  local -a eval_cmd=(
    "$ACCELERATE_BIN" launch
    --config_file "$ACCELERATE_CONFIG"
    --num_processes "$NUM_PROCESSES"
    --main_process_port "$MAIN_PROCESS_PORT"
    inference_rats40k.py
    --checkpoint "$CHECKPOINT"
    --data_path "$DATA_PATH"
    --eval_split TSAD_test
    --llm_model_path "$LLM_MODEL_PATH"
    --patch_len 16
    --stride 16
    --input_len 128
    --d_model 512
    --n_heads 8
    --e_layers 4
    --it_d_model 3584
    --it_n_heads 28
    --it_layers 2
    --prefix_num 8
    --adapter_type itformer
    --batch_size "$EVAL_BATCH"
    --max_new_tokens "$EVAL_MAX_NEW_TOKENS"
    --num_workers "$EVAL_NUM_WORKERS"
    --output_dir "$EVAL_OUTPUT_DIR"
    "${extra_args[@]}"
  )

  echo "============================================================"
  echo "ITFormer RATs40K checkpoint evaluation"
  echo "Checkpoint: $CHECKPOINT"
  echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"
  echo "Processes: $NUM_PROCESSES"
  echo "Per-GPU eval batch: $EVAL_BATCH"
  echo "Output: $EVAL_OUTPUT_DIR"
  echo "Log: $LOG_FILE"
  echo "============================================================"

  "${eval_cmd[@]}"
  echo "ITFormer RATs40K evaluation completed: $EVAL_OUTPUT_DIR"
}

run_eval