"""
Core SGWT functions for BioGSP package.

Includes kernel functions, filter computation, forward and inverse transforms.
Uses PyTorch for accelerated matrix operations where beneficial.
"""

import numpy as np
import torch
from typing import Dict, List, Optional, Tuple, Union, Callable

from .utils import gft, igft, get_device


def sgwt_get_kernels(kernel_type: str = "heat") -> Dict[str, Callable]:
    """
    Get a unified kernel family (low-pass and band-pass) by kernel_type.
    
    Returns a pair of functions implementing the scaling (low-pass) and 
    wavelet (band-pass) kernels for a given kernel family.
    
    Parameters
    ----------
    kernel_type : str
        Kernel family name: "mexican_hat", "meyer", or "heat".
        
    Returns
    -------
    dict
        Dictionary with 'scaling' and 'wavelet' functions.
        Each function takes (x, scale_param) as arguments.
        
    Examples
    --------
    >>> kernels = sgwt_get_kernels("heat")
    >>> scaling_val = kernels["scaling"](0.5, 1.0)
    >>> wavelet_val = kernels["wavelet"](0.5, 1.0)
    """
    if kernel_type == "mexican_hat":
        def scaling_fun(x, scale_param):
            t = np.asarray(x) / scale_param
            return np.exp(-0.5 * t ** 2)
        
        def wavelet_fun(x, scale_param):
            t = np.asarray(x) / scale_param
            return (t ** 2) * np.exp(-0.5 * t ** 2)
        
        return {"scaling": scaling_fun, "wavelet": wavelet_fun}
    
    elif kernel_type == "meyer":
        def scaling_fun(x, scale_param):
            x = np.asarray(x)
            a = 0.5 * scale_param
            b = 1.0 * scale_param
            
            result = np.zeros_like(x, dtype=np.float64)
            
            # x <= a: 1
            mask1 = x <= a
            result[mask1] = 1.0
            
            # a < x < b: cos(pi/2 * s) where s = (x-a)/(b-a)
            mask2 = (x > a) & (x < b)
            s = (x[mask2] - a) / (b - a)
            result[mask2] = np.cos(np.pi / 2 * s)
            
            # x >= b: 0 (already initialized)
            
            return result
        
        def wavelet_fun(x, scale_param):
            x = np.asarray(x)
            a = 0.5 * scale_param
            b = 1.0 * scale_param
            c = 2.0 * scale_param
            
            result = np.zeros_like(x, dtype=np.float64)
            
            # a < x <= b: sin(pi/2 * s)
            mask1 = (x > a) & (x <= b)
            s1 = (x[mask1] - a) / (b - a)
            result[mask1] = np.sin(np.pi / 2 * s1)
            
            # b < x < c: cos(pi/2 * s)
            mask2 = (x > b) & (x < c)
            s2 = (x[mask2] - b) / (c - b)
            result[mask2] = np.cos(np.pi / 2 * s2)
            
            return result
        
        return {"scaling": scaling_fun, "wavelet": wavelet_fun}
    
    elif kernel_type == "heat":
        def scaling_fun(x, scale_param):
            t = np.asarray(x) / scale_param
            return np.exp(-t)
        
        def wavelet_fun(x, scale_param):
            t = np.asarray(x) / scale_param
            return t * np.exp(-t)
        
        return {"scaling": scaling_fun, "wavelet": wavelet_fun}
    
    else:
        raise ValueError(f"Kernel type '{kernel_type}' not supported. "
                        "Use 'mexican_hat', 'meyer', or 'heat'.")


def compute_sgwt_filters(
    eigenvalues: np.ndarray,
    scales: np.ndarray,
    lmax: Optional[float] = None,
    kernel_type: str = "heat"
) -> List[np.ndarray]:
    """
    Compute SGWT filters.
    
    Compute wavelet and scaling function coefficients in the spectral domain.
    
    Parameters
    ----------
    eigenvalues : np.ndarray
        Eigenvalues of the graph Laplacian.
    scales : np.ndarray
        Vector of scales for the wavelets.
    lmax : float, optional
        Maximum eigenvalue. If None, computed from eigenvalues.
    kernel_type : str
        Kernel family: "mexican_hat", "meyer", or "heat".
        
    Returns
    -------
    list of np.ndarray
        List of filters (scaling function + wavelets).
        First element is scaling function, rest are wavelets at each scale.
        
    Examples
    --------
    >>> eigenvals = np.array([0, 0.1, 0.5, 1.0, 1.5])
    >>> scales = np.array([2, 1, 0.5])
    >>> filters = compute_sgwt_filters(eigenvals, scales)
    """
    eigenvalues = np.asarray(eigenvalues, dtype=np.float64)
    scales = np.asarray(scales, dtype=np.float64)
    
    if lmax is None:
        lmax = np.max(eigenvalues) * 0.95  # Avoid numerical issues at lambda_max
    
    J = len(scales)  # Number of scales
    
    # Get unified kernel family
    kernels = sgwt_get_kernels(kernel_type)
    
    # Initialize filter bank: 1 scaling function + J wavelets
    filters = []
    
    # Scaling function (low-pass filter)
    scaling_filter = kernels["scaling"](eigenvalues, scales[0])
    filters.append(scaling_filter)
    
    # Wavelet functions at different scales
    for j in range(J):
        wavelet_filter = kernels["wavelet"](eigenvalues, scales[j])
        filters.append(wavelet_filter)
    
    return filters


def sgwt_auto_scales(
    lmax: float,
    J: int = 5,
    scaling_factor: float = 2.0
) -> np.ndarray:
    """
    Generate automatic scales for SGWT.
    
    Generate logarithmically spaced scales for SGWT.
    
    Parameters
    ----------
    lmax : float
        Maximum eigenvalue.
    J : int
        Number of scales.
    scaling_factor : float
        Scaling factor between consecutive scales.
        
    Returns
    -------
    np.ndarray
        Vector of scales.
        
    Examples
    --------
    >>> scales = sgwt_auto_scales(lmax=2.0, J=5, scaling_factor=2)
    """
    return lmax / (scaling_factor ** np.arange(J))


def sgwt_forward(
    signal: np.ndarray,
    eigenvectors: np.ndarray,
    eigenvalues: np.ndarray,
    scales: np.ndarray,
    lmax: Optional[float] = None,
    kernel_type: str = "heat",
    use_torch: bool = True,
    device: Optional[torch.device] = None
) -> Dict:
    """
    Forward SGWT transform (single or batch).
    
    Transform signal(s) to spectral domain and apply SGWT filters.
    Handles both single signals (vector) and multiple signals (matrix).
    
    Parameters
    ----------
    signal : np.ndarray
        Input signal vector OR matrix (n_vertices x n_signals).
    eigenvectors : np.ndarray
        Eigenvectors of the graph Laplacian.
    eigenvalues : np.ndarray
        Eigenvalues of the graph Laplacian.
    scales : np.ndarray
        Vector of scales for the wavelets.
    lmax : float, optional
        Maximum eigenvalue.
    kernel_type : str
        Kernel family: "mexican_hat", "meyer", or "heat".
    use_torch : bool
        Whether to use PyTorch for computation.
    device : torch.device, optional
        Device for PyTorch computations.
        
    Returns
    -------
    dict
        Dictionary containing:
        - 'fourier_coefficients': dict with 'original' and 'filtered' coefficients
        - 'filters': filter bank used
        
    Examples
    --------
    >>> signal = np.random.randn(100)
    >>> eigenvectors = np.random.randn(100, 50)
    >>> eigenvalues = np.sort(np.random.rand(50))
    >>> scales = np.array([2, 1, 0.5])
    >>> result = sgwt_forward(signal, eigenvectors, eigenvalues, scales)
    """
    if device is None:
        device = get_device()
    
    signal = np.asarray(signal, dtype=np.float64)
    eigenvectors = np.asarray(eigenvectors, dtype=np.float64)
    eigenvalues = np.asarray(eigenvalues, dtype=np.float64)
    scales = np.asarray(scales, dtype=np.float64)
    
    # Validate dimensions
    if signal.ndim == 1:
        if eigenvectors.shape[0] != len(signal):
            raise ValueError("Number of vertices in signal must match eigenvectors")
        signal = signal.reshape(-1, 1)
        was_1d = True
    else:
        if eigenvectors.shape[0] != signal.shape[0]:
            raise ValueError("Number of vertices in signals matrix must match eigenvectors")
        was_1d = False
    
    # Compute filters
    filters = compute_sgwt_filters(eigenvalues, scales, lmax, kernel_type)
    
    # Transform signal(s) to spectral domain using GFT
    signal_hat = gft(signal, eigenvectors, use_torch=use_torch, device=device)
    if signal_hat.ndim == 1:
        signal_hat = signal_hat.reshape(-1, 1)
    
    # Store original and filtered Fourier coefficients
    filtered_coeffs = {}
    
    # Apply filters
    # First filter is scaling function
    filtered_coeffs["scaling"] = signal_hat * filters[0].reshape(-1, 1)
    
    # Remaining filters are wavelets
    for i, filt in enumerate(filters[1:], start=1):
        filtered_coeffs[f"wavelet_scale_{i}"] = signal_hat * filt.reshape(-1, 1)
    
    # Squeeze if input was 1D
    original = signal_hat.squeeze() if was_1d else signal_hat
    for key in filtered_coeffs:
        if was_1d:
            filtered_coeffs[key] = filtered_coeffs[key].squeeze()
    
    return {
        "fourier_coefficients": {
            "original": original,
            "filtered": filtered_coeffs,
        },
        "filters": filters,
    }


def sgwt_inverse(
    sgwt_decomp: Dict,
    eigenvectors: np.ndarray,
    original_signal: Optional[np.ndarray] = None,
    use_torch: bool = True,
    device: Optional[torch.device] = None
) -> Dict:
    """
    Inverse SGWT transform (single or batch).
    
    Reconstruct signal(s) from filtered Fourier coefficients using inverse GFT.
    
    Parameters
    ----------
    sgwt_decomp : dict
        SGWT decomposition object from sgwt_forward.
    eigenvectors : np.ndarray
        Eigenvectors of the graph Laplacian.
    original_signal : np.ndarray, optional
        Original signal for error calculation.
    use_torch : bool
        Whether to use PyTorch for computation.
    device : torch.device, optional
        Device for PyTorch computations.
        
    Returns
    -------
    dict
        Dictionary containing:
        - 'vertex_approximations': dict with inverse-transformed signals
        - 'reconstructed_signal': full reconstructed signal
        - 'reconstruction_error': RMSE if original_signal provided
        
    Examples
    --------
    >>> # After forward transform
    >>> inverse_result = sgwt_inverse(sgwt_decomp, eigenvectors, original_signal)
    """
    if device is None:
        device = get_device()
    
    eigenvectors = np.asarray(eigenvectors, dtype=np.float64)
    filtered_fourier = sgwt_decomp["fourier_coefficients"]["filtered"]
    
    # Perform inverse GFT on all filtered coefficients
    vertex_approximations = {}
    
    for name, coeffs in filtered_fourier.items():
        coeffs = np.asarray(coeffs, dtype=np.float64)
        vertex_approx = igft(coeffs, eigenvectors, use_torch=use_torch, device=device)
        
        # Rename for clarity
        if name == "scaling":
            new_name = "low_pass"
        elif name.startswith("wavelet_scale_"):
            scale_num = name.replace("wavelet_scale_", "")
            new_name = f"wavelet_{scale_num}"
        else:
            new_name = name
        
        vertex_approximations[new_name] = vertex_approx
    
    # Reconstruct signal - sum of all components
    reconstructed = None
    for approx in vertex_approximations.values():
        if reconstructed is None:
            reconstructed = approx.copy()
        else:
            reconstructed = reconstructed + approx
    
    # Calculate reconstruction error if original signal provided
    reconstruction_error = None
    if original_signal is not None:
        original_signal = np.asarray(original_signal, dtype=np.float64)
        
        if original_signal.ndim != reconstructed.ndim:
            if reconstructed.ndim == 1:
                original_signal = original_signal.flatten()
            else:
                original_signal = original_signal.reshape(reconstructed.shape)
        
        if original_signal.ndim == 1:
            reconstruction_error = np.sqrt(np.mean((original_signal - reconstructed) ** 2))
        else:
            reconstruction_error = np.sqrt(np.mean((original_signal - reconstructed) ** 2, axis=0))
    
    return {
        "vertex_approximations": vertex_approximations,
        "reconstructed_signal": reconstructed,
        "reconstruction_error": reconstruction_error,
    }


def compare_kernel_families(
    x_range: Tuple[float, float] = (0, 3),
    scale_param: float = 1.0,
    plot_results: bool = True
) -> Dict[str, np.ndarray]:
    """
    Compare different kernel families.
    
    Visualize and compare different kernel families (both scaling and wavelet filters).
    
    Parameters
    ----------
    x_range : tuple
        Range of x values to evaluate.
    scale_param : float
        Scale parameter for all functions.
    plot_results : bool
        Whether to plot the comparison.
        
    Returns
    -------
    dict
        Dictionary with x values and kernel values for each family.
        
    Examples
    --------
    >>> comparison = compare_kernel_families()
    >>> comparison = compare_kernel_families(x_range=(0, 5), scale_param=1.5)
    """
    x_vals = np.linspace(x_range[0], x_range[1], 200)
    
    kernel_types = ["mexican_hat", "meyer", "heat"]
    
    results = {"x": x_vals}
    
    for ktype in kernel_types:
        kernels = sgwt_get_kernels(ktype)
        results[f"{ktype}_scaling"] = kernels["scaling"](x_vals, scale_param)
        results[f"{ktype}_wavelet"] = kernels["wavelet"](x_vals, scale_param)
    
    if plot_results:
        try:
            import matplotlib.pyplot as plt
            
            fig, axes = plt.subplots(1, 2, figsize=(12, 5))
            
            # Plot scaling functions
            ax1 = axes[0]
            for ktype in kernel_types:
                ax1.plot(x_vals, results[f"{ktype}_scaling"], label=ktype.replace("_", " ").title(), linewidth=2)
            ax1.set_xlabel("Eigenvalue")
            ax1.set_ylabel("Scaling Function Value")
            ax1.set_title("SGWT Scaling Functions")
            ax1.legend()
            ax1.grid(True, alpha=0.3)
            
            # Plot wavelet functions
            ax2 = axes[1]
            for ktype in kernel_types:
                ax2.plot(x_vals, results[f"{ktype}_wavelet"], label=ktype.replace("_", " ").title(), linewidth=2)
            ax2.set_xlabel("Eigenvalue")
            ax2.set_ylabel("Wavelet Function Value")
            ax2.set_title("SGWT Wavelet Functions")
            ax2.legend()
            ax2.grid(True, alpha=0.3)
            
            plt.tight_layout()
            plt.show()
            
        except ImportError:
            print("matplotlib not available for plotting")
    
    return results
