# Use NVIDIA L4T base image for Jetson devices with JetPack 6.2
# JetPack 6.2 uses L4T 36.4 with CUDA 12.6
FROM nvcr.io/nvidia/l4t-base:r36.4.0

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV CUDA_HOME=/usr/local/cuda
ENV PATH=${CUDA_HOME}/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
ENV OPENBLAS_CORETYPE=ARMV8

# Install system dependencies (ARM64-compatible for JetPack 6.2)
RUN apt-get update && apt-get install -y \
    python3.10 \
    python3.10-dev \
    python3-pip \
    python3.10-venv \
    git \
    wget \
    curl \
    build-essential \
    cmake \
    ninja-build \
    libopencv-dev \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender-dev \
    libgomp1 \
    libgl1-mesa-glx \
    libglib2.0-0 \
    libblas-dev \
    liblapack-dev \
    libatlas-base-dev \
    gfortran \
    libhdf5-serial-dev \
    hdf5-tools \
    libhdf5-dev \
    zlib1g-dev \
    zip \
    libjpeg8-dev \
    liblapack3 \
    libopenblas-base \
    libopenmpi-dev \
    vim \
    && rm -rf /var/lib/apt/lists/*

# Create symbolic links for python (JetPack 6.2 uses Python 3.10)
RUN ln -sf /usr/bin/python3.10 /usr/bin/python3 && \
    ln -sf /usr/bin/python3.10 /usr/bin/python

# Upgrade pip and install build tools
RUN python3 -m pip install --upgrade pip setuptools wheel

# Install Cython first (required for some packages)
RUN pip install Cython==3.1.2

# Install JetPack components (CUDA 12.6, cuDNN, TensorRT for JetPack 6.2)
RUN apt-get update && apt-get install -y nvidia-jetpack && rm -rf /var/lib/apt/lists/*

# Install PyTorch for Jetson JetPack 6.2 (CUDA 12.6 compatible)
# Using NVIDIA's pre-built wheels optimized for Jetson
RUN pip install torch==2.7.0 torchvision==0.22.0 \
    --index-url=https://pypi.jetson-ai-lab.dev/jp6/cu126/

# Install cuSPARSELt if needed (required for PyTorch 2.4+ on Jetson)
RUN wget https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-sbsa/libcusparse_lt-linux-sbsa-0.5.2.1-archive.tar.xz && \
    tar xf libcusparse_lt-linux-sbsa-0.5.2.1-archive.tar.xz && \
    cp -a libcusparse_lt-linux-sbsa-0.5.2.1-archive/include/* /usr/local/cuda/include/ && \
    cp -a libcusparse_lt-linux-sbsa-0.5.2.1-archive/lib/* /usr/local/cuda/lib64/ && \
    rm -rf libcusparse_lt-*

# Install ninja (required for detectron2 compilation)
RUN pip install ninja==1.11.1.4

# Install fvcore and other detectron2 dependencies
RUN pip install fvcore==0.1.5.post20221221 \
    opencv-python==4.11.0.86 \
    matplotlib==3.9.4

# Install COCO API (pycocotools) - using the Windows-compatible version you used
RUN git clone https://github.com/philferriere/cocoapi.git && \
    cd cocoapi/PythonAPI && \
    python setup.py build_ext --inplace && \
    python setup.py install && \
    cd / && rm -rf cocoapi

# Install detectron2 from source with ARM64 optimizations for JetPack 6.2
# CUDA 12.6 and compute capability for Orin (8.7)
ENV FORCE_CUDA=1
ENV TORCH_CUDA_ARCH_LIST="8.7"
ENV CUDA_HOME=/usr/local/cuda-12.6
RUN pip install --no-build-isolation 'git+https://github.com/facebookresearch/detectron2.git'

# Install remaining requirements
COPY requirements.txt /tmp/requirements.txt
RUN pip install -r /tmp/requirements.txt

# Create working directory
WORKDIR /app

# Copy your application files
COPY FINAL.py /app/
COPY speedtrack.py /app/

# Create necessary directories
RUN mkdir -p /app/xpctest /app/output_json /app/ins /app/xpc

# Set up runtime environment
ENV PYTHONPATH=/app:$PYTHONPATH

# Expose any ports if needed (uncomment if you add a web interface)
# EXPOSE 8000

# Default command
CMD ["python", "FINAL.py"]