#!/bin/bash

set -e

echo "====================================="
echo "TensorFlow GPU WSL Setup Script"
echo "====================================="

PROJECT_DIR="$HOME/CNN"
VENV_NAME="tf_gpu"

echo "Updating system packages..."
sudo apt update
sudo apt install -y python3.10 python3.10-venv python3-pip

echo ""
echo "Creating project directory if needed..."
mkdir -p $PROJECT_DIR

cd $PROJECT_DIR

echo ""
echo "Creating Python virtual environment..."
python3.10 -m venv $VENV_NAME

echo ""
echo "Activating virtual environment..."
source $VENV_NAME/bin/activate

echo ""
echo "Upgrading pip..."
pip install --upgrade pip

echo ""
echo "Installing TensorFlow GPU build..."
pip install tensorflow[and-cuda]==2.19

echo ""
echo "Adding recommended TensorFlow environment variables..."

BASHRC="$HOME/.bashrc"

if ! grep -q TF_FORCE_GPU_ALLOW_GROWTH $BASHRC; then
echo "" >> $BASHRC
echo "# TensorFlow GPU settings" >> $BASHRC
echo "export TF_FORCE_GPU_ALLOW_GROWTH=true" >> $BASHRC
echo "export TF_CPP_MIN_LOG_LEVEL=2" >> $BASHRC
fi

echo ""
echo "Testing TensorFlow GPU detection..."

python <<EOF
import tensorflow as tf
print("TensorFlow version:", tf.__version__)
print("GPUs detected:", tf.config.list_physical_devices('GPU'))
EOF

echo ""
echo "====================================="
echo "Setup complete"
echo "====================================="
echo ""
echo "To activate the environment later:"
echo ""
echo "cd ~/CNN"
echo "source venv/bin/activate"