#!/bin/bash
# Unified activation script for all suno environments

# Function to list available environments
list_envs() {
    echo "Available Suno environments:"
    echo ""
    printf "%-20s %-15s %-12s %-12s\n" "Environment" "PyTorch" "flash_attn" "flash_attn_3"
    printf "%-20s %-15s %-12s %-12s\n" "--------------------" "---------------" "------------" "------------"

    # Get all suno* environments
    conda env list | grep -E "^suno" | awk '{print $1}' | while read env; do
        # Get versions
        result=$(conda run -n "$env" python -c "
torch_ver = 'N/A'
flash_attn_ver = 'N/A'
flash_attn_3_ver = 'N/A'

try:
    import torch
    torch_ver = torch.__version__
except:
    pass

try:
    import flash_attn
    flash_attn_ver = flash_attn.__version__
except:
    pass

try:
    import flash_attn_interface
    flash_attn_3_ver = 'YES'
except:
    pass

print(f'{torch_ver}|{flash_attn_ver}|{flash_attn_3_ver}')
" 2>/dev/null || echo "N/A|N/A|N/A")

        IFS='|' read -r torch_ver fa_ver fa3_ver <<< "$result"
        printf "%-20s %-15s %-12s %-12s\n" "$env" "$torch_ver" "$fa_ver" "$fa3_ver"
    done
    echo ""
}

# Check if an environment name was provided
if [ -z "$1" ]; then
    echo "Usage: source activate_suno.sh <env_name>"
    echo ""
    list_envs
    return 1
fi

ENV_NAME="$1"

# Source conda
source "$(conda info --base)/etc/profile.d/conda.sh"

# Check if environment exists
if ! conda env list | grep -q "^${ENV_NAME} "; then
    echo "Error: Environment '${ENV_NAME}' not found"
    echo ""
    list_envs
    return 1
fi

# Activate environment
conda activate "${ENV_NAME}"

# Set environment variables
export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}
export OMP_NUM_THREADS=1
export TRITON_CACHE_DIR=/mnt/localdisk/.triton_cache_$USER

echo "Environment ${ENV_NAME} activated!"
echo "Python: $(which python)"
echo "PyTorch: $(python -c 'import torch; print(torch.__version__)' 2>/dev/null || echo 'Not available')"
echo "CUDA available: $(python -c 'import torch; print(torch.cuda.is_available())' 2>/dev/null || echo 'Unknown')"

# Check for flash attention
FA_VERSION=$(python -c 'import flash_attn; print(flash_attn.__version__)' 2>/dev/null || echo 'N/A')
if [ "$FA_VERSION" != "N/A" ]; then
    echo "Flash Attention: $FA_VERSION"
fi

FA3_CHECK=$(python -c 'import flash_attn_interface; print("YES")' 2>/dev/null || echo 'N/A')
if [ "$FA3_CHECK" = "YES" ]; then
    echo "Flash Attention v3: Available"
fi