#!/usr/bin/env bash
# Train the prereq_op_selector 4-way specialist.
#
# Classes:
#   0 = standard       [validate-intake, resolve-dependencies]
#   1 = needs_review   [validate-intake, architect-review]
#   2 = needs_approval [validate-intake, manual-approval]
#   3 = full_gates     [validate-intake, architect-review, manual-approval]
#
# Usage: bash specialists/scripts/train_prereq_op.sh
#        (run from whetstone_DSL root)

set -euo pipefail

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

TRAIN_DATA="$ROOT/specialists/data/generated/prereq_op_train.tsv"
EVAL_DATA="$ROOT/specialists/data/generated/prereq_op_eval.tsv"
OUT_DIR="/mnt/storage/fabricate_runs/whetstone_prereq_op"
HISTORY="$ROOT/specialists/runs/prereq_op_history.json"
RUN_LOG="$ROOT/specialists/runs/prereq_op_training.log"

mkdir -p "$ROOT/specialists/runs" "$OUT_DIR"

if [[ ! -f "$TRAIN_DATA" ]]; then
  echo "Missing training data. Run:" >&2
  echo "  python3 specialists/scripts/gen_prereq_op_data.py" >&2
  exit 1
fi

echo "Training prereq_op_selector specialist..."
echo "  train: $TRAIN_DATA"
echo "  eval:  $EVAL_DATA"
echo "  out:   $OUT_DIR"
echo "  log:   $RUN_LOG"
echo ""

cd "$FABRICATE_ROOT"

.venv/bin/python3 run_grokking_until.py \
  --dataset        "$TRAIN_DATA" \
  --eval-dataset   "$EVAL_DATA" \
  --out-dir        "$OUT_DIR" \
  --labels         "standard,needs_review,needs_approval,full_gates" \
  --lr             0.001 \
  --weight-decay   0.01 \
  --max-steps      5000 \
  --grok-loss-threshold   0.10 \
  --grok-acc-jump         15.0 \
  --stop-after-grokking-blocks 3 \
  --history-out    "$HISTORY" \
  2>&1 | tee "$RUN_LOG"

echo ""
echo "Done. Checkpoint at: $OUT_DIR/checkpoint.bin"