#!/usr/bin/env bash
# Train a model via Argo Workflows.
#
# Usage:
#   ./scripts/argo-train.sh dqn                          # defaults: HEAD, H100, 50 epochs
#   ./scripts/argo-train.sh dqn --sha abc1234            # specific commit
#   ./scripts/argo-train.sh dqn --epochs 100 --trials 40 # override training params
#   ./scripts/argo-train.sh ppo --gpu-pool ci-training    # L40S instead of H100
#   ./scripts/argo-train.sh dqn --baseline                # skip hyperopt
#   ./scripts/argo-train.sh dqn --watch                  # follow logs
#
# Supported models:
#   RL:          dqn, ppo
#   Supervised:  tft, mamba2, tggn, tlob, liquid, kan, xlstm, diffusion
#
# Requires: argo CLI
set -euo pipefail

SHA="HEAD"
BRANCH="main"
TRIALS=""
EPOCHS=""
GPU_POOL=""
SYMBOL=""
WATCH=false
BASELINE=false
CAPITAL=""
SANITIZER="none"

usage() {
  cat <<EOF
Usage: $(basename "$0") <model> [OPTIONS]

Models:
  dqn, ppo                                          (RL)
  tft, mamba2, tggn, tlob, liquid, kan, xlstm       (supervised)

Options:
  --sha <commit>       Git commit SHA (default: HEAD)
  --branch <branch>    Git branch (default: main)
  --trials <n>         Hyperopt trials (default: 20)
  --epochs <n>         Training epochs (default: 50)
  --gpu-pool <pool>    GPU node pool (default: ci-training-h100)
  --symbol <sym>       Trading symbol (default: ES.FUT)
  --capital <n>        Initial capital (default: 35000)
  --baseline           Skip hyperopt (trials=0)
  --sanitizer <tool>   Run under compute-sanitizer (memcheck|racecheck|synccheck)
  --watch              Follow workflow logs
  -h, --help           Show this help
EOF
  exit 0
}

[[ $# -eq 0 ]] && { echo "Error: model argument required"; usage; }

MODEL="$1"; shift

case "$MODEL" in
  dqn|ppo|tft|mamba2|tggn|tlob|liquid|kan|xlstm|diffusion) ;;
  *) echo "Error: unknown model '$MODEL'"; usage ;;
esac

while [[ $# -gt 0 ]]; do
  case $1 in
    --sha)       SHA="$2"; shift 2 ;;
    --branch)    BRANCH="$2"; shift 2 ;;
    --trials)    TRIALS="$2"; shift 2 ;;
    --epochs)    EPOCHS="$2"; shift 2 ;;
    --gpu-pool)  GPU_POOL="$2"; shift 2 ;;
    --symbol)    SYMBOL="$2"; shift 2 ;;
    --capital)   CAPITAL="$2"; shift 2 ;;
    --baseline)  BASELINE=true; shift ;;
    --sanitizer) SANITIZER="$2"; shift 2 ;;
    --watch)     WATCH=true; shift ;;
    -h|--help)   usage ;;
    *)           echo "Unknown option: $1"; usage ;;
  esac
done

# Auto-derive cuda-compute-cap from GPU pool — cubins must match device sm_XX.
# Default pool is ci-training-h100 (sm_90). Override for other architectures:
#   ci-training-h100*  → sm_90 (Hopper)
#   ci-training-l40s   → sm_89 (Ada Lovelace)
case "${GPU_POOL:-ci-training-h100}" in
  *l40s*)     CUDA_COMPUTE_CAP="89" ;;
  *h100*|*)   CUDA_COMPUTE_CAP="90" ;;  # default Hopper
esac

CMD="argo submit -n foxhunt --from=wftmpl/train"
CMD="$CMD -p commit-sha=$SHA"
CMD="$CMD -p git-branch=$BRANCH"
CMD="$CMD -p model=$MODEL"
CMD="$CMD -p cuda-compute-cap=$CUDA_COMPUTE_CAP"

[[ -n "$TRIALS" ]]    && CMD="$CMD -p hyperopt-trials=$TRIALS"
[[ -n "$EPOCHS" ]]    && CMD="$CMD -p train-epochs=$EPOCHS"
[[ -n "$GPU_POOL" ]]  && CMD="$CMD -p gpu-pool=$GPU_POOL"
[[ -n "$SYMBOL" ]]    && CMD="$CMD -p symbol=$SYMBOL"
[[ -n "$CAPITAL" ]]   && CMD="$CMD -p initial-capital=$CAPITAL"
[[ "$SANITIZER" != "none" ]] && CMD="$CMD -p sanitizer=$SANITIZER"

$BASELINE && CMD="$CMD -p hyperopt-trials=0"
$WATCH    && CMD="$CMD --watch"

echo "Submitting $MODEL training workflow..."
echo "  sha:    $SHA"
echo "  branch: $BRANCH"
echo "  model:  $MODEL"
$BASELINE              && echo "  mode:   baseline (no hyperopt)"
[[ -n "$TRIALS" ]]     && echo "  trials: $TRIALS"
[[ -n "$EPOCHS" ]]     && echo "  epochs: $EPOCHS"
[[ -n "$GPU_POOL" ]]   && echo "  gpu:    $GPU_POOL"
echo "  sm:     $CUDA_COMPUTE_CAP"
echo ""
eval "$CMD"