# Base Image
FROM nvidia/cuda:11.3.1-cudnn8-runtime-ubuntu20.04

# System Dependencies and Cleanup
RUN apt-get update -y && \
    DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y tzdata && \
    apt-get install -y software-properties-common ffmpeg libsm6 libxext6 libhdf5-serial-dev netcdf-bin libnetcdf-dev && \
    add-apt-repository ppa:ubuntugis/ubuntugis-unstable && \
    apt-get update && \
    apt-get install -y curl build-essential gdal-bin libgdal-dev libpq-dev python3-gdal python3-pip apt-transport-https ca-certificates gnupg && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

# Copy only the necessary files
COPY requirements.txt /micrograd/requirements.txt

# Install Python Packages
RUN pip install --no-cache-dir -r /micrograd/requirements.txt

# Set Working Directory and Prepare App
WORKDIR /micrograd
COPY micrograd /micrograd/micrograd
COPY test /micrograd/test
COPY app.py /micrograd/app.py

RUN mkdir -p /root/.cache/torch/hub/checkpoints/

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