FROM mlebench-env

ENV OMP_NUM_THREADS=8 \
    MKL_NUM_THREADS=8 \
    OPENBLAS_NUM_THREADS=8 \
    NUMEXPR_NUM_THREADS=8

RUN apt-get update && apt-get install -y --no-install-recommends \
    htop \
    iotop \
    nmon \
    procps \
    vim \
    less \
    tree \
 && rm -rf /var/lib/apt/lists/*

# where to put submission.csv, will be extracted
ARG SUBMISSION_DIR
ENV SUBMISSION_DIR=${SUBMISSION_DIR}
# where to put any logs, will be extracted
ARG LOGS_DIR
ENV LOGS_DIR=${LOGS_DIR}
# where to put any code, will be extracted
ARG CODE_DIR
ENV CODE_DIR=${CODE_DIR}
# where to put any other agent-specific files, will not be necessarily extracted
ARG AGENT_DIR
ENV AGENT_DIR=${AGENT_DIR}

RUN mkdir ${LOGS_DIR} ${CODE_DIR} ${AGENT_DIR}

ARG CONDA_ENV_NAME=agent
ARG REQUIREMENTS=${AGENT_DIR}/requirements.txt

# copy just the requirements file, so that we can cache conda separately from the agent files
COPY requirements.txt ${AGENT_DIR}/requirements.txt
COPY aideml ${AGENT_DIR}/aideml
COPY llm-mcts ${AGENT_DIR}/llm-mcts

# Requirements for opencv
RUN apt-get update && apt-get install ffmpeg libsm6 libxext6  -y
RUN apt update && apt install graphviz graphviz-dev -y

# create conda environment and install the requirements to it
RUN conda run -n ${CONDA_ENV_NAME} pip install -r ${AGENT_DIR}/requirements.txt && \
    conda clean -afy
RUN conda run -n ${CONDA_ENV_NAME} pip install ${AGENT_DIR}/aideml ${AGENT_DIR}/llm-mcts && \
    conda clean -afy

RUN conda run -n ${CONDA_ENV_NAME} pip install pygraphviz openai==1.56.1 httpx==0.27.2 --force-reinstall

# put all the agent files in the expected location
COPY . ${AGENT_DIR}
