#!/usr/bin/env bash
# SUM Parts - WSL2 conda env setup (no sudo required)
#
# NOTE: Anaconda's "defaults" channels (repo.anaconda.com/pkgs/*) require accepting
# a Terms of Service that carries commercial-license obligations for larger orgs.
# This script deliberately avoids them entirely: conda-forge for packages,
# nvidia channel for the CUDA toolkit, both with --override-channels.
set -euo pipefail

CONDA_ROOT="$HOME/miniconda3"
ENV_NAME="sumparts"

export PATH="$CONDA_ROOT/bin:$PATH"
source "$CONDA_ROOT/etc/profile.d/conda.sh"

echo "=== [0/5] pin channels to conda-forge (drop anaconda defaults) ==="
conda config --remove channels defaults 2>/dev/null || true
conda config --add channels conda-forge
conda config --set channel_priority strict

echo "=== [1/5] create env: $ENV_NAME (python 3.10) ==="
conda create -n "$ENV_NAME" -y --override-channels -c conda-forge python=3.10

conda activate "$ENV_NAME"

echo "=== [2/5] CUDA Toolkit 11.8 (nvidia channel, no sudo) ==="
conda install -y --override-channels -c "nvidia/label/cuda-11.8.0" cuda-toolkit

echo "=== [3/5] PyTorch 2.0.1 + cu118 (pip wheels) ==="
pip install --no-cache-dir \
    torch==2.0.1+cu118 torchvision==0.15.2+cu118 \
    --index-url https://download.pytorch.org/whl/cu118

echo "=== [4/5] pin numpy<2 (torch 2.0.x is ABI-incompatible with numpy 2) ==="
pip install --no-cache-dir "numpy<2"

echo "=== [5/5] verify ==="
which nvcc
nvcc --version | tail -2
python - <<'PY'
import torch, numpy
print("torch  :", torch.__version__)
print("cuda   :", torch.version.cuda)
print("avail  :", torch.cuda.is_available())
print("device :", torch.cuda.get_device_name(0) if torch.cuda.is_available() else None)
print("numpy  :", numpy.__version__)
PY

echo "ENV DONE"