#!/usr/bin/env bash
# Train the confidence_tier 3-way specialist.
# Classes: high (conf>=80) / medium (60-79) / low (<60)
# Input: "conflicts=N ambiguity=M deps=D prereqs=P queueready=yes/no"
set -euo pipefail

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

TRAIN_DATA="$ROOT/specialists/data/generated/confidence_tier_train.tsv"
EVAL_DATA="$ROOT/specialists/data/generated/confidence_tier_eval.tsv"
OUT_DIR="/mnt/storage/fabricate_runs/whetstone_confidence_tier"
HISTORY="$ROOT/specialists/runs/confidence_tier_history.json"
RUN_LOG="$ROOT/specialists/runs/confidence_tier_training.log"

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

if [[ ! -f "$TRAIN_DATA" ]]; then
  echo "Generating training data..." >&2
  python3 "$ROOT/specialists/scripts/gen_confidence_tier_data.py" \
    --train-out "$TRAIN_DATA" --eval-out "$EVAL_DATA"
fi

cd "$FABRICATE_ROOT"

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