FROM python:3.10-slim

# CRITICAL: Set legacy Keras before any TensorFlow imports
# Required for MEGNet compatibility with TensorFlow 2.16+
ENV TF_USE_LEGACY_KERAS=1

WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y \
    build-essential \
    libgomp1 \
    git \
    && rm -rf /var/lib/apt/lists/*

# CRITICAL: Install CPU-only versions to save space and avoid GPU conflicts
# Install PyTorch CPU first (from special index)
RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu

# Install orb-models (required for ORBFeaturizer)
RUN pip install --no-cache-dir orb-models

# Install MatterVial and its dependencies
# Then pin TensorFlow to 2.15.x (last version compatible with MEGNet)
RUN pip install --no-cache-dir \
    "mattervial @ git+https://github.com/rogeriog/MatterVial.git" \
    pymatgen \
    scikit-learn \
    pandas \
    fastapi \
    uvicorn \
    python-multipart

# CRITICAL: Force TensorFlow 2.15 and compatible Keras AFTER all other installs
# TensorFlow 2.16+ uses Keras 3.x which breaks MEGNet's Trainer.compile()
# tf_keras is required when TF_USE_LEGACY_KERAS=1
RUN pip install --no-cache-dir --force-reinstall \
    "tensorflow-cpu>=2.15,<2.16" \
    "keras>=2.15,<2.16" \
    "tf_keras>=2.15,<2.16" \
    "tensorboard>=2.15,<2.16" \
    "ml-dtypes>=0.3.1,<0.4"

COPY . .

# HF Spaces requires port 7860
EXPOSE 7860

# Launch with generous timeout for model loading
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860", "--timeout-keep-alive", "120"]