"""
PySpinW Backends - Computational backends for matrix operations.

Supports:
- PyTorch (default if available, GPU-accelerated)
- NumPy + Numba (CPU-optimized with JIT)
- JAX (experimental, for TPU/GPU)
"""

from .numpy_backend import NumPyBackend

# Try to import PyTorch directly to check availability
PYTORCH_AVAILABLE = False
try:
    import torch
    PYTORCH_AVAILABLE = True
    from .pytorch_backend import PyTorchBackend
except ImportError:
    PyTorchBackend = None

# Default backend
_current_backend = None

def get_backend(name: str = 'auto'):
    """Get a computational backend by name."""
    global _current_backend
    
    if name == 'auto':
        # Prefer PyTorch if available, otherwise NumPy
        if PYTORCH_AVAILABLE:
            name = 'pytorch'
        else:
            name = 'numpy'
    
    if name == 'pytorch':
        if not PYTORCH_AVAILABLE:
            raise ImportError("PyTorch not installed. Install with: pip install torch")
        _current_backend = PyTorchBackend()
    elif name == 'numpy' or name == 'numba':
        _current_backend = NumPyBackend(use_numba=(name == 'numba'))
    elif name == 'jax':
        try:
            from .jax_backend import JAXBackend
            _current_backend = JAXBackend()
        except ImportError:
            raise ImportError("JAX not installed. Install with: pip install jax jaxlib")
    else:
        raise ValueError(f"Unknown backend: {name}")
    
    return _current_backend

def current_backend():
    """Get the current backend (auto-selects best available)."""
    global _current_backend
    if _current_backend is None:
        _current_backend = get_backend('auto')
    return _current_backend

__all__ = ['get_backend', 'current_backend', 'NumPyBackend', 'PYTORCH_AVAILABLE']

if PYTORCH_AVAILABLE:
    __all__.append('PyTorchBackend')
