# thanks to Paul DeCoursey (optimuspaul@github.com) for inspiration
# see: https://github.com/MVIG-SJTU/AlphaPose/issues/1057

#FROM nvidia/cuda:11.6.0-cudnn8-devel-ubuntu22.04
#TODO: fix deprecation with either of these
#FROM nvidia/cuda:11.6.0-devel-ubuntu20.04
FROM nvidia/cuda:11.6.2-devel-ubuntu20.04
#FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu22.04
#FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04

RUN apt clean

RUN DEBIAN_FRONTEND=noninteractive apt update && DEBIAN_FRONTEND=noninteractive apt install -y --allow-change-held-packages \
    ffmpeg \
    git \
    build-essential \
    ninja-build \
    cuda-libraries-11-6 \
    mesa-common-dev \
    libosmesa6 libosmesa6-dev \
    libgles2-mesa-dev \
    libglu1-mesa-dev \
    libgles2-mesa-dev \
    libcublas-11-6 \
    libcublas-dev-11-6 \
    libcusparse-dev-11-6 \
    cuda-nvcc-11-6 \
    libcusolver-dev-11-6 \
    cuda-nvrtc-dev-11-6 \
    libcurand-dev-11-6 \
    cuda-nvml-dev-11-6 \
    libcufft-dev-11-6 \
    cuda-toolkit-11-6 \
    nvidia-cuda-toolkit \
    libyaml-dev

ENV CUDA_HOME='/usr/local/cuda'

RUN DEBIAN_FRONTEND=noninteractive apt -y install python3 python3-pip

RUN pip install "torch==1.12.1+cu116" "torchvision==0.13.1+cu116" "torchaudio==0.12.1" --extra-index-url https://download.pytorch.org/whl/cu116
RUN pip install opencv-contrib-python-headless
RUN pip install pip --upgrade

RUN pip install cython pycocotools

RUN pip install "cython_bbox==0.1.3"
RUN pip install "numpy==1.22.4"
RUN apt install -y python3-tk

RUN python3 -c "import torch; print(torch.version.cuda)" 

RUN mkdir /build
RUN cd /build && git clone https://github.com/HaoyiZhu/HalpeCOCOAPI.git
RUN cd /build && git clone https://github.com/MVIG-SJTU/AlphaPose.git

RUN cd /build/HalpeCOCOAPI/PythonAPI && python3 setup.py build develop --user

WORKDIR /build/AlphaPose

# from setup.py; see if it avoids installation issues
RUN pip install six 
RUN pip install terminaltables scipy matplotlib visdom \
        tqdm tensorboardx easydict \
        pyyaml halpecocotools \
        munkres timm==0.1.20 natsort

RUN TORCH_CUDA_ARCH_LIST="6.1;7.5;8.6" python3 setup.py build develop --user

RUN apt install bc -y && \
    pip install boto3 && \
    python3 -c "import torchvision.models as tm; tm.resnet152(pretrained=True)"


# Resnet 50 is needed for some of the AlphaPose models
# TODO: alternatively make this a build param/tag for resnet50/152
#       see: https://stackoverflow.com/questions/33935807/how-to-define-a-variable-in-a-dockerfile#37639938
RUN python3 -c "import torchvision.models as tm; tm.resnet50(pretrained=True)"
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 2