"""
NumPy backend with optional Numba JIT compilation for CPU-optimized calculations.
"""

import numpy as np
from scipy import linalg as la
from typing import Tuple, Optional
from .base import BaseBackend

# Try to import numba
try:
    from numba import jit, prange
    NUMBA_AVAILABLE = True
except ImportError:
    NUMBA_AVAILABLE = False
    # Dummy decorator
    def jit(*args, **kwargs):
        def decorator(func):
            return func
        return decorator
    prange = range


class NumPyBackend(BaseBackend):
    """
    NumPy backend with optional Numba JIT acceleration.
    
    Parameters
    ----------
    use_numba : bool
        Whether to use Numba JIT compilation for hot loops
    """
    
    def __init__(self, use_numba: bool = True):
        self.use_numba = use_numba and NUMBA_AVAILABLE
        self.dtype = np.float64
        self.complex_dtype = np.complex128
    
    @property
    def name(self) -> str:
        suffix = " + numba" if self.use_numba else ""
        return f"numpy{suffix}"
    
    def to_array(self, data, dtype=None):
        dtype = dtype or self.dtype
        return np.asarray(data, dtype=dtype)
    
    def to_complex(self, data):
        return np.asarray(data, dtype=self.complex_dtype)
    
    def to_numpy(self, data) -> np.ndarray:
        return np.asarray(data)
    
    def zeros(self, shape, dtype=None):
        return np.zeros(shape, dtype=dtype or self.dtype)
    
    def zeros_complex(self, shape):
        return np.zeros(shape, dtype=self.complex_dtype)
    
    def ones(self, shape, dtype=None):
        return np.ones(shape, dtype=dtype or self.dtype)
    
    def eye(self, n: int, dtype=None):
        return np.eye(n, dtype=dtype or self.dtype)
    
    def eye_complex(self, n: int):
        return np.eye(n, dtype=self.complex_dtype)
    
    def arange(self, start, stop=None, step=1, dtype=None):
        dtype = dtype or self.dtype
        if stop is None:
            return np.arange(start, step=step, dtype=dtype)
        return np.arange(start, stop, step, dtype=dtype)
    
    def linspace(self, start, stop, num: int):
        return np.linspace(start, stop, num, dtype=self.dtype)
    
    def concatenate(self, arrays, axis=0):
        return np.concatenate(arrays, axis=axis)
    
    def stack(self, arrays, axis=0):
        return np.stack(arrays, axis=axis)
    
    def reshape(self, a, shape):
        return np.reshape(a, shape)
    
    def transpose(self, a, axes=None):
        return np.transpose(a, axes)
    
    # Linear algebra
    def matmul(self, a, b):
        return np.matmul(a, b)
    
    def eigh(self, a) -> Tuple:
        return np.linalg.eigh(a)
    
    def eig(self, a) -> Tuple:
        return np.linalg.eig(a)
    
    def inv(self, a):
        return np.linalg.inv(a)
    
    def det(self, a):
        return np.linalg.det(a)
    
    def svd(self, a, full_matrices=True) -> Tuple:
        return np.linalg.svd(a, full_matrices=full_matrices)
    
    def norm(self, a, ord=None, axis=None):
        return np.linalg.norm(a, ord=ord, axis=axis)
    
    def cross(self, a, b, axis=-1):
        return np.cross(a, b, axis=axis)
    
    def dot(self, a, b):
        return np.dot(a, b)
    
    def einsum(self, subscripts, *operands):
        return np.einsum(subscripts, *operands)
    
    # Element-wise operations
    def exp(self, a):
        return np.exp(a)
    
    def sin(self, a):
        return np.sin(a)
    
    def cos(self, a):
        return np.cos(a)
    
    def sqrt(self, a):
        return np.sqrt(a)
    
    def conj(self, a):
        return np.conj(a)
    
    def real(self, a):
        return np.real(a)
    
    def imag(self, a):
        return np.imag(a)
    
    def abs(self, a):
        return np.abs(a)
    
    def sign(self, a):
        return np.sign(a)
    
    # Reductions
    def sum(self, a, axis=None, keepdims=False):
        return np.sum(a, axis=axis, keepdims=keepdims)
    
    def mean(self, a, axis=None, keepdims=False):
        return np.mean(a, axis=axis, keepdims=keepdims)
    
    def max(self, a, axis=None, keepdims=False):
        return np.max(a, axis=axis, keepdims=keepdims)
    
    def min(self, a, axis=None, keepdims=False):
        return np.min(a, axis=axis, keepdims=keepdims)
    
    def cholesky(self, a):
        return np.linalg.cholesky(a)
    
    def bogoliubov_diagonalize(self, H, n_modes: int) -> Tuple:
        """Standard Bogoliubov diagonalization."""
        A = H[:n_modes, :n_modes]
        B = H[:n_modes, n_modes:]
        
        K = np.matmul(A - B, A + B)
        eigenvalues, eigenvectors = np.linalg.eigh(K)
        omega = np.sqrt(np.abs(eigenvalues))
        
        return omega, eigenvectors
    
    def colpa_diagonalize(self, H, n_modes: int) -> Tuple:
        """
        Colpa's method for para-unitary diagonalization.
        
        Uses scipy.linalg for better numerical stability.
        """
        n2 = 2 * n_modes
        
        # Build para-unit matrix g
        g = np.zeros((n2, n2), dtype=self.complex_dtype)
        g[:n_modes, :n_modes] = np.eye(n_modes, dtype=self.complex_dtype)
        g[n_modes:, n_modes:] = -np.eye(n_modes, dtype=self.complex_dtype)
        
        # Hermitize
        H_herm = (H + np.conj(H.T)) / 2
        
        # Form K = g @ H
        K = np.matmul(g, H_herm)
        
        # Eigenvalue decomposition
        eigenvalues, eigenvectors = la.eig(K)
        
        # Sort by real part
        idx = np.argsort(np.real(eigenvalues))[::-1]
        eigenvalues = eigenvalues[idx]
        eigenvectors = eigenvectors[:, idx]
        
        omega = np.real(eigenvalues[:n_modes])
        T = eigenvectors
        
        return omega, T
    
    def generalized_bogoliubov(self, H, n_modes: int, hermit: bool = True) -> Tuple:
        """General Bogoliubov transformation for spin wave Hamiltonians."""
        single_q = H.ndim == 2
        if single_q:
            H = H[np.newaxis, ...]
        
        n_q = H.shape[0]
        
        omega_all = []
        V_all = []
        
        if self.use_numba:
            omega_all, V_all = _bogoliubov_loop_numba(H, n_modes, hermit)
        else:
            for iq in range(n_q):
                Hq = H[iq]
                if hermit:
                    omega, V = self.colpa_diagonalize(Hq, n_modes)
                else:
                    eigenvalues, V = np.linalg.eig(Hq)
                    omega = np.real(eigenvalues[:n_modes])
                omega_all.append(omega)
                V_all.append(V)
            omega_all = np.stack(omega_all, axis=0)
            V_all = np.stack(V_all, axis=0)
        
        if single_q:
            omega_all = omega_all[0]
            V_all = V_all[0]
        
        return omega_all, V_all


# Numba-accelerated loops
if NUMBA_AVAILABLE:
    @jit(nopython=True, parallel=True)
    def _bogoliubov_loop_numba(H, n_modes, hermit):
        """JIT-compiled loop over Q-points for Bogoliubov diagonalization."""
        n_q = H.shape[0]
        n2 = 2 * n_modes
        
        omega_all = np.zeros((n_q, n_modes), dtype=np.float64)
        V_all = np.zeros((n_q, n2, n2), dtype=np.complex128)
        
        for iq in prange(n_q):
            # Note: numpy linalg calls inside numba are limited
            # This is a simplified version - actual implementation
            # would need to call scipy outside numba
            pass
        
        return omega_all, V_all
else:
    def _bogoliubov_loop_numba(H, n_modes, hermit):
        """Fallback without numba."""
        return None, None
