#!/usr/bin/env bash
# STEP 1 of docs/DECODE_FIX_AND_BPTT.md: short fine-tune from fieldsat 4996k with the
# byte-tick decode-label fix active and --efference-every 0 (removes the copy rule the
# off-by-one efference taught).  Exact fieldsat recipe otherwise.  New dir, new suffix.
#   env START=4996000 END_TARGET=5012000 CHUNK=8000 bash scripts/train_decfix.sh
set -eu
cd /workspace/exports/atom-ai
START=${START:-4996000}; END_TARGET=${END_TARGET:-5012000}; CHUNK=${CHUNK:-8000}
LR=${LR:-1.5e-5}; SLR=${SLR:-3.6e-5}; BASE=496000
OUT=checkpoints/byte_tick_decfix; SUF=d32_ms1M_lrd_b4828_fieldsat_decfix
LOGD=logs/decfix; mkdir -p "$OUT" "$LOGD"
export PYTHONPATH=. OMP_NUM_THREADS=${THREADS:-1} MKL_NUM_THREADS=${THREADS:-1}
CUR=$START
RESUME=checkpoints/byte_tick_fieldsat/atom_native_step_${START}_d32_ms1M_lrd_b4828_fieldsat.pt
while [ "$CUR" -lt "$END_TARGET" ]; do
  END=$((CUR + CHUNK)); NAME="atom_native_step_${END}_${SUF}.pt"
  if [ ! -f "$OUT/$NAME" ]; then
    echo "$(date -Is) chunk $CUR->$END resume=$RESUME"
    .venv/bin/python -u tools/run_atom_native.py \
      --stream --data-glob 'data/corpus_fr_medium*.txt' --loop-shards --chunk-bytes 65536 \
      --stream-skip-packets $((CUR - BASE)) --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 0.5 --field-leak 0.3 > "$LOGD/train_${CUR}_${END}_decfix.log" 2>&1
    cp "$OUT/run_metrics.json" "$OUT/run_metrics_${END}_${SUF}.json"
  fi
  RESUME="$OUT/$NAME"; CUR=$END
done
echo "$(date -Is) done"