#!/bin/bash
# DCU ENVIRONMENT Activation Script + Dependency Installation
# Run with bash `source activate-q2.sh`
# By KsuserKqy
# 2025-10-18

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
BOLD_RED='\033[1;31m'
NC='\033[0m'

if ! command -v conda &>/dev/null; then
    echo -e "[$RED x $NC] Conda not found. Please install Conda first."
    return 1 2>/dev/null || exit 1
fi

CONDA_BASE=$(conda info --base 2>/dev/null)
if [ -f "$CONDA_BASE/etc/profile.d/conda.sh" ]; then
    source "$CONDA_BASE/etc/profile.d/conda.sh"
    conda activate dcu-q2
    echo -e "[$GREEN √ $NC] Conda environment 'dcu-q2' activated."
else
    echo -e "[$RED x $NC] Could not find conda.sh for initialization."
    return 1 2>/dev/null || exit 1
fi

echo -e "[$YELLOW → $NC] Installing Python packages ..."

# Step 1: install llmcompressor without dependencies
pip install llmcompressor --no-deps -i https://pypi.tuna.tsinghua.edu.cn/simple
if [ $? -ne 0 ]; then
    echo -e "[$RED x $NC] Failed to install llmcompressor"
    return 1 2>/dev/null || exit 1
fi

# Step 2: install DCU-compatible torch
pip install https://download.sourcefind.cn:65024/file/4/pytorch/DAS1.6/torch-2.5.1+das.opt1.dtk25041-cp310-cp310-manylinux_2_28_x86_64.whl
pip install llmcompressor --no-deps
if [ $? -ne 0 ]; then
    echo -e "[$RED x $NC] Failed to install custom torch"
    return 1 2>/dev/null || exit 1
fi

# Step 3: install remaining libraries
pip install \
    modelscope \
    datasets==4.1.1 \
    pandas \
    huggingface_hub \
    multiprocess \
    httpx \
    xxhash \
    transformers==4.56.2 \
    loguru \
    accelerate \
    compressed_tensors \
    pillow==11.3.0 \
    nvidia-ml-py \
    numpy==1.26.4 \
    -i https://pypi.tuna.tsinghua.edu.cn/simple

if [ $? -eq 0 ]; then
    echo -e "[$GREEN √ $NC] All Python packages installed successfully."
else
    echo -e "[$RED x $NC] Some packages failed to install."
fi

echo -e "[$GREEN √ $NC] Environment setup completed."
