FROM nvidia/cuda:12.6.0-cudnn-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV CUDA_HOME=/usr/local/cuda
ENV PATH=${CUDA_HOME}/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}

# Install system dependencies
RUN apt-get update && apt-get install -y \
    python3.11 \
    python3.11-dev \
    python3-pip \
    git \
    wget \
    curl \
    ffmpeg \
    libavcodec-dev \
    libavformat-dev \
    libsndfile1 \
    sox \
    libsox-dev \
    openjdk-11-jdk \
    && rm -rf /var/lib/apt/lists/*

# Install Miniconda
RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh && \
    bash /tmp/miniconda.sh -b -p /opt/conda && \
    rm /tmp/miniconda.sh

ENV PATH=/opt/conda/bin:${PATH}

# Create conda environment
RUN conda create -n pvmd python=3.11 -y && \
    conda clean -afy

# Make RUN commands use the new environment
SHELL ["conda", "run", "-n", "pvmd", "/bin/bash", "-c"]

# Install PyTorch with CUDA
RUN pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

# Install ML dependencies
RUN pip install --no-cache-dir \
    transformers \
    datasets \
    evaluate \
    accelerate \
    pyspark \
    soundfile \
    librosa \
    jiwer \
    tensorboard \
    wandb \
    huggingface-hub \
    scipy \
    pandas \
    numpy

# Set working directory
WORKDIR /workspace

# Copy project files
COPY . /workspace/

# Default command
CMD ["/bin/bash"]