# STARRY Python ML Services Dockerfile
# Multi-stage build for PyTorch + TensorFlow services

# ============================================================
# Stage 1: Base image with CUDA support
# ============================================================
# Use CUDA 12.1 runtime - PyTorch wheel includes cudnn
FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04 AS base

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1

# Install system dependencies
RUN apt-get update && apt-get install -y \
    python3.11 \
    python3.11-dev \
    python3-pip \
    libgl1-mesa-glx \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender-dev \
    git \
    && rm -rf /var/lib/apt/lists/*

# Set Python 3.11 as default
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1 \
    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1

# Upgrade pip
RUN python -m pip install --upgrade pip

# ============================================================
# Stage 2: PyTorch services (layout, mask, semantic, gauge, loc)
# ============================================================
FROM base AS pytorch-services

WORKDIR /app

# Install PyTorch with CUDA support
# Using cu121 for compatibility with CUDA driver 12.4
RUN pip install --no-cache-dir \
    "numpy>=1.26.0,<2.0.0" \
    torch==2.5.1 \
    torchvision==0.20.1 \
    --index-url https://download.pytorch.org/whl/cu121

# Install common dependencies
RUN pip install --no-cache-dir \
    "opencv-python-headless<4.11" \
    Pillow>=8.0.0 \
    PyYAML>=5.4.0 \
    pyzmq>=22.0.0 \
    msgpack>=1.0.0 \
    dill \
    scipy \
    imgaug \
    scikit-image \
    python-dotenv \
    fs \
    tqdm \
    einops \
    lmdb

# Source code should be mounted as volume at runtime:
# -v /path/to/deep-starry:/app/deep-starry:ro
ENV PYTHONPATH=/app/deep-starry

# Default command (override with docker-compose)
CMD ["python", "--help"]


# ============================================================
# Stage 3: TensorFlow services (ocr, brackets)
# ============================================================
FROM base AS tensorflow-services

WORKDIR /app

# Install TensorFlow with legacy Keras support
RUN pip install --no-cache-dir \
    "numpy==1.26.4" \
    tensorflow==2.20.0 \
    tf_keras==2.20.1 \
    "opencv-python-headless<4.11" \
    Pillow>=8.0.0 \
    PyYAML>=5.4.0 \
    pyzmq>=22.0.0 \
    msgpack>=1.0.0 \
    zhon \
    nltk \
    distance \
    anyconfig \
    munch \
    tensorboardX \
    scipy \
    scikit-image \
    python-dotenv

# Set legacy Keras environment
ENV TF_USE_LEGACY_KERAS=1

# Source code should be mounted as volume at runtime:
# -v /path/to/starry-ocr:/app/starry-ocr:ro
ENV PYTHONPATH=/app/starry-ocr

# Default command
CMD ["python", "--help"]


# ============================================================
# Stage 4: All-in-one image (for convenience)
# ============================================================
FROM base AS all-in-one

WORKDIR /app

# Install all dependencies (larger image but simpler deployment)
# Using cu121 for compatibility with CUDA driver 12.4
# Note: numpy<2.0 required for imgaug compatibility
RUN pip install --no-cache-dir \
    "numpy>=1.26.0,<2.0.0" \
    torch==2.5.1 \
    torchvision==0.20.1 \
    --index-url https://download.pytorch.org/whl/cu121

RUN pip install --no-cache-dir \
    "numpy>=1.26.0,<2.0.0" \
    tensorflow==2.20.0 \
    tf_keras==2.20.1 \
    "opencv-python-headless<4.11" \
    Pillow>=8.0.0 \
    PyYAML>=5.4.0 \
    pyzmq>=22.0.0 \
    msgpack>=1.0.0 \
    dill \
    scipy \
    imgaug \
    scikit-image \
    zhon \
    nltk \
    distance \
    anyconfig \
    munch \
    tensorboardX \
    python-dotenv \
    pyclipper \
    shapely \
    polygon3 \
    Polygon3 \
    tqdm \
    fs \
    einops \
    lmdb

ENV TF_USE_LEGACY_KERAS=1

# Source code should be mounted as volumes at runtime:
# -v /path/to/deep-starry:/app/deep-starry:ro
# -v /path/to/starry-ocr:/app/starry-ocr:ro
ENV PYTHONPATH=/app/deep-starry:/app/starry-ocr

# Default working directory
WORKDIR /app

CMD ["python", "--help"]