#!/usr/bin/env bash
# STEP 3 of docs/DECODE_FIX_AND_BPTT.md: long full-corpus run with truncated BPTT through the
# field (--field-bptt K), resumable, 8k-step chunks, each chunk a NEW checkpoint in a NEW dir.
# Re-running resumes from the newest checkpoint in $OUT (an interrupted chunk is re-run).
#
# Env: INIT=<ckpt to start from>  START=<its step>  END_TARGET  CHUNK=8000
#      K=64 LEAK=0.05 LEARN=1 (learnable per-mode leak) LR=1e-3 SLR=1e-3 LADROP=0.5
#      OUT=checkpoints/byte_tick_bptt  SUF=d32_ms1M_fieldsat_decfix_bptt  MEASURE_EVERY=2
# Stream position: SKIP = (CUR - 496000) mod 1000452 (one corpus epoch = 1,000,453 bytes = 1,000,452
# byte-tick transitions): same text position as the legacy skip without replaying 4.5M packets
# (the 256-period phase feature differs by a constant offset; irrelevant for this new branch).
# Launch:  nohup env INIT=... START=... bash scripts/train_bptt.sh >> logs/bptt/master.log 2>&1 &
#
# Opt-in variables (defaults = exact previous behaviour):
#   DATA_GLOB='data/corpus_fr_medium*.txt'  corpus glob (mixed corpus: 'data/corpus_fr_mix_[0-9]*.txt')
#   EPOCH=1000452        corpus length in byte-tick transitions (mix: total bytes from the manifest)
#   SKIP_BYTES=0         1 = position the stream with --stream-skip-bytes (seek, no replay) instead of
#                        --stream-skip-packets (replay costs ~5 min per 4.5M packets on a 7 MB corpus)
#   TSV_SCRIPT=$A/tsv_bptt.py   progress-row script (mix: docs/artifacts/mix_long_memory/tsv_mix.py)
#   MEASURE_ASYNC=0      1 = run the progress row in the background (1 thread, pinned to MEASURE_CPUS),
#                        at most MEASURE_MAX concurrent, so training does not wait for it
#   CPUS=""              taskset core list for the training process ("" = no pinning)
#   EXTRA_ARGS=""        extra trainer flags (e.g. a new opt-in lever)
set -u
cd /workspace/exports/atom-ai
INIT=${INIT:?INIT ckpt}; START=${START:?START step}; END_TARGET=${END_TARGET:-$((START + 2000000))}
CHUNK=${CHUNK:-8000}; K=${K:-64}; LEAK=${LEAK:-0.05}; LEARN=${LEARN:-1}; LR=${LR:-1e-3}; SLR=${SLR:-$LR}
LADROP=${LADROP:-0.5}; MEASURE_EVERY=${MEASURE_EVERY:-2}
OUT=${OUT:-checkpoints/byte_tick_bptt}; SUF=${SUF:-d32_ms1M_fieldsat_decfix_bptt}
A=docs/artifacts/decfix_bptt; TSV=${TSV:-$A/bptt_progress.tsv}; LOGD=${LOGD:-logs/bptt}
EPOCH=${EPOCH:-1000452}; BASE=${BASE:-496000}; PY=.venv/bin/python
DATA_GLOB=${DATA_GLOB:-data/corpus_fr_medium*.txt}; SKIP_BYTES=${SKIP_BYTES:-0}; TSV_SCRIPT=${TSV_SCRIPT:-$A/tsv_bptt.py}
MEASURE_ASYNC=${MEASURE_ASYNC:-0}; MEASURE_CPUS=${MEASURE_CPUS:-5-7}; MEASURE_MAX=${MEASURE_MAX:-2}; CPUS=${CPUS:-}
EXTRA_ARGS=${EXTRA_ARGS:-}
[ -x "$PY" ] || { echo "$(date -Is) $PY missing: run scripts/ensure_venv.sh"; exit 1; }
PIN=(); [ -n "$CPUS" ] && PIN=(taskset -c "$CPUS")
SKIP_FLAG=--stream-skip-packets; [ "$SKIP_BYTES" = "1" ] && SKIP_FLAG=--stream-skip-bytes
export PYTHONPATH=. OMP_NUM_THREADS=${THREADS:-1} MKL_NUM_THREADS=${THREADS:-1}
mkdir -p "$OUT" "$LOGD"
exec 9>"$LOGD/driver.lock"; flock -n 9 || { echo "$(date -Is) another driver holds $LOGD/driver.lock"; exit 1; }
echo "=== $(date -Is) bptt driver pid=$$ INIT=$INIT START=$START K=$K LEAK=$LEAK LEARN=$LEARN LR=$LR SLR=$SLR LADROP=$LADROP OUT=$OUT ==="
LEARN_FLAG=""; [ "$LEARN" = "1" ] && LEARN_FLAG="--field-leak-learnable"
latest() { ls "$OUT"/atom_native_step_*_${SUF}.pt 2>/dev/null | sed -n 's/.*step_\([0-9]*\)_.*/\1/p' | sort -n | tail -1; }
CUR=$(latest); if [ -z "$CUR" ]; then CUR=$START; RESUME=$INIT; else RESUME="$OUT/atom_native_step_${CUR}_${SUF}.pt"; fi
N=0
while [ "$CUR" -lt "$END_TARGET" ]; do
  END=$((CUR + CHUNK)); NAME="atom_native_step_${END}_${SUF}.pt"; SKIP=$(( (CUR - BASE) % EPOCH ))
  LOG="$LOGD/train_${CUR}_${END}_bptt.log"; T0=$(date +%s)
  echo "$(date -Is) chunk $CUR->$END skip=$SKIP resume=$RESUME"
  "${PIN[@]}" "$PY" -u tools/run_atom_native.py \
    --stream --data-glob "$DATA_GLOB" --loop-shards --chunk-bytes 65536 \
    $SKIP_FLAG "$SKIP" --resume "$RESUME" --output-dir "$OUT" \
    --checkpoint-name "$NAME" --steps "$CHUNK" \
    --d-model 32 --n-modes 32 --n-atoms-max 64 --max-span-bytes 1 \
    --field-obligatory-hard --no-enable-merge --no-payload-copy --last-atom-readout \
    --efference-every 0 --learning-rate "$LR" --surface-learning-rate "$SLR" --seed 20260913 \
    --atom-flush-every 64 --episode-length 8000 --episode-reset --slow-every 1 \
    --field-loss-weight 0.05 --field-contrast-weight 0 --allow-parallel-train \
    --last-atom-dropout "$LADROP" --field-leak "$LEAK" --field-bptt "$K" $LEARN_FLAG $EXTRA_ARGS > "$LOG" 2>&1
  RC=$?
  if [ "$RC" -ne 0 ] || [ ! -f "$OUT/$NAME" ]; then echo "$(date -Is) chunk FAILED rc=$RC see $LOG"; exit 1; fi
  cp "$OUT/run_metrics.json" "$OUT/run_metrics_${END}_${SUF}.json"
  echo "$(date -Is) chunk $CUR->$END saved wall=$(( $(date +%s) - T0 ))s"
  RESUME="$OUT/$NAME"; CUR=$END; N=$((N + 1))
  if [ $((N % MEASURE_EVERY)) -eq 0 ]; then
    if [ "$MEASURE_ASYNC" = "1" ]; then
      NRUN=$(pgrep -f -- "$TSV_SCRIPT $TSV" | wc -l)
      if [ "$NRUN" -lt "$MEASURE_MAX" ]; then
        OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 nohup taskset -c "$MEASURE_CPUS" "$PY" -u $TSV_SCRIPT "$TSV" "$OUT/$NAME" \
          "$OUT/run_metrics_${END}_${SUF}.json" >> "$LOGD/measure.log" 2>&1 < /dev/null 9>&- &
      else
        echo "$(date -Is) measure skipped at $END ($NRUN running)"
      fi
    else
      "$PY" -u $TSV_SCRIPT "$TSV" "$OUT/$NAME" "$OUT/run_metrics_${END}_${SUF}.json" >> "$LOGD/measure.log" 2>&1 \
        && tail -n 1 "$TSV" | cut -c1-400
    fi
  fi
done
echo "$(date -Is) DONE at $CUR"