"""
Utility functions for BioGSP package.

Includes graph construction, Laplacian computation, eigendecomposition,
Graph Fourier Transform, and other helper functions.
Uses PyTorch for accelerated matrix operations.
"""

import numpy as np
import torch
from scipy.sparse import csr_matrix, diags, eye as sparse_eye
from scipy.sparse.linalg import eigsh
from sklearn.neighbors import NearestNeighbors
from typing import Tuple, Optional, Union, List, Dict, Any


def get_device(use_mps: bool = False) -> torch.device:
    """
    Get the best available device (CUDA, MPS, or CPU).
    
    Note: MPS doesn't support float64, so it's disabled by default.
    """
    if torch.cuda.is_available():
        return torch.device("cuda")
    elif use_mps and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def cal_laplacian(
    W: Union[np.ndarray, csr_matrix],
    laplacian_type: str = "normalized"
) -> csr_matrix:
    """
    Calculate Graph Laplacian Matrix.
    
    Compute unnormalized, normalized, or random-walk Laplacian from an adjacency matrix.
    
    Parameters
    ----------
    W : np.ndarray or scipy.sparse matrix
        Square adjacency matrix (can be dense or sparse).
    laplacian_type : str
        Type of Laplacian to compute: "unnormalized", "normalized", or "randomwalk".
        
    Returns
    -------
    scipy.sparse.csr_matrix
        Laplacian matrix.
        
    Examples
    --------
    >>> W = np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]])
    >>> L = cal_laplacian(W, laplacian_type="normalized")
    """
    if laplacian_type not in ["unnormalized", "normalized", "randomwalk"]:
        raise ValueError("laplacian_type must be 'unnormalized', 'normalized', or 'randomwalk'")
    
    # Convert to sparse if needed
    if not isinstance(W, csr_matrix):
        W = csr_matrix(W)
    
    n = W.shape[0]
    
    # Compute row degrees
    deg = np.array(W.sum(axis=1)).flatten()
    deg[deg == 0] = 1  # Avoid division by zero
    
    # Create diagonal matrices
    D = diags(deg, format="csr")
    
    if laplacian_type == "unnormalized":
        L = D - W
    elif laplacian_type == "normalized":
        D_half_inv = diags(1.0 / np.sqrt(deg), format="csr")
        I = sparse_eye(n, format="csr")
        L = I - D_half_inv @ W @ D_half_inv
    else:  # randomwalk
        D_inv = diags(1.0 / deg, format="csr")
        I = sparse_eye(n, format="csr")
        L = I - D_inv @ W
    
    return L


def fast_decomposition_lap(
    laplacian_mat: Union[np.ndarray, csr_matrix],
    k_eigen: int = 25,
    which: str = "SM",
    sigma: Optional[float] = None,
    use_torch: bool = True,
    device: Optional[torch.device] = None
) -> Dict[str, np.ndarray]:
    """
    Fast eigendecomposition of Laplacian matrix using PyTorch or SciPy.
    
    Parameters
    ----------
    laplacian_mat : np.ndarray or scipy.sparse matrix
        Laplacian matrix.
    k_eigen : int
        Number of eigenvalues to compute.
    which : str
        Which eigenvalues to compute ("SM" for smallest, "LM" for largest).
    sigma : float, optional
        Shift parameter for eigenvalue computation.
    use_torch : bool
        Whether to use PyTorch for eigendecomposition (faster for dense matrices).
    device : torch.device, optional
        Device to use for PyTorch computations.
        
    Returns
    -------
    dict
        Dictionary with 'evalues' (eigenvalues) and 'evectors' (eigenvectors).
        
    Examples
    --------
    >>> L = np.array([[2, -1, -1], [-1, 2, -1], [-1, -1, 2]])
    >>> decomp = fast_decomposition_lap(L, k_eigen=2)
    """
    if device is None:
        device = get_device()
    
    n = laplacian_mat.shape[0]
    k_eigen = min(k_eigen, n - 1)  # Ensure k_eigen doesn't exceed matrix size
    
    # For sparse matrices or when k_eigen is small relative to n, use SciPy
    if isinstance(laplacian_mat, csr_matrix) or k_eigen < n // 2:
        if not isinstance(laplacian_mat, csr_matrix):
            laplacian_mat = csr_matrix(laplacian_mat)
        
        try:
            eigenvalues, eigenvectors = eigsh(
                laplacian_mat, 
                k=k_eigen, 
                which=which,
                sigma=sigma,
                return_eigenvectors=True
            )
        except Exception:
            # Fall back to dense computation if sparse fails
            L_dense = laplacian_mat.toarray() if isinstance(laplacian_mat, csr_matrix) else laplacian_mat
            eigenvalues, eigenvectors = np.linalg.eigh(L_dense)
            if which == "SM":
                idx = np.argsort(eigenvalues)[:k_eigen]
            else:
                idx = np.argsort(eigenvalues)[-k_eigen:][::-1]
            eigenvalues = eigenvalues[idx]
            eigenvectors = eigenvectors[:, idx]
    elif use_torch:
        # Use PyTorch for dense eigendecomposition
        L_dense = laplacian_mat.toarray() if isinstance(laplacian_mat, csr_matrix) else laplacian_mat
        # Use float32 if on MPS (doesn't support float64)
        dtype = torch.float32 if device.type == "mps" else torch.float64
        L_tensor = torch.tensor(L_dense, dtype=dtype, device=device)
        
        # PyTorch eigendecomposition
        eigenvalues_all, eigenvectors_all = torch.linalg.eigh(L_tensor)
        
        # Select eigenvalues based on which
        if which == "SM":
            idx = torch.argsort(eigenvalues_all)[:k_eigen]
        else:  # LM
            idx = torch.argsort(eigenvalues_all)[-k_eigen:].flip(0)
        
        eigenvalues = eigenvalues_all[idx].cpu().numpy()
        eigenvectors = eigenvectors_all[:, idx].cpu().numpy()
    else:
        # Use NumPy
        L_dense = laplacian_mat.toarray() if isinstance(laplacian_mat, csr_matrix) else laplacian_mat
        eigenvalues_all, eigenvectors_all = np.linalg.eigh(L_dense)
        
        if which == "SM":
            idx = np.argsort(eigenvalues_all)[:k_eigen]
        else:
            idx = np.argsort(eigenvalues_all)[-k_eigen:][::-1]
        
        eigenvalues = eigenvalues_all[idx]
        eigenvectors = eigenvectors_all[:, idx]
    
    # Ensure eigenvalues are sorted in ascending order
    sort_idx = np.argsort(eigenvalues)
    eigenvalues = eigenvalues[sort_idx]
    eigenvectors = eigenvectors[:, sort_idx]
    
    return {"evalues": eigenvalues, "evectors": eigenvectors}


def gft(
    signal: Union[np.ndarray, torch.Tensor],
    U: Union[np.ndarray, torch.Tensor],
    use_torch: bool = True,
    device: Optional[torch.device] = None
) -> np.ndarray:
    """
    Graph Fourier Transform.
    
    Compute the Graph Fourier Transform (GFT) of a signal using Laplacian eigenvectors.
    
    Parameters
    ----------
    signal : np.ndarray or torch.Tensor
        Input signal (vector or matrix where each column is a signal).
    U : np.ndarray or torch.Tensor
        Matrix of eigenvectors (n_vertices x k_eigen).
    use_torch : bool
        Whether to use PyTorch for computation.
    device : torch.device, optional
        Device to use for PyTorch computations.
        
    Returns
    -------
    np.ndarray
        Transformed signal in the spectral domain.
        
    Examples
    --------
    >>> signal = np.random.randn(100)
    >>> U = np.random.randn(100, 50)  # eigenvectors
    >>> signal_hat = gft(signal, U)
    """
    if device is None:
        device = get_device()
    
    # Convert to numpy if needed
    if isinstance(signal, torch.Tensor):
        signal = signal.cpu().numpy()
    if isinstance(U, torch.Tensor):
        U = U.cpu().numpy()
    
    # Ensure signal is 2D
    if signal.ndim == 1:
        signal = signal.reshape(-1, 1)
    
    if use_torch:
        # Use PyTorch for matrix multiplication
        # Use float32 if on MPS (doesn't support float64)
        dtype = torch.float32 if device.type == "mps" else torch.float64
        signal_tensor = torch.tensor(signal, dtype=dtype, device=device)
        U_tensor = torch.tensor(U, dtype=dtype, device=device)
        
        # GFT: U^T @ signal
        result = (U_tensor.T @ signal_tensor).cpu().numpy().astype(np.float64)
    else:
        # Use NumPy
        result = U.T @ signal
    
    # Squeeze if input was 1D
    if result.shape[1] == 1:
        result = result.squeeze()
    
    return result


def igft(
    fourier_coeffs: Union[np.ndarray, torch.Tensor],
    U: Union[np.ndarray, torch.Tensor],
    use_torch: bool = True,
    device: Optional[torch.device] = None
) -> np.ndarray:
    """
    Inverse Graph Fourier Transform.
    
    Compute the Inverse Graph Fourier Transform (IGFT) of spectral coefficients.
    
    Parameters
    ----------
    fourier_coeffs : np.ndarray or torch.Tensor
        Input Fourier coefficients (vector or matrix).
    U : np.ndarray or torch.Tensor
        Matrix of eigenvectors.
    use_torch : bool
        Whether to use PyTorch for computation.
    device : torch.device, optional
        Device to use for PyTorch computations.
        
    Returns
    -------
    np.ndarray
        Reconstructed signal in the vertex domain.
    """
    if device is None:
        device = get_device()
    
    # Convert to numpy if needed
    if isinstance(fourier_coeffs, torch.Tensor):
        fourier_coeffs = fourier_coeffs.cpu().numpy()
    if isinstance(U, torch.Tensor):
        U = U.cpu().numpy()
    
    # Track if input was 1D
    was_1d = fourier_coeffs.ndim == 1
    if was_1d:
        fourier_coeffs = fourier_coeffs.reshape(-1, 1)
    
    if use_torch:
        # Use PyTorch for matrix multiplication
        # Use float32 if on MPS (doesn't support float64)
        dtype = torch.float32 if device.type == "mps" else torch.float64
        coeffs_tensor = torch.tensor(fourier_coeffs, dtype=dtype, device=device)
        U_tensor = torch.tensor(U, dtype=dtype, device=device)
        
        # IGFT: U @ fourier_coeffs
        result = (U_tensor @ coeffs_tensor).cpu().numpy().astype(np.float64)
    else:
        # Use NumPy
        result = U @ fourier_coeffs
    
    # Squeeze if input was 1D
    if was_1d:
        result = result.squeeze()
    
    return result


def cosine_similarity(
    x: np.ndarray,
    y: np.ndarray,
    eps: float = 1e-12
) -> float:
    """
    Calculate cosine similarity between two vectors.
    
    Parameters
    ----------
    x : np.ndarray
        First vector.
    y : np.ndarray
        Second vector.
    eps : float
        Small value for numerical stability.
        
    Returns
    -------
    float
        Cosine similarity value (between -1 and 1).
        
    Examples
    --------
    >>> x = np.array([1, 2, 3])
    >>> y = np.array([2, 3, 4])
    >>> similarity = cosine_similarity(x, y)
    """
    x = np.asarray(x, dtype=np.float64).flatten()
    y = np.asarray(y, dtype=np.float64).flatten()
    
    if len(x) != len(y):
        raise ValueError("Vectors must have the same length")
    
    # Calculate magnitudes
    mag_x = np.sqrt(np.sum(x ** 2))
    mag_y = np.sqrt(np.sum(y ** 2))
    
    # Handle zero magnitude cases
    if mag_x < eps and mag_y < eps:
        return 0.0
    if mag_x < eps or mag_y < eps:
        return 0.0
    
    # Calculate cosine similarity
    dot_product = np.dot(x, y)
    cosine_sim = dot_product / max(mag_x * mag_y, eps)
    
    # Clamp to [-1, 1]
    return float(np.clip(cosine_sim, -1.0, 1.0))


def find_knee_point(y: np.ndarray, sensitivity: float = 1.0) -> int:
    """
    Find knee point in a curve using the maximum curvature method.
    
    Parameters
    ----------
    y : np.ndarray
        Numeric vector of y values.
    sensitivity : float
        Sensitivity parameter (not used in this implementation).
        
    Returns
    -------
    int
        Index of the knee point (0-based).
        
    Examples
    --------
    >>> y = np.array([1, 2, 3, 10, 11, 12])
    >>> knee_idx = find_knee_point(y)
    """
    y = np.asarray(y, dtype=np.float64)
    
    if len(y) < 3:
        return 0
    
    # Normalize data to [0, 1]
    x = np.arange(len(y))
    y_norm = (y - y.min()) / (y.max() - y.min() + 1e-12)
    x_norm = (x - x.min()) / (x.max() - x.min() + 1e-12)
    
    # Calculate distances from line connecting first and last points
    x1, y1 = x_norm[0], y_norm[0]
    x2, y2 = x_norm[-1], y_norm[-1]
    
    distances = np.zeros(len(y))
    for i in range(1, len(y) - 1):
        xi, yi = x_norm[i], y_norm[i]
        # Point-to-line distance
        distances[i] = abs((y2 - y1) * xi - (x2 - x1) * yi + x2 * y1 - y2 * x1) / \
                       np.sqrt((y2 - y1) ** 2 + (x2 - x1) ** 2 + 1e-12)
    
    return int(np.argmax(distances))


def build_knn_graph(
    coords: np.ndarray,
    k: int = 25
) -> csr_matrix:
    """
    Build k-nearest neighbor graph from coordinates.
    
    Parameters
    ----------
    coords : np.ndarray
        Coordinate matrix (n_points x n_dims).
    k : int
        Number of nearest neighbors.
        
    Returns
    -------
    scipy.sparse.csr_matrix
        Adjacency matrix of the k-NN graph.
    """
    n_points = coords.shape[0]
    k = min(k, n_points - 1)
    
    # Find k nearest neighbors
    nn = NearestNeighbors(n_neighbors=k + 1, algorithm="auto")
    nn.fit(coords)
    distances, indices = nn.kneighbors(coords)
    
    # Build adjacency matrix
    rows = []
    cols = []
    
    for i in range(n_points):
        for j in indices[i, 1:]:  # Skip self (first neighbor)
            rows.append(i)
            cols.append(j)
    
    # Make symmetric
    data = np.ones(len(rows))
    A = csr_matrix((data, (rows, cols)), shape=(n_points, n_points))
    A = A + A.T
    A.data = np.ones_like(A.data)  # Binary adjacency
    
    return A


def check_kband(
    data: np.ndarray,
    coords: np.ndarray,
    signals: List[str],
    signal_data: Dict[str, np.ndarray],
    alpha: float = 0.05,
    k: int = 25,
    laplacian_type: str = "normalized",
    verbose: bool = True
) -> Dict[str, Any]:
    """
    Check K-band limited property of signals.
    
    Analyze whether signals are k-band limited by comparing low-frequency 
    and high-frequency Fourier coefficients.
    
    Parameters
    ----------
    data : np.ndarray
        Coordinate data.
    coords : np.ndarray
        Coordinates for graph construction.
    signals : list of str
        Signal names to analyze.
    signal_data : dict
        Dictionary mapping signal names to signal arrays.
    alpha : float
        Significance level for Wilcoxon test.
    k : int
        Number of nearest neighbors for graph construction.
    laplacian_type : str
        Type of Laplacian.
    verbose : bool
        Whether to print progress messages.
        
    Returns
    -------
    dict
        Results including is_kband_limited flag and per-signal results.
    """
    from scipy.stats import mannwhitneyu
    
    n_nodes = coords.shape[0]
    k_eigen = int(4 * np.sqrt(n_nodes))
    
    if verbose:
        print(f"Analyzing k-band limited property for {len(signals)} signals")
        print(f"Number of nodes: {n_nodes}")
        print(f"Building graph with k = {k} neighbors...")
    
    # Build graph
    A = build_knn_graph(coords, k)
    
    # Compute Laplacian
    if verbose:
        print("Computing Laplacian matrix...")
    L = cal_laplacian(A, laplacian_type)
    
    if verbose:
        print(f"Computing {k_eigen} low and high frequency eigenvalues/eigenvectors")
    
    # Low-frequency eigendecomposition
    if verbose:
        print("Computing low-frequency eigendecomposition...")
    low_decomp = fast_decomposition_lap(L, k_eigen=k_eigen, which="SM")
    
    # High-frequency eigendecomposition
    if verbose:
        print("Computing high-frequency eigendecomposition...")
    high_decomp = fast_decomposition_lap(L, k_eigen=k_eigen, which="LM")
    
    # Find knee points
    if verbose:
        print("Finding knee points in eigenvalue spectra...")
    knee_point_low = find_knee_point(low_decomp["evalues"])
    knee_point_high = find_knee_point(high_decomp["evalues"])
    
    if verbose:
        print(f"Low-frequency knee point at index: {knee_point_low}")
        print(f"High-frequency knee point at index: {knee_point_high}")
    
    # Select eigenvectors up to knee points
    low_freq_eigenvecs = low_decomp["evectors"][:, :knee_point_low + 1]
    high_freq_eigenvecs = high_decomp["evectors"][:, :knee_point_high + 1]
    
    # Concatenate
    combined_eigenvecs = np.hstack([low_freq_eigenvecs, high_freq_eigenvecs])
    
    if verbose:
        print(f"Selected {knee_point_low + 1} low-frequency and {knee_point_high + 1} high-frequency components")
    
    # Analyze each signal
    signal_results = {}
    all_kband_limited = True
    
    for sig in signals:
        if verbose:
            print(f"Analyzing signal: {sig}")
        
        signal_vec = signal_data[sig]
        
        # Compute Fourier coefficients
        fourier_coeffs = gft(signal_vec, combined_eigenvecs)
        
        # Split into low and high frequency
        low_freq_fc = fourier_coeffs[:knee_point_low + 1]
        high_freq_fc = fourier_coeffs[knee_point_low + 1:]
        
        # Remove DC component
        low_freq_fc_no_dc = low_freq_fc[1:] if len(low_freq_fc) > 1 else np.array([])
        
        # Perform statistical test
        if len(low_freq_fc_no_dc) > 0 and len(high_freq_fc) > 0:
            low_filtered = np.abs(low_freq_fc_no_dc)
            low_filtered = low_filtered[low_filtered > np.mean(low_filtered)]
            
            high_filtered = np.abs(high_freq_fc)
            high_filtered = high_filtered[high_filtered > np.mean(high_filtered)]
            
            if len(low_filtered) > 0 and len(high_filtered) > 0:
                stat, p_value = mannwhitneyu(low_filtered, high_filtered, alternative="greater")
                p_value_adjusted = p_value * len(signals)
                is_significant = p_value_adjusted < alpha
            else:
                p_value_adjusted = np.nan
                is_significant = False
        else:
            p_value_adjusted = np.nan
            is_significant = False
        
        signal_results[sig] = {
            "low_freq_fc": low_freq_fc,
            "high_freq_fc": high_freq_fc,
            "low_freq_fc_no_dc": low_freq_fc_no_dc,
            "p_value_adjusted": p_value_adjusted,
            "is_kband_limited": is_significant,
        }
        
        if not is_significant:
            all_kband_limited = False
        
        if verbose:
            print(f"  Wilcoxon test p-value: {p_value_adjusted:.6f}")
            print(f"  K-band limited: {is_significant}")
    
    if verbose:
        print("\n=== Summary ===")
        print(f"All signals k-band limited: {all_kband_limited}")
    
    return {
        "is_kband_limited": all_kband_limited,
        "knee_point_low": knee_point_low,
        "knee_point_high": knee_point_high,
        "signal_results": signal_results,
    }
