#!/usr/bin/env bash
# =============================================================================
# decide_l3.sh -- the decision gate: does training on documents LONGER than the
# 16,384 window help, or only on documents that FIT inside it?
#
#   l1  b1 4-8k    docs fit inside the window
#   l2  b2 8-16k   docs fit (at the boundary)
#   l3  b3 16-32k  docs EXCEED the window -> trained on truncated slices
#   armA_c  flux, median 826 -- the control
#
# If l3 <= l2, truncated long documents add nothing and the right move is to
# EXTEND the window (16k -> 32k/64k) before feeding longer data. If l3 > l2,
# truncation still carries signal and more length is worth buying.
#
# SCOPE: niah_single only, at 4096/8192/16360. The doc-length effect was
# resolvable at 8k and 16k and not at 1-4k, so those are the decisive cells;
# 4096 is kept for the trend. val_backward is omitted -- it has sat at 0-4%
# (chance ~2.9%) in every arm measured, so it would double the cost for nothing.
#
# HAYSTACK CAVEAT, stated up front: b2's val shard is the filler, which is
# IN-DISTRIBUTION for l2 and out-of-distribution for l1/l3/armA_c. That biases
# the comparison IN L2's FAVOUR -- exactly the arm l3 is being measured against.
# So "l3 beats l2" would be robust, while "l3 loses to l2" is ambiguous and
# needs the b3-val cross-check (cross_check_l3.sh) before it can be believed.
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"
ROOT="$(cd ../.. && pwd)"
export LD_LIBRARY_PATH="$ROOT/Tensor-Implementations/lib:$ROOT/Profiler/lib:${LD_LIBRARY_PATH:-}"
export OMPI_ALLOW_RUN_AS_ROOT=1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
PY=/root/Adhi/BluTrain/envv/bin/python
CK=checkpoints_bluscriptcp/rank_0
LOG=decide_logs; mkdir -p "$LOG"
M="$LOG/master.log"
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }

$PY -c "import tiktoken" 2>/dev/null || { say "FATAL tiktoken missing"; exit 1; }

EXTRA="armA_c=blumodelcp_run28_step_400.ckpt \
l1=lenbucket_l1_blumodelcp_run1_step_400.ckpt \
l2=lenbucket_l2_blumodelcp_run2_step_400.ckpt \
l3=blumodelcp_run1_step_400.ckpt"
for kv in $EXTRA; do
  [ -f "$CK/${kv#*=}" ] || { say "FATAL missing $CK/${kv#*=}"; exit 1; }
done
MODELS="armA_c l1 l2 l3"
say "=== l3 decision eval ==="
say "models: $MODELS   haystack: b2val (in-distribution for l2 -- biased toward l2)"

for LEN in 4096 8192 16360; do
  say "PROMPT_LEN=$LEN  niah_single  N=100"
  GEN_DATA_ROOT=lendata/b3_16384_32768 \
  HAY="$PWD/lendata/haystack/b2val_haystack.bin" \
  PY=$PY OUT=decide_len_$LEN TASKS="niah_single" CKPTS="$MODELS" N=100 \
  PROMPT_LEN=$LEN MAX_NEW=24 T=16384 DEPTHS="0.0,0.25,0.5,0.75,1.0" \
  YARN_ORIG=4096 GPUS=0,1 JOBS_PER_GPU=2 BF16=1 GREEDY=1 \
  CKPT_EXTRA="$EXTRA" ./run_longeval.sh >> "$LOG/eval_$LEN.log" 2>&1 \
    || say "  WARN nonzero at $LEN"
  [ -s "decide_len_$LEN/summary.csv" ] && say "  -> $(sed -n 1p decide_len_$LEN/summary.csv) / $(sed -n 2p decide_len_$LEN/summary.csv)" \
                                       || say "  -> NO RESULT at $LEN"
done
say "=== DECISION EVAL DONE ==="