# Use your specific research base image
FROM dongjun57/ctm-docker:latest

# Install system dependencies
RUN apt update && apt -y upgrade && \
    apt -y install curl git vim tmux wget software-properties-common \
    python3 python3-dev python3-pip python3-setuptools python3-wheel \
    libopenmpi-dev

# Removed flash-attn and xformers for CPU compatibility
RUN pip install packaging && \
    pip install tensorflow torch torchvision torchaudio \
    blobfile tqdm numpy scipy pandas Cython piq==0.7.0 \
    joblib==0.14.0 albumentations==0.4.3 lmdb pillow \
    clip@git+https://github.com/openai/CLIP.git \
    mpi4py nvidia-ml-py3 timm==0.4.12 legacy dill \
    jupyter jupyterlab notebook gdown matplotlib

# --- Hugging Face Specific Setup ---
# Create a user to avoid root permissions issues
RUN useradd -m -u 1000 user
USER user
ENV HOME=/home/user \
    PATH=/home/user/.local/bin:$PATH

WORKDIR $HOME/app

# --- FIX: CLONE FROM GITHUB INSTEAD OF COPYING LOCAL EMPTY FILES ---
RUN git clone https://github.com/ChiehHsinJesseLai/ctm-notebook.git .

COPY --chown=user model095000.pt .

COPY --chown=user cifar10-32x32.npz .

COPY --chown=user Eval-ctm-notebook.ipynb .

COPY --chown=user model.py .

# Hugging Face runs on port 7860
EXPOSE 7860

# Start JupyterLab (added the notebook_dir flag to be safe)
CMD ["jupyter", "lab", "--ip=0.0.0.0", "--port=7860", "--no-browser", "--NotebookApp.token=''", "--NotebookApp.password=''", "--NotebookApp.allow_origin='*'", "--NotebookApp.notebook_dir='/home/user/app'", "--ServerApp.disable_check_xsrf=True", "--NotebookApp.tornado_settings={'headers': {'Content-Security-Policy': 'frame-ancestors *'}}"]