#!/usr/bin/env bash
# ==============================================================================
# ⚡ ULTRON Systems // Turnkey Colab CLI Deployment Runner
# Automates environment setup, Google Drive mounting, dependencies, and execution
# ==============================================================================

set -e

echo "=================================================================="
echo "⚡ ULTRON // CONVOLUTIONAL TRANSFORMER COLAB CLI RUNNER"
echo "=================================================================="

# 1. Detect GPU telemetry
if command -v nvidia-smi &> /dev/null; then
    echo "⚡ [ULTRON] NVIDIA GPU Detected:"
    nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
else
    echo "⚠️ [ULTRON] No CUDA GPU detected! Running on CPU mode."
fi

# 2. Install optimized dependencies (pyarrow, fastparquet, huggingface_hub)
echo "⚡ [ULTRON] Verifying and installing essential runtime dependencies..."
pip install -q --upgrade pip
pip install -q pyarrow fastparquet huggingface_hub pandas numpy torch torchvision torchaudio

# Authenticate Hugging Face Hub if token exists
if [ -n "$HF_TOKEN" ]; then
    echo "⚡ [ULTRON] Authenticating Hugging Face Hub with HF_TOKEN..."
    python3 -c "import os, huggingface_hub; huggingface_hub.login(token=os.environ.get('HF_TOKEN'), add_to_git_credential=True)" 2>/dev/null || true
fi

# 3. Mount Google Drive for persistent checkpoints (if running inside Google Colab)
python3 -c "
try:
    from google.colab import drive
    print('⚡ [ULTRON] Mounting Google Drive to /content/drive ...')
    drive.mount('/content/drive', force_remount=False)
    print('✅ [ULTRON] Google Drive mounted successfully.')
except ImportError:
    print('⚡ [ULTRON] Non-Colab environment detected. Skipping Drive mount.')
"

# 4. Create persistent checkpoint directories
mkdir -p /content/drive/MyDrive/ohlcv_checkpoints 2>/dev/null || mkdir -p ./checkpoints
mkdir -p /tmp/ohlcv_cache

# 5. Prevent Colab Disconnect / Idle Keepalive (background task)
python3 -c "
import threading, time
def keepalive():
    while True:
        time.sleep(300)
        # Touch a keepalive heartbeat
        with open('/tmp/ultron_heartbeat.log', 'a') as f:
            f.write(f'Keepalive heartbeat at {time.ctime()}\n')
threading.Thread(target=keepalive, daemon=True).start()
" &

# 6. Launch Training Loop
echo "⚡ [ULTRON] Launching autonomous training loop..."
python3 train.py \
    --start_year 2010 \
    --end_year 2025 \
    --batch_size 64 \
    --grad_accum 4 \
    --lr 0.0003 \
    --epochs 5