#!/bin/bash
set -e

echo "🔧 Starting environment setup..."

# Miniconda install path
CONDA_DIR=/opt/conda

# Download and install Miniconda to /opt/conda
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh
bash /tmp/miniconda.sh -b -u -p $CONDA_DIR
rm /tmp/miniconda.sh
export PATH="$CONDA_DIR/bin:$PATH"

# Create environment
conda update -y -n base -c defaults conda
conda env create -f /cosmos-predict1.yaml
source $CONDA_DIR/bin/activate cosmos-predict1

# Install Python dependencies
pip install --no-cache-dir -r /requirements.txt
pip install transformer-engine[pytorch]==1.12.0
pip install git+https://github.com/microsoft/MoGe.git

# Build and install Apex
git clone https://github.com/NVIDIA/apex.git
cd apex
CUDA_HOME=$CONDA_DIR pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \
  --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" .
cd ..
rm -rf apex

echo "✅ Environment setup complete."