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

RUN_ID="${RUN_ID:-ssm-improvement-sweep-$(date +%Y%m%d-%H%M%S)}"
DATA_PATH="${DATA_PATH:-/home/student/Data/TaoData/pretrain.jsonl}"
TOKENIZER_PATH="${TOKENIZER_PATH:-/home/student/YouZheng/tokenizers/taodata_pilot_8k/tokenizer.model}"
SSM_REPO_PATH="${SSM_REPO_PATH:-/home/student/YouZheng/gamma_ssm_repo}"
PYTHON_BIN="${PYTHON_BIN:-/home/student/.venv/bin/python}"
JOB_ROOT="${JOB_ROOT:-/home/student/YouZheng/jobs/taotern}"
REMOTE_REPO="${REMOTE_REPO:-$(pwd)}"

# This is intentionally below the eventual 200M scale. It is a model-selection
# sweep that can run unattended and resume completed benchmark cases.
MAX_TOKENS="${MAX_TOKENS:-100000000}"
MAX_RECORDS="${MAX_RECORDS:-150000}"
TRAIN_STEPS="${TRAIN_STEPS:-5000}"
EVAL_BATCHES="${EVAL_BATCHES:-96}"
BATCH_SIZES="${BATCH_SIZES:-32,64}"
SEQ_LEN="${SEQ_LEN:-512}"

JOB_COMMAND="PYTHONPATH=$REMOTE_REPO/src:$SSM_REPO_PATH $PYTHON_BIN scripts/benchmark_taonet_real_tokens.py \
  --data-path $DATA_PATH \
  --text-field text \
  --tokenizer-type sentencepiece \
  --tokenizer-path $TOKENIZER_PATH \
  --max-records $MAX_RECORDS \
  --max-tokens $MAX_TOKENS \
  --eval-fraction 0.1 \
  --architectures taonet,taonet_ssm,taonet_hybrid \
  --batch-sizes $BATCH_SIZES \
  --seq-len $SEQ_LEN \
  --hidden-dim 256 \
  --num-layers 4 \
  --num-heads 4 \
  --d-latent-kv 192 \
  --hidden-dim-ff 1024 \
  --ssm-core dplr \
  --ssm-hidden-dims 16,32 \
  --ssm-mixer-dims 128,256 \
  --ssm-num-lanes-list 1,2 \
  --ssm-lane-combine channel \
  --ssm-lane-modes full,split \
  --ssm-split-mixes none,hadamard \
  --ssm-rank 1 \
  --ssm-kernel-mode conv \
  --no-ssm-finite-tail-correction \
  --ssm-gate-types channel \
  --hybrid-patterns attention_first,ssm_first,single_ssm_middle,single_ssm_late \
  --dtype bf16 \
  --device cuda \
  --warmup 2 \
  --repeats 3 \
  --backward \
  --train-steps $TRAIN_STEPS \
  --learning-rate 0.0008 \
  --weight-decay 0.01 \
  --eval-batches $EVAL_BATCHES \
  --ssm-local-shift \
  --ssm-local-shift-per-channel \
  --ssm-local-shift-init 0.1 \
  --output-dir \"\$REPOBRIDGE_OUTPUT_DIR\" \
  --resume-completed \
  --incremental-output \
  --save-case-checkpoints \
  --checkpoint-dir \"\$TAOTERN_CHECKPOINT_DIR\""

export RUN_ID JOB_ROOT JOB_COMMAND
export OUTPUT_DIR="${OUTPUT_DIR:-$JOB_ROOT/$RUN_ID/outputs}"
export CHECKPOINT_DIR="${CHECKPOINT_DIR:-$JOB_ROOT/$RUN_ID/checkpoints}"

bash scripts/remote/submit_detached_job.sh