"""
Main SGWT class and workflow functions for BioGSP package.

Provides the SGWT class that encapsulates the complete workflow:
initSGWT -> runSpecGraph -> runSGWT -> runSGCC
"""

import numpy as np
import pandas as pd
import torch
from typing import Dict, List, Optional, Union, Any
from dataclasses import dataclass, field

from .utils import (
    cal_laplacian,
    fast_decomposition_lap,
    build_knn_graph,
    cosine_similarity,
    get_device,
)
from .core import (
    sgwt_auto_scales,
    sgwt_forward,
    sgwt_inverse,
    compute_sgwt_filters,
)


@dataclass
class SGWTData:
    """Container for SGWT data."""
    data: pd.DataFrame
    x_col: str
    y_col: str
    signals: List[str]


@dataclass
class SGWTGraph:
    """Container for SGWT graph data."""
    adjacency_matrix: Any
    laplacian_matrix: Any
    eigenvalues: np.ndarray
    eigenvectors: np.ndarray


@dataclass
class SGWTParameters:
    """Container for SGWT parameters."""
    scales: Optional[np.ndarray] = None
    J: int = 5
    scaling_factor: float = 2.0
    kernel_type: str = "heat"
    k: Optional[int] = None
    laplacian_type: Optional[str] = None


class SGWT:
    """
    Spectral Graph Wavelet Transform (SGWT) object.
    
    Main class for performing SGWT analysis on spatial data.
    Workflow: init -> run_spec_graph -> run_sgwt -> run_sgcc
    
    Parameters
    ----------
    data : pd.DataFrame
        Data frame containing spatial coordinates and signal data.
    x_col : str
        Column name for X coordinates.
    y_col : str
        Column name for Y coordinates.
    signals : list of str, optional
        Signal column names to analyze. If None, all non-coordinate columns are used.
    scales : np.ndarray, optional
        Vector of scales for wavelets. If None, auto-generated.
    J : int
        Number of scales to generate if scales is None.
    scaling_factor : float
        Scaling factor between consecutive scales.
    kernel_type : str
        Kernel family: "mexican_hat", "meyer", or "heat".
        
    Attributes
    ----------
    Data : SGWTData
        Container for input data.
    Graph : SGWTGraph or None
        Container for graph data (after run_spec_graph).
    Forward : dict or None
        Forward transform results (after run_sgwt).
    Inverse : dict or None
        Inverse transform results (after run_sgwt).
    Parameters : SGWTParameters
        Container for parameters.
        
    Examples
    --------
    >>> import pandas as pd
    >>> import numpy as np
    >>> data = pd.DataFrame({
    ...     'x': np.random.rand(100),
    ...     'y': np.random.rand(100),
    ...     'signal1': np.random.randn(100),
    ...     'signal2': np.random.randn(100)
    ... })
    >>> sg = SGWT(data, signals=['signal1', 'signal2'])
    >>> sg.run_spec_graph(k=15)
    >>> sg.run_sgwt()
    >>> similarity = sg.run_sgcc('signal1', 'signal2')
    """
    
    def __init__(
        self,
        data: pd.DataFrame,
        x_col: str = "x",
        y_col: str = "y",
        signals: Optional[List[str]] = None,
        scales: Optional[np.ndarray] = None,
        J: int = 5,
        scaling_factor: float = 2.0,
        kernel_type: str = "heat",
    ):
        """Initialize SGWT object."""
        # Validate input
        if data is None:
            raise ValueError("data must be provided")
        
        if not isinstance(data, pd.DataFrame):
            data = pd.DataFrame(data)
        
        if x_col not in data.columns or y_col not in data.columns:
            raise ValueError(f"Data must contain columns: {x_col} and {y_col}")
        
        # Auto-detect signals if not provided
        if signals is None:
            signals = [col for col in data.columns if col not in [x_col, y_col]]
            if len(signals) == 0:
                raise ValueError("No signal columns found in data")
        
        # Validate signal columns exist
        missing_signals = [s for s in signals if s not in data.columns]
        if missing_signals:
            raise ValueError(f"Signal columns not found in data: {', '.join(missing_signals)}")
        
        # Initialize data container
        self.Data = SGWTData(
            data=data.copy(),
            x_col=x_col,
            y_col=y_col,
            signals=signals,
        )
        
        # Initialize other containers
        self.Graph: Optional[SGWTGraph] = None
        self.Forward: Optional[Dict] = None
        self.Inverse: Optional[Dict] = None
        
        # Initialize parameters
        self.Parameters = SGWTParameters(
            scales=scales,
            J=J,
            scaling_factor=scaling_factor,
            kernel_type=kernel_type,
        )
    
    def run_spec_graph(
        self,
        k: int = 25,
        laplacian_type: str = "normalized",
        length_eigenvalue: Optional[int] = None,
        verbose: bool = True,
        use_torch: bool = True,
    ) -> "SGWT":
        """
        Build spectral graph for SGWT object.
        
        Generate Graph slot information including adjacency matrix,
        Laplacian matrix, eigenvalues, and eigenvectors.
        
        Parameters
        ----------
        k : int
            Number of nearest neighbors for graph construction.
        laplacian_type : str
            Type of Laplacian: "unnormalized", "normalized", or "randomwalk".
        length_eigenvalue : int, optional
            Number of eigenvalues/eigenvectors to compute. If None, uses full length.
        verbose : bool
            Whether to print progress messages.
        use_torch : bool
            Whether to use PyTorch for eigendecomposition.
            
        Returns
        -------
        SGWT
            Updated SGWT object with Graph slot populated.
            
        Examples
        --------
        >>> sg.run_spec_graph(k=30, laplacian_type="normalized")
        """
        # Extract data
        data = self.Data.data
        x_col = self.Data.x_col
        y_col = self.Data.y_col
        
        # Set default length_eigenvalue
        if length_eigenvalue is None:
            length_eigenvalue = len(data)
        
        if verbose:
            print("Building graph from spatial coordinates...")
        
        # Get coordinates
        coords = data[[x_col, y_col]].values
        
        # Build k-nearest neighbor graph
        A = build_knn_graph(coords, k)
        
        if verbose:
            print("Computing Laplacian and eigendecomposition...")
        
        # Compute Laplacian matrix
        L = cal_laplacian(A, laplacian_type)
        
        # Eigendecomposition
        decomp = fast_decomposition_lap(L, k_eigen=length_eigenvalue, which="SM", use_torch=use_torch)
        
        # Update Graph slot
        self.Graph = SGWTGraph(
            adjacency_matrix=A,
            laplacian_matrix=L,
            eigenvalues=decomp["evalues"],
            eigenvectors=decomp["evectors"],
        )
        
        # Update parameters
        self.Parameters.k = k
        self.Parameters.laplacian_type = laplacian_type
        
        # Auto-generate scales if not provided
        if self.Parameters.scales is None:
            lmax = np.max(decomp["evalues"]) * 0.95
            self.Parameters.scales = sgwt_auto_scales(
                lmax, 
                self.Parameters.J, 
                self.Parameters.scaling_factor
            )
            if verbose:
                scales_str = ", ".join([f"{s:.4f}" for s in self.Parameters.scales])
                print(f"Auto-generated scales: {scales_str}")
        
        if verbose:
            print("Graph construction completed.")
        
        return self
    
    def run_sgwt(
        self,
        use_batch: bool = True,
        verbose: bool = True,
        use_torch: bool = True,
    ) -> "SGWT":
        """
        Run SGWT forward and inverse transforms for all signals.
        
        Perform SGWT analysis on all signals in the SGWT object.
        Uses batch processing for multiple signals when possible.
        
        Parameters
        ----------
        use_batch : bool
            Whether to use batch processing for multiple signals.
        verbose : bool
            Whether to print progress messages.
        use_torch : bool
            Whether to use PyTorch for computation.
            
        Returns
        -------
        SGWT
            Updated SGWT object with Forward and Inverse slots populated.
            
        Examples
        --------
        >>> sg.run_sgwt()
        """
        # Validate input
        if self.Graph is None:
            raise ValueError("Graph slot is empty. Run run_spec_graph() first.")
        
        if self.Parameters.scales is None:
            raise ValueError("Scales not found. Make sure to run run_spec_graph() before run_sgwt().")
        
        # Extract components
        eigenvalues = self.Graph.eigenvalues
        eigenvectors = self.Graph.eigenvectors
        params = self.Parameters
        signals = self.Data.signals
        data = self.Data.data
        
        n_signals = len(signals)
        
        if verbose:
            print(f"Performing SGWT analysis for {n_signals} signals...")
        
        device = get_device() if use_torch else None
        
        if use_batch and n_signals > 1:
            # Batch processing for multiple signals
            if verbose:
                print("Using batch processing for efficiency...")
            
            # Create signal matrix (n_vertices x n_signals)
            signals_matrix = data[signals].values
            
            # Batch forward transform
            batch_forward = sgwt_forward(
                signals_matrix,
                eigenvectors,
                eigenvalues,
                params.scales,
                kernel_type=params.kernel_type,
                use_torch=use_torch,
                device=device,
            )
            
            # Batch inverse transform
            batch_inverse = sgwt_inverse(
                batch_forward,
                eigenvectors,
                signals_matrix,
                use_torch=use_torch,
                device=device,
            )
            
            # Split results back into individual signals
            forward_list = {}
            inverse_list = {}
            
            for i, sig in enumerate(signals):
                # Extract individual signal results from batch
                filtered = {}
                for key, val in batch_forward["fourier_coefficients"]["filtered"].items():
                    if val.ndim > 1:
                        filtered[key] = val[:, i]
                    else:
                        filtered[key] = val
                
                original = batch_forward["fourier_coefficients"]["original"]
                if original.ndim > 1:
                    original = original[:, i]
                
                forward_list[sig] = {
                    "fourier_coefficients": {
                        "original": original,
                        "filtered": filtered,
                    },
                    "filters": batch_forward["filters"],
                }
                
                # Extract individual vertex approximations
                coefficients_individual = {}
                for comp_name, comp_val in batch_inverse["vertex_approximations"].items():
                    if comp_val.ndim > 1:
                        coefficients_individual[comp_name] = comp_val[:, i]
                    else:
                        coefficients_individual[comp_name] = comp_val
                
                reconstructed = batch_inverse["reconstructed_signal"]
                if reconstructed.ndim > 1:
                    reconstructed = reconstructed[:, i]
                
                error = batch_inverse["reconstruction_error"]
                if error is not None and hasattr(error, "__len__") and len(error) > 1:
                    error = error[i]
                
                inverse_list[sig] = {
                    "vertex_approximations": coefficients_individual,
                    "reconstructed_signal": reconstructed,
                    "reconstruction_error": error,
                }
        else:
            # Individual processing
            if verbose and n_signals > 1:
                print("Using individual processing...")
            
            forward_list = {}
            inverse_list = {}
            
            for sig in signals:
                if verbose:
                    print(f"Processing signal: {sig}")
                
                # Extract signal vector
                sig_vec = data[sig].values
                
                # Forward transform
                fwd = sgwt_forward(
                    sig_vec,
                    eigenvectors,
                    eigenvalues,
                    params.scales,
                    kernel_type=params.kernel_type,
                    use_torch=use_torch,
                    device=device,
                )
                forward_list[sig] = fwd
                
                # Inverse transform
                inv = sgwt_inverse(fwd, eigenvectors, sig_vec, use_torch=use_torch, device=device)
                inverse_list[sig] = inv
        
        # Update SGWT object
        self.Forward = forward_list
        self.Inverse = inverse_list
        
        if verbose:
            print("SGWT analysis completed.")
        
        return self
    
    def run_sgcc(
        self,
        signal1: Union[str, "SGWT"],
        signal2: Optional[Union[str, "SGWT"]] = None,
        eps: float = 1e-12,
        validate: bool = True,
        return_parts: bool = True,
        low_only: bool = False,
    ) -> Dict:
        """
        Run SGCC weighted similarity analysis in Fourier domain.
        
        Calculate energy-normalized weighted similarity between two signals
        using Fourier domain coefficients directly.
        
        Parameters
        ----------
        signal1 : str or SGWT
            Signal name or SGWT object.
        signal2 : str or SGWT, optional
            Signal name or SGWT object.
        eps : float
            Small value for numerical stability.
        validate : bool
            Whether to check consistency.
        return_parts : bool
            Whether to return detailed components.
        low_only : bool
            Whether to compute only low-frequency similarity.
            
        Returns
        -------
        dict or float
            Similarity analysis results or just the similarity value.
            
        Examples
        --------
        >>> similarity = sg.run_sgcc("signal1", "signal2")
        """
        def get_decomp(x, sg=None):
            if isinstance(x, str):
                if sg is None or sg.Forward is None:
                    raise ValueError("SGWT object with Forward results required when using signal names")
                if x not in sg.Forward:
                    raise ValueError(f"Signal '{x}' not found in SGWT Forward results")
                return sg.Forward[x]
            elif isinstance(x, SGWT):
                if x.Forward is None or len(x.Forward) == 0:
                    raise ValueError("SGWT object must have Forward results")
                return list(x.Forward.values())[0]
            elif isinstance(x, dict) and "fourier_coefficients" in x:
                return x
            else:
                raise ValueError("Invalid input type for signal")
        
        # Get decompositions
        A = get_decomp(signal1, self)
        B = get_decomp(signal2, self) if signal2 is not None else get_decomp(signal1, self)
        
        # Extract filtered Fourier coefficients
        fourier_a = A["fourier_coefficients"]["filtered"]
        fourier_b = B["fourier_coefficients"]["filtered"]
        
        if fourier_a is None or fourier_b is None:
            raise ValueError("Fourier coefficients not found in Forward results")
        
        # Get scaling coefficients and exclude DC component
        f_low_a = np.asarray(fourier_a["scaling"], dtype=np.float64)
        f_low_b = np.asarray(fourier_b["scaling"], dtype=np.float64)
        
        if len(f_low_a) > 1:
            f_low_a = f_low_a[1:]
        if len(f_low_b) > 1:
            f_low_b = f_low_b[1:]
        
        # Handle non-finite values
        f_low_a = np.nan_to_num(f_low_a, nan=0.0, posinf=0.0, neginf=0.0)
        f_low_b = np.nan_to_num(f_low_b, nan=0.0, posinf=0.0, neginf=0.0)
        
        # Ensure same length
        min_length = min(len(f_low_a), len(f_low_b))
        if len(f_low_a) != len(f_low_b) and validate:
            print(f"Warning: Scaling coefficients have different lengths: {len(f_low_a)} vs {len(f_low_b)}. "
                  f"Truncating to minimum length: {min_length}")
        f_low_a = f_low_a[:min_length]
        f_low_b = f_low_b[:min_length]
        
        # Energies in Fourier domain
        E_low_a = np.sum(f_low_a ** 2)
        E_low_b = np.sum(f_low_b ** 2)
        
        # Low-frequency cosine similarity
        c_low = cosine_similarity(f_low_a, f_low_b, eps)
        
        # Short-circuit for low-only
        if low_only:
            if return_parts:
                return {
                    "c_low": c_low,
                    "c_nonlow": np.nan,
                    "w_low": 1.0,
                    "w_NL": 0.0,
                    "S": c_low,
                    "E_low_a": E_low_a,
                    "E_NL_a": np.nan,
                    "E_low_b": E_low_b,
                    "E_NL_b": np.nan,
                    "n": len(f_low_a),
                    "J": np.nan,
                }
            return c_low
        
        # Collect wavelet Fourier coefficients
        wavelet_names_a = [k for k in fourier_a.keys() if k.startswith("wavelet_scale_")]
        wavelet_names_b = [k for k in fourier_b.keys() if k.startswith("wavelet_scale_")]
        
        if len(wavelet_names_a) == 0 or len(wavelet_names_b) == 0:
            raise ValueError("No wavelet Fourier coefficients found")
        
        # Order by scale index
        wavelet_names_a = sorted(wavelet_names_a, key=lambda x: int(x.replace("wavelet_scale_", "")))
        wavelet_names_b = sorted(wavelet_names_b, key=lambda x: int(x.replace("wavelet_scale_", "")))
        
        # Handle different numbers of scales
        min_scales = min(len(wavelet_names_a), len(wavelet_names_b))
        if len(wavelet_names_a) != len(wavelet_names_b) and validate:
            print(f"Warning: Different numbers of wavelet scales: {len(wavelet_names_a)} vs {len(wavelet_names_b)}. "
                  f"Using first {min_scales} scales.")
        wavelet_names_a = wavelet_names_a[:min_scales]
        wavelet_names_b = wavelet_names_b[:min_scales]
        
        # Extract and flatten wavelet coefficients
        f_wave_a = np.concatenate([np.asarray(fourier_a[name]).flatten() for name in wavelet_names_a])
        f_wave_b = np.concatenate([np.asarray(fourier_b[name]).flatten() for name in wavelet_names_b])
        
        # Handle non-finite values
        f_wave_a = np.nan_to_num(f_wave_a, nan=0.0, posinf=0.0, neginf=0.0)
        f_wave_b = np.nan_to_num(f_wave_b, nan=0.0, posinf=0.0, neginf=0.0)
        
        # Ensure same length
        min_wave_length = min(len(f_wave_a), len(f_wave_b))
        if len(f_wave_a) != len(f_wave_b) and validate:
            print(f"Warning: Wavelet coefficients have different lengths: {len(f_wave_a)} vs {len(f_wave_b)}. "
                  f"Truncating to minimum length: {min_wave_length}")
        f_wave_a = f_wave_a[:min_wave_length]
        f_wave_b = f_wave_b[:min_wave_length]
        
        # Energies for wavelet components
        E_NL_a = np.sum(f_wave_a ** 2)
        E_NL_b = np.sum(f_wave_b ** 2)
        
        J = len(wavelet_names_a)
        n = len(f_low_a)
        
        # Non-low cosine similarity
        c_nonlow = cosine_similarity(f_wave_a, f_wave_b, eps)
        
        # Energy-based macro weights
        w_low_a = E_low_a / (E_low_a + E_NL_a + eps)
        w_low_b = E_low_b / (E_low_b + E_NL_b + eps)
        w_low = np.clip(0.5 * (w_low_a + w_low_b), 0, 1)
        w_NL = 1 - w_low
        S = w_low * c_low + w_NL * c_nonlow
        
        if return_parts:
            return {
                "c_low": c_low,
                "c_nonlow": c_nonlow,
                "w_low": w_low,
                "w_NL": w_NL,
                "S": S,
                "E_low_a": E_low_a,
                "E_NL_a": E_NL_a,
                "E_low_b": E_low_b,
                "E_NL_b": E_NL_b,
                "n": n,
                "J": J,
            }
        return S
    
    def energy_analysis(self, signal_name: Optional[str] = None) -> pd.DataFrame:
        """
        Analyze SGWT energy distribution across scales in Fourier domain.
        
        Parameters
        ----------
        signal_name : str, optional
            Name of signal to analyze. If None, uses first signal.
            
        Returns
        -------
        pd.DataFrame
            Energy analysis results.
            
        Examples
        --------
        >>> energy_df = sg.energy_analysis("signal1")
        """
        if self.Forward is None:
            raise ValueError("Forward results not found. Run run_sgwt() first.")
        
        if signal_name is None:
            signal_name = list(self.Forward.keys())[0]
        
        if signal_name not in self.Forward:
            raise ValueError(f"Signal '{signal_name}' not found in Forward results")
        
        forward_result = self.Forward[signal_name]
        fourier_coeffs = forward_result["fourier_coefficients"]["filtered"]
        scales = self.Parameters.scales
        
        energies = []
        scale_names = []
        scale_values = []
        
        # Scaling (low-pass) energy - exclude DC component
        if "scaling" in fourier_coeffs:
            scaling_coeffs = np.asarray(fourier_coeffs["scaling"])
            if len(scaling_coeffs) > 1:
                scaling_coeffs = scaling_coeffs[1:]
            scaling_energy = np.sum(np.abs(scaling_coeffs) ** 2)
            
            energies.append(scaling_energy)
            scale_names.append("low_pass")
            scale_values.append(scales[0])
        
        # Wavelet energies
        wavelet_names = sorted([k for k in fourier_coeffs.keys() if k.startswith("wavelet_scale_")],
                              key=lambda x: int(x.replace("wavelet_scale_", "")))
        
        for name in wavelet_names:
            wavelet_coeffs = np.asarray(fourier_coeffs[name])
            if len(wavelet_coeffs) > 1:
                wavelet_coeffs = wavelet_coeffs[1:]
            wavelet_energy = np.sum(np.abs(wavelet_coeffs) ** 2)
            
            scale_idx = int(name.replace("wavelet_scale_", "")) - 1
            energies.append(wavelet_energy)
            scale_names.append(f"wavelet_{scale_idx + 1}")
            scale_values.append(scales[scale_idx] if scale_idx < len(scales) else np.nan)
        
        # Calculate energy ratios
        total_energy = sum(energies)
        energy_ratios = [e / total_energy if total_energy > 0 else 0 for e in energies]
        
        return pd.DataFrame({
            "scale": scale_names,
            "energy": energies,
            "energy_ratio": energy_ratios,
            "scale_value": scale_values,
            "signal": signal_name,
        })
    
    def __repr__(self) -> str:
        """String representation of SGWT object."""
        lines = ["SGWT Object", "==========="]
        
        # Data information
        if self.Data is not None:
            lines.append("Data:")
            lines.append(f"  Dimensions: {len(self.Data.data)} x {len(self.Data.data.columns)}")
            lines.append(f"  Coordinates: {self.Data.x_col}, {self.Data.y_col}")
            lines.append(f"  Signals: {', '.join(self.Data.signals)}")
        
        # Parameters
        if self.Parameters is not None:
            lines.append("\nParameters:")
            lines.append(f"  k (neighbors): {self.Parameters.k}")
            lines.append(f"  J (scales): {self.Parameters.J}")
            lines.append(f"  Kernel type: {self.Parameters.kernel_type}")
            lines.append(f"  Laplacian type: {self.Parameters.laplacian_type}")
            if self.Parameters.scales is not None:
                scales_str = ", ".join([f"{s:.4f}" for s in self.Parameters.scales])
                lines.append(f"  Scales: {scales_str}")
        
        # Status
        lines.append("\nStatus:")
        lines.append(f"  Graph computed: {self.Graph is not None}")
        lines.append(f"  Forward computed: {self.Forward is not None}")
        lines.append(f"  Inverse computed: {self.Inverse is not None}")
        
        # Reconstruction errors
        if self.Inverse is not None:
            lines.append("\nReconstruction Errors:")
            for sig, inv in self.Inverse.items():
                err = inv.get("reconstruction_error")
                if err is not None:
                    if hasattr(err, "__len__"):
                        err = float(np.mean(err))
                    lines.append(f"  {sig}: {err:.6f}")
        
        return "\n".join(lines)
    
    def _repr_html_(self) -> str:
        """HTML representation for Jupyter notebooks."""
        return f"<pre>{self.__repr__()}</pre>"


# Convenience function to match R API
def init_sgwt(
    data: Union[pd.DataFrame, np.ndarray, Dict],
    x_col: str = "x",
    y_col: str = "y",
    signals: Optional[List[str]] = None,
    scales: Optional[np.ndarray] = None,
    J: int = 5,
    scaling_factor: float = 2.0,
    kernel_type: str = "heat",
) -> SGWT:
    """
    Initialize SGWT object (convenience function matching R API).
    
    Parameters
    ----------
    data : pd.DataFrame, np.ndarray, or dict
        Data containing spatial coordinates and signals.
    x_col : str
        Column name for X coordinates.
    y_col : str
        Column name for Y coordinates.
    signals : list of str, optional
        Signal column names.
    scales : np.ndarray, optional
        Wavelet scales.
    J : int
        Number of scales.
    scaling_factor : float
        Scaling factor between scales.
    kernel_type : str
        Kernel family.
        
    Returns
    -------
    SGWT
        Initialized SGWT object.
    """
    if isinstance(data, np.ndarray):
        data = pd.DataFrame(data)
    elif isinstance(data, dict):
        data = pd.DataFrame(data)
    
    return SGWT(
        data=data,
        x_col=x_col,
        y_col=y_col,
        signals=signals,
        scales=scales,
        J=J,
        scaling_factor=scaling_factor,
        kernel_type=kernel_type,
    )
