FROM docker.io/nvidia/cuda:11.8.0-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update
RUN apt-get install -y wget sudo pkg-config libopencv-dev unzip python3-pip

RUN pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0

RUN mkdir /go
ENV GOPATH=/go

RUN wget https://go.dev/dl/go1.22.2.linux-amd64.tar.gz
RUN tar -xvf go1.22.2.linux-amd64.tar.gz -C /usr/local
ENV PATH=$PATH:/usr/local/go/bin

RUN mkdir /app
WORKDIR /app

ADD  go.mod .
ADD  go.sum .
ADD  main.go .
ADD  logic logic

RUN go install || true

WORKDIR /root

RUN  wget https://github.com/sugarme/gotch/releases/download/v0.9.0/setup-libtorch.sh
RUN chmod +x setup-libtorch.sh
ENV CUDA_VER=11.8
ENV GOTCH_VER=v0.9.1
RUN bash setup-libtorch.sh
ENV GOTCH_LIBTORCH="/usr/local/lib/libtorch"
ENV LIBRARY_PATH="$LIBRARY_PATH:$GOTCH_LIBTORCH/lib"
ENV export CPATH="$CPATH:$GOTCH_LIBTORCH/lib:$GOTCH_LIBTORCH/include:$GOTCH_LIBTORCH/include/torch/csrc/api/include"
ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:$GOTCH_LIBTORCH/lib:/usr/lib64-nvidia:/usr/local/cuda-${CUDA_VERSION}/lib64"
RUN wget https://github.com/sugarme/gotch/releases/download/v0.9.0/setup-gotch.sh
RUN chmod +x setup-gotch.sh
RUN echo 'root ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
RUN bash setup-gotch.sh

RUN ln -s /usr/local/lib/libtorch/include/torch/csrc /usr/local/lib/libtorch/include/torch/csrc/api/include/torch
RUN mkdir -p /go/pkg/mod/github.com/sugarme/gotch@v0.9.1/libtch/libtorch/include/torch/csrc/api
RUN find /usr/local/lib/libtorch/include -maxdepth 1 -type d | tail -n +2 | grep -ve 'torch$' | xargs -I{} ln -s {} /go/pkg/mod/github.com/sugarme/gotch@v0.9.1/libtch/libtorch/include
RUN ln -s /usr/local/lib/libtorch/include/torch/csrc/api/include /go/pkg/mod/github.com/sugarme/gotch@v0.9.1/libtch/libtorch/include/torch/csrc/api/include
RUN find /usr/local/lib/libtorch/include/torch -maxdepth 1 -type f | xargs -I{} ln -s {} /go/pkg/mod/github.com/sugarme/gotch@v0.9.1/libtch/libtorch/include/torch
RUN ln -s /usr/local/lib/libtorch/lib/libcudnn.so.8 /usr/local/lib/libcudnn.so

WORKDIR /app

ADD . .
RUN go install || true

CMD ["bash", "-c", "go run ."]