"""
Visualization functions for BioGSP package.

Plotting functions for SGWT decomposition, kernels, and simulated patterns.
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Optional, Tuple, Union, Any

try:
    import matplotlib.pyplot as plt
    import matplotlib.colors as mcolors
    from matplotlib.gridspec import GridSpec
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False

try:
    import seaborn as sns
    HAS_SEABORN = True
except ImportError:
    HAS_SEABORN = False


def _check_matplotlib():
    """Check if matplotlib is available."""
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib is required for visualization. Install with: pip install matplotlib")


def plot_sgwt_decomposition(
    sg: Any,
    signal_name: Optional[str] = None,
    plot_scales: Optional[List[int]] = None,
    ncol: int = 3,
    figsize: Optional[Tuple[float, float]] = None,
    cmap: str = "viridis",
) -> Any:
    """
    Plot SGWT decomposition results.
    
    Visualize SGWT decomposition components including original signal,
    scaling function, wavelet coefficients, and reconstructed signal.
    
    Parameters
    ----------
    sg : SGWT
        SGWT object with Forward and Inverse results computed.
    signal_name : str, optional
        Name of signal to plot. If None, uses first signal.
    plot_scales : list of int, optional
        Which wavelet scales to plot. Default: first 4.
    ncol : int
        Number of columns in the plot layout.
    figsize : tuple, optional
        Figure size.
    cmap : str
        Colormap name.
        
    Returns
    -------
    matplotlib.figure.Figure
        Figure with combined plots.
        
    Examples
    --------
    >>> fig = plot_sgwt_decomposition(sg, signal_name="signal1")
    >>> plt.show()
    """
    _check_matplotlib()
    
    # Validate input
    if sg.Forward is None or sg.Inverse is None:
        raise ValueError("SGWT object must have Forward and Inverse results computed")
    
    # Default to first signal
    if signal_name is None:
        signal_name = list(sg.Forward.keys())[0]
    
    if signal_name not in sg.Forward:
        raise ValueError(f"Signal '{signal_name}' not found in SGWT results")
    
    # Get results
    inverse_result = sg.Inverse[signal_name]
    coefficients = inverse_result["vertex_approximations"]
    
    # Default scales to plot
    if plot_scales is None:
        n_wavelets = len([k for k in coefficients.keys() if k.startswith("wavelet_")])
        plot_scales = list(range(1, min(5, n_wavelets + 1)))
    
    # Prepare data
    data = sg.Data.data
    x_col = sg.Data.x_col
    y_col = sg.Data.y_col
    
    # Count plots
    n_plots = 2 + len(plot_scales) + 1  # original + scaling + wavelets + reconstructed
    nrow = int(np.ceil(n_plots / ncol))
    
    if figsize is None:
        figsize = (4 * ncol, 4 * nrow)
    
    fig, axes = plt.subplots(nrow, ncol, figsize=figsize)
    axes = axes.flatten() if n_plots > 1 else [axes]
    
    plot_idx = 0
    
    # Plot original signal
    ax = axes[plot_idx]
    scatter = ax.scatter(data[x_col], data[y_col], c=data[signal_name], cmap=cmap, s=10)
    ax.set_title(f"Original Signal: {signal_name}")
    ax.set_aspect("equal")
    ax.axis("off")
    plot_idx += 1
    
    # Plot scaling function (low-pass)
    ax = axes[plot_idx]
    scatter = ax.scatter(data[x_col], data[y_col], c=np.real(coefficients["low_pass"]), cmap=cmap, s=10)
    ax.set_title("Low-pass (Scaling)")
    ax.set_aspect("equal")
    ax.axis("off")
    plot_idx += 1
    
    # Plot wavelet coefficients
    for i in plot_scales:
        wavelet_name = f"wavelet_{i}"
        if wavelet_name in coefficients:
            ax = axes[plot_idx]
            scatter = ax.scatter(data[x_col], data[y_col], c=np.real(coefficients[wavelet_name]), cmap=cmap, s=10)
            ax.set_title(f"Band-pass Scale {i}")
            ax.set_aspect("equal")
            ax.axis("off")
            plot_idx += 1
    
    # Plot reconstructed signal
    ax = axes[plot_idx]
    scatter = ax.scatter(data[x_col], data[y_col], c=inverse_result["reconstructed_signal"], cmap=cmap, s=10)
    error = inverse_result.get("reconstruction_error")
    if error is not None:
        if hasattr(error, "__len__"):
            error = float(np.mean(error))
        ax.set_title(f"Reconstructed (RMSE: {error:.4f})")
    else:
        ax.set_title("Reconstructed")
    ax.set_aspect("equal")
    ax.axis("off")
    plot_idx += 1
    
    # Hide unused axes
    for i in range(plot_idx, len(axes)):
        axes[i].axis("off")
    
    plt.tight_layout()
    return fig


def plot_fourier_modes(
    sg: Any,
    mode_type: str = "both",
    n_modes: int = 6,
    ncol: int = 3,
    point_size: float = 10,
    figsize: Optional[Tuple[float, float]] = None,
    cmap: str = "plasma",
) -> Any:
    """
    Plot Fourier modes (eigenvectors) from SGWT object.
    
    Parameters
    ----------
    sg : SGWT
        SGWT object with Graph slot computed.
    mode_type : str
        Type of modes: "low", "high", or "both".
    n_modes : int
        Number of modes to plot for each type.
    ncol : int
        Number of columns in plot layout.
    point_size : float
        Size of points.
    figsize : tuple, optional
        Figure size.
    cmap : str
        Colormap name.
        
    Returns
    -------
    matplotlib.figure.Figure
        Figure with Fourier mode plots.
    """
    _check_matplotlib()
    
    if sg.Graph is None:
        raise ValueError("SGWT object must have Graph slot computed. Run run_spec_graph() first.")
    
    # Extract components
    data = sg.Data.data
    x_col = sg.Data.x_col
    y_col = sg.Data.y_col
    eigenvalues = sg.Graph.eigenvalues
    eigenvectors = sg.Graph.eigenvectors
    
    if mode_type not in ["low", "high", "both"]:
        raise ValueError("mode_type must be 'low', 'high', or 'both'")
    
    n_total = len(eigenvalues)
    n_modes = min(n_modes, n_total // 2)
    
    # Determine which modes to plot
    modes_to_plot = []
    
    if mode_type in ["low", "both"]:
        # Skip DC component (first mode)
        low_indices = list(range(1, n_modes + 1))
        modes_to_plot.extend([("Low Freq", i, eigenvalues[i], eigenvectors[:, i]) for i in low_indices if i < n_total])
    
    if mode_type in ["high", "both"]:
        high_indices = list(range(n_total - n_modes, n_total))
        modes_to_plot.extend([("High Freq", i, eigenvalues[i], eigenvectors[:, i]) for i in high_indices if i >= 0])
    
    n_plots = len(modes_to_plot)
    if n_plots == 0:
        raise ValueError("No plots to create. Check parameters.")
    
    nrow = int(np.ceil(n_plots / ncol))
    
    if figsize is None:
        figsize = (4 * ncol, 4 * nrow)
    
    fig, axes = plt.subplots(nrow, ncol, figsize=figsize)
    if n_plots > 1:
        axes = axes.flatten()
    else:
        axes = [axes]
    
    for idx, (mode_label, mode_idx, eigenval, mode_vec) in enumerate(modes_to_plot):
        ax = axes[idx]
        scatter = ax.scatter(data[x_col], data[y_col], c=mode_vec, cmap=cmap, s=point_size)
        ax.set_title(f"{mode_label} {mode_idx}\nλ = {eigenval:.4f}")
        ax.set_aspect("equal")
        ax.axis("off")
    
    # Hide unused axes
    for i in range(len(modes_to_plot), len(axes)):
        axes[i].axis("off")
    
    # Add main title
    title_map = {
        "low": f"Low-Frequency Fourier Modes (n = {n_modes})",
        "high": f"High-Frequency Fourier Modes (n = {n_modes})",
        "both": f"Fourier Modes: Low & High Frequency (n = {n_modes} each)",
    }
    fig.suptitle(title_map[mode_type], fontsize=14, fontweight="bold")
    
    plt.tight_layout()
    return fig


def visualize_sgwt_kernels(
    eigenvalues: np.ndarray,
    scales: Optional[np.ndarray] = None,
    J: int = 4,
    scaling_factor: float = 2.0,
    kernel_type: str = "heat",
    lmax: Optional[float] = None,
    eigenvalue_range: Optional[Tuple[float, float]] = None,
    resolution: int = 1000,
    figsize: Tuple[float, float] = (12, 8),
) -> Dict:
    """
    Visualize SGWT kernels and scaling functions.
    
    Parameters
    ----------
    eigenvalues : np.ndarray
        Eigenvalues from graph Laplacian.
    scales : np.ndarray, optional
        Wavelet scales. If None, auto-generated.
    J : int
        Number of scales if auto-generating.
    scaling_factor : float
        Scaling factor between scales.
    kernel_type : str
        Kernel type.
    lmax : float, optional
        Maximum eigenvalue.
    eigenvalue_range : tuple, optional
        Range for plotting.
    resolution : int
        Number of points for smooth curves.
    figsize : tuple
        Figure size.
        
    Returns
    -------
    dict
        Dictionary containing plot and filter values.
    """
    _check_matplotlib()
    from .core import sgwt_auto_scales, compute_sgwt_filters
    
    eigenvalues = np.asarray(eigenvalues)
    
    if lmax is None:
        lmax = np.max(eigenvalues) * 0.95
    
    if scales is None:
        scales = sgwt_auto_scales(lmax, J, scaling_factor)
    
    if eigenvalue_range is None:
        eigenvalue_range = (0, np.max(eigenvalues))
    
    # Create smooth eigenvalue sequence
    lambda_smooth = np.linspace(eigenvalue_range[0], eigenvalue_range[1], resolution)
    
    # Compute filters
    filters_smooth = compute_sgwt_filters(lambda_smooth, scales, lmax, kernel_type)
    
    # Create figure
    fig, axes = plt.subplots(2, 1, figsize=figsize, gridspec_kw={"height_ratios": [3, 1]})
    
    # Plot kernels
    ax1 = axes[0]
    colors = plt.cm.viridis(np.linspace(0, 1, len(filters_smooth)))
    colors[0] = [0.9, 0.3, 0.2, 1.0]  # Red for scaling
    
    ax1.plot(lambda_smooth, filters_smooth[0], color=colors[0], linewidth=2, label="Scaling Function")
    for j, filt in enumerate(filters_smooth[1:], 1):
        ax1.plot(lambda_smooth, filt, color=colors[j], linewidth=2, label=f"Wavelet Scale {j}")
    
    # Add vertical lines for eigenvalues (sample)
    n_sample = min(50, len(eigenvalues))
    sample_idx = np.linspace(0, len(eigenvalues) - 1, n_sample, dtype=int)
    for ev in eigenvalues[sample_idx]:
        ax1.axvline(ev, alpha=0.3, color="gray", linewidth=0.3)
    
    ax1.set_xlabel("Eigenvalue (λ)")
    ax1.set_ylabel("Filter Response")
    ax1.set_title(f"SGWT Filter Bank: Scaling Function and Wavelet Kernels\n"
                  f"Kernel Type: {kernel_type} | J = {len(scales)} | Scaling Factor = {scaling_factor}")
    ax1.legend(loc="upper right", ncol=2)
    ax1.grid(True, alpha=0.3)
    
    # Plot eigenvalue histogram
    ax2 = axes[1]
    ax2.hist(eigenvalues, bins=50, color="steelblue", alpha=0.7, edgecolor="white")
    ax2.set_xlabel("Eigenvalue (λ)")
    ax2.set_ylabel("Count")
    ax2.set_title("Eigenvalue Distribution")
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    # Compute filters for actual eigenvalues
    filters_actual = compute_sgwt_filters(eigenvalues, scales, lmax, kernel_type)
    
    return {
        "plot": fig,
        "filters_smooth": filters_smooth,
        "filters_actual": filters_actual,
        "lambda_smooth": lambda_smooth,
        "scales": scales,
        "parameters": {
            "J": len(scales),
            "scaling_factor": scaling_factor,
            "kernel_type": kernel_type,
            "lmax": lmax,
        },
    }


def visualize_similarity_xy(
    similarity_results: Union[Dict, List[Dict]],
    point_size: float = 50,
    point_color: str = "steelblue",
    add_diagonal: bool = True,
    add_axes_lines: bool = True,
    title: str = "Low-frequency vs Non-low-frequency Similarity",
    show_labels: bool = False,
    figsize: Tuple[float, float] = (8, 8),
) -> Any:
    """
    Visualize similarity in low vs non-low frequency space.
    
    Parameters
    ----------
    similarity_results : dict or list of dict
        Similarity results from run_sgcc.
    point_size : float
        Size of points.
    point_color : str
        Color of points.
    add_diagonal : bool
        Whether to add diagonal lines.
    add_axes_lines : bool
        Whether to add axis lines at 0.
    title : str
        Plot title.
    show_labels : bool
        Whether to show point labels.
    figsize : tuple
        Figure size.
        
    Returns
    -------
    matplotlib.figure.Figure
        Figure with similarity plot.
    """
    _check_matplotlib()
    
    # Handle single result vs list
    if isinstance(similarity_results, dict) and "c_low" in similarity_results:
        similarity_results = {"result": similarity_results}
    elif isinstance(similarity_results, list):
        similarity_results = {f"result_{i}": r for i, r in enumerate(similarity_results)}
    
    # Extract data
    c_low_vals = []
    c_nonlow_vals = []
    labels = []
    
    for name, result in similarity_results.items():
        if "c_low" in result and "c_nonlow" in result:
            c_low_vals.append(result["c_low"] * result.get("w_low", 1.0))
            c_nonlow_vals.append(result["c_nonlow"] * result.get("w_NL", 1.0))
            labels.append(name)
    
    if len(c_low_vals) == 0:
        raise ValueError("No valid similarity results found")
    
    fig, ax = plt.subplots(figsize=figsize)
    
    ax.scatter(c_low_vals, c_nonlow_vals, s=point_size, c=point_color, alpha=0.7)
    
    if add_axes_lines:
        ax.axhline(0, linestyle="--", color="gray", alpha=0.7)
        ax.axvline(0, linestyle="--", color="gray", alpha=0.7)
    
    if add_diagonal:
        ax.plot([-1, 1], [-1, 1], linestyle=":", color="gray", alpha=0.7)
        ax.plot([-1, 1], [1, -1], linestyle=":", color="gray", alpha=0.7)
    
    if show_labels:
        for x, y, label in zip(c_low_vals, c_nonlow_vals, labels):
            ax.annotate(label, (x, y), fontsize=8, ha="left", va="bottom")
    
    ax.set_xlim(-1, 1)
    ax.set_ylim(-1, 1)
    ax.set_xlabel("Low-frequency Similarity (c_low)")
    ax.set_ylabel("Non-low-frequency Similarity (c_nonlow)")
    ax.set_title(title)
    ax.set_aspect("equal")
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    return fig


def visualize_multiscale(
    sim_data: Dict[str, pd.DataFrame],
    Ra_seq: np.ndarray,
    n_steps: int,
    bg_color: str = "#E5E5E5",
    signal1_color: str = "#16964a",
    signal2_color: str = "#2958a8",
    figsize: Optional[Tuple[float, float]] = None,
) -> Any:
    """
    Visualize multi-scale concentric ring patterns.
    
    Parameters
    ----------
    sim_data : dict
        Output from simulate_multiscale.
    Ra_seq : np.ndarray
        Ra values used in simulation.
    n_steps : int
        Number of steps.
    bg_color : str
        Background color.
    signal1_color : str
        Color for signal_1.
    signal2_color : str
        Color for signal_2.
    figsize : tuple, optional
        Figure size.
        
    Returns
    -------
    matplotlib.figure.Figure
        Figure with pattern visualizations.
    """
    _check_matplotlib()
    
    Ra_seq = np.asarray(Ra_seq)
    n_Ra = len(Ra_seq)
    
    if figsize is None:
        figsize = (2 * n_steps, 2 * n_Ra)
    
    fig, axes = plt.subplots(n_Ra, n_steps, figsize=figsize)
    
    for i, Ra in enumerate(Ra_seq):
        for j in range(1, n_steps + 1):
            key = f"Ra_{Ra}_Step_{j}"
            if key in sim_data:
                df = sim_data[key]
                ax = axes[i, j - 1] if n_Ra > 1 else axes[j - 1]
                
                colors = np.where(df["signal_1"] == 1, signal1_color,
                                 np.where(df["signal_2"] == 1, signal2_color, bg_color))
                
                ax.scatter(df["X"], df["Y"], c=colors, s=1, marker="s")
                ax.set_aspect("equal")
                ax.axis("off")
                
                if j == 1:
                    ax.set_ylabel(f"Ra={Ra}", fontsize=8)
                if i == 0:
                    ax.set_title(f"Step {j}", fontsize=8)
    
    plt.tight_layout()
    return fig


def visualize_stripe_patterns(
    sim_data: Dict[str, pd.DataFrame],
    gap_seq: List[float],
    width_seq: List[float],
    theta_seq: List[float],
    bg_color: str = "gray",
    signal1_color: str = "#1f6f8b",
    signal2_color: str = "#e67e22",
    overlap_color: str = "#7a4dbf",
    figsize: Optional[Tuple[float, float]] = None,
) -> Any:
    """
    Visualize stripe pattern simulation results.
    
    Parameters
    ----------
    sim_data : dict
        Output from simulate_stripe_patterns.
    gap_seq : list
        Gap values.
    width_seq : list
        Width values.
    theta_seq : list
        Theta values.
    bg_color : str
        Background color.
    signal1_color : str
        Color for signal_1.
    signal2_color : str
        Color for signal_2.
    overlap_color : str
        Color for overlap.
    figsize : tuple, optional
        Figure size.
        
    Returns
    -------
    matplotlib.figure.Figure
        Figure with pattern visualizations.
    """
    _check_matplotlib()
    
    n_patterns = len(gap_seq) * len(width_seq) * len(theta_seq)
    ncol = len(width_seq)
    nrow = int(np.ceil(n_patterns / ncol))
    
    if figsize is None:
        figsize = (3 * ncol, 3 * nrow)
    
    fig, axes = plt.subplots(nrow, ncol, figsize=figsize)
    axes = axes.flatten() if n_patterns > 1 else [axes]
    
    idx = 0
    for gap in gap_seq:
        for width in width_seq:
            for theta in theta_seq:
                key = f"gap_{gap}_w_{width}_th_{theta}"
                if key in sim_data and idx < len(axes):
                    df = sim_data[key]
                    ax = axes[idx]
                    
                    # Determine colors
                    colors = np.full(len(df), bg_color)
                    mask1 = (df["signal_1"] == 1) & (df["signal_2"] == 0)
                    mask2 = (df["signal_2"] == 1) & (df["signal_1"] == 0)
                    mask_overlap = (df["signal_1"] == 1) & (df["signal_2"] == 1)
                    
                    colors[mask1] = signal1_color
                    colors[mask2] = signal2_color
                    colors[mask_overlap] = overlap_color
                    
                    ax.scatter(df["X"], df["Y"], c=colors, s=1, marker="s")
                    ax.set_aspect("equal")
                    ax.axis("off")
                    ax.set_title(f"g={gap}, w={width}, θ={theta}", fontsize=8)
                    
                    idx += 1
    
    # Hide unused axes
    for i in range(idx, len(axes)):
        axes[i].axis("off")
    
    plt.tight_layout()
    return fig


def visualize_moving_circles(
    sim_data: Dict[str, pd.DataFrame],
    bg_color: str = "#E5E5E5",
    signal1_color: str = "#16964a",
    signal2_color: str = "#2958a8",
    figsize: Optional[Tuple[float, float]] = None,
) -> Any:
    """
    Visualize moving circles patterns.
    
    Parameters
    ----------
    sim_data : dict
        Output from simulate_moving_circles.
    bg_color : str
        Background color.
    signal1_color : str
        Color for signal_1.
    signal2_color : str
        Color for signal_2.
    figsize : tuple, optional
        Figure size.
        
    Returns
    -------
    matplotlib.figure.Figure
        Figure with pattern visualizations.
    """
    _check_matplotlib()
    
    # Extract unique Ra and Step values
    Ra_values = sorted(set(df["Ra"].iloc[0] for df in sim_data.values()))
    Step_values = sorted(set(df["Step"].iloc[0] for df in sim_data.values()))
    
    n_Ra = len(Ra_values)
    n_steps = len(Step_values)
    
    if figsize is None:
        figsize = (2 * n_steps, 2 * n_Ra)
    
    fig, axes = plt.subplots(n_Ra, n_steps, figsize=figsize)
    
    for i, Ra in enumerate(Ra_values):
        for j, step in enumerate(Step_values):
            key = f"Ra_{Ra}_Step_{step}"
            if key in sim_data:
                df = sim_data[key]
                ax = axes[i, j] if n_Ra > 1 else (axes[j] if n_steps > 1 else axes)
                
                colors = np.where(df["signal_1"] == 1, signal1_color,
                                 np.where(df["signal_2"] == 1, signal2_color, bg_color))
                
                ax.scatter(df["X"], df["Y"], c=colors, s=1, marker="s")
                ax.set_aspect("equal")
                ax.axis("off")
    
    plt.tight_layout()
    return fig


def visualize_checkerboard(
    df: pd.DataFrame,
    color1: str = "black",
    color2: str = "white",
    figsize: Tuple[float, float] = (6, 6),
) -> Any:
    """
    Visualize checkerboard pattern.
    
    Parameters
    ----------
    df : pd.DataFrame
        DataFrame from simulate_checkerboard.
    color1 : str
        Color for signal_1.
    color2 : str
        Color for signal_2.
    figsize : tuple
        Figure size.
        
    Returns
    -------
    matplotlib.figure.Figure
        Figure with checkerboard visualization.
    """
    _check_matplotlib()
    
    fig, ax = plt.subplots(figsize=figsize)
    
    colors = np.where(df["signal_1"] == 1, color1, color2)
    ax.scatter(df["X"], df["Y"], c=colors, s=10, marker="s")
    ax.set_aspect("equal")
    ax.axis("off")
    
    plt.tight_layout()
    return fig
