#!/bin/bash
# Anti-UAV PyTorch Complete Environment Setup Script
# Creates new virtual environment with all dependencies

set -e  # Exit on any error

echo "🚀 Anti-UAV PyTorch Complete Environment Setup"
echo "=============================================="

# Check if conda is available
if ! command -v conda &> /dev/null; then
    echo "❌ Conda not found. Please install Anaconda or Miniconda first."
    exit 1
fi

# Environment name
ENV_NAME="anti_uav_complete"
PYTHON_VERSION="3.10"

echo "📋 Setup Configuration:"
echo "  Environment: $ENV_NAME"
echo "  Python: $PYTHON_VERSION"
echo "  CUDA: 12.1 (compatible with PyTorch 2.3.0)"

# Remove existing environment if it exists
echo "🧹 Cleaning up existing environment..."
conda remove -n $ENV_NAME --all -y 2>/dev/null || true

# Create new environment
echo "🔨 Creating new conda environment..."
conda create -n $ENV_NAME python=$PYTHON_VERSION -y

# Activate environment
echo "⚡ Activating environment: $ENV_NAME"
source $(conda info --base)/etc/profile.d/conda.sh
conda activate $ENV_NAME

# Verify Python version
echo "🐍 Python version:"
python --version

# Install PyTorch with CUDA 12.1 (compatible version)
echo "🔥 Installing PyTorch 2.3.0 + CUDA 12.1..."
pip install torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 --index-url https://download.pytorch.org/whl/cu121

# Install Ultralytics YOLOv8
echo "🎯 Installing YOLOv8 (Ultralytics)..."
pip install ultralytics==8.2.0

# Install core dependencies
echo "📦 Installing core dependencies..."
pip install \
    numpy \
    opencv-python \
    Pillow \
    matplotlib \
    seaborn \
    pandas \
    tqdm \
    pyyaml \
    scipy

# Install hyperparameter optimization
echo "⚙️ Installing HPO dependencies..."
pip install \
    optuna \
    ray[tune] \
    bayesian-optimization \
    hyperopt \
    scikit-optimize

# Install model export and optimization
echo "🔄 Installing model export dependencies..."
pip install \
    onnx \
    onnxruntime \
    onnxruntime-gpu

# Install API and web service dependencies
echo "🌐 Installing API service dependencies..."
pip install \
    fastapi \
    uvicorn[standard] \
    websockets \
    python-multipart \
    redis \
    celery

# Install advanced ML dependencies
echo "🧠 Installing advanced ML dependencies..."
pip install \
    transformers \
    timm \
    einops \
    albumentations \
    imgaug

# Install monitoring and logging
echo "📊 Installing monitoring dependencies..."
pip install \
    wandb \
    tensorboard \
    psutil \
    gpustat

# Install evaluation and visualization
echo "📈 Installing evaluation dependencies..."
pip install \
    scikit-learn \
    plotly \
    jupyter \
    ipywidgets

# Install development and testing tools
echo "🛠️ Installing development tools..."
pip install \
    pytest \
    black \
    flake8 \
    isort \
    pre-commit

# Try to install TensorRT (may fail if not available)
echo "🚄 Attempting TensorRT installation..."
pip install tensorrt || echo "⚠️ TensorRT installation failed - continuing without it"

# Verify installations
echo "✅ Verifying critical installations..."

# Test PyTorch and CUDA
python -c "
import torch
print(f'PyTorch version: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'CUDA version: {torch.version.cuda}')
    print(f'GPU count: {torch.cuda.device_count()}')
    for i in range(torch.cuda.device_count()):
        print(f'  GPU {i}: {torch.cuda.get_device_name(i)}')
else:
    print('⚠️ CUDA not available - CPU mode only')
"

# Test Ultralytics
python -c "
from ultralytics import YOLO
print('✅ Ultralytics YOLO imported successfully')
"

# Test key ML libraries
python -c "
import optuna, ray, fastapi, transformers, onnx
print('✅ All key ML libraries imported successfully')
"

# Display final package list
echo "📋 Installed packages:"
pip list | grep -E "(torch|ultralytics|optuna|ray|fastapi|transformers|onnx|wandb)"

echo ""
echo "🎉 Environment setup completed successfully!"
echo ""
echo "🚀 To activate the environment:"
echo "   conda activate $ENV_NAME"
echo ""
echo "📝 Next steps:"
echo "   1. Activate the environment: conda activate $ENV_NAME"
echo "   2. Navigate to project: cd /home/zty/Anti-UAV/anti_uav_pytorch"
echo "   3. Run verification tests"
echo ""
echo "✅ Environment: $ENV_NAME is ready for production use!"