#!/usr/bin/env bash
# Submit the alpha-perception workflow.
#
# Trains the stacked Mamba2 -> CfC -> heads perception model on MBP-10
# from the training-data PVC. Emits alpha_train_summary.json with
# per-horizon validation AUC to the feature-cache PVC.
#
# Defaults match the validated synthetic-overfit smoke config; cluster
# runs can override for sweep work.
#
# Usage:
#   ./scripts/argo-alpha-perception.sh                   # current HEAD on ml-alpha-phase-a
#   ./scripts/argo-alpha-perception.sh --branch main
#   ./scripts/argo-alpha-perception.sh --epochs 1 --n-train-seqs 1000  # quick smoke
#   ./scripts/argo-alpha-perception.sh --watch

set -euo pipefail

SHA=HEAD
BRANCH=$(git symbolic-ref --short HEAD 2>/dev/null || echo "ml-alpha-phase-a")
GPU_POOL=ci-training-l40s
EPOCHS=5
SEQ_LEN=32
MAMBA2_STATE_DIM=16
LR_CFC=3e-3
LR_MAMBA2=1e-3
N_TRAIN_SEQS=8000
N_VAL_SEQS=1000
SEED=16962
BATCH_SIZE=1
AUTO_HORIZON_WEIGHTS=false
EARLY_STOP_METRIC=mean_auc
EARLY_STOP_PATIENCE=5
CV_FOLD=0
CV_N_FOLDS=1
CV_TRAIN_WINDOW=0
WATCH=false

usage() {
  cat <<EOF
Usage: $0 [OPTIONS]
  --sha <commit>            Git SHA (default: HEAD on the current branch)
  --branch <branch>         Git branch (default: $BRANCH)
  --gpu-pool <pool>         GPU pool (default: $GPU_POOL)
  --epochs <n>              Training epochs (default: $EPOCHS)
  --seq-len <n>             Snapshots per sequence window (default: $SEQ_LEN)
  --mamba2-state-dim <n>    Mamba2 SSM state dim (default: $MAMBA2_STATE_DIM)
  --lr-cfc <f>              CfC learning rate (default: $LR_CFC)
  --lr-mamba2 <f>           Mamba2 learning rate (default: $LR_MAMBA2)
  --n-train-seqs <n>        Train sequences per epoch (default: $N_TRAIN_SEQS)
  --n-val-seqs <n>          Val sequences per epoch (default: $N_VAL_SEQS)
  --seed <n>                Random seed (default: $SEED)
  --cv-fold <k>             CV fold index (default: $CV_FOLD)
  --cv-n-folds <N>          Total CV folds (default: $CV_N_FOLDS)
  --cv-train-window <W>     Files per train window (default: $CV_TRAIN_WINDOW = auto)
  --watch                   Follow logs via argo watch
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --sha)               SHA="$2"; shift 2 ;;
    --branch)            BRANCH="$2"; shift 2 ;;
    --gpu-pool)          GPU_POOL="$2"; shift 2 ;;
    --epochs)            EPOCHS="$2"; shift 2 ;;
    --seq-len)           SEQ_LEN="$2"; shift 2 ;;
    --mamba2-state-dim)  MAMBA2_STATE_DIM="$2"; shift 2 ;;
    --lr-cfc)            LR_CFC="$2"; shift 2 ;;
    --lr-mamba2)         LR_MAMBA2="$2"; shift 2 ;;
    --n-train-seqs)      N_TRAIN_SEQS="$2"; shift 2 ;;
    --n-val-seqs)        N_VAL_SEQS="$2"; shift 2 ;;
    --seed)              SEED="$2"; shift 2 ;;
    --batch-size)        BATCH_SIZE="$2"; shift 2 ;;
    --auto-horizon-weights) AUTO_HORIZON_WEIGHTS=true; shift ;;
    --early-stop-metric)  EARLY_STOP_METRIC="$2"; shift 2 ;;
    --early-stop-patience) EARLY_STOP_PATIENCE="$2"; shift 2 ;;
    --cv-fold)           CV_FOLD="$2"; shift 2 ;;
    --cv-n-folds)        CV_N_FOLDS="$2"; shift 2 ;;
    --cv-train-window)   CV_TRAIN_WINDOW="$2"; shift 2 ;;
    --watch)             WATCH=true; shift ;;
    -h|--help)           usage; exit 0 ;;
    *) echo "Unknown option: $1"; usage; exit 1 ;;
  esac
done

case "$GPU_POOL" in
  ci-training-l40s)
    SM=89 ;;
  ci-training-h100)
    SM=90 ;;
  *)
    echo "Unknown gpu-pool: $GPU_POOL"
    exit 1 ;;
esac

# Resolve commit-sha=HEAD to an actual SHA locally before submission.
# The in-cluster check-cache pod uses this SHA to look up the binary
# cache without needing git inside the alpine pod.
if [ "$SHA" = "HEAD" ]; then
  echo "Resolving HEAD for branch $BRANCH..."
  if ! git rev-parse --git-dir >/dev/null 2>&1; then
    echo "ERROR: not in a git repo; cannot resolve HEAD"
    exit 1
  fi
  git fetch --quiet origin "$BRANCH" 2>/dev/null || true
  SHA=$(git rev-parse "origin/$BRANCH" 2>/dev/null \
        || git rev-parse "$BRANCH" 2>/dev/null \
        || git rev-parse HEAD)
  echo "  resolved: $SHA"
fi

echo "Submitting alpha-perception workflow..."
echo "  branch:           $BRANCH"
echo "  sha:              $SHA"
echo "  gpu-pool:         $GPU_POOL (sm_$SM)"
echo "  epochs:           $EPOCHS"
echo "  seq-len:          $SEQ_LEN"
echo "  mamba2-state-dim: $MAMBA2_STATE_DIM"
echo "  lr-cfc:           $LR_CFC"
echo "  lr-mamba2:        $LR_MAMBA2"
echo "  n-train-seqs:     $N_TRAIN_SEQS"
echo "  n-val-seqs:       $N_VAL_SEQS"
echo "  seed:             $SEED"

WATCH_FLAG=""
if [[ "$WATCH" == "true" ]]; then
  WATCH_FLAG="--watch"
fi

argo submit -n foxhunt --from=wftmpl/alpha-perception \
  -p commit-sha="$SHA" \
  -p git-branch="$BRANCH" \
  -p cuda-compute-cap="$SM" \
  -p gpu-pool="$GPU_POOL" \
  -p epochs="$EPOCHS" \
  -p seq-len="$SEQ_LEN" \
  -p mamba2-state-dim="$MAMBA2_STATE_DIM" \
  -p lr-cfc="$LR_CFC" \
  -p lr-mamba2="$LR_MAMBA2" \
  -p n-train-seqs="$N_TRAIN_SEQS" \
  -p n-val-seqs="$N_VAL_SEQS" \
  -p seed="$SEED" \
  -p batch-size="$BATCH_SIZE" \
  -p auto-horizon-weights="$AUTO_HORIZON_WEIGHTS" \
  -p early-stop-metric="$EARLY_STOP_METRIC" \
  -p early-stop-patience="$EARLY_STOP_PATIENCE" \
  -p cv-fold="$CV_FOLD" \
  -p cv-n-folds="$CV_N_FOLDS" \
  -p cv-train-window="$CV_TRAIN_WINDOW" \
  $WATCH_FLAG