#!/bin/bash
# DQN Retrain Deployment to Runpod
# Deploys a Runpod pod to retrain DQN with fixed reward function and monitoring

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Get script directory
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
cd "$SCRIPT_DIR"

echo -e "${GREEN}====================================================================${NC}"
echo -e "${GREEN}DQN Retrain Deployment Script${NC}"
echo -e "${GREEN}====================================================================${NC}"

# 1. Check prerequisites
echo -e "\n${YELLOW}Step 1: Checking prerequisites...${NC}"

# Check if .venv is activated
if [[ -z "$VIRTUAL_ENV" ]]; then
    echo -e "${YELLOW}Activating virtual environment...${NC}"
    source .venv/bin/activate
fi

# Set PYTHONPATH for runpod module
export PYTHONPATH="${SCRIPT_DIR}:${PYTHONPATH}"

# Verify runpod module is available
if ! python3 -c "import runpod" 2>/dev/null; then
    echo -e "${RED}ERROR: runpod module not found${NC}"
    echo "Install dependencies: pip install -r runpod/requirements.txt"
    exit 1
fi
echo -e "${GREEN}✓ Virtual environment and runpod module OK${NC}"

# Check if code compiles
echo -e "\n${YELLOW}Step 2: Verifying DQN training code compiles...${NC}"
echo "(This will take a moment...)"
if ! cargo build -p ml --example train_dqn --release 2>&1 | tail -5; then
    echo -e "${RED}ERROR: DQN training code failed to compile${NC}"
    exit 1
fi
echo -e "${GREEN}✓ DQN training code compiles successfully${NC}"

# 2. Define training command for Runpod
# IMPORTANT:
# - Docker image has train_dqn binary in /usr/local/bin/
# - train_dqn supports parquet via --parquet-file argument
# - Path /runpod-volume/ is the volume mount point
# - Data file: /runpod-volume/test_data/ES_FUT_180d.parquet
# - We override the Docker CMD to run train_dqn with custom args
TRAINING_COMMAND="train_dqn --parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet --epochs 100 --min-epochs-before-stopping 50 --learning-rate 0.0001 --batch-size 32 --gamma 0.9626 --epsilon-start 0.3 --epsilon-end 0.05 --epsilon-decay 0.995 --buffer-size 104346 --min-replay-size 500 --checkpoint-frequency 10 --output-dir /runpod-volume/ml_training/dqn_fixed_reward --checkpoint-dir /runpod-volume/ml_training/dqn_fixed_reward/checkpoints --verbose"

echo -e "\n${YELLOW}Step 3: Deployment Configuration${NC}"
echo "  GPU Type: RTX A4000 (16GB VRAM, \$0.25/hr)"
echo "  Docker Image: jgrusewski/foxhunt:latest"
echo "  Training Command: $TRAINING_COMMAND"
echo "  Expected Duration: ~1-2 hours"
echo "  Expected Cost: ~\$0.25-\$0.50"
echo ""
echo "Monitoring features:"
echo "  - Real-time log streaming from S3"
echo "  - Automatic validation of:"
echo "    * Reward variance > 0.1"
echo "    * Action diversity ~30-35% each"
echo "    * Q-value balance across BUY/SELL/HOLD"
echo ""

# 3. Ask for confirmation
read -p "Deploy DQN retrain to Runpod? (y/n): " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo -e "${YELLOW}Deployment cancelled.${NC}"
    exit 0
fi

# 4. Deploy pod using runpod_deploy.py
echo -e "\n${YELLOW}Step 4: Deploying Runpod pod...${NC}"

python3 scripts/runpod_deploy.py \
    --gpu-type "RTX A4000" \
    --image "jgrusewski/foxhunt:latest" \
    --command "$TRAINING_COMMAND" \
    --container-disk 50 \
    --monitor \
    --timeout 3h

# Script will automatically:
# - Find available RTX A4000 in EUR-IS-1
# - Deploy pod with volume mounted at /runpod-volume/
# - Stream training logs in real-time
# - Show reward/action/Q-value metrics

echo -e "\n${GREEN}====================================================================${NC}"
echo -e "${GREEN}Deployment completed!${NC}"
echo -e "${GREEN}====================================================================${NC}"
echo -e "\nNext steps:"
echo "  1. Monitor logs above for:"
echo "     - Reward std > 0.1 (healthy variance)"
echo "     - Action distribution: BUY ~30-35%, SELL ~30-35%, HOLD ~30-35%"
echo "     - Q-value balance (BUY/SELL/HOLD similar magnitudes)"
echo "  2. Check S3 for saved checkpoints:"
echo "     aws s3 ls s3://se3zdnb5o4/ml_training/dqn_fixed_reward/ --profile runpod --recursive"
echo "  3. Pod will auto-terminate after 3h timeout or manual termination:"
echo "     curl -X POST -H \"Authorization: Bearer \$RUNPOD_API_KEY\" \\"
echo "       https://rest.runpod.io/v1/pods/<pod_id>/terminate"
echo ""