# # Use the pre-built PyTorch Docker image with CUDA
# FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-devel

# # Install build deps
# RUN apt-get update -y && apt-get install -y wget git vim ninja-build \
#     && apt-get clean && rm -f /var/lib/apt/lists/*_*

# # Set up environment variables
# ENV CUDA_HOME=/usr/local/cuda
# # ENV CONDA_DIR=/opt/conda
# # ENV PATH=$CONDA_DIR/bin:$PATH

# # Set CUDA architecture flags
# ENV TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6"

# WORKDIR /app

# # Clone the InstantSplat repository and download the pre-trained model
# RUN git clone --recursive https://github.com/NVlabs/InstantSplat.git . && \
#     git submodule update --init --recursive && \
#     cd submodules/dust3r/ && \
#     mkdir -p checkpoints/ && \
#     wget https://download.europe.naverlabs.com/ComputerVision/DUSt3R/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth -P checkpoints/
    
# RUN pip install submodules/simple-knn && \
#     sed -i '154s/p_view.z <= 0.2f/p_view.z <= 0.001f/' submodules/diff-gaussian-rasterization/cuda_rasterizer/auxiliary.h && \
#     pip install submodules/diff-gaussian-rasterization && \
#     cd submodules/dust3r/croco/models/curope/ && \
#     python setup.py build_ext --inplace

# # Create the conda environment and install dependencies
# # RUN conda update -n base -c defaults conda

# # RUN conda create -n instantsplat python=3.11 cmake=3.14.0 -y && \
# #     echo "source activate instantsplat" > ~/.bashrc && \
# #     /bin/bash -c "source activate instantsplat" && \
# #     conda install pytorch torchvision pytorch-cuda=11.8 -c pytorch -c nvidia -y && \
# #     cd InstantSplat && \
# #     pip install -r requirements.txt && \
# #     pip install submodules/simple-knn && \
# #     sed -i '154s/p_view.z <= 0.2f/p_view.z <= 0.001f/' submodules/diff-gaussian-rasterization/cuda_rasterizer/auxiliary.h && \
# #     pip install submodules/diff-gaussian-rasterization

# RUN pip install -r requirements.txt
# RUN pip install imageio[ffmpeg]

# # Set the working directory
# RUN useradd -m -u 1000 user
# RUN chown -R user:user /app
# USER user

# COPY --chown=user . /app
# # COPY --chown=user ./requirements.txt requirements.txt
# # COPY ./requirements.txt requirements.txt

# ENV PATH="/home/user/.local/bin:$PATH"

# # COPY --chown=user ./run_train_infer.sh run_train_infer.sh
# # COPY ./run_train_infer.sh run_train_infer.sh
# RUN chmod +x run_train_infer.sh


# # Default command to run when the container starts
# # CMD ["/bin/bash", "run_train_infer.sh"]

# # Copy the current directory contents into the container at $HOME/app setting the owner to the user

# CMD ["python", "app.py"]


FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-devel

ARG DEBIAN_FRONTEND=noninteractive

ARG FORCE_CUDA=1
ENV FORCE_CUDA=${FORCE_CUDA}

ENV NVIDIA_VISIBLE_DEVICES \
    ${NVIDIA_VISIBLE_DEVICES:-all}
ENV NVIDIA_DRIVER_CAPABILITIES \
    ${NVIDIA_DRIVER_CAPABILITIES:+$NVIDIA_DRIVER_CAPABILITIES,}graphics
    
RUN apt-get update -y && apt-get install -y git wget

WORKDIR /
RUN git clone --recursive https://github.com/NVlabs/InstantSplat.git &&\
    cd InstantSplat &&\
    git submodule update --init --recursive &&\
    cd submodules/dust3r/ &&\
    mkdir -p checkpoints/ &&\
    wget https://download.europe.naverlabs.com/ComputerVision/DUSt3R/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth -P checkpoints/

WORKDIR /InstantSplat

RUN pip3 install -r requirements.txt

ARG TORCH_CUDA_ARCH_LIST="3.5;5.0;6.0;6.1;7.0;7.5;8.0;8.6+PTX"

RUN pip3 install submodules/simple-knn
RUN sed -i 's/p_view.z <= 0.2f/p_view.z <= 0.001f/g' /InstantSplat/submodules/diff-gaussian-rasterization/cuda_rasterizer/auxiliary.h
RUN pip3 install submodules/diff-gaussian-rasterization

RUN cd submodules/dust3r/croco/models/curope/ &&\
    python3 setup.py build_ext --inplace
    
RUN pip3 install plyfile
RUN pip3 install imageio[ffmpeg]
RUN apt-get install -y libgl1-mesa-dev libglib2.0-0
RUN apt-get install -y tree


ENV CUDA_VISIBLE_DEVICES=0
ENV MKL_THREADING_LAYER=GNU
ENV MKL_SERVICE_FORCE_INTEL=1

RUN useradd -m -u 1000 user
RUN chown -R user:user /InstantSplat
USER user

COPY --chown=user . /InstantSplat


CMD ["python", "app.py"]