#!/bin/bash
# Runpod 225-Feature Training Deployment Script
# GPU: 16GB VRAM (RTX 4000 Ada or A4000)
# Features: 225 (201 Wave C + 24 Wave D)
# Training Time: ~15-20 minutes
# Cost: ~$0.05-$0.10

set -euo pipefail

# Configuration
GPU_TYPE="RTX 4000 Ada"
VRAM_MIN=16
COST_CEILING="0.30"
TRAINING_TIMEOUT=1200
FEATURE_COUNT=225
CONTAINER_IMAGE="runpod/pytorch:2.0.1-py3.10-cuda11.8.0-devel"

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

echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}Runpod 225-Feature Training Deployment${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""

# Step 1: Verify prerequisites
echo -e "${YELLOW}Step 1/8: Verifying prerequisites...${NC}"

# Check API key
if ! runpodctl config 2>&1 | grep -q "apiKey"; then
    echo -e "${RED}❌ Runpod API key not configured!${NC}"
    echo -e "${YELLOW}Configure with: runpodctl config --apiKey YOUR_API_KEY${NC}"
    exit 1
fi
echo -e "${GREEN}✅ API key configured${NC}"

# Check training data
if [ ! -f "test_data/ES_FUT_180d.parquet" ]; then
    echo -e "${RED}❌ Training data not found: test_data/ES_FUT_180d.parquet${NC}"
    exit 1
fi
DATA_SIZE=$(du -h test_data/ES_FUT_180d.parquet | cut -f1)
echo -e "${GREEN}✅ Training data found: ${DATA_SIZE}${NC}"

# Verify 225 features
if ! grep -q "features: \[f64; 225\]" ml/src/features/unified.rs; then
    echo -e "${RED}❌ 225 features not configured in ml/src/features/unified.rs!${NC}"
    exit 1
fi
echo -e "${GREEN}✅ 225 features configured${NC}"

# Step 2: Build release binary
echo -e "\n${YELLOW}Step 2/8: Building release binary...${NC}"
echo "This may take 5-6 minutes..."

cargo build --release -p ml --example train_tft_parquet --features cuda 2>&1 | tee /tmp/build.log | grep -E "Compiling|Finished" || true

if [ ! -f "target/release/examples/train_tft_parquet" ]; then
    echo -e "${RED}❌ Build failed! Check /tmp/build.log${NC}"
    exit 1
fi
BINARY_SIZE=$(du -h target/release/examples/train_tft_parquet | cut -f1)
echo -e "${GREEN}✅ Binary built: ${BINARY_SIZE}${NC}"

# Step 3: List available GPUs
echo -e "\n${YELLOW}Step 3/8: Checking available GPUs...${NC}"
echo "Looking for 16GB+ GPUs under \$${COST_CEILING}/hr..."

runpodctl get gpus --filter "vram>=16" 2>&1 | grep -E "RTX 4000|A4000|RTX 4090|A5000" | head -5 || {
    echo -e "${YELLOW}⚠️  No pre-filtered results, listing all GPUs...${NC}"
    runpodctl get gpus 2>&1 | head -20
}

# Step 4: Create pod
echo -e "\n${YELLOW}Step 4/8: Creating Runpod pod...${NC}"
echo "GPU: ${GPU_TYPE}, VRAM: ${VRAM_MIN}GB, Max Cost: \$${COST_CEILING}/hr"

POD_CREATE_OUTPUT=$(runpodctl create pod \
    --name "foxhunt-tft-225-training-$(date +%s)" \
    --gpuType "RTX 4000 Ada" \
    --minVram ${VRAM_MIN} \
    --maxCost ${COST_CEILING} \
    --containerDiskSize 50 \
    --volumeSize 50 \
    --imageName "${CONTAINER_IMAGE}" \
    --env "FEATURE_COUNT=${FEATURE_COUNT}" \
    --env "CUDA_VISIBLE_DEVICES=0" 2>&1)

POD_ID=$(echo "$POD_CREATE_OUTPUT" | jq -r '.id' 2>/dev/null || echo "$POD_CREATE_OUTPUT" | grep -oP 'pod-[a-z0-9]+' | head -1)

if [ -z "$POD_ID" ]; then
    echo -e "${RED}❌ Failed to create pod!${NC}"
    echo "$POD_CREATE_OUTPUT"
    exit 1
fi

echo -e "${GREEN}✅ Pod created: ${POD_ID}${NC}"
echo "$POD_ID" > /tmp/runpod_pod_id.txt

# Step 5: Wait for pod ready
echo -e "\n${YELLOW}Step 5/8: Waiting for pod initialization...${NC}"
echo "This usually takes 30-60 seconds..."

WAIT_COUNT=0
MAX_WAIT=120
while [ $WAIT_COUNT -lt $MAX_WAIT ]; do
    POD_STATUS=$(runpodctl get pod "${POD_ID}" 2>&1 | grep -oP 'status: \K[A-Z]+' || echo "UNKNOWN")

    if [ "$POD_STATUS" = "RUNNING" ]; then
        echo -e "${GREEN}✅ Pod is running!${NC}"
        break
    fi

    echo -n "."
    sleep 5
    WAIT_COUNT=$((WAIT_COUNT + 5))
done

if [ $WAIT_COUNT -ge $MAX_WAIT ]; then
    echo -e "\n${RED}❌ Pod failed to start within ${MAX_WAIT} seconds${NC}"
    runpodctl remove pod "${POD_ID}"
    exit 1
fi

# Step 6: Upload training data and binary
echo -e "\n${YELLOW}Step 6/8: Uploading training data and binary...${NC}"

# Create directories on pod
runpodctl exec "${POD_ID}" -- mkdir -p /workspace/data /workspace/models

# Upload training data
echo "Uploading training data (${DATA_SIZE})..."
runpodctl send "${POD_ID}" test_data/ES_FUT_180d.parquet /workspace/data/ 2>&1 | grep -E "Success|Error" || echo "Upload in progress..."

# Upload binary
echo "Uploading training binary (${BINARY_SIZE})..."
runpodctl send "${POD_ID}" target/release/examples/train_tft_parquet /workspace/ 2>&1 | grep -E "Success|Error" || echo "Upload in progress..."

# Make binary executable
runpodctl exec "${POD_ID}" -- chmod +x /workspace/train_tft_parquet

echo -e "${GREEN}✅ Files uploaded${NC}"

# Step 7: Execute training
echo -e "\n${YELLOW}Step 7/8: Starting training...${NC}"
echo "Training configuration:"
echo "  • Features: 225 (201 Wave C + 24 Wave D)"
echo "  • Epochs: 50"
echo "  • Batch size: 32"
echo "  • Data: ES.FUT 180 days"
echo "  • Estimated time: 15-20 minutes"
echo ""

TRAINING_START=$(date +%s)

# Execute training with comprehensive logging
runpodctl exec "${POD_ID}" -- /workspace/train_tft_parquet \
    --parquet-file /workspace/data/ES_FUT_180d.parquet \
    --epochs 50 \
    --batch-size 32 \
    --learning-rate 0.001 \
    --lookback-window 60 \
    --forecast-horizon 10 \
    --hidden-dim 256 \
    --num-attention-heads 8 \
    2>&1 | tee /tmp/training_output.log

TRAINING_END=$(date +%s)
TRAINING_DURATION=$((TRAINING_END - TRAINING_START))
TRAINING_MINUTES=$((TRAINING_DURATION / 60))
TRAINING_SECONDS=$((TRAINING_DURATION % 60))

echo -e "${GREEN}✅ Training completed in ${TRAINING_MINUTES}m ${TRAINING_SECONDS}s${NC}"

# Step 8: Download trained model
echo -e "\n${YELLOW}Step 8/8: Downloading trained model...${NC}"

mkdir -p models/runpod_trained

# List models on pod
echo "Available models:"
runpodctl exec "${POD_ID}" -- ls -lh /workspace/*.safetensors 2>/dev/null || {
    echo -e "${YELLOW}⚠️  No .safetensors found, checking models directory...${NC}"
    runpodctl exec "${POD_ID}" -- ls -lh /workspace/models/ 2>/dev/null || echo "No models found"
}

# Download all model files
runpodctl receive "${POD_ID}" "/workspace/*.safetensors" models/runpod_trained/ 2>&1 || {
    echo -e "${YELLOW}⚠️  Trying alternative model location...${NC}"
    runpodctl receive "${POD_ID}" "/workspace/models/*.safetensors" models/runpod_trained/ 2>&1
}

# Also download training logs if available
runpodctl receive "${POD_ID}" "/workspace/*.log" models/runpod_trained/ 2>&1 || echo "No logs to download"

if ls models/runpod_trained/*.safetensors 1>/dev/null 2>&1; then
    MODEL_SIZE=$(du -h models/runpod_trained/*.safetensors | head -1 | cut -f1)
    echo -e "${GREEN}✅ Model downloaded: ${MODEL_SIZE}${NC}"
else
    echo -e "${RED}❌ No model files downloaded!${NC}"
fi

# Step 9: Calculate cost
echo -e "\n${YELLOW}Calculating training cost...${NC}"

POD_INFO=$(runpodctl get pod "${POD_ID}" 2>&1)
GPU_COST=$(echo "$POD_INFO" | grep -oP 'costPerHr: \K[0-9.]+' || echo "0.25")
COST_HOURS=$(echo "scale=4; ${TRAINING_DURATION} / 3600" | bc)
TOTAL_COST=$(echo "scale=4; ${GPU_COST} * ${COST_HOURS}" | bc)

echo -e "${BLUE}Cost Breakdown:${NC}"
echo "  • GPU Rate: \$${GPU_COST}/hr"
echo "  • Training Time: ${TRAINING_MINUTES}m ${TRAINING_SECONDS}s"
echo "  • Total Cost: \$${TOTAL_COST}"

# Step 10: Terminate pod
echo -e "\n${YELLOW}Terminating pod...${NC}"
runpodctl remove pod "${POD_ID}" 2>&1 | grep -E "Success|Removed" || echo "Pod termination in progress..."
echo -e "${GREEN}✅ Pod terminated (no ongoing charges)${NC}"

# Summary
echo ""
echo -e "${BLUE}========================================${NC}"
echo -e "${GREEN}Training Deployment Complete!${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""
echo "Results:"
echo "  • Pod ID: ${POD_ID}"
echo "  • Training Time: ${TRAINING_MINUTES}m ${TRAINING_SECONDS}s"
echo "  • Total Cost: \$${TOTAL_COST}"
echo "  • Model Location: models/runpod_trained/"
echo "  • Training Log: /tmp/training_output.log"
echo ""
echo "Next Steps:"
echo "  1. Run backtesting: ./scripts/backtest_runpod_225.sh"
echo "  2. Validate Wave D targets (Sharpe ≥2.0, Win Rate ≥60%, Drawdown ≤15%)"
echo "  3. Train remaining models (DQN, PPO, MAMBA-2)"
echo ""