#!/bin/bash
# MAMBA-2 Fixed Binary Deployment Monitor
# Pod ID: 8e6o2r2snavgzf
# Expected completion: 2025-10-27 10:21 UTC (~81 minutes from 09:00)

set -euo pipefail

POD_ID="8e6o2r2snavgzf"
CHECKPOINT_DIR="/runpod-volume/models/mamba2_FIXED_sgd_bs512_lr5e4_shuffle_50ep"

echo "=================================="
echo "MAMBA-2 FIXED BINARY MONITOR"
echo "=================================="
echo "Pod ID: $POD_ID"
echo "GPU: RTX 4090 (24GB VRAM)"
echo "Cost: \$0.59/hr"
echo "Expected runtime: ~81 minutes"
echo "=================================="
echo ""

# Load RunPod credentials
if [ ! -f ".env.runpod" ]; then
    echo "ERROR: .env.runpod not found"
    exit 1
fi

source .env.runpod

if [ -z "$RUNPOD_API_KEY" ]; then
    echo "ERROR: RUNPOD_API_KEY not set"
    exit 1
fi

# Function to check pod status
check_status() {
    echo "Checking pod status..."
    curl -s -H "Authorization: Bearer $RUNPOD_API_KEY" \
        "https://rest.runpod.io/v1/pods/$POD_ID" | python3 -m json.tool
    echo ""
}

# Function to download results
download_results() {
    echo "Downloading results from S3..."
    aws s3 sync "s3://se3zdnb5o4/models/mamba2_FIXED_sgd_bs512_lr5e4_shuffle_50ep" \
        "./local_models/mamba2_FIXED" \
        --profile runpod \
        --endpoint-url https://s3api-eur-is-1.runpod.io
    echo ""
    echo "Results downloaded to ./local_models/mamba2_FIXED/"
    ls -lh "./local_models/mamba2_FIXED/"
}

# Function to verify training success
verify_training() {
    echo "Verifying training success..."

    if [ ! -d "./local_models/mamba2_FIXED" ]; then
        echo "ERROR: Results not downloaded yet. Run with 'download' first."
        return 1
    fi

    # Check for model checkpoint
    if [ -f "./local_models/mamba2_FIXED/mamba2_model_epoch_50.safetensors" ]; then
        echo "✅ Model checkpoint found"
        ls -lh "./local_models/mamba2_FIXED/mamba2_model_epoch_50.safetensors"
    else
        echo "❌ Model checkpoint NOT found"
    fi

    # Check for metrics
    if [ -f "./local_models/mamba2_FIXED/training_metrics.json" ]; then
        echo "✅ Training metrics found"
        cat "./local_models/mamba2_FIXED/training_metrics.json"
    else
        echo "❌ Training metrics NOT found"
    fi

    # Check for loss history
    if [ -f "./local_models/mamba2_FIXED/loss_history.csv" ]; then
        echo "✅ Loss history found"
        echo "Last 10 epochs:"
        tail -n 10 "./local_models/mamba2_FIXED/loss_history.csv"
    else
        echo "❌ Loss history NOT found"
    fi

    # Check training log for key indicators
    if [ -f "./local_models/mamba2_FIXED/training.log" ]; then
        echo "✅ Training log found"
        echo ""
        echo "Checking for success indicators..."

        # Check optimizer
        if grep -q "Optimizer: SGD" "./local_models/mamba2_FIXED/training.log"; then
            echo "✅ SGD optimizer confirmed (not Adam)"
        else
            echo "❌ SGD optimizer NOT found (check for Adam)"
        fi

        # Check for zero gradients
        if grep -q "grad: 0.0000" "./local_models/mamba2_FIXED/training.log"; then
            echo "❌ Zero gradients detected (P0 fix failed)"
        else
            echo "✅ No zero gradients detected"
        fi

        # Check for E11 spike
        if grep -q "E11" "./local_models/mamba2_FIXED/training.log" || \
           grep -q "1e+11" "./local_models/mamba2_FIXED/training.log"; then
            echo "❌ E11 spike detected (numerical instability)"
        else
            echo "✅ No E11 spike detected"
        fi

    else
        echo "❌ Training log NOT found"
    fi
}

# Main menu
case "${1:-status}" in
    status)
        check_status
        ;;
    download)
        download_results
        ;;
    verify)
        verify_training
        ;;
    ssh)
        echo "SSH to pod $POD_ID..."
        echo "ssh root@$POD_ID.ssh.runpod.io"
        echo ""
        echo "Once connected, check training status:"
        echo "  cd $CHECKPOINT_DIR"
        echo "  tail -f training.log"
        ;;
    jupyter)
        echo "Jupyter URL: https://$POD_ID-8888.proxy.runpod.net"
        echo ""
        echo "Navigate to: $CHECKPOINT_DIR/training.log"
        ;;
    all)
        check_status
        echo ""
        echo "=================================="
        read -p "Download results? (y/n) " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            download_results
            echo ""
            verify_training
        fi
        ;;
    *)
        echo "Usage: $0 {status|download|verify|ssh|jupyter|all}"
        echo ""
        echo "Commands:"
        echo "  status   - Check pod status via API"
        echo "  download - Download results from S3"
        echo "  verify   - Verify training success (requires download first)"
        echo "  ssh      - Show SSH command"
        echo "  jupyter  - Show Jupyter URL"
        echo "  all      - Status + download + verify (interactive)"
        exit 1
        ;;
esac