#!/usr/bin/env bash
# Train all 4 specialists with a larger model (hidden_dim=512, layers=2, heads=4).
# ~10-12x the capacity of the default (hidden_dim=128, layers=1, heads=1).
# Uses combined sprint-plan + project data for prereq_op.
#
# Usage: bash specialists/scripts/train_all_large.sh
#        (run from whetstone_DSL root)

set -euo pipefail

ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
FABRICATE_ROOT="/home/bill/Documents/WhetstoneAI_Fabricate"
RUN_BASE="/mnt/storage/fabricate_runs"

# Model size — ~16x the default 128/1/1
# Note: gpu_transformer engine requires heads=1
HIDDEN=512
LAYERS=4
HEADS=1
MAX_STEPS=20000

mkdir -p "$ROOT/specialists/runs"

cd "$FABRICATE_ROOT"

run_specialist() {
  local name="$1"
  local train="$2"
  local eval="$3"
  local labels="$4"
  local out="$RUN_BASE/${name}_large"
  local history="$ROOT/specialists/runs/${name}_large_history.json"
  local log="$ROOT/specialists/runs/${name}_large.log"

  mkdir -p "$out"
  echo "=== $name (large) ===" | tee "$log"
  echo "  train:  $train" | tee -a "$log"
  echo "  eval:   $eval" | tee -a "$log"
  echo "  model:  hidden=$HIDDEN layers=$LAYERS heads=$HEADS" | tee -a "$log"
  echo "" | tee -a "$log"

  .venv/bin/python3 run_grokking_until.py \
    --dataset              "$train" \
    --eval-dataset         "$eval" \
    --out-dir              "$out" \
    --labels               "$labels" \
    --lr                   0.0005 \
    --weight-decay         0.01 \
    --max-steps            "$MAX_STEPS" \
    --hidden-dim           "$HIDDEN" \
    --layers               "$LAYERS" \
    --heads                "$HEADS" \
    --batch-size           256 \
    --grok-loss-threshold  0.05 \
    --grok-acc-jump        10.0 \
    --stop-after-grokking-blocks 5 \
    --morph-interval       99999 \
    --reset \
    --history-out          "$history" \
    2>&1 | tee -a "$log"

  echo "" | tee -a "$log"
  echo "Done: $out/checkpoint.bin" | tee -a "$log"
}

# Run sequentially — one model at a time gets full GPU, avoids memory contention.
# prereq_op — combined sprint-plan + project data (text-labeled)
run_specialist "whetstone_prereq_op" \
  "$ROOT/specialists/data/combined/prereq_op_train.tsv" \
  "$ROOT/specialists/data/combined/prereq_op_eval.tsv" \
  "standard,needs_review,needs_approval,full_gates"

# worker_type
run_specialist "whetstone_worker_type" \
  "$ROOT/specialists/data/generated/worker_type_train.tsv" \
  "$ROOT/specialists/data/generated/worker_type_eval.tsv" \
  "implementer,reviewer,architect,qa"

# verification_type
run_specialist "whetstone_verification_type" \
  "$ROOT/specialists/data/generated/verification_type_train.tsv" \
  "$ROOT/specialists/data/generated/verification_type_eval.tsv" \
  "unit,integration,schema,smoke,docs"

# automatability
run_specialist "whetstone_automatability" \
  "$ROOT/specialists/data/generated/automatability_train.tsv" \
  "$ROOT/specialists/data/generated/automatability_eval.tsv" \
  "deterministic,template,specialist,slm,llm,human"
echo ""
echo "All large-model training runs complete."