# Use the nvidia cuda12.2.0 base image with Ubuntu 20.04
FROM nvidia/cuda:12.2.0-devel-ubuntu20.04

# Update and upgrade Ubuntu packages
RUN apt-get update && apt-get upgrade -y && \
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    software-properties-common git tmux vim-tiny python3-pip wget htop \
    openssh-server silversearcher-ag sqlite3 libsqlite3-dev

# Install Mambaforge python 3.10.12
ENV MAMBAFORGE_VERSION=23.1.0-4
ENV MAMBAFORGE_FILE=Mambaforge-${MAMBAFORGE_VERSION}-Linux-x86_64.sh
RUN wget -q https://github.com/conda-forge/miniforge/releases/download/${MAMBAFORGE_VERSION}/${MAMBAFORGE_FILE} && \
    chmod +x ${MAMBAFORGE_FILE} && \
    ./${MAMBAFORGE_FILE} -b -p $HOME/mambaforge && \
    rm ${MAMBAFORGE_FILE}
# Set path to conda
ENV PATH /root/mambaforge/bin:$PATH

# Python Env Commands: ray + pytorch1.13
# pytorch 1.13: https://pytorch.org/get-started/previous-versions/
# https://github.com/pytorch/pytorch/wiki/PyTorch-Versions
RUN /bin/bash -c "source $HOME/mambaforge/etc/profile.d/conda.sh && mamba install -y ipython ipdb -c conda-forge -c anaconda"
RUN /bin/bash -c "source $HOME/mambaforge/etc/profile.d/conda.sh && mamba install -y tensorboard"
RUN /bin/bash -c "source $HOME/mambaforge/etc/profile.d/conda.sh && mamba install -y pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda==11.8 -c pytorch -c nvidia"

# Initialize mamba in your shell.
RUN /bin/bash -c "source $HOME/mambaforge/etc/profile.d/conda.sh && mamba init"

# Python packages
RUN /bin/bash -c "source $HOME/mambaforge/etc/profile.d/conda.sh && \
    pip install pettingzoo"
RUN /bin/bash -c "source $HOME/mambaforge/etc/profile.d/conda.sh && \
    mamba install -y tianshou -c conda-forge"

# Installing mujoco old version
# the necessary libraries and packages
RUN apt-get update && apt-get install -y \
    libosmesa6-dev \
    libgl1-mesa-glx \
    libglfw3 \
    patchelf \
    xz-utils \
    && rm -rf /var/lib/apt/lists/*
# Install mujoco-py and other Python packages using the provided pip mirror
RUN /bin/bash -c "source $HOME/mambaforge/etc/profile.d/conda.sh && \
    pip install -U  'mujoco-py<2.2,>=2.1'"
RUN /bin/bash -c "source $HOME/mambaforge/etc/profile.d/conda.sh && \
    pip install 'cython<3'"
COPY ./mujoco210 /root/.mujoco/mujoco210
RUN echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/root/.mujoco/mujoco210/bin' >> /root/.bashrc
ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/root/.mujoco/mujoco210/bin"
RUN python -c "import mujoco_py"

# MISC python packages
RUN /bin/bash -c "source $HOME/mambaforge/etc/profile.d/conda.sh && \
    mamba install -y -c conda-forge optuna optuna-dashboard ax-platform einops"
RUN /bin/bash -c "source $HOME/mambaforge/etc/profile.d/conda.sh && \
    pip install aim"

# Install ray
RUN /bin/bash -c "source $HOME/mambaforge/etc/profile.d/conda.sh && \
    mamba install -y -c conda-forge ray-all==2.11.0"

# Copy your script into the container
# Add /usr/bin to PATH and create it if it doesn't exist
RUN mkdir -p /root/bin && echo 'export PATH="$PATH:/root/bin"' >> ~/.bashrc

# Clean up
RUN apt-get clean && \
    rm -rf /var/lib/apt/lists/*

WORKDIR /root/myCodeLab
