"""
PyTorch backend implementation for GPU-accelerated spin wave calculations.
"""

import numpy as np
from typing import Tuple, Optional, Union
from .base import BaseBackend, BackendError

try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    torch = None

# Placeholder class if torch not available
if not TORCH_AVAILABLE:
    class PyTorchBackend:
        """Placeholder when PyTorch is not installed."""
        def __init__(self, *args, **kwargs):
            raise ImportError("PyTorch not installed. Install with: pip install torch")
else:
    class PyTorchBackend(BaseBackend):
        """
        PyTorch backend for GPU-accelerated matrix operations.
        
        Parameters
        ----------
        device : str
            Device to use ('cuda', 'cpu', or 'mps')
        dtype : torch.dtype
            Default floating point type (default: torch.float64)
        """
        
        def __init__(self, device: str = None, dtype = None):
            if dtype is None:
                dtype = torch.float64
                
            if device is None:
                if torch.cuda.is_available():
                    device = 'cuda'
                elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                    device = 'mps'
                else:
                    device = 'cpu'
            
            self.device = torch.device(device)
            self.dtype = dtype
            self.complex_dtype = torch.complex128 if dtype == torch.float64 else torch.complex64
        
        @property
        def name(self) -> str:
            return f"pytorch ({self.device})"
        
        def to_array(self, data, dtype=None):
            if dtype is None:
                dtype = self.dtype
            if isinstance(data, torch.Tensor):
                return data.to(device=self.device, dtype=dtype)
            return torch.tensor(data, device=self.device, dtype=dtype)
        
        def to_complex(self, data):
            """Convert to complex tensor."""
            if isinstance(data, torch.Tensor):
                if data.is_complex():
                    return data.to(self.device)
                return data.to(dtype=self.complex_dtype, device=self.device)
            return torch.tensor(data, dtype=self.complex_dtype, device=self.device)
        
        def to_numpy(self, data) -> np.ndarray:
            if isinstance(data, torch.Tensor):
                return data.detach().cpu().numpy()
            return np.asarray(data)
        
        def zeros(self, shape, dtype=None):
            dtype = dtype or self.dtype
            return torch.zeros(shape, dtype=dtype, device=self.device)
        
        def zeros_complex(self, shape):
            return torch.zeros(shape, dtype=self.complex_dtype, device=self.device)
        
        def ones(self, shape, dtype=None):
            dtype = dtype or self.dtype
            return torch.ones(shape, dtype=dtype, device=self.device)
        
        def eye(self, n: int, dtype=None):
            dtype = dtype or self.dtype
            return torch.eye(n, dtype=dtype, device=self.device)
        
        def eye_complex(self, n: int):
            return torch.eye(n, dtype=self.complex_dtype, device=self.device)
        
        def arange(self, start, stop=None, step=1, dtype=None):
            dtype = dtype or self.dtype
            if stop is None:
                return torch.arange(start, step=step, dtype=dtype, device=self.device)
            return torch.arange(start, stop, step=step, dtype=dtype, device=self.device)
        
        def linspace(self, start, stop, num: int):
            return torch.linspace(start, stop, num, dtype=self.dtype, device=self.device)
        
        def concatenate(self, arrays, axis=0):
            return torch.cat(arrays, dim=axis)
        
        def stack(self, arrays, axis=0):
            return torch.stack(arrays, dim=axis)
        
        def reshape(self, a, shape):
            return a.reshape(shape)
        
        def transpose(self, a, axes=None):
            if axes is None:
                return a.T
            return a.permute(axes)
        
        # Linear algebra
        def matmul(self, a, b):
            return torch.matmul(a, b)
        
        def eigh(self, a) -> Tuple:
            return torch.linalg.eigh(a)
        
        def eig(self, a) -> Tuple:
            return torch.linalg.eig(a)
        
        def inv(self, a):
            return torch.linalg.inv(a)
        
        def det(self, a):
            return torch.linalg.det(a)
        
        def svd(self, a, full_matrices=True) -> Tuple:
            return torch.linalg.svd(a, full_matrices=full_matrices)
        
        def norm(self, a, ord=None, axis=None):
            return torch.linalg.norm(a, ord=ord, dim=axis)
        
        def cross(self, a, b, axis=-1):
            return torch.cross(a, b, dim=axis)
        
        def dot(self, a, b):
            return torch.dot(a.flatten(), b.flatten()) if a.dim() == 1 else torch.tensordot(a, b, dims=1)
        
        def einsum(self, subscripts, *operands):
            return torch.einsum(subscripts, *operands)
        
        # Element-wise operations
        def exp(self, a):
            return torch.exp(a)
        
        def sin(self, a):
            return torch.sin(a)
        
        def cos(self, a):
            return torch.cos(a)
        
        def sqrt(self, a):
            return torch.sqrt(a)
        
        def conj(self, a):
            return torch.conj(a)
        
        def real(self, a):
            return torch.real(a)
        
        def imag(self, a):
            return torch.imag(a)
        
        def abs(self, a):
            return torch.abs(a)
        
        def sign(self, a):
            return torch.sign(a)
        
        # Reductions
        def sum(self, a, axis=None, keepdims=False):
            if axis is None:
                return a.sum()
            return a.sum(dim=axis, keepdim=keepdims)
        
        def mean(self, a, axis=None, keepdims=False):
            if axis is None:
                return a.mean()
            return a.mean(dim=axis, keepdim=keepdims)
        
        def max(self, a, axis=None, keepdims=False):
            if axis is None:
                return a.max()
            return a.max(dim=axis, keepdim=keepdims).values
        
        def min(self, a, axis=None, keepdims=False):
            if axis is None:
                return a.min()
            return a.min(dim=axis, keepdim=keepdims).values
        
        def cholesky(self, a):
            """Cholesky decomposition."""
            return torch.linalg.cholesky(a)
        
        def bogoliubov_diagonalize(self, H, n_modes: int) -> Tuple:
            """
            Standard Bogoliubov diagonalization for positive definite Hamiltonians.
            """
            # Extract blocks
            A = H[:n_modes, :n_modes]
            B = H[:n_modes, n_modes:]
            
            # Form the matrix to diagonalize
            K = self.matmul(A - B, A + B)
            
            # Eigendecomposition
            eigenvalues, eigenvectors = self.eigh(K)
            
            # Energies are square roots
            omega = self.sqrt(self.abs(eigenvalues))
            
            return omega, eigenvectors
        
        def colpa_diagonalize(self, H, n_modes: int) -> Tuple:
            """
            Colpa's method for para-unitary diagonalization.
            
            This solves the generalized eigenvalue problem for bosonic systems
            following J.H.P. Colpa, Physica 93A (1978) 327.
            """
            n2 = 2 * n_modes
            
            # Build the para-unit matrix g = diag(I, -I)
            g = self.zeros_complex((n2, n2))
            g[:n_modes, :n_modes] = self.eye_complex(n_modes)
            g[n_modes:, n_modes:] = -self.eye_complex(n_modes)
            
            # Make H Hermitian if not already
            H_herm = (H + self.conj(H.T)) / 2
            
            # Form the matrix K = g @ H
            K = self.matmul(g, H_herm)
            
            # Eigenvalue decomposition
            eigenvalues, eigenvectors = self.eig(K)
            
            # Sort by real part of eigenvalues
            idx = torch.argsort(self.real(eigenvalues), descending=True)
            eigenvalues = eigenvalues[idx]
            eigenvectors = eigenvectors[:, idx]
            
            # Extract positive eigenvalues (spin wave energies)
            omega = self.real(eigenvalues[:n_modes])
            
            # Normalize transformation matrix for para-unitarity
            T = eigenvectors
            
            return omega, T
        
        def generalized_bogoliubov(self, H, n_modes: int, hermit: bool = True) -> Tuple:
            """
            General Bogoliubov transformation for spin wave Hamiltonians.
            """
            # Handle single q-point case
            single_q = H.dim() == 2
            if single_q:
                H = H.unsqueeze(0)
            
            n_q = H.shape[0]
            n2 = 2 * n_modes
            
            omega_all = []
            V_all = []
            
            for iq in range(n_q):
                Hq = H[iq]
                
                if hermit:
                    omega, V = self.colpa_diagonalize(Hq, n_modes)
                else:
                    eigenvalues, V = self.eig(Hq)
                    omega = self.real(eigenvalues[:n_modes])
                
                omega_all.append(omega)
                V_all.append(V)
            
            omega_out = self.stack(omega_all, axis=0)
            V_out = self.stack(V_all, axis=0)
            
            if single_q:
                omega_out = omega_out.squeeze(0)
                V_out = V_out.squeeze(0)
            
            return omega_out, V_out
