#!/bin/bash
# MAMBA-2 Hyperopt Monitoring Script
# Pod ID: w4srx0tgm5hfgu
# Created: 2025-10-28

set -e

POD_ID="w4srx0tgm5hfgu"
S3_BUCKET="s3://se3zdnb5o4"
S3_ENDPOINT="https://s3api-eur-is-1.runpod.io"
AWS_PROFILE="runpod"

echo "========================================"
echo "MAMBA-2 Hyperopt Monitor"
echo "========================================"
echo "Pod ID: $POD_ID"
echo "GPU: RTX A4000 (16GB VRAM)"
echo "Cost: \$0.25/hr"
echo "Expected Duration: ~60 minutes"
echo "========================================"
echo ""

# Function to check S3 results
check_s3_results() {
    echo "📦 Checking S3 for results..."
    aws s3 ls "$S3_BUCKET/models/" \
        --profile "$AWS_PROFILE" \
        --endpoint-url "$S3_ENDPOINT" \
        --recursive \
        --human-readable \
        | grep "mamba2_hyperopt" || echo "  (No results yet)"
    echo ""
}

# Function to show monitoring commands
show_commands() {
    echo "📝 Monitoring Commands:"
    echo ""
    echo "1. SSH into pod:"
    echo "   ssh root@$POD_ID.ssh.runpod.io"
    echo ""
    echo "2. View logs (inside pod):"
    echo "   docker logs -f \$(docker ps -q)"
    echo ""
    echo "3. Monitor GPU (inside pod):"
    echo "   watch -n 1 nvidia-smi"
    echo ""
    echo "4. Check process (inside pod):"
    echo "   ps aux | grep hyperopt_mamba2_demo"
    echo ""
    echo "5. Check outputs (inside pod):"
    echo "   ls -lh /runpod-volume/models/"
    echo ""
    echo "6. Runpod Console:"
    echo "   https://www.runpod.io/console/pods"
    echo ""
    echo "7. Jupyter Notebook:"
    echo "   https://$POD_ID-8888.proxy.runpod.net"
    echo ""
}

# Function to download results
download_results() {
    echo "📥 Downloading results from S3..."
    RESULTS_DIR="./runpod_hyperopt_results_$(date +%Y%m%d_%H%M%S)"
    mkdir -p "$RESULTS_DIR"

    aws s3 sync "$S3_BUCKET/models/" "$RESULTS_DIR/" \
        --profile "$AWS_PROFILE" \
        --endpoint-url "$S3_ENDPOINT" \
        --exclude "*" \
        --include "mamba2_hyperopt*" \
        --no-progress

    echo "  ✅ Results downloaded to: $RESULTS_DIR"
    ls -lh "$RESULTS_DIR/"
    echo ""
}

# Function to estimate completion time
estimate_completion() {
    echo "⏱️  Estimated Timeline:"
    echo "  - Initialization: ~2-3 minutes"
    echo "  - Trial 1: ~5 minutes"
    echo "  - Trial 10: ~20 minutes"
    echo "  - Trial 20: ~40 minutes"
    echo "  - Trial 30 (Complete): ~60 minutes"
    echo "  - Auto-Termination: ~61 minutes"
    echo ""
    echo "💰 Cost Estimate: \$0.19-\$0.38 (45-90 minutes)"
    echo ""
}

# Main menu
case "${1:-help}" in
    status|s)
        check_s3_results
        estimate_completion
        ;;
    commands|c)
        show_commands
        ;;
    download|d)
        download_results
        ;;
    watch|w)
        echo "🔄 Watching for results (checks every 30 seconds)..."
        echo "   Press Ctrl+C to stop"
        echo ""
        while true; do
            check_s3_results
            sleep 30
        done
        ;;
    help|h|*)
        echo "Usage: $0 [command]"
        echo ""
        echo "Commands:"
        echo "  status (s)     - Check S3 for results and show timeline"
        echo "  commands (c)   - Show monitoring commands"
        echo "  download (d)   - Download results from S3"
        echo "  watch (w)      - Watch for results (checks every 30s)"
        echo "  help (h)       - Show this help"
        echo ""
        echo "Examples:"
        echo "  $0 status      # Check current status"
        echo "  $0 watch       # Monitor in real-time"
        echo "  $0 download    # Download completed results"
        echo ""
        ;;
esac