#!/usr/bin/env bash
set -euo pipefail

# GPU Training Launcher -- detects GPU and selects local or cloud path
#
# Usage:
#   ./scripts/train_launcher.sh --model dqn [--data path/to/data] [--cloud] [--epochs N]
#   ./scripts/train_launcher.sh --model ppo --data data/databento/ES.FUT/ohlcv-1m/ --epochs 50
#   ./scripts/train_launcher.sh --model tft --cloud
#
# Models: dqn, ppo, tft, mamba2, cfc
# Flags:
#   --cloud    Force cloud training even if local GPU is available
#   --epochs   Number of training epochs (default: 20)
#   --data     Path to training data (parquet file or directory)
#   --lr       Learning rate (default: 0.0001)
#   --output   Output directory for trained models (default: ml/trained_models)

MODEL=""
DATA_PATH=""
CLOUD=false
EPOCHS=20
LEARNING_RATE=0.0001
OUTPUT_DIR="ml/trained_models"
EXTRA_ARGS=()

while [[ $# -gt 0 ]]; do
    case "$1" in
        --model)   MODEL="$2"; shift 2 ;;
        --data)    DATA_PATH="$2"; shift 2 ;;
        --cloud)   CLOUD=true; shift ;;
        --epochs)  EPOCHS="$2"; shift 2 ;;
        --lr)      LEARNING_RATE="$2"; shift 2 ;;
        --output)  OUTPUT_DIR="$2"; shift 2 ;;
        --help|-h) usage; exit 0 ;;
        *)         EXTRA_ARGS+=("$1"); shift ;;
    esac
done

usage() {
    cat <<'USAGE'
GPU Training Launcher -- detects GPU and routes to local or cloud

Usage:
  ./scripts/train_launcher.sh --model <model> [OPTIONS]

Models: dqn, ppo, tft, mamba2, cfc

Options:
  --model <name>   Model to train (required)
  --data <path>    Training data parquet file or directory
  --epochs <N>     Number of training epochs (default: 20)
  --lr <rate>      Learning rate (default: 0.0001)
  --output <dir>   Output directory (default: ml/trained_models)
  --cloud          Force cloud training
  --help           Show this help

Examples:
  ./scripts/train_launcher.sh --model dqn --epochs 100
  ./scripts/train_launcher.sh --model ppo --data test_data/ES_FUT_180d.parquet --epochs 50
  ./scripts/train_launcher.sh --model tft --cloud
USAGE
}

if [[ -z "$MODEL" ]]; then
    echo "Error: --model is required (dqn, ppo, tft, mamba2, cfc)"
    echo "Run with --help for usage information"
    exit 1
fi

VALID_MODELS="dqn ppo tft mamba2 cfc"
if ! echo "$VALID_MODELS" | grep -qw "$MODEL"; then
    echo "Error: Invalid model '$MODEL'. Valid: $VALID_MODELS"
    exit 1
fi

# Map model names to cargo example targets
# These are the actual ml/examples/train_*.rs binaries
declare -A EXAMPLE_MAP
EXAMPLE_MAP[dqn]="train_dqn"
EXAMPLE_MAP[ppo]="train_ppo_parquet"
EXAMPLE_MAP[tft]="train_tft_parquet"
EXAMPLE_MAP[mamba2]="train_mamba2_parquet"
EXAMPLE_MAP[cfc]="train_liquid_dbn"

EXAMPLE_NAME="${EXAMPLE_MAP[$MODEL]}"

# ---------------------------------------------------------------------------
# GPU Detection
# ---------------------------------------------------------------------------
GPU_DETECTED=false
VRAM_MB=0
GPU_NAME="none"

if command -v nvidia-smi &>/dev/null; then
    VRAM_MB=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 || echo "0")
    if [[ "$VRAM_MB" -gt 0 ]]; then
        GPU_DETECTED=true
        GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1 || echo "unknown")
        echo "Local GPU detected: $GPU_NAME (${VRAM_MB}MB VRAM)"
    fi
else
    echo "No NVIDIA GPU detected (nvidia-smi not found)"
fi

# ---------------------------------------------------------------------------
# Batch size selection based on VRAM and model
# ---------------------------------------------------------------------------
# These values are empirically tested; PPO is constrained on 4GB (max 230).
# See scripts/measure_vram.sh for VRAM profiling methodology.
declare -A BS_TIER_LOW    # < 4 GB
declare -A BS_TIER_MED    # 4-8 GB  (e.g. RTX 3050 Ti 4GB)
declare -A BS_TIER_HIGH   # 8-16 GB (e.g. RTX 3070 8GB)
declare -A BS_TIER_MAX    # 16+ GB  (e.g. L4 24GB)

BS_TIER_LOW[dqn]=32;   BS_TIER_MED[dqn]=128;  BS_TIER_HIGH[dqn]=256;  BS_TIER_MAX[dqn]=512
BS_TIER_LOW[ppo]=32;   BS_TIER_MED[ppo]=230;  BS_TIER_HIGH[ppo]=512;  BS_TIER_MAX[ppo]=1024
BS_TIER_LOW[tft]=8;    BS_TIER_MED[tft]=32;   BS_TIER_HIGH[tft]=64;   BS_TIER_MAX[tft]=256
BS_TIER_LOW[mamba2]=8; BS_TIER_MED[mamba2]=64; BS_TIER_HIGH[mamba2]=128; BS_TIER_MAX[mamba2]=512
BS_TIER_LOW[cfc]=32;   BS_TIER_MED[cfc]=128;  BS_TIER_HIGH[cfc]=256;  BS_TIER_MAX[cfc]=512

BATCH_SIZE=64
if [[ "$GPU_DETECTED" == true ]]; then
    if [[ "$VRAM_MB" -lt 4096 ]]; then
        BATCH_SIZE="${BS_TIER_LOW[$MODEL]}"
    elif [[ "$VRAM_MB" -lt 8192 ]]; then
        BATCH_SIZE="${BS_TIER_MED[$MODEL]}"
    elif [[ "$VRAM_MB" -lt 16384 ]]; then
        BATCH_SIZE="${BS_TIER_HIGH[$MODEL]}"
    else
        BATCH_SIZE="${BS_TIER_MAX[$MODEL]}"
    fi
fi

echo ""
echo "Configuration:"
echo "  Model:         $MODEL (example: $EXAMPLE_NAME)"
echo "  Epochs:        $EPOCHS"
echo "  Batch size:    $BATCH_SIZE"
echo "  Learning rate: $LEARNING_RATE"
echo "  Output:        $OUTPUT_DIR"
if [[ -n "$DATA_PATH" ]]; then
    echo "  Data:          $DATA_PATH"
fi

# ---------------------------------------------------------------------------
# Build the cargo command arguments
# ---------------------------------------------------------------------------
build_cargo_args() {
    local args=()
    args+=(--epochs "$EPOCHS")
    args+=(--batch-size "$BATCH_SIZE")
    args+=(--learning-rate "$LEARNING_RATE")
    args+=(--output-dir "$OUTPUT_DIR")

    if [[ -n "$DATA_PATH" ]]; then
        args+=(--parquet-file "$DATA_PATH")
    fi

    # Pass through any extra arguments
    if [[ ${#EXTRA_ARGS[@]} -gt 0 ]]; then
        args+=("${EXTRA_ARGS[@]}")
    fi

    echo "${args[@]}"
}

CARGO_EXTRA=$(build_cargo_args)

# ---------------------------------------------------------------------------
# Route to local or cloud
# ---------------------------------------------------------------------------
CLOUD_HOST="${FOXHUNT_GPU_HOST:-gpu.fxhnt.ai}"
CLOUD_DIR="${FOXHUNT_CLOUD_DIR:-/opt/foxhunt}"

need_cloud() {
    # Cloud is needed when: forced, no GPU, or insufficient VRAM (< 4GB)
    [[ "$CLOUD" == true ]] || [[ "$GPU_DETECTED" == false ]] || [[ "$VRAM_MB" -lt 4096 ]]
}

if need_cloud; then
    echo ""
    echo "==> Routing to cloud training at $CLOUD_HOST"
    echo ""

    # Build the remote command
    REMOTE_CMD="cd $CLOUD_DIR && git pull --ff-only && SQLX_OFFLINE=true cargo run --release -p ml --example $EXAMPLE_NAME --features cuda -- $CARGO_EXTRA"

    echo "To run on cloud GPU:"
    echo ""
    echo "  ssh training@$CLOUD_HOST '$REMOTE_CMD'"
    echo ""
    echo "To sync model artifacts back:"
    echo ""
    echo "  rsync -avz training@$CLOUD_HOST:$CLOUD_DIR/checkpoints/ ./checkpoints/"
    echo "  rsync -avz training@$CLOUD_HOST:$CLOUD_DIR/ml/trained_models/ ./ml/trained_models/"
    echo ""

    # If we have SSH access, offer to run it directly
    if [[ "$CLOUD" == true ]]; then
        read -rp "Execute on cloud now? [y/N] " CONFIRM
        if [[ "$CONFIRM" =~ ^[Yy]$ ]]; then
            echo "Connecting to $CLOUD_HOST..."
            # shellcheck disable=SC2029
            ssh "training@$CLOUD_HOST" "$REMOTE_CMD"
            echo ""
            echo "Syncing model artifacts..."
            rsync -avz "training@$CLOUD_HOST:$CLOUD_DIR/$OUTPUT_DIR/" "./$OUTPUT_DIR/"
            rsync -avz "training@$CLOUD_HOST:$CLOUD_DIR/checkpoints/" "./checkpoints/"
            echo "Done."
        else
            echo "Skipped. Run the SSH command above manually."
        fi
    fi
else
    echo ""
    echo "==> Running local training on $GPU_NAME"
    echo ""

    mkdir -p "$OUTPUT_DIR"

    export SQLX_OFFLINE=true

    # shellcheck disable=SC2086
    exec cargo run --release -p ml --example "$EXAMPLE_NAME" --features cuda -- $CARGO_EXTRA
fi