# Install Mamba
conda install -n base mamba -c conda-forge -y

# Create environment with Mamba
mamba create -n bio_env python=3.10 -y
mamba activate bio_env

# Install packages
mamba install -c pytorch -c nvidia \
  pytorch=2.3.0 torchvision=0.18.0 torchaudio=2.3.0 \
  pytorch-cuda=12.1 cuda-toolkit=12.1.0 \
  cuda-nvcc=12.1.105 -y

mamba install -c conda-forge \
  numpy pandas scipy scikit-learn matplotlib seaborn \
  tqdm networkx pyyaml biopython omegaconf pydantic \
  pytest wandb loguru gdown tensorboard torchmetrics -y

pip install pymongo hydra-core torcheval typing biopandas 

pip install torch-scatter torch-sparse \
    torch-geometric -f https://data.pyg.org/whl/torch-2.3.0+cu121.html

pip install torch-cluster -f https://data.pyg.org/whl/torch-2.3.0+cu121.html

pip install ase


import torch, torch_geometric

print(f"PyTorch: {torch.__version__}")
print(f"CUDA: {torch.version.cuda}, Available: {torch.cuda.is_available()}")
print(f"PyG: {torch_geometric.__version__}")

# Test GPU tensor
if torch.cuda.is_available():
    print(f"GPU test: {torch.randn(2,4).cuda()}")
else:
    print("WARNING: CUDA not available!")
