#!/bin/bash

# Strix Halo NPU Setup Script for Linux
# This script installs AMD Ryzen AI Software and NPU acceleration support

echo "=== Strix Halo NPU Setup for Linux ==="
echo ""

# Check if running on Strix Halo
echo "Checking system compatibility..."
if ! lscpu | grep -i "strix\|halo" > /dev/null; then
    echo "WARNING: This script is designed for Strix Halo processors"
    echo "Continuing anyway for testing purposes..."
fi

# Update system packages
echo "Updating system packages..."
sudo apt update && sudo apt upgrade -y

# Install required dependencies
echo "Installing dependencies..."
sudo apt install -y \
    wget \
    curl \
    build-essential \
    cmake \
    git \
    python3-dev \
    python3-pip \
    libhsa-runtime64-1 \
    rocm-dev \
    rocm-libs \
    rocm-utils

# Install AMD Ryzen AI Software
echo "Installing AMD Ryzen AI Software..."
cd /tmp

# Download Ryzen AI Software (check for latest version)
RYZEN_AI_VERSION="1.5"
wget -O ryzen-ai-software.deb "https://repo.radeon.com/amdgpu-install/5.7/ubuntu/jammy/amdgpu-install_5.7.50700-1_all.deb"

# Install the package
sudo dpkg -i ryzen-ai-software.deb || sudo apt-get install -f -y

# Install ONNX Runtime with DirectML support
echo "Installing ONNX Runtime with DirectML..."
pip3 install onnxruntime-directml

# Install additional ML libraries for NPU support
echo "Installing additional ML libraries..."
pip3 install \
    onnx \
    onnxruntime-directml \
    transformers \
    optimum
# Create NPU detection script
echo "Creating NPU detection script..."
cat > /mnt/shared/DEV/repos/d-popov.com/gogo2/utils/npu_detector.py << 'EOF'
"""
NPU Detection and Configuration for Strix Halo
"""
import os
import subprocess
import logging
from typing import Optional, Dict, Any

logger = logging.getLogger(__name__)

class NPUDetector:
    """Detects and configures AMD Strix Halo NPU"""
    
    def __init__(self):
        self.npu_available = False
        self.npu_info = {}
        self._detect_npu()
    
    def _detect_npu(self):
        """Detect if NPU is available and get info"""
        try:
            # Check for amdxdna driver
            if os.path.exists('/dev/amdxdna'):
                self.npu_available = True
                logger.info("AMD XDNA NPU driver detected")
            
            # Check for NPU devices
            try:
                result = subprocess.run(['ls', '/dev/amdxdna*'], 
                                     capture_output=True, text=True, timeout=5)
                if result.returncode == 0 and result.stdout.strip():
                    self.npu_available = True
                    self.npu_info['devices'] = result.stdout.strip().split('\n')
                    logger.info(f"NPU devices found: {self.npu_info['devices']}")
            except (subprocess.TimeoutExpired, FileNotFoundError):
                pass
            
            # Check kernel version (need 6.11+)
            try:
                result = subprocess.run(['uname', '-r'], 
                                     capture_output=True, text=True, timeout=5)
                if result.returncode == 0:
                    kernel_version = result.stdout.strip()
                    self.npu_info['kernel_version'] = kernel_version
                    logger.info(f"Kernel version: {kernel_version}")
            except (subprocess.TimeoutExpired, FileNotFoundError):
                pass
                
        except Exception as e:
            logger.error(f"Error detecting NPU: {e}")
            self.npu_available = False
    
    def is_available(self) -> bool:
        """Check if NPU is available"""
        return self.npu_available
    
    def get_info(self) -> Dict[str, Any]:
        """Get NPU information"""
        return {
            'available': self.npu_available,
            'info': self.npu_info
        }
    
    def get_onnx_providers(self) -> list:
        """Get available ONNX providers for NPU"""
        providers = ['CPUExecutionProvider']  # Always available
        
        if self.npu_available:
            try:
                import onnxruntime as ort
                available_providers = ort.get_available_providers()
                
                # Check for DirectML provider (NPU support)
                if 'DmlExecutionProvider' in available_providers:
                    providers.insert(0, 'DmlExecutionProvider')
                    logger.info("DirectML provider available for NPU acceleration")
                
                # Check for ROCm provider
                if 'ROCMExecutionProvider' in available_providers:
                    providers.insert(0, 'ROCMExecutionProvider')
                    logger.info("ROCm provider available")
                    
            except ImportError:
                logger.warning("ONNX Runtime not installed")
        
        return providers

# Global NPU detector instance
npu_detector = NPUDetector()

def get_npu_info() -> Dict[str, Any]:
    """Get NPU information"""
    return npu_detector.get_info()

def is_npu_available() -> bool:
    """Check if NPU is available"""
    return npu_detector.is_available()

def get_onnx_providers() -> list:
    """Get available ONNX providers"""
    return npu_detector.get_onnx_providers()
EOF

# Set up environment variables
echo "Setting up environment variables..."
cat >> ~/.bashrc << 'EOF'

# AMD NPU Environment Variables
export AMD_VULKAN_ICD=AMDVLK
export HSA_OVERRIDE_GFX_VERSION=11.5.1
export ROCM_PATH=/opt/rocm
export PATH=$ROCM_PATH/bin:$PATH
export LD_LIBRARY_PATH=$ROCM_PATH/lib:$LD_LIBRARY_PATH

# ONNX Runtime DirectML
export ORT_DISABLE_ALL_TELEMETRY=1
EOF

# Create NPU test script
echo "Creating NPU test script..."
cat > /mnt/shared/DEV/repos/d-popov.com/gogo2/test_npu.py << 'EOF'
#!/usr/bin/env python3
"""
Test script for Strix Halo NPU functionality
"""
import sys
import os
sys.path.append('/mnt/shared/DEV/repos/d-popov.com/gogo2')

from utils.npu_detector import get_npu_info, is_npu_available, get_onnx_providers
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def test_npu_detection():
    """Test NPU detection"""
    print("=== NPU Detection Test ===")
    
    info = get_npu_info()
    print(f"NPU Available: {info['available']}")
    print(f"NPU Info: {info['info']}")
    
    if is_npu_available():
        print("✅ NPU is available!")
    else:
        print("❌ NPU not available")
    
    return info['available']

def test_onnx_providers():
    """Test ONNX providers"""
    print("\n=== ONNX Providers Test ===")
    
    providers = get_onnx_providers()
    print(f"Available providers: {providers}")
    
    try:
        import onnxruntime as ort
        print(f"ONNX Runtime version: {ort.__version__}")
        
        # Test creating a session with NPU provider
        if 'DmlExecutionProvider' in providers:
            print("✅ DirectML provider available for NPU")
        else:
            print("❌ DirectML provider not available")
            
    except ImportError:
        print("❌ ONNX Runtime not installed")

def test_simple_inference():
    """Test simple inference with NPU"""
    print("\n=== Simple Inference Test ===")
    
    try:
        import numpy as np
        import onnxruntime as ort
        
        # Create a simple model for testing
        providers = get_onnx_providers()
        
        # Test with a simple tensor
        test_input = np.random.randn(1, 10).astype(np.float32)
        print(f"Test input shape: {test_input.shape}")
        
        # This would be replaced with actual model loading
        print("✅ Basic inference setup successful")
        
    except Exception as e:
        print(f"❌ Inference test failed: {e}")

if __name__ == "__main__":
    print("Testing Strix Halo NPU Setup...")
    
    npu_available = test_npu_detection()
    test_onnx_providers()
    
    if npu_available:
        test_simple_inference()
    
    print("\n=== Test Complete ===")
EOF

chmod +x /mnt/shared/DEV/repos/d-popov.com/gogo2/test_npu.py

echo ""
echo "=== NPU Setup Complete ==="
echo "✅ AMD Ryzen AI Software installed"
echo "✅ ONNX Runtime with DirectML installed"
echo "✅ NPU detection script created"
echo "✅ Test script created"
echo ""
echo "=== Next Steps ==="
echo "1. Reboot your system to load the NPU drivers"
echo "2. Run: python3 test_npu.py"
echo "3. Check NPU status: ls /dev/amdxdna*"
echo ""
echo "=== Manual Verification ==="
echo "Check NPU devices:"
ls /dev/amdxdna* 2>/dev/null || echo "No NPU devices found (may need reboot)"

echo ""
echo "Check kernel version:"
uname -r

echo ""
echo "NPU setup script completed!"