# Copyright 2024-2026 Gentoo Authors
# Distributed under the terms of the GNU General Public License v2
# ROCm implementation of torchvision

EAPI=8

PYTHON_COMPAT=( python3_{11..14} )
DISTUTILS_SINGLE_IMPL=1
DISTUTILS_USE_PEP517=setuptools
DISTUTILS_EXT=1
ROCM_SKIP_GLOBALS=1
inherit distutils-r1 multiprocessing rocm

MY_PN=torchvision
MYPN=vision
MYP=${MYPN}-${PV}

DESCRIPTION="Datasets, transforms and models for computer vision (ROCm backend)"
HOMEPAGE="https://github.com/pytorch/vision"
SRC_URI="https://github.com/pytorch/${MYPN}/archive/refs/tags/v${PV}.tar.gz
	-> torchvision-${PV}.tar.gz"

S="${WORKDIR}"/${MYP}

LICENSE="BSD"
SLOT="0"
KEYWORDS="~amd64"
IUSE="+ffmpeg +jpeg +png +webp"
RESTRICT="test"

REQUIRED_USE="${PYTHON_REQUIRED_USE}"

CAFFE2_PREFIX="/usr/lib/caffe2/rocm"
PYTORCH_PREFIX="/usr/lib/pytorch/rocm"
TORCHVISION_PREFIX="/usr/lib/torchvision/rocm"

RDEPEND="
	${PYTHON_DEPS}
	$(python_gen_cond_dep '
		dev-python/numpy[${PYTHON_USEDEP}]
		dev-python/pillow[${PYTHON_USEDEP}]
	')
	jpeg? ( media-libs/libjpeg-turbo:= )
	png? ( media-libs/libpng:= )
	webp? ( media-libs/libwebp )
	ffmpeg? ( media-video/ffmpeg )
	~sci-ml/caffe2-rocm-2.12.0[${PYTHON_SINGLE_USEDEP}]
	~sci-ml/pytorch-rocm-2.12.0[${PYTHON_SINGLE_USEDEP}]
	app-eselect/eselect-torchvision
"

DEPEND="${RDEPEND}"

src_prepare() {
	distutils-r1_src_prepare
}

src_configure() {
	rocm_add_sandbox -w
	distutils-r1_src_configure
}

python_compile() {
	addpredict /dev/kfd
	addpredict /dev/random

	# Point torch imports to ROCm backend
	export PYTHONPATH="${PYTORCH_PREFIX}:${PYTHONPATH}"
	export LD_LIBRARY_PATH="${CAFFE2_PREFIX}/lib64:${LD_LIBRARY_PATH}"

	export FORCE_CUDA=1
	export TORCHVISION_USE_PNG=$(usex png 1 0)
	export TORCHVISION_USE_JPEG=$(usex jpeg 1 0)
	export TORCHVISION_USE_WEBP=$(usex webp 1 0)
	export TORCHVISION_USE_FFMPEG=$(usex ffmpeg 1 0)
	export TORCHVISION_USE_NVJPEG=0
	export TORCHVISION_USE_VIDEO_CODEC=0

	# ROCm caffe2 headers must be found before eselect symlinks in /usr/include
	export CPLUS_INCLUDE_PATH="${CAFFE2_PREFIX}/include:${CPLUS_INCLUDE_PATH}"
	export C_INCLUDE_PATH="${CAFFE2_PREFIX}/include:${C_INCLUDE_PATH}"

	MAX_JOBS="$(get_makeopts_jobs)" \
		distutils-r1_python_compile -j1
}

python_install() {
	local dest="${ED}${TORCHVISION_PREFIX}"
	mkdir -p "${dest}" || die

	export PYTHONPATH="${PYTORCH_PREFIX}:${PYTHONPATH}"
	export LD_LIBRARY_PATH="${CAFFE2_PREFIX}/lib64:${LD_LIBRARY_PATH}"

	distutils-r1_python_install

	# Move from site-packages to our prefix
	local sitedir="$(python_get_sitedir)"
	if [[ -d "${ED}${sitedir}/torchvision" ]]; then
		mv "${ED}${sitedir}/torchvision" "${dest}/torchvision" || die
	fi

	# Move dist-info
	local distinfo
	for distinfo in "${ED}${sitedir}"/torchvision-*.dist-info; do
		[[ -d "${distinfo}" ]] && mv "${distinfo}" "${dest}/" || true
	done

	# Fix RPATH on _C.so to find ROCm caffe2 libs
	local cso="${dest}/torchvision/_C.so"
	if [[ -f "${cso}" ]]; then
		patchelf --set-rpath "${CAFFE2_PREFIX}/lib64:\$ORIGIN" "${cso}" || \
			ewarn "patchelf failed on _C.so — RPATH may need manual fix"
	fi
}

pkg_postinst() {
	local active
	active=$(eselect torchvision show 2>/dev/null)
	if [[ "${active}" == "(unset)" || -z "${active}" ]]; then
		eselect torchvision set rocm
		elog "torchvision backend set to: rocm"
	fi
}