# Copyright (c) 2021-2024, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

# Docker file for aarch64 based Jetson device
ARG BASE_IMAGE="nvcr.io/nvidia/l4t-cuda:12.2.12-devel"
FROM ${BASE_IMAGE}

# Store list of packages (must be first)
RUN mkdir -p /opt/nvidia/isaac_ros_dev_base && dpkg-query -W | sort > /opt/nvidia/isaac_ros_dev_base/aarch64-start-packages.csv

# Disable terminal interaction for apt
ENV DEBIAN_FRONTEND=noninteractive
ENV SHELL /bin/bash
SHELL ["/bin/bash", "-c"]

# Ensure we have universe
RUN --mount=type=cache,target=/var/cache/apt \
apt-get update && apt-get install -y \
        software-properties-common \
&& add-apt-repository universe \
&& apt-get update

# Fundamentals
RUN --mount=type=cache,target=/var/cache/apt \
apt-get update && apt-get install -y \
    apt-utils \
    bash-completion \
    build-essential \
    ca-certificates \
    curl \
    git \
    git-lfs \
    gnupg2 \
    iputils-ping \
    libgoogle-glog-dev \
    locales \
    lsb-release \
    software-properties-common \
    sudo \
    tar \
    unzip \
    vim \
    wget \
    zlib1g-dev

# Add Isaac apt repository
RUN --mount=type=cache,target=/var/cache/apt \
    wget -qO - https://isaac.download.nvidia.com/isaac-ros/repos.key | apt-key add - && \
    grep -qxF "deb https://isaac.download.nvidia.com/isaac-ros/release-3 $(lsb_release -cs) release-3.0" /etc/apt/sources.list || \
    echo "deb https://isaac.download.nvidia.com/isaac-ros/release-3 $(lsb_release -cs) release-3.0" | tee -a /etc/apt/sources.list \
    && apt-get update

# Setup Jetson debian repositories
RUN --mount=type=cache,target=/var/cache/apt \
    apt-key adv --fetch-keys https://repo.download.nvidia.com/jetson/jetson-ota-public.asc \
    && apt-key adv --fetch-keys http://l4t-repo.nvidia.com/jetson-ota-internal.key \
    && echo 'deb https://repo.download.nvidia.com/jetson/common r36.3 main' > /etc/apt/sources.list.d/nvidia-l4t-apt-source.list \
    && echo 'deb https://repo.download.nvidia.com/jetson/t234 r36.3 main' >> /etc/apt/sources.list.d/nvidia-l4t-apt-source.list \
    && apt-get update

# Python basics
RUN --mount=type=cache,target=/var/cache/apt \
apt-get update && apt-get install -y \
    python3-dev \
    python3-distutils \
    python3-flake8 \
    python3-pip \
    python3-pytest-cov \
    python3-venv \
    python3-zmq \
    python3.10 \
    python3.10-venv

# Set Python3 as default
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 1

# Core dev libraries
RUN --mount=type=cache,target=/var/cache/apt \
apt-get update && apt-get install -y \
    ffmpeg \
    gfortran \
    graphicsmagick-libmagick-dev-compat \
    jq \
    kmod \
    lcov \
    libasio-dev \
    libassimp-dev \
    libatlas-base-dev \
    libblas3 \
    libatlas3-base \
    libboost-all-dev \
    libboost-dev \
    libceres-dev \
    libbullet-dev \
    libcunit1-dev \
    libffi7 \
    libfreetype6 \
    libgraphicsmagick++1-dev \
    libhidapi-libusb0 \
    libinput10 \
    libjpeg8 \
    liblapack3 \
    libmnl0 \
    libmnl-dev \
    libncurses5-dev \
    libode-dev \
    libopenblas0 \
    libopencv-dev=4.5.4+dfsg-9ubuntu4 \
    libopenmpi3 \
    libpcap-dev \
    libpcl-dev \
    libsuitesparse-dev \
    libtinyxml2-dev \
    libturbojpeg \
    linuxptp \
    libunwind8 \
    libv4l-0 \
    libx264-dev \
    libxaw7-dev \
    libyaml-cpp-dev \
    llvm-14 \
    nlohmann-json3-dev \
    python3-opencv=4.5.4+dfsg-9ubuntu4 \
    python3-scipy

# Additional Python dependencies
RUN python3 -m pip install -U \
    Cython \
    pymongo \
    wheel \
    scikit-learn \
    ninja \
    networkx \
    numpy \
    numpy-quaternion \
    pyyaml \
    setuptools_scm>=6.2 \
    trimesh \
    yourdfpy>=0.0.53 \
    warp-lang>=0.9.0 \
    scipy>=1.7.0 \
    tqdm \
    importlib_resources


# Update environment
RUN update-alternatives --install /usr/bin/llvm-config llvm-config /usr/bin/llvm-config-14 14
ENV LD_LIBRARY_PATH="/opt/nvidia/vpi3/lib64:${LD_LIBRARY_PATH}"
ENV LD_LIBRARY_PATH="/usr/lib/aarch64-linux-gnu/tegra:${LD_LIBRARY_PATH}"
ENV LD_LIBRARY_PATH="/usr/local/cuda-12.2/targets/aarch64-linux/lib:${LD_LIBRARY_PATH}"
ENV LD_LIBRARY_PATH="/usr/lib/aarch64-linux-gnu/tegra-egl:${LD_LIBRARY_PATH}"
ENV LD_LIBRARY_PATH="/usr/lib/aarch64-linux-gnu/tegra/weston:${LD_LIBRARY_PATH}"
ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/lib/aarch64-linux-gnu-host"
ENV PATH="/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/src/tensorrt/bin:${PATH}"

# Install CUDA packages
RUN --mount=type=cache,target=/var/cache/apt \
apt-get update && apt-get install -y --no-install-recommends \
    cuda-cudart-12-2 \
    cuda-libraries-12-2 \
    cuda-nvml-dev-12-2 \
    cuda-sanitizer-12-2 \
    cuda-toolkit-12-2 \
    libcublas-12-2 \
    libcudnn8 \
    libcusparse-12-2 \
    libnpp-12-2

# Install TensorRT and VPI
RUN --mount=type=cache,target=/var/cache/apt \
mkdir -p /lib/firmware && \
apt-get update && apt-get install -y \
    libnvvpi3 \
    tensorrt \
    vpi3-dev

# Install Tao converter
RUN mkdir -p /opt/nvidia/tao && cd /opt/nvidia/tao && \
    wget --content-disposition 'https://api.ngc.nvidia.com/v2/resources/org/nvidia/team/tao/tao-converter/v5.1.0_jp6.0_aarch64/files?redirect=true&path=tao-converter' -O tao-converter && \
    chmod 755 tao-converter

ENV PATH="${PATH}:/opt/nvidia/tao"
ENV TRT_LIB_PATH="/usr/lib/aarch64-linux-gnu"
ENV TRT_INCLUDE_PATH="/usr/include/aarch64-linux-gnu"

# PyTorch (NV CUDA edition)
# https://docs.nvidia.com/deeplearning/frameworks/install-pytorch-jetson-platform/index.html
RUN python3 -m pip install --no-cache \
        https://developer.download.nvidia.com/compute/redist/jp/v60dp/pytorch/torch-2.2.0a0+6a974be.nv23.11-cp310-cp310-linux_aarch64.whl

# Install Triton server from https://github.com/triton-inference-server/server/releases/tag/v2.40.0
RUN --mount=type=cache,target=/var/cache/apt \
apt-get update && apt-get install -y --no-install-recommends \
    libb64-0d \
    libre2-9 \
    rapidjson-dev \
    libopenblas-dev \
    libarchive-dev

RUN  --mount=type=cache,target=/var/cache/apt \
    cd /opt \
    && wget https://github.com/triton-inference-server/server/releases/download/v2.40.0/tritonserver2.40.0-igpu.tar.gz \
    && tar -xzvf tritonserver2.40.0-igpu.tar.gz \
    && chmod 644 /opt/tritonserver/backends/tensorflow/libtensorflow_cc.so.2 \
    && chmod 644 /opt/tritonserver/backends/tensorflow/libtensorflow_framework.so.2 \
    && rm tritonserver2.40.0-igpu.tar.gz

ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/tritonserver/lib"

# Install boost version >= 1.78 for boost::span
# Current libboost-dev apt packages are < 1.78, so install from tar.gz
RUN --mount=type=cache,target=/var/cache/apt \
    wget -O /tmp/boost.tar.gz \
    https://boostorg.jfrog.io/artifactory/main/release/1.80.0/source/boost_1_80_0.tar.gz \
    && (cd /tmp && tar xzf boost.tar.gz) \
    && cd /tmp/boost_1_80_0 \
    && ./bootstrap.sh --prefix=/usr \
    && ./b2 install \
    && rm -rf /tmp/boost*

# Install CV-CUDA
RUN --mount=type=cache,target=/var/cache/apt \
    cd /tmp && \
    wget https://github.com/CVCUDA/CV-CUDA/releases/download/v0.5.0-beta/nvcv-lib-0.5.0_beta_DP-cuda12-aarch64-linux.deb && \
    dpkg -i nvcv-lib-0.5.0_beta_DP-cuda12-aarch64-linux.deb && \
    wget https://github.com/CVCUDA/CV-CUDA/releases/download/v0.5.0-beta/nvcv-dev-0.5.0_beta_DP-cuda12-aarch64-linux.deb && \
    dpkg -i nvcv-dev-0.5.0_beta_DP-cuda12-aarch64-linux.deb

# Add MQTT binaries and libraries
RUN --mount=type=cache,target=/var/cache/apt \
apt-add-repository ppa:mosquitto-dev/mosquitto-ppa \
&& apt-get update && apt-get install -y \
        mosquitto \
        mosquitto-clients

# Install jtop
RUN python3 -m pip install -U \
    jetson-stats

# Store list of packages (must be last)
RUN mkdir -p /opt/nvidia/isaac_ros_dev_base && dpkg-query -W | sort > /opt/nvidia/isaac_ros_dev_base/aarch64-end-packages.csv