# 1. Use NVIDIA CUDA Devel Base (Has nvcc compiler)
FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04

# 2. Set workspace variables
ENV DEBIAN_FRONTEND=noninteractive
ENV PATH="/opt/conda/envs/smi-ssed-env/bin:/opt/conda/bin:$PATH"
ENV CUDA_HOME="/usr/local/cuda"
# Force build for custom kernels
ENV CAUSAL_CONV1D_FORCE_BUILD=TRUE
ENV MAMBA_FORCE_BUILD=TRUE
ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6"

# 3. Install System Dependencies
RUN apt-get update && apt-get install -y \
    wget \
    git \
    build-essential \
    && rm -rf /var/lib/apt/lists/*

# 4. Install Miniforge
RUN wget --quiet https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh -O ~/miniforge.sh && \
    /bin/bash ~/miniforge.sh -b -p /opt/conda && \
    rm ~/miniforge.sh

# 5. Create  workspace
RUN pip install torch==2.7.1 --index-url https://download.pytorch.org/whl/cu118

# 6. Create and move to the 'workspace' folder
WORKDIR /workspace

# 7. Clone the mamba repository here
RUN git clone https://github.com/state-spaces/mamba.git
WORKDIR /workspace/mamba
RUN git checkout v1.2.0
RUN pip install . --no-build-isolation
# 8. Clone the  ibm-research/materials.smi_ssedrepository here
WORKDIR /workspace
RUN apt-get update && apt-get install -y libxrender1 libxext6 libsm6
RUN pip install pandas==2.3.1 rdkit "transformers==4.55.0" h5py==3.13.0
WORKDIR /workspace

RUN git clone https://huggingface.co/ibm-research/materials.smi_ssed
RUN sed -i "s|torch.load(\(.*ckpt_path.*\))|torch.load(\1, weights_only=False)|" /workspace/materials.smi_ssed/smi_ssed/inference/smi_ssed/load.py
# 9. Update PYTHONPATH
# We append the inference path so python can find the scripts
WORKDIR /workspace
RUN git clone https://huggingface.co/SuvenduK/ChemicalDice

ENV PYTHONPATH="/workspace/ChemicalDice:/workspace/materials.smi_ssed/smi_ssed/inference"


WORKDIR /app

# 7. Default command
CMD ["/bin/bash"]