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

if [[ $# -ne 4 ]]; then
  echo "Usage: $0 METHOD MODEL_PATH GPU BENCHMARK" >&2
  exit 2
fi

METHOD="$1"
MODEL_PATH="$2"
GPU="$3"
BENCHMARK="$4"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
REPO="$(cd "$RECIPE/.." && pwd)"
EVAL_ROOT="${EVAL_ROOT:-$REPO/outputs/eval}"
OUT="$EVAL_ROOT/$METHOD"
LOG="$EVAL_ROOT/logs/$METHOD"
TOKENIZER="${VANILLA_DIR:-$REPO}"
HELMET_ROWS="$RECIPE/data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl"
MRCR_ROWS="$RECIPE/data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl"
RULER_LIMIT="${RULER_LIMIT:-20}"
BABILONG_LIMIT="${BABILONG_LIMIT:-50}"
AHA_GATE_HARD_THRESHOLD="${AHA_GATE_HARD_THRESHOLD:-0.5}"
mkdir -p "$OUT" "$LOG"

if [[ -f "$OUT/${BENCHMARK}.DONE" ]]; then
  exit 0
fi

export CUDA_VISIBLE_DEVICES="$GPU"
export TOKENIZERS_PARALLELISM=false
export PYTHONUNBUFFERED=1
export AHA_SPARSITY_STATS_PATH="$OUT/sparsity_${BENCHMARK}.json"
export AHA_GATE_HARD_THRESHOLD
export AHA_FORCE_FULL_DECODE=0
export AHA_FORCE_FULL_PREFILL_TAIL=0
export AHA_FORCE_FULL_HEADS=
export AHA_FORCE_LOW_ALPHA_FULL_HEADS=0
export AHA_DUO_PREFILL_FULL=0

model_args="pretrained=${MODEL_PATH},trust_remote_code=True,dtype=bfloat16,max_length=16384,attn_implementation=sdpa"
split_a="niah_single_1,niah_single_3,niah_multikey_2,niah_multiquery,ruler_vt,ruler_fwe,ruler_qa_hotpot"
split_b="niah_single_2,niah_multikey_1,niah_multikey_3,niah_multivalue,ruler_cwe,ruler_qa_squad"

case "$BENCHMARK" in
  ruler_a|ruler_b)
    tasks="$split_a"
    [[ "$BENCHMARK" == ruler_b ]] && tasks="$split_b"
    python -m lm_eval \
      --model hf --model_args "$model_args" \
      --tasks "$tasks" --metadata '{"max_seq_lengths":[8192]}' \
      --batch_size 1 --limit "$RULER_LIMIT" --log_samples \
      --output_path "$OUT/ruler8k_splits/${BENCHMARK#ruler_}/lm_eval" \
      >"$LOG/${BENCHMARK}.log" 2>&1
    ;;
  babilong)
    python -m lm_eval \
      --model hf --model_args "$model_args" \
      --tasks babilong_longctx --metadata '{"max_seq_lengths":"8k"}' \
      --num_fewshot 2 --batch_size 1 --limit "$BABILONG_LIMIT" --log_samples \
      --output_path "$OUT/babilong8k_qa1_qa5_n50" \
      >"$LOG/babilong.log" 2>&1
    ;;
  helmet)
    python "$SCRIPT_DIR/eval_qwen_external_longctx_pilot.py" \
      --benchmark helmet_icl --model-path "$MODEL_PATH" --tokenizer-path "$TOKENIZER" \
      --attn-implementation sdpa --rows "$HELMET_ROWS" \
      --method "$METHOD" --output "$OUT/helmet_icl8k_n50.jsonl" \
      >"$LOG/helmet.log" 2>&1
    ;;
  mrcr)
    python "$SCRIPT_DIR/eval_qwen_external_longctx_pilot.py" \
      --benchmark mrcr --model-path "$MODEL_PATH" --tokenizer-path "$TOKENIZER" \
      --attn-implementation sdpa --rows "$MRCR_ROWS" \
      --method "$METHOD" --output "$OUT/mrcr_8k_2_4_8needle_n10.jsonl" \
      >"$LOG/mrcr.log" 2>&1
    ;;
  *) echo "Unknown benchmark: $BENCHMARK" >&2; exit 2 ;;
esac

touch "$OUT/${BENCHMARK}.DONE"