#!/bin/bash

# If any command fails, the script will exit immediately.
set -e

ENV_NAME=molos

CUDA_TOOLKIT_VERSION=12.4.1
GCC_VERSION=12.4.0
GXX_VERSION=12.4.0
PYTHON_VERSION=3.12

# If you are using Lmod or Environment Module, and Miniconda3 is pre-installed on the system, you can load Miniconda3 with uncommenting the following lines.
# module purge
# module load miniconda3

if conda info --envs | grep -q "^$ENV_NAME[[:space:]]"; then
    echo "${ENV_NAME} is built, removing it to rebuild environment."
    conda remove -y --name $ENV_NAME --all
fi

echo "Building conda environment ${ENV_NAME} with python ${PYTHON_VERSION}."
conda create -y --name $ENV_NAME python=$PYTHON_VERSION

# Initialize conda for bash shell.
eval "$($CONDA_PREFIX/bin/conda 'shell.bash' 'hook')"

echo "Activating conda environment ${ENV_NAME}."
conda activate $ENV_NAME

echo "Installing Compilers."
conda install -y -c conda-forge compilers

echo "Installing GCC ${GCC_VERSION} and G++ ${GXX_VERSION}."
conda install -y -c conda-forge gcc_linux-64=${GCC_VERSION} gxx_linux-64=${GXX_VERSION}

echo "Installing CUDA Toolkit ${CUDA_TOOLKIT_VERSION}."
conda install -y nvidia/label/cuda-${CUDA_TOOLKIT_VERSION}::cuda-toolkit

echo "Removing all PyTorch extension caches."
rm -rf "$HOME/.cache/torch_extensions"

echo "Installing Python packages."
pip install -r requirements.txt --no-cache-dir

echo "Successfully built conda environment ${ENV_NAME} with python ${PYTHON_VERSION}."
