#!/usr/bin/env bash
# Validation-only screening for tasks on which additional depth is useful to BP.
# The test generator is not evaluated. Only tasks passing the BP depth-gain gate
# should advance to local-learning and lesion experiments.
#
# Usage:
#   depth_task_validation_sweep.sh <gpu> <dataset> [depths] [modes] [width] [model_seeds] [task_seeds]
set -eu

cd "$(dirname "$0")/.."
GPU="${1:?GPU index required}"
DATASET="${2:?dataset required: teacher, hierarchical, or tentmap}"
DEPTHS="${3:-1 2 4 6 8 12}"
MODES="${4:-bp}"
WIDTH_OVERRIDE="${5:-}"
MODEL_SEEDS="${6:-0}"
TASK_SEEDS="${7:-0}"
PYTHON="${PYTHON:-/home/yurenh2/miniconda3/envs/ep_pascal/bin/python3}"

case "$DATASET" in
  teacher)
    WIDTH="${WIDTH_OVERRIDE:-64}"
    ACT=tanh
    RESIDUAL=1
    EPOCHS=25
    ETA=0.03
    TASK_ARGS="--task_n_in 128 --task_classes 10 --teacher_depth 8 --teacher_width 64 --teacher_residual 1"
    ;;
  hierarchical)
    WIDTH="${WIDTH_OVERRIDE:-64}"
    ACT=tanh
    RESIDUAL=1
    EPOCHS=25
    ETA=0.03
    TASK_ARGS="--task_levels 6 --task_classes 10"
    ;;
  tentmap)
    WIDTH="${WIDTH_OVERRIDE:-16}"
    ACT=relu
    RESIDUAL=0
    EPOCHS=40
    ETA=0.02
    TASK_ARGS="--task_levels 7 --task_n_in 4"
    ;;
  *)
    echo "unknown depth task: $DATASET" >&2
    exit 2
    ;;
esac

for task_seed in $TASK_SEEDS; do
  for model_seed in $MODEL_SEEDS; do
    for mode in $MODES; do
      for depth in $DEPTHS; do
        tag="depthtask_dev_v1_${DATASET}_${mode}_w${WIDTH}_d${depth}_t${task_seed}_s${model_seed}"
        out="results/${tag}.json"
        if [ -s "$out" ]; then
          echo "[$tag] exists; skipping"
          continue
        fi
        CUDA_VISIBLE_DEVICES="$GPU" "$PYTHON" experiments/run.py \
          --mode "$mode" --dataset "$DATASET" --device cuda \
          --depth "$depth" --width "$WIDTH" --act "$ACT" --residual "$RESIDUAL" \
          --epochs "$EPOCHS" --batch_size 128 --eta "$ETA" --momentum 0.9 \
          --eta_A 0.02 --eta_P 0.002 \
          --pert_sigma 0.01 --pert_every 4 --pert_ndirs 16 --pert_mode simultaneous \
          --task_train_examples 50000 --task_test_examples 10000 \
          --task_seed "$task_seed" $TASK_ARGS \
          --val_examples 5000 --split_seed 2027 --eval_split validation \
          --diagnostics alignment --seed "$model_seed" --log_every 200 --tag "$tag"
      done
    done
  done
done