#!/bin/bash

# PyTorch Wheel Installation Script
# This script installs PyTorch and all compatible dependencies from local wheels

set -e  # Exit on any error

WHEEL_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
echo "Installing PyTorch and dependencies from: $WHEEL_DIR"

# Check if we're in the correct directory
if [ ! -f "$WHEEL_DIR/torch-2.7.1-cp310-cp310-manylinux_2_28_x86_64.whl" ]; then
    echo "Error: PyTorch wheel not found in $WHEEL_DIR"
    echo "Please ensure you're running this script from the pytorch_wheel directory"
    exit 1
fi

echo "Step 1: Installing core Python packages..."
pip install "$WHEEL_DIR"/setuptools-*.whl "$WHEEL_DIR"/wheel-*.whl "$WHEEL_DIR"/pip-*.whl || true

echo "Step 2: Installing basic dependencies..."
pip install "$WHEEL_DIR"/six-*.whl "$WHEEL_DIR"/typing_extensions-*.whl "$WHEEL_DIR"/packaging-*.whl

echo "Step 3: Installing NumPy (compatible version)..."
pip install "$WHEEL_DIR"/numpy-1.26.4-*.whl

echo "Step 4: Installing core utilities..."
pip install "$WHEEL_DIR"/tqdm-*.whl "$WHEEL_DIR"/requests-*.whl "$WHEEL_DIR"/filelock-*.whl "$WHEEL_DIR"/regex-*.whl

echo "Step 5: Installing NVIDIA CUDA libraries..."
pip install "$WHEEL_DIR"/nvidia_*.whl

echo "Step 6: Installing PyTorch ecosystem..."
pip install "$WHEEL_DIR"/sympy-*.whl "$WHEEL_DIR"/mpmath-*.whl "$WHEEL_DIR"/networkx-*.whl "$WHEEL_DIR"/triton-*.whl

echo "Step 7: Installing PyTorch..."
pip install "$WHEEL_DIR"/torch-*.whl "$WHEEL_DIR"/torchaudio-*.whl

echo "Step 8: Installing HuggingFace and tokenizers..."
pip install "$WHEEL_DIR"/safetensors-*.whl "$WHEEL_DIR"/tokenizers-*.whl "$WHEEL_DIR"/huggingface_hub-*.whl "$WHEEL_DIR"/fsspec-*.whl

echo "Step 9: Installing ML libraries..."
pip install "$WHEEL_DIR"/transformers-*.whl "$WHEEL_DIR"/accelerate-*.whl "$WHEEL_DIR"/peft-*.whl "$WHEEL_DIR"/datasets-*.whl

echo "Step 10: Installing scientific libraries..."
pip install "$WHEEL_DIR"/scipy-*.whl "$WHEEL_DIR"/librosa-*.whl "$WHEEL_DIR"/opencv_python-*.whl

echo "Step 11: Installing JSON/Schema handling..."
pip install "$WHEEL_DIR"/rpds_py-*.whl "$WHEEL_DIR"/referencing-*.whl "$WHEEL_DIR"/jsonschema_specifications-*.whl "$WHEEL_DIR"/jsonschema-*.whl

echo "Step 12: Installing system utilities..."
pip install "$WHEEL_DIR"/attrs-*.whl "$WHEEL_DIR"/distro-*.whl "$WHEEL_DIR"/tomli-*.whl "$WHEEL_DIR"/platformdirs-*.whl "$WHEEL_DIR"/tornado-*.whl "$WHEEL_DIR"/monotonic-*.whl

echo "Step 13: Installing additional utilities..."
pip install "$WHEEL_DIR"/babel-*.whl "$WHEEL_DIR"/fs-*.whl "$WHEEL_DIR"/appdirs-*.whl "$WHEEL_DIR"/more_itertools-*.whl

echo "Step 14: Installing PyTorch Lightning..."
pip install "$WHEEL_DIR"/lightning_utilities-*.whl "$WHEEL_DIR"/pytorch_lightning-*.whl

echo "Step 15: Installing Jupyter support..."
pip install "$WHEEL_DIR"/nbconvert-*.whl

echo "Step 16: Installing OpenTelemetry (compatible versions)..."
pip install "$WHEEL_DIR"/opentelemetry_semantic_conventions-*.whl "$WHEEL_DIR"/opentelemetry_api-*.whl "$WHEEL_DIR"/opentelemetry_sdk-*.whl "$WHEEL_DIR"/opentelemetry_util_http-*.whl "$WHEEL_DIR"/opentelemetry_instrumentation-*.whl "$WHEEL_DIR"/opentelemetry_instrumentation_asgi-*.whl

echo "Step 17: Installing data handling..."
pip install "$WHEEL_DIR"/pydantic_settings-*.whl "$WHEEL_DIR"/dataclasses_json-*.whl "$WHEEL_DIR"/pypdf-*.whl

echo "Step 18: Installing API clients..."
pip install "$WHEEL_DIR"/httpx-*.whl "$WHEEL_DIR"/openai-*.whl

echo "Step 19: Installing CLI utilities..."
pip install "$WHEEL_DIR"/decorator-*.whl "$WHEEL_DIR"/typer-*.whl

echo "Step 20: Installing remaining packages..."
pip install "$WHEEL_DIR"/keyring-*.whl "$WHEEL_DIR"/virtualenv-*.whl "$WHEEL_DIR"/protobuf-*.whl "$WHEEL_DIR"/mako-*.whl "$WHEEL_DIR"/markdown-*.whl "$WHEEL_DIR"/lxml-*.whl "$WHEEL_DIR"/chardet-*.whl "$WHEEL_DIR"/tensorboard-*.whl

echo ""
echo "Installation completed successfully!"
echo ""
echo "Testing imports..."
python -c "
import torch
import numpy
import transformers
import accelerate
import peft
import datasets
import scipy
import librosa
print('✅ All major packages imported successfully!')
print(f'PyTorch version: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'CUDA device count: {torch.cuda.device_count()}')
"

echo ""
echo "🎉 PyTorch installation with all compatible dependencies completed!"
echo "No more dependency conflicts should occur."