with import <nixpkgs> {
  config.sudaSupport = true;
  config.allowUnfree = true;
};
let
  python3 = pkgs.python311;
in
mkShell {
  name = "tensorflow-cuda-shell";
  buildInputs = with python3.pkgs; [
    pip
    numpy
    setuptools
    virtualenv
    tensorflow tensorboard
  ];
  shellHook = ''
    export CUDA_PATH=${pkgs.cudatoolkit}
    export LD_LIBRARY_PATH=${pkgs.stdenv.cc.cc.lib}/lib:${pkgs.cudaPackages_11.cudatoolkit}/lib:${pkgs.cudaPackages_11.cudnn}/lib:${pkgs.cudaPackages_11.cudatoolkit.lib}/lib:$LD_LIBRARY_PATH
    export EXTRA_LDFLAGS="-L/lib -L${pkgs.linuxPackages.nvidia_x11}/lib"
    export EXTRA_CCFLAGS="-I/usr/include"
    alias pip="PIP_PREFIX='$(pwd)/_build/pip_packages' TMPDIR='$HOME' \pip"
    export PYTHONPATH="$(pwd)/_build/pip_packages/lib/python3.7/site-packages:$PYTHONPATH"
    export PATH="$(pwd)/_build/pip_packages/bin:$PATH"
    unset SOURCE_DATE_EPOCH

    # Use bash instead of zsh
    export SHELL=${pkgs.bash}/bin/bash

    # set up a virtuial environment
    export VIRTUAL_ENV=$PWD/venv
    if [ ! -d "$VIRTUAL_ENV" ]; then
      virtualenv "$VIRTUAL_ENV"
      source "$VIRTUAL_ENV/bin/activate"
      python -m pip install --upgrade pip
      # python -m pip install -r requirements.txt
    else
      source "$VIRTUAL_ENV/bin/activate"
    fi
  '';
}