# CASCADE Hyperlattice - GPU Docker Image
# 🔮 GLASS BOX AI CAPSULE - MAXIMUM POWER MODE
# For heavy GPU compute on Hugging Face Spaces
#
# Uses NVIDIA CUDA base for GPU acceleration
# Includes: PyTorch, Transformers, DreamerV3 RSSM

FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04

# Set environment
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1

WORKDIR /app

# System dependencies
RUN apt-get update && apt-get install -y \
    build-essential \
    curl \
    git \
    libgl1-mesa-glx \
    libglib2.0-0 \
    python3.11 \
    python3-pip \
    python3.11-venv \
    && rm -rf /var/lib/apt/lists/*

# Use python3.11 as default
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1 \
    && update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1

# Upgrade pip
RUN pip install --no-cache-dir --upgrade pip setuptools wheel

# Install PyTorch with CUDA support
RUN pip install --no-cache-dir \
    torch>=2.1.0 \
    torchvision>=0.16.0 \
    --index-url https://download.pytorch.org/whl/cu121

# Install core dependencies
COPY requirements.txt ./
RUN pip install --no-cache-dir -r requirements.txt

# Install GPU-specific packages
RUN pip install --no-cache-dir \
    transformers>=4.36.0 \
    accelerate>=0.25.0 \
    bitsandbytes>=0.41.0 \
    peft>=0.7.0 \
    jax[cuda12_pip]>=0.4.20 \
    flax>=0.7.5

# Copy source code
COPY src/ ./src/

# Expose ports
# 8501 - Streamlit main app
# 9090 - Rerun web viewer for 4D viz
# 7860 - Gradio (optional)
EXPOSE 8501
EXPOSE 9090
EXPOSE 7860

# Health check
HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health

# Environment for GPU
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility

# Run with GPU support
CMD ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0", "--server.enableCORS=false", "--server.enableXsrfProtection=false"]