#!/bin/bash
#
# YOLO-TS Separate Environment Setup Script
#
# This script automates the creation of a dedicated conda environment
# for YOLO-TS with its custom Ultralytics modules.
#
# Usage:
#   ./scripts/setup_yolo_ts_env.sh
#

set -e  # Exit on error

echo "========================================"
echo "YOLO-TS Separate Environment Setup"
echo "========================================"
echo ""

# Find conda
if [ -n "$CONDA_EXE" ]; then
    CONDA="$CONDA_EXE"
elif [ -f "$HOME/anaconda3/bin/conda" ]; then
    CONDA="$HOME/anaconda3/bin/conda"
elif [ -f "$HOME/miniconda3/bin/conda" ]; then
    CONDA="$HOME/miniconda3/bin/conda"
elif command -v conda &> /dev/null; then
    CONDA="conda"
else
    echo "❌ Error: conda not found."
    echo "Please install Anaconda or Miniconda first:"
    echo "  https://docs.conda.io/en/latest/miniconda.html"
    exit 1
fi

echo "✓ Found conda: $CONDA"
echo ""

# Check if yolots environment already exists
if "$CONDA" env list | grep -q "^yolots "; then
    echo "⚠️  'yolots' environment already exists."
    read -p "Do you want to remove and recreate it? (y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        echo "Removing existing environment..."
        "$CONDA" env remove -n yolots -y
    else
        echo "Skipping environment creation."
        exit 0
    fi
fi

# Create conda environment
echo "Creating 'yolots' conda environment with Python 3.11..."
"$CONDA" create -n yolots python=3.11 -y

echo ""
echo "✓ Environment created"
echo ""

# Install PyTorch
echo "Installing PyTorch with CUDA 12.1 support..."
"$CONDA" run -n yolots pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

echo ""
echo "✓ PyTorch installed"
echo ""

# Clone YOLO-TS repository
YOLOTS_DIR="/tmp/YOLO-TS"
if [ -d "$YOLOTS_DIR" ]; then
    echo "YOLO-TS repository already exists at $YOLOTS_DIR"
    read -p "Do you want to update it? (y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        echo "Updating YOLO-TS repository..."
        cd "$YOLOTS_DIR"
        git pull
    fi
else
    echo "Cloning YOLO-TS repository to $YOLOTS_DIR..."
    git clone https://github.com/Heqiang-Huang/YOLO-TS.git "$YOLOTS_DIR"
fi

echo ""
echo "✓ YOLO-TS repository ready"
echo ""

# Install YOLO-TS dependencies
echo "Installing YOLO-TS dependencies..."
cd "$YOLOTS_DIR"
"$CONDA" run -n yolots pip install -r requirements.txt

echo ""
echo "✓ Dependencies installed"
echo ""

# Install custom Ultralytics by copying to site-packages
echo "Installing custom Ultralytics with YOLO-TS modules..."
SITE_PACKAGES=$("$CONDA" run -n yolots python -c "import site; print(site.getsitepackages()[0])")
cp -r "$YOLOTS_DIR/ultralytics" "$SITE_PACKAGES/"

echo ""
echo "✓ Custom Ultralytics installed"
echo ""

# Test the installation
echo "Testing installation..."
"$CONDA" run -n yolots python -c "
from ultralytics import YOLO
from ultralytics.nn.modules.block import SimFusion_3in
print('✓ Custom modules available')
print('✓ YOLO-TS environment ready!')
" 2>/dev/null && SUCCESS=true || SUCCESS=false

echo ""

if [ "$SUCCESS" = true ]; then
    echo "========================================"
    echo "✅ Setup Complete!"
    echo "========================================"
    echo ""
    echo "YOLO-TS environment 'yolots' is ready to use."
    echo ""
    echo "Next steps:"
    echo "1. Download or train a YOLO-TS model"
    echo "2. Place model in: models/traffic-sign-detection/YOLO-TS/weights/"
    echo "3. Test with: ./scripts/yolo_ts_inference.sh <image_path>"
    echo ""
    echo "For more information, see: docs/YOLO_TS_SEPARATE_ENV_SETUP.md"
    echo ""
    echo "To activate the environment manually:"
    echo "  conda activate yolots"
    echo ""
else
    echo "========================================"
    echo "⚠️  Setup completed with warnings"
    echo "========================================"
    echo ""
    echo "The environment was created but testing failed."
    echo "This might be normal if you need to restart or install additional dependencies."
    echo ""
    echo "Try manually:"
    echo "  conda activate yolots"
    echo "  cd $YOLOTS_DIR"
    echo "  pip install -e ."
    echo ""
fi
