#!/bin/bash
# Multi-Model Hyperopt RunPod Deployment Script
# Provides easy deployment interface for all 4 hyperopt models

set -e

# Activate virtual environment
source .venv/bin/activate

# Set PYTHONPATH to include custom runpod module
export PYTHONPATH=/home/jgrusewski/Work/foxhunt:$PYTHONPATH

# Default configuration
GPU_TYPE="${GPU_TYPE:-RTX A4000}"
IMAGE="${IMAGE:-jgrusewski/foxhunt:latest}"
TRIALS="${TRIALS:-50}"
EPOCHS="${EPOCHS:-50}"
TIMEOUT="${TIMEOUT:-120m}"
BATCH_SIZE_MAX="${BATCH_SIZE_MAX:-96}"
PARQUET_FILE="/runpod-volume/test_data/ES_FUT_180d.parquet"
BASE_DIR="/runpod-volume/ml_training"

# Show menu if no argument provided
if [ $# -eq 0 ]; then
    echo "======================================================================"
    echo "Multi-Model Hyperopt RunPod Deployment"
    echo "======================================================================"
    echo "Usage: $0 <model> [options]"
    echo ""
    echo "Models:"
    echo "  mamba2    - MAMBA-2 hyperparameter optimization"
    echo "  dqn       - DQN hyperparameter optimization"
    echo "  ppo       - PPO hyperparameter optimization"
    echo "  tft       - TFT hyperparameter optimization"
    echo ""
    echo "Environment Variables (optional):"
    echo "  GPU_TYPE         - GPU type (default: RTX A4000)"
    echo "  IMAGE            - Docker image (default: jgrusewski/foxhunt:latest)"
    echo "  TRIALS           - Number of trials (default: 50)"
    echo "  EPOCHS           - Epochs per trial (default: 50)"
    echo "  TIMEOUT          - Max monitoring time (default: 120m)"
    echo "  BATCH_SIZE_MAX   - Max batch size (default: 96)"
    echo ""
    echo "Examples:"
    echo "  $0 mamba2"
    echo "  GPU_TYPE='RTX 4090' TRIALS=100 $0 dqn"
    echo "  $0 tft"
    echo "======================================================================"
    exit 1
fi

MODEL=$1

# Build command based on model
case $MODEL in
    mamba2)
        COMMAND="hyperopt_mamba2_demo \
          --parquet-file ${PARQUET_FILE} \
          --base-dir ${BASE_DIR} \
          --trials ${TRIALS} \
          --epochs ${EPOCHS} \
          --batch-size-max ${BATCH_SIZE_MAX} \
          --early-stopping-patience 5"
        ;;
    dqn)
        COMMAND="hyperopt_dqn_demo \
          --parquet-file ${PARQUET_FILE} \
          --base-dir ${BASE_DIR} \
          --trials ${TRIALS} \
          --epochs ${EPOCHS} \
          --batch-size-max ${BATCH_SIZE_MAX} \
          --early-stopping-patience 5"
        ;;
    ppo)
        COMMAND="hyperopt_ppo_demo \
          --parquet-file ${PARQUET_FILE} \
          --base-dir ${BASE_DIR} \
          --trials ${TRIALS} \
          --epochs ${EPOCHS} \
          --batch-size-max ${BATCH_SIZE_MAX} \
          --early-stopping-patience 5"
        ;;
    tft)
        COMMAND="hyperopt_tft_demo \
          --parquet-file ${PARQUET_FILE} \
          --base-dir ${BASE_DIR} \
          --trials ${TRIALS} \
          --epochs ${EPOCHS} \
          --batch-size-max ${BATCH_SIZE_MAX} \
          --early-stopping-patience 5"
        ;;
    *)
        echo "ERROR: Unknown model '${MODEL}'"
        echo "Valid models: mamba2, dqn, ppo, tft"
        exit 1
        ;;
esac

echo "======================================================================"
echo "${MODEL^^} Hyperopt RunPod Deployment"
echo "======================================================================"
echo "GPU Type: ${GPU_TYPE}"
echo "Docker Image: ${IMAGE}"
echo "Trials: ${TRIALS}"
echo "Epochs per Trial: ${EPOCHS}"
echo "Batch Size Max: ${BATCH_SIZE_MAX}"
echo "Max Monitoring: ${TIMEOUT}"
echo "Command: ${COMMAND}"
echo "======================================================================"
echo ""

# Deploy pod with monitoring and auto-stop
python3 scripts/runpod_deploy.py \
  --gpu-type "${GPU_TYPE}" \
  --image "${IMAGE}" \
  --command "${COMMAND}" \
  --monitor \
  --auto-stop \
  --timeout "${TIMEOUT}" \
  --monitor-interval 15

echo ""
echo "======================================================================"
echo "Deployment Complete!"
echo "======================================================================"
echo "Results will be saved to: ${BASE_DIR}/"
echo "Check S3 bucket for outputs: s3://se3zdnb5o4/ml_training/"
echo "======================================================================"