#!/usr/bin/env fish

# Quick Setup Script for Arch Linux + RTX 2050
# Road Anomaly Detection Training

echo "🚗 Setting up Road Anomaly Detection training environment..."
echo "   Target: NVIDIA RTX 2050 (4GB VRAM)"
echo ""

# =============================================================================
# RTX 2050 CUDA Environment Variables (4GB VRAM optimization)
# =============================================================================
set -gx CUDA_VISIBLE_DEVICES 0
set -gx PYTORCH_CUDA_ALLOC_CONF "max_split_size_mb:512"
set -gx TF_FORCE_GPU_ALLOW_GROWTH "true"

# Ensure `uv` (env manager) is installed and create/activate a virtual environment
if not type -q uv
	echo "Installing uv (preferred via pipx)..."
	if type -q pipx
		pipx install uv || pip install --user uv
	else if type -q pip
		pip install --user uv
	else
		echo "Error: neither pipx nor pip found. Install pip or pipx and retry."
		exit 1
	end
end

# Create virtualenv with uv (creates .venv by default)
if not test -d .venv
	uv venv
else
	echo ".venv already exists; skipping creation."
end

# Activate the fish activation script if available
if test -f .venv/bin/activate.fish
	source .venv/bin/activate.fish
else if test -f .venv/bin/activate
	source .venv/bin/activate
end

# Install PyTorch with CUDA into the uv-managed environment
uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# Install project dependencies into the uv-managed environment
uv pip install ultralytics opencv-python matplotlib albumentations numpy pandas seaborn pillow tqdm PyYAML tensorboard onnx onnxruntime

# Install TFLite export dependencies
uv pip install tensorflow

# Verify setup
echo ""
echo "Verifying installation..."
python -c "import torch; print('✓ PyTorch:', torch.__version__); print('✓ CUDA Available:', torch.cuda.is_available()); print('✓ GPU:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'None')"
python -c "from ultralytics import YOLO; print('✓ Ultralytics YOLOv8 ready')"

echo ""
echo "═══════════════════════════════════════════════════════════"
echo "✓ Setup complete! Ready to train."
echo "═══════════════════════════════════════════════════════════"
echo ""
echo "Next steps:"
echo "  1. Verify dataset:  python verify_dataset.py"
echo "  2. Train model:     python train_road_anomaly_model.py"
echo "  3. Run inference:   python inference.py --model <model.pt> --source <image/video>"
echo "  4. Package for RPi: python package_for_rpi.py"
echo ""
echo "Monitor GPU during training:"
echo "  watch -n 1 nvidia-smi"
echo ""