#!/bin/bash
# Complete installation script for Mirel with MX format support on H200

echo "Installing Mirel dependencies for GPT-OSS with MX format support..."

# Upgrade pip first
pip install --upgrade pip

# Install main requirements WITH SPECIFIC VERSIONS for ZeroGPU compatibility
pip install huggingface_hub>=0.34.0
pip install transformers>=4.55.0
pip install accelerate>=0.33.0
pip install torch==2.4.0  # SPECIFIC VERSION for ZeroGPU - DO NOT use nightly!
pip install gradio>=5.42.0
pip install spaces

# Install LoRA/PEFT support
pip install peft>=0.11.0
pip install bitsandbytes>=0.43.1

# Install Harmony format
pip install openai-harmony

# Install standard Triton (MX will fallback to bf16)
pip install triton>=3.0.0

# Note: triton_kernels with ragged_tma is not available in stable releases
# The model will fall back to bf16 mode which is fine for inference
echo "Note: MX format requires bleeding-edge Triton features not available in stable releases."
echo "The model will use bf16 mode instead, which works fine but uses more memory."

# Optional but recommended
pip install safetensors>=0.4.0
pip install sentencepiece>=0.2.0
pip install protobuf>=3.20.0
pip install "numpy<2.0.0"

# Verify critical imports
echo "Verifying installation..."
python -c "
import sys
errors = []

try:
    import torch
    print(f'✓ PyTorch {torch.__version__}')
    # Check CUDA availability without initializing it (for ZeroGPU)
    print(f'  CUDA available: Will be checked at runtime')
except ImportError as e:
    errors.append(f'✗ PyTorch: {e}')

try:
    import transformers
    print(f'✓ Transformers {transformers.__version__}')
except ImportError as e:
    errors.append(f'✗ Transformers: {e}')

try:
    import peft
    print(f'✓ PEFT {peft.__version__}')
except ImportError as e:
    errors.append(f'✗ PEFT: {e}')

try:
    import triton
    print(f'✓ Triton {triton.__version__}')
except ImportError as e:
    errors.append(f'✗ Triton: {e}')

try:
    import openai_harmony
    print('✓ OpenAI Harmony')
except ImportError as e:
    errors.append(f'✗ OpenAI Harmony: {e}')

try:
    import gradio
    print(f'✓ Gradio {gradio.__version__}')
except ImportError as e:
    errors.append(f'✗ Gradio: {e}')

try:
    import spaces
    print('✓ Spaces (ZeroGPU support)')
except ImportError as e:
    errors.append(f'✗ Spaces: {e}')

if errors:
    print('\n❌ Installation issues found:')
    for error in errors:
        print(f'  {error}')
    sys.exit(1)
else:
    print('\n✅ All dependencies installed successfully!')
    print('Note: Model will run in bf16 mode (MX format requires unreleased Triton features)')
"

echo "Installation complete!"

# Optional but recommended
pip install safetensors>=0.4.0
pip install sentencepiece>=0.2.0
pip install protobuf>=3.20.0
pip install "numpy<2.0.0"

# Verify critical imports
echo "Verifying installation..."
python -c "
import sys
errors = []

try:
    import torch
    print(f'✓ PyTorch {torch.__version__}')
except ImportError as e:
    errors.append(f'✗ PyTorch: {e}')

try:
    import transformers
    print(f'✓ Transformers {transformers.__version__}')
except ImportError as e:
    errors.append(f'✗ Transformers: {e}')

try:
    import peft
    print(f'✓ PEFT {peft.__version__}')
except ImportError as e:
    errors.append(f'✗ PEFT: {e}')

try:
    import triton
    print(f'✓ Triton {triton.__version__}')
except ImportError as e:
    errors.append(f'✗ Triton: {e}')

try:
    import triton_kernels
    print('✓ Triton Kernels (MX format support)')
except ImportError as e:
    errors.append(f'✗ Triton Kernels (CRITICAL): {e}')
    print('⚠️  WARNING: MX format will NOT work without triton_kernels!')

try:
    import openai_harmony
    print('✓ OpenAI Harmony')
except ImportError as e:
    errors.append(f'✗ OpenAI Harmony: {e}')

try:
    import gradio
    print(f'✓ Gradio {gradio.__version__}')
except ImportError as e:
    errors.append(f'✗ Gradio: {e}')

if errors:
    print('\n❌ Installation issues found:')
    for error in errors:
        print(f'  {error}')
    sys.exit(1)
else:
    print('\n✅ All dependencies installed successfully!')
    print('You can now run the Mirel app with MX format support on H200 GPUs')
"

echo "Installation complete!"