import logging
import multiprocessing as mp
import numbers
import os
from dataclasses import dataclass
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
from canns_ripser import ripser
from matplotlib import animation, cm, gridspec
from numpy.exceptions import AxisError

# from ripser import ripser
from scipy import signal
from scipy.ndimage import (
    _nd_image,
    _ni_support,
    binary_closing,
    gaussian_filter,
    gaussian_filter1d,
)
from scipy.ndimage._filters import _invalid_origin
from scipy.sparse import coo_matrix
from scipy.sparse.linalg import lsmr
from scipy.spatial.distance import pdist, squareform
from scipy.stats import binned_statistic_2d, multivariate_normal
from sklearn import preprocessing
from tqdm import tqdm

# Import PlotConfig for unified plotting
from ..plotting import PlotConfig


# ==================== Configuration Classes ====================
@dataclass
class SpikeEmbeddingConfig:
    """Configuration for spike train embedding."""

    res: int = 100000
    dt: int = 1000
    sigma: int = 5000
    smooth: bool = True
    speed_filter: bool = True
    min_speed: float = 2.5


@dataclass
class TDAConfig:
    """Configuration for Topological Data Analysis."""

    dim: int = 6
    num_times: int = 5
    active_times: int = 15000
    k: int = 1000
    n_points: int = 1200
    metric: str = "cosine"
    nbs: int = 800
    maxdim: int = 1
    coeff: int = 47
    show: bool = True
    do_shuffle: bool = False
    num_shuffles: int = 1000
    progress_bar: bool = True


@dataclass
class CANN2DPlotConfig(PlotConfig):
    """Specialized PlotConfig for CANN2D visualizations."""

    # 3D projection specific parameters
    zlabel: str = "Component 3"
    dpi: int = 300

    # Torus animation specific parameters
    numangsint: int = 51
    r1: float = 1.5  # Major radius
    r2: float = 1.0  # Minor radius
    window_size: int = 300
    frame_step: int = 5
    n_frames: int = 20

    @classmethod
    def for_projection_3d(cls, **kwargs) -> "CANN2DPlotConfig":
        """Create configuration for 3D projection plots."""
        defaults = {
            "title": "3D Data Projection",
            "xlabel": "Component 1",
            "ylabel": "Component 2",
            "zlabel": "Component 3",
            "figsize": (10, 8),
            "dpi": 300,
        }
        defaults.update(kwargs)
        return cls.for_static_plot(**defaults)

    @classmethod
    def for_torus_animation(cls, **kwargs) -> "CANN2DPlotConfig":
        """Create configuration for 3D torus bump animations."""
        defaults = {
            "title": "3D Bump on Torus",
            "figsize": (8, 8),
            "fps": 5,
            "repeat": True,
            "show_progress_bar": True,
            "numangsint": 51,
            "r1": 1.5,
            "r2": 1.0,
            "window_size": 300,
            "frame_step": 5,
            "n_frames": 20,
        }
        defaults.update(kwargs)
        time_steps = kwargs.get("time_steps_per_second", 1000)
        config = cls.for_animation(time_steps, **defaults)
        # Add torus-specific attributes
        config.numangsint = defaults["numangsint"]
        config.r1 = defaults["r1"]
        config.r2 = defaults["r2"]
        config.window_size = defaults["window_size"]
        config.frame_step = defaults["frame_step"]
        config.n_frames = defaults["n_frames"]
        return config


# ==================== Constants ====================
class Constants:
    """Constants used throughout CANN2D analysis."""

    DEFAULT_FIGSIZE = (10, 8)
    DEFAULT_DPI = 300
    GAUSSIAN_SIGMA_FACTOR = 100
    SPEED_CONVERSION_FACTOR = 100
    TIME_CONVERSION_FACTOR = 0.01
    MULTIPROCESSING_CORES = 4


# ==================== Custom Exceptions ====================
class CANN2DError(Exception):
    """Base exception for CANN2D analysis errors."""

    pass


class DataLoadError(CANN2DError):
    """Raised when data loading fails."""

    pass


class ProcessingError(CANN2DError):
    """Raised when data processing fails."""

    pass


try:
    from numba import jit, njit, prange

    HAS_NUMBA = True
except ImportError:
    HAS_NUMBA = False
    print(
        "Using numba for FAST CANN2D analysis, now using pure numpy implementation.",
        "Try numba by `pip install numba` to speed up the process.",
    )

    # Create dummy decorators if numba is not available
    def jit(*args, **kwargs):
        def decorator(func):
            return func

        return decorator

    def njit(*args, **kwargs):
        def decorator(func):
            return func

        return decorator

    def prange(x):
        return range(x)


def embed_spike_trains(spike_trains, config: SpikeEmbeddingConfig | None = None, **kwargs):
    """
    Load and preprocess spike train data from npz file.

    This function converts raw spike times into a time-binned spike matrix,
    optionally applying Gaussian smoothing and filtering based on animal movement speed.

    Parameters:
        spike_trains : dict containing 'spike', 't', and optionally 'x', 'y'.
        config : SpikeEmbeddingConfig, optional configuration object
        **kwargs : backward compatibility parameters

    Returns:
        spikes_bin (ndarray): Binned and optionally smoothed spike matrix of shape (T, N).
        xx (ndarray, optional): X coordinates (if speed_filter=True).
        yy (ndarray, optional): Y coordinates (if speed_filter=True).
        tt (ndarray, optional): Time points (if speed_filter=True).
    """
    # Handle backward compatibility and configuration
    if config is None:
        config = SpikeEmbeddingConfig(
            res=kwargs.get("res", 100000),
            dt=kwargs.get("dt", 1000),
            sigma=kwargs.get("sigma", 5000),
            smooth=kwargs.get("smooth0", True),
            speed_filter=kwargs.get("speed0", True),
            min_speed=kwargs.get("min_speed", 2.5),
        )

    try:
        # Step 1: Extract and filter spike data
        spikes_filtered = _extract_spike_data(spike_trains, config)

        # Step 2: Create time bins
        time_bins = _create_time_bins(spike_trains["t"], config)

        # Step 3: Bin spike data
        spikes_bin = _bin_spike_data(spikes_filtered, time_bins, config)

        # Step 4: Apply temporal smoothing if requested
        if config.smooth:
            spikes_bin = _apply_temporal_smoothing(spikes_bin, config)

        # Step 5: Apply speed filtering if requested
        if config.speed_filter:
            return _apply_speed_filtering(spikes_bin, spike_trains, config)

        return spikes_bin

    except Exception as e:
        raise ProcessingError(f"Failed to embed spike trains: {e}") from e


def _extract_spike_data(
    spike_trains: dict[str, Any], config: SpikeEmbeddingConfig
) -> dict[int, np.ndarray]:
    """Extract and filter spike data within time window."""
    try:
        # Handle different spike data formats
        spike_data = spike_trains["spike"]
        if hasattr(spike_data, "item") and callable(spike_data.item):
            # numpy array with .item() method (from npz file)
            spikes_all = spike_data[()]
        elif isinstance(spike_data, dict):
            # Already a dictionary
            spikes_all = spike_data
        elif isinstance(spike_data, list | np.ndarray):
            # List or array format
            spikes_all = spike_data
        else:
            # Try direct access
            spikes_all = spike_data

        t = spike_trains["t"]

        min_time0 = np.min(t)
        max_time0 = np.max(t)

        # Extract spike intervals for each cell
        if isinstance(spikes_all, dict):
            # Dictionary format
            spikes = {}
            for i, key in enumerate(spikes_all.keys()):
                s = np.array(spikes_all[key])
                spikes[i] = s[(s >= min_time0) & (s < max_time0)]
        else:
            # List/array format
            cell_inds = np.arange(len(spikes_all))
            spikes = {}

            for i, m in enumerate(cell_inds):
                s = np.array(spikes_all[m]) if len(spikes_all[m]) > 0 else np.array([])
                # Filter spikes within time window
                if len(s) > 0:
                    spikes[i] = s[(s >= min_time0) & (s < max_time0)]
                else:
                    spikes[i] = np.array([])

        return spikes

    except KeyError as e:
        raise DataLoadError(f"Missing required data key: {e}") from e
    except Exception as e:
        raise ProcessingError(f"Error extracting spike data: {e}") from e


def _create_time_bins(t: np.ndarray, config: SpikeEmbeddingConfig) -> np.ndarray:
    """Create time bins for spike discretization."""
    min_time0 = np.min(t)
    max_time0 = np.max(t)

    min_time = min_time0 * config.res
    max_time = max_time0 * config.res

    return np.arange(np.floor(min_time), np.ceil(max_time) + 1, config.dt)


def _bin_spike_data(
    spikes: dict[int, np.ndarray], time_bins: np.ndarray, config: SpikeEmbeddingConfig
) -> np.ndarray:
    """Convert spike times to binned spike matrix."""
    min_time = time_bins[0]
    max_time = time_bins[-1]

    spikes_bin = np.zeros((len(time_bins), len(spikes)), dtype=int)

    for n in spikes:
        spike_times = np.array(spikes[n] * config.res - min_time, dtype=int)
        # Filter valid spike times
        spike_times = spike_times[(spike_times < (max_time - min_time)) & (spike_times > 0)]
        spike_times = np.array(spike_times / config.dt, int)

        # Bin spikes
        for j in spike_times:
            if j < len(time_bins):
                spikes_bin[j, n] += 1

    return spikes_bin


def _apply_temporal_smoothing(spikes_bin: np.ndarray, config: SpikeEmbeddingConfig) -> np.ndarray:
    """Apply Gaussian temporal smoothing to spike matrix."""
    # Calculate smoothing parameters (legacy implementation used custom kernel)
    # Current implementation uses scipy's gaussian_filter1d for better performance

    # Apply smoothing (simplified version - could be further optimized)
    smoothed = np.zeros((spikes_bin.shape[0], spikes_bin.shape[1]))

    # Use scipy's gaussian_filter1d for better performance

    sigma_bins = config.sigma / config.dt

    for n in range(spikes_bin.shape[1]):
        smoothed[:, n] = gaussian_filter1d(
            spikes_bin[:, n].astype(float), sigma=sigma_bins, mode="constant"
        )

    # Normalize
    normalization_factor = 1 / np.sqrt(2 * np.pi * (config.sigma / config.res) ** 2)
    return smoothed * normalization_factor


def _apply_speed_filtering(
    spikes_bin: np.ndarray, spike_trains: dict[str, Any], config: SpikeEmbeddingConfig
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Apply speed-based filtering to spike data."""
    try:
        xx, yy, tt_pos, speed = _load_pos(
            spike_trains["t"], spike_trains["x"], spike_trains["y"], res=config.res, dt=config.dt
        )

        valid = speed > config.min_speed

        return (spikes_bin[valid, :], xx[valid], yy[valid], tt_pos[valid])

    except KeyError as e:
        raise DataLoadError(f"Missing position data for speed filtering: {e}") from e
    except Exception as e:
        raise ProcessingError(f"Error in speed filtering: {e}") from e


def plot_projection(
    reduce_func,
    embed_data,
    config: CANN2DPlotConfig | None = None,
    title="Projection (3D)",
    xlabel="Component 1",
    ylabel="Component 2",
    zlabel="Component 3",
    save_path=None,
    show=True,
    dpi=300,
    figsize=(10, 8),
    **kwargs,
):
    """
    Plot a 3D projection of the embedded data.

    Parameters:
        reduce_func (callable): Function to reduce the dimensionality of the data.
        embed_data (ndarray): Data to be projected.
        config (PlotConfig, optional): Configuration object for unified plotting parameters
        **kwargs: backward compatibility parameters
        title (str): Title of the plot.
        xlabel (str): Label for the x-axis.
        ylabel (str): Label for the y-axis.
        zlabel (str): Label for the z-axis.
        save_path (str, optional): Path to save the plot. If None, plot will not be saved.
        show (bool): Whether to display the plot.
        dpi (int): Dots per inch for saving the figure.
        figsize (tuple): Size of the figure.

    Returns:
        fig: The created figure object.
    """

    # Handle backward compatibility and configuration
    if config is None:
        config = CANN2DPlotConfig.for_projection_3d(
            title=title,
            xlabel=xlabel,
            ylabel=ylabel,
            zlabel=zlabel,
            save_path=save_path,
            show=show,
            figsize=figsize,
            dpi=dpi,
            **kwargs,
        )

    reduced_data = reduce_func(embed_data[::5])

    fig = plt.figure(figsize=config.figsize)
    ax = fig.add_subplot(111, projection="3d")
    ax.scatter(reduced_data[:, 0], reduced_data[:, 1], reduced_data[:, 2], s=1, alpha=0.5)

    ax.set_title(config.title)
    ax.set_xlabel(config.xlabel)
    ax.set_ylabel(config.ylabel)
    ax.set_zlabel(config.zlabel)

    if config.save_path is None and config.show is None:
        raise ValueError("Either save path or show must be provided.")
    if config.save_path:
        plt.savefig(config.save_path, dpi=config.dpi)
    if config.show:
        plt.show()

    plt.close(fig)

    return fig


def tda_vis(embed_data: np.ndarray, config: TDAConfig | None = None, **kwargs) -> dict[str, Any]:
    """
    Topological Data Analysis visualization with optional shuffle testing.

    Parameters:
        embed_data : ndarray
            Embedded spike train data.
        config : TDAConfig, optional
            Configuration object with all TDA parameters
        **kwargs : backward compatibility parameters

    Returns:
        dict : Dictionary containing:
            - persistence: persistence diagrams from real data
            - indstemp: indices of sampled points
            - movetimes: selected time points
            - n_points: number of sampled points
            - shuffle_max: shuffle analysis results (if do_shuffle=True, otherwise None)
    """
    # Handle backward compatibility and configuration
    if config is None:
        config = TDAConfig(
            dim=kwargs.get("dim", 6),
            num_times=kwargs.get("num_times", 5),
            active_times=kwargs.get("active_times", 15000),
            k=kwargs.get("k", 1000),
            n_points=kwargs.get("n_points", 1200),
            metric=kwargs.get("metric", "cosine"),
            nbs=kwargs.get("nbs", 800),
            maxdim=kwargs.get("maxdim", 1),
            coeff=kwargs.get("coeff", 47),
            show=kwargs.get("show", True),
            do_shuffle=kwargs.get("do_shuffle", False),
            num_shuffles=kwargs.get("num_shuffles", 1000),
            progress_bar=kwargs.get("progress_bar", True),
        )

    try:
        # Compute persistent homology for real data
        print("Computing persistent homology for real data...")
        real_persistence = _compute_real_persistence(embed_data, config)

        # Perform shuffle analysis if requested
        shuffle_max = None
        if config.do_shuffle:
            shuffle_max = _perform_shuffle_analysis(embed_data, config)

        # Visualization
        _handle_visualization(real_persistence["persistence"], shuffle_max, config)

        # Return results as dictionary
        return {
            "persistence": real_persistence["persistence"],
            "indstemp": real_persistence["indstemp"],
            "movetimes": real_persistence["movetimes"],
            "n_points": real_persistence["n_points"],
            "shuffle_max": shuffle_max,
        }

    except Exception as e:
        raise ProcessingError(f"TDA analysis failed: {e}") from e


def _compute_real_persistence(embed_data: np.ndarray, config: TDAConfig) -> dict[str, Any]:
    """Compute persistent homology for real data with progress tracking."""

    logging.info("Processing real data - Starting TDA analysis (5 steps)")

    # Step 1: Time point downsampling
    logging.info("Step 1/5: Time point downsampling")
    times_cube = _downsample_timepoints(embed_data, config.num_times)

    # Step 2: Select most active time points
    logging.info("Step 2/5: Selecting active time points")
    movetimes = _select_active_timepoints(embed_data, times_cube, config.active_times)

    # Step 3: PCA dimensionality reduction
    logging.info("Step 3/5: PCA dimensionality reduction")
    dimred = _apply_pca_reduction(embed_data, movetimes, config.dim)

    # Step 4: Point cloud sampling (denoising)
    logging.info("Step 4/5: Point cloud denoising")
    indstemp = _apply_denoising(dimred, config)

    # Step 5: Compute persistent homology
    logging.info("Step 5/5: Computing persistent homology")
    persistence = _compute_persistence_homology(dimred, indstemp, config)

    logging.info("TDA analysis completed successfully")

    # Return all necessary data in dictionary format
    return {
        "persistence": persistence,
        "indstemp": indstemp,
        "movetimes": movetimes,
        "n_points": config.n_points,
    }


def _downsample_timepoints(embed_data: np.ndarray, num_times: int) -> np.ndarray:
    """Downsample timepoints for computational efficiency."""
    return np.arange(0, embed_data.shape[0], num_times)


def _select_active_timepoints(
    embed_data: np.ndarray, times_cube: np.ndarray, active_times: int
) -> np.ndarray:
    """Select most active timepoints based on total activity."""
    activity_scores = np.sum(embed_data[times_cube, :], 1)
    # Match external TDAvis: sort indices first, then map to times_cube
    movetimes = np.sort(np.argsort(activity_scores)[-active_times:])
    return times_cube[movetimes]


def _apply_pca_reduction(embed_data: np.ndarray, movetimes: np.ndarray, dim: int) -> np.ndarray:
    """Apply PCA dimensionality reduction."""
    scaled_data = preprocessing.scale(embed_data[movetimes, :])
    dimred, *_ = _pca(scaled_data, dim=dim)
    return dimred


def _apply_denoising(dimred: np.ndarray, config: TDAConfig) -> np.ndarray:
    """Apply point cloud denoising."""
    indstemp, *_ = _sample_denoising(
        dimred,
        k=config.k,
        num_sample=config.n_points,
        omega=1,  # Match external TDAvis: uses 1, not default 0.2
        metric=config.metric,
    )
    return indstemp


def _compute_persistence_homology(
    dimred: np.ndarray, indstemp: np.ndarray, config: TDAConfig
) -> dict[str, Any]:
    """Compute persistent homology using ripser."""
    d = _second_build(dimred, indstemp, metric=config.metric, nbs=config.nbs)
    np.fill_diagonal(d, 0)

    return ripser(
        d,
        maxdim=config.maxdim,
        coeff=config.coeff,
        do_cocycles=True,
        distance_matrix=True,
        progress_bar=config.progress_bar,
    )


def _perform_shuffle_analysis(embed_data: np.ndarray, config: TDAConfig) -> dict[int, Any]:
    """Perform shuffle analysis with progress tracking."""
    print(f"\nStarting shuffle analysis with {config.num_shuffles} iterations...")

    # Create parameters dict for shuffle analysis
    shuffle_params = {
        "dim": config.dim,
        "num_times": config.num_times,
        "active_times": config.active_times,
        "k": config.k,
        "n_points": config.n_points,
        "metric": config.metric,
        "nbs": config.nbs,
        "maxdim": config.maxdim,
        "coeff": config.coeff,
    }

    shuffle_max = _run_shuffle_analysis(
        embed_data,
        num_shuffles=config.num_shuffles,
        num_cores=Constants.MULTIPROCESSING_CORES,
        progress_bar=config.progress_bar,
        **shuffle_params,
    )

    # Print shuffle analysis summary
    _print_shuffle_summary(shuffle_max)

    return shuffle_max


def _print_shuffle_summary(shuffle_max: dict[int, Any]) -> None:
    """Print summary of shuffle analysis results."""
    print("\nSummary of shuffle-based analysis:")
    for dim_idx in [0, 1, 2]:
        if shuffle_max and dim_idx in shuffle_max and shuffle_max[dim_idx]:
            values = shuffle_max[dim_idx]
            print(
                f"H{dim_idx}: {len(values)} valid iterations | "
                f"Mean maximum persistence: {np.mean(values):.4f} | "
                f"99.9th percentile: {np.percentile(values, 99.9):.4f}"
            )


def _handle_visualization(
    real_persistence: dict[str, Any], shuffle_max: dict[int, Any] | None, config: TDAConfig
) -> None:
    """Handle visualization based on configuration."""
    if config.show:
        if config.do_shuffle and shuffle_max is not None:
            _plot_barcode_with_shuffle(real_persistence, shuffle_max)
        else:
            _plot_barcode(real_persistence)
        plt.show()
    else:
        plt.close()


def _load_pos(t, x, y, res=100000, dt=1000):
    """
    Compute animal position and speed from spike data file.

    Interpolates animal positions to match spike time bins and computes smoothed velocity vectors and speed.

    Parameters:
        t (ndarray): Time points of the spikes (in seconds).
        x (ndarray): X coordinates of the animal's position.
        y (ndarray): Y coordinates of the animal's position.
        res (int): Time scaling factor to align with spike resolution.
        dt (int): Temporal bin size in microseconds.

    Returns:
        xx (ndarray): Interpolated x positions.
        yy (ndarray): Interpolated y positions.
        tt (ndarray): Corresponding time points (in seconds).
        speed (ndarray): Speed at each time point (in cm/s).
    """

    min_time0 = np.min(t)
    max_time0 = np.max(t)

    times = np.where((t >= min_time0) & (t < max_time0))
    x = x[times]
    y = y[times]
    t = t[times]

    min_time = min_time0 * res
    max_time = max_time0 * res

    tt = np.arange(np.floor(min_time), np.ceil(max_time) + 1, dt) / res

    idt = np.concatenate(([0], np.digitize(t[1:-1], tt[:]) - 1, [len(tt) + 1]))
    idtt = np.digitize(np.arange(len(tt)), idt) - 1

    idx = np.concatenate((np.unique(idtt), [np.max(idtt) + 1]))
    divisor = np.bincount(idtt)
    steps = 1.0 / divisor[divisor > 0]
    N = np.max(divisor)
    ranges = np.multiply(np.arange(N)[np.newaxis, :], steps[:, np.newaxis])
    ranges[ranges >= 1] = np.nan

    rangesx = x[idx[:-1], np.newaxis] + np.multiply(
        ranges, (x[idx[1:]] - x[idx[:-1]])[:, np.newaxis]
    )
    xx = rangesx[~np.isnan(ranges)]

    rangesy = y[idx[:-1], np.newaxis] + np.multiply(
        ranges, (y[idx[1:]] - y[idx[:-1]])[:, np.newaxis]
    )
    yy = rangesy[~np.isnan(ranges)]

    xxs = _gaussian_filter1d(xx - np.min(xx), sigma=100)
    yys = _gaussian_filter1d(yy - np.min(yy), sigma=100)
    dx = (xxs[1:] - xxs[:-1]) * 100
    dy = (yys[1:] - yys[:-1]) * 100
    speed = np.sqrt(dx**2 + dy**2) / 0.01
    speed = np.concatenate(([speed[0]], speed))
    return xx, yy, tt, speed


def _gaussian_filter1d(
    input,
    sigma,
    axis=-1,
    order=0,
    output=None,
    mode="reflect",
    cval=0.0,
    truncate=4.0,
    *,
    radius=None,
):
    """1-D Gaussian filter.

    Parameters
    ----------
    %(input)s
    sigma : scalar
        standard deviation for Gaussian kernel
    %(axis)s
    order : int, optional
        An order of 0 corresponds to convolution with a Gaussian
        kernel. A positive order corresponds to convolution with
        that derivative of a Gaussian.
    %(output)s
    %(mode_reflect)s
    %(cval)s
    truncate : float, optional
        Truncate the filter at this many standard deviations.
        Default is 4.0.
    radius : None or int, optional
        Radius of the Gaussian kernel. If specified, the size of
        the kernel will be ``2*radius + 1``, and `truncate` is ignored.
        Default is None.

    Returns
    -------
    gaussian_filter1d : ndarray

    Notes
    -----
    The Gaussian kernel will have size ``2*radius + 1`` along each axis. If
    `radius` is None, a default ``radius = round(truncate * sigma)`` will be
    used.

    Examples
    --------
    >>> from scipy.ndimage import gaussian_filter1d
    >>> import numpy as np
    >>> gaussian_filter1d([1.0, 2.0, 3.0, 4.0, 5.0], 1)
    array([ 1.42704095,  2.06782203,  3.        ,  3.93217797,  4.57295905])
    >>> _gaussian_filter1d([1.0, 2.0, 3.0, 4.0, 5.0], 4)
    array([ 2.91948343,  2.95023502,  3.        ,  3.04976498,  3.08051657])
    >>> import matplotlib.pyplot as plt
    >>> rng = np.random.default_rng()
    >>> x = rng.standard_normal(101).cumsum()
    >>> y3 = _gaussian_filter1d(x, 3)
    >>> y6 = _gaussian_filter1d(x, 6)
    >>> plt.plot(x, 'k', label='original data')
    >>> plt.plot(y3, '--', label='filtered, sigma=3')
    >>> plt.plot(y6, ':', label='filtered, sigma=6')
    >>> plt.legend()
    >>> plt.grid()
    >>> plt.show()

    """
    sd = float(sigma)
    # make the radius of the filter equal to truncate standard deviations
    lw = int(truncate * sd + 0.5)
    if radius is not None:
        lw = radius
    if not isinstance(lw, numbers.Integral) or lw < 0:
        raise ValueError("Radius must be a nonnegative integer.")
    # Since we are calling correlate, not convolve, revert the kernel
    weights = _gaussian_kernel1d(sigma, order, lw)[::-1]
    return _correlate1d(input, weights, axis, output, mode, cval, 0)


def _gaussian_kernel1d(sigma, order, radius):
    """
    Computes a 1-D Gaussian convolution kernel.
    """
    if order < 0:
        raise ValueError("order must be non-negative")
    exponent_range = np.arange(order + 1)
    sigma2 = sigma * sigma
    x = np.arange(-radius, radius + 1)
    phi_x = np.exp(-0.5 / sigma2 * x**2)
    phi_x = phi_x / phi_x.sum()

    if order == 0:
        return phi_x
    else:
        # f(x) = q(x) * phi(x) = q(x) * exp(p(x))
        # f'(x) = (q'(x) + q(x) * p'(x)) * phi(x)
        # p'(x) = -1 / sigma ** 2
        # Implement q'(x) + q(x) * p'(x) as a matrix operator and apply to the
        # coefficients of q(x)
        q = np.zeros(order + 1)
        q[0] = 1
        D = np.diag(exponent_range[1:], 1)  # D @ q(x) = q'(x)
        P = np.diag(np.ones(order) / -sigma2, -1)  # P @ q(x) = q(x) * p'(x)
        Q_deriv = D + P
        for _ in range(order):
            q = Q_deriv.dot(q)
        q = (x[:, None] ** exponent_range).dot(q)
        return q * phi_x


def _correlate1d(input, weights, axis=-1, output=None, mode="reflect", cval=0.0, origin=0):
    """Calculate a 1-D correlation along the given axis.

    The lines of the array along the given axis are correlated with the
    given weights.

    Parameters
    ----------
    %(input)s
    weights : array
        1-D sequence of numbers.
    %(axis)s
    %(output)s
    %(mode_reflect)s
    %(cval)s
    %(origin)s

    Returns
    -------
    result : ndarray
        Correlation result. Has the same shape as `input`.

    Examples
    --------
    >>> from scipy.ndimage import correlate1d
    >>> correlate1d([2, 8, 0, 4, 1, 9, 9, 0], weights=[1, 3])
    array([ 8, 26,  8, 12,  7, 28, 36,  9])
    """
    input = np.asarray(input)
    weights = np.asarray(weights)
    complex_input = input.dtype.kind == "c"
    complex_weights = weights.dtype.kind == "c"
    if complex_input or complex_weights:
        if complex_weights:
            weights = weights.conj()
            weights = weights.astype(np.complex128, copy=False)
        kwargs = dict(axis=axis, mode=mode, origin=origin)
        output = _ni_support._get_output(output, input, complex_output=True)
        return _complex_via_real_components(_correlate1d, input, weights, output, cval, **kwargs)

    output = _ni_support._get_output(output, input)
    weights = np.asarray(weights, dtype=np.float64)
    if weights.ndim != 1 or weights.shape[0] < 1:
        raise RuntimeError("no filter weights given")
    if not weights.flags.contiguous:
        weights = weights.copy()
    axis = _normalize_axis_index(axis, input.ndim)
    if _invalid_origin(origin, len(weights)):
        raise ValueError(
            "Invalid origin; origin must satisfy "
            "-(len(weights) // 2) <= origin <= "
            "(len(weights)-1) // 2"
        )
    mode = _ni_support._extend_mode_to_code(mode)
    _nd_image.correlate1d(input, weights, axis, output, mode, cval, origin)
    return output


def _complex_via_real_components(func, input, weights, output, cval, **kwargs):
    """Complex convolution via a linear combination of real convolutions."""
    complex_input = input.dtype.kind == "c"
    complex_weights = weights.dtype.kind == "c"
    if complex_input and complex_weights:
        # real component of the output
        func(input.real, weights.real, output=output.real, cval=np.real(cval), **kwargs)
        output.real -= func(input.imag, weights.imag, output=None, cval=np.imag(cval), **kwargs)
        # imaginary component of the output
        func(input.real, weights.imag, output=output.imag, cval=np.real(cval), **kwargs)
        output.imag += func(input.imag, weights.real, output=None, cval=np.imag(cval), **kwargs)
    elif complex_input:
        func(input.real, weights, output=output.real, cval=np.real(cval), **kwargs)
        func(input.imag, weights, output=output.imag, cval=np.imag(cval), **kwargs)
    else:
        if np.iscomplexobj(cval):
            raise ValueError("Cannot provide a complex-valued cval when the input is real.")
        func(input, weights.real, output=output.real, cval=cval, **kwargs)
        func(input, weights.imag, output=output.imag, cval=cval, **kwargs)
    return output


def _normalize_axis_index(axis, ndim):
    # Check if `axis` is in the correct range and normalize it
    if axis < -ndim or axis >= ndim:
        msg = f"axis {axis} is out of bounds for array of dimension {ndim}"
        raise AxisError(msg)

    if axis < 0:
        axis = axis + ndim
    return axis


def _compute_persistence(
    sspikes,
    dim=6,
    num_times=5,
    active_times=15000,
    k=1000,
    n_points=1200,
    metric="cosine",
    nbs=800,
    maxdim=1,
    coeff=47,
    progress_bar=True,
):
    # Time point downsampling
    times_cube = np.arange(0, sspikes.shape[0], num_times)

    # Select most active time points
    movetimes = np.sort(np.argsort(np.sum(sspikes[times_cube, :], 1))[-active_times:])
    movetimes = times_cube[movetimes]

    # PCA dimensionality reduction
    scaled_data = preprocessing.scale(sspikes[movetimes, :])
    dimred, *_ = _pca(scaled_data, dim=dim)

    # Point cloud sampling (denoising)
    indstemp, *_ = _sample_denoising(dimred, k, n_points, 1, metric)

    # Build distance matrix
    d = _second_build(dimred, indstemp, metric=metric, nbs=nbs)
    np.fill_diagonal(d, 0)

    # Compute persistent homology
    persistence = ripser(
        d,
        maxdim=maxdim,
        coeff=coeff,
        do_cocycles=True,
        distance_matrix=True,
        progress_bar=progress_bar,
    )

    return persistence


def _pca(data, dim=2):
    """
    Perform PCA (Principal Component Analysis) for dimensionality reduction.

    Parameters:
        data (ndarray): Input data matrix of shape (N_samples, N_features).
        dim (int): Target dimension for PCA projection.

    Returns:
        components (ndarray): Projected data of shape (N_samples, dim).
        var_exp (list): Variance explained by each principal component.
        evals (ndarray): Eigenvalues corresponding to the selected components.
    """
    if dim < 2:
        return data, [0]
    m, n = data.shape
    # mean center the data
    # data -= data.mean(axis=0)
    # calculate the covariance matrix
    R = np.cov(data, rowvar=False)
    # calculate eigenvectors & eigenvalues of the covariance matrix
    # use 'eigh' rather than 'eig' since R is symmetric,
    # the performance gain is substantial
    evals, evecs = np.linalg.eig(R)
    # sort eigenvalue in decreasing order
    idx = np.argsort(evals)[::-1]
    evecs = evecs[:, idx]
    # sort eigenvectors according to same index
    evals = evals[idx]
    # select the first n eigenvectors (n is desired dimension
    # of rescaled data array, or dims_rescaled_data)
    evecs = evecs[:, :dim]
    # carry out the transformation on the data using eigenvectors
    # and return the re-scaled data, eigenvalues, and eigenvectors

    tot = np.sum(evals)
    var_exp = [(i / tot) * 100 for i in sorted(evals[:dim], reverse=True)]
    components = np.dot(evecs.T, data.T).T
    return components, var_exp, evals[:dim]


def _sample_denoising(data, k=10, num_sample=500, omega=0.2, metric="euclidean"):
    """
    Perform denoising and greedy sampling based on mutual k-NN graph.

    Parameters:
        data (ndarray): High-dimensional point cloud data.
        k (int): Number of neighbors for local density estimation.
        num_sample (int): Number of samples to retain.
        omega (float): Suppression factor during greedy sampling.
        metric (str): Distance metric used for kNN ('euclidean', 'cosine', etc).

    Returns:
        inds (ndarray): Indices of sampled points.
        d (ndarray): Pairwise similarity matrix of sampled points.
        Fs (ndarray): Sampling scores at each step.
    """
    if HAS_NUMBA:
        return _sample_denoising_numba(data, k, num_sample, omega, metric)
    else:
        return _sample_denoising_numpy(data, k, num_sample, omega, metric)


def _sample_denoising_numpy(data, k=10, num_sample=500, omega=0.2, metric="euclidean"):
    """Original numpy implementation for fallback."""
    n = data.shape[0]
    X = squareform(pdist(data, metric))
    knn_indices = np.argsort(X)[:, :k]
    knn_dists = X[np.arange(X.shape[0])[:, None], knn_indices].copy()

    sigmas, rhos = _smooth_knn_dist(knn_dists, k, local_connectivity=0)
    rows, cols, vals = _compute_membership_strengths(knn_indices, knn_dists, sigmas, rhos)
    result = coo_matrix((vals, (rows, cols)), shape=(n, n))
    result.eliminate_zeros()
    transpose = result.transpose()
    prod_matrix = result.multiply(transpose)
    result = result + transpose - prod_matrix
    result.eliminate_zeros()
    X = result.toarray()
    F = np.sum(X, 1)
    Fs = np.zeros(num_sample)
    Fs[0] = np.max(F)
    i = np.argmax(F)
    inds_all = np.arange(n)
    inds_left = inds_all > -1
    inds_left[i] = False
    inds = np.zeros(num_sample, dtype=int)
    inds[0] = i
    for j in np.arange(1, num_sample):
        F -= omega * X[i, :]
        Fmax = np.argmax(F[inds_left])
        # Exactly match external TDAvis implementation (including the indexing logic)
        Fs[j] = F[Fmax]
        i = inds_all[inds_left][Fmax]

        inds_left[i] = False
        inds[j] = i
    d = np.zeros((num_sample, num_sample))

    for j, i in enumerate(inds):
        d[j, :] = X[i, inds]
    return inds, d, Fs


def _sample_denoising_numba(data, k=10, num_sample=500, omega=0.2, metric="euclidean"):
    """Optimized numba implementation."""
    n = data.shape[0]
    X = squareform(pdist(data, metric))
    knn_indices = np.argsort(X)[:, :k]
    knn_dists = X[np.arange(X.shape[0])[:, None], knn_indices].copy()

    sigmas, rhos = _smooth_knn_dist(knn_dists, k, local_connectivity=0)
    rows, cols, vals = _compute_membership_strengths(knn_indices, knn_dists, sigmas, rhos)

    # Build symmetric adjacency matrix using optimized function
    X_adj = _build_adjacency_matrix_numba(rows, cols, vals, n)

    # Greedy sampling using optimized function
    inds, Fs = _greedy_sampling_numba(X_adj, num_sample, omega)

    # Build final distance matrix
    d = _build_distance_matrix_numba(X_adj, inds)

    return inds, d, Fs


@njit(fastmath=True)
def _build_adjacency_matrix_numba(rows, cols, vals, n):
    """Build symmetric adjacency matrix efficiently with numba.

    This matches the scipy sparse matrix operations:
    result = result + transpose - prod_matrix
    where prod_matrix = result.multiply(transpose)
    """
    # Initialize matrices
    X = np.zeros((n, n), dtype=np.float64)
    X_T = np.zeros((n, n), dtype=np.float64)

    # Build adjacency matrix and its transpose simultaneously
    for i in range(len(rows)):
        X[rows[i], cols[i]] = vals[i]
        X_T[cols[i], rows[i]] = vals[i]  # Transpose

    # Apply the symmetrization formula: A = A + A^T - A ⊙ A^T
    # This matches scipy's: result + transpose - prod_matrix
    for i in range(n):
        for j in range(n):
            prod_val = X[i, j] * X_T[i, j]  # Element-wise multiplication
            X[i, j] = X[i, j] + X_T[i, j] - prod_val

    return X


@njit(fastmath=True)
def _greedy_sampling_numba(X, num_sample, omega):
    """Optimized greedy sampling with numba."""
    n = X.shape[0]
    F = np.sum(X, axis=1)
    Fs = np.zeros(num_sample)
    inds = np.zeros(num_sample, dtype=np.int64)
    inds_left = np.ones(n, dtype=np.bool_)

    # Initialize with maximum F
    i = np.argmax(F)
    Fs[0] = F[i]
    inds[0] = i
    inds_left[i] = False

    # Greedy sampling loop
    for j in range(1, num_sample):
        # Update F values
        for k in range(n):
            F[k] -= omega * X[i, k]

        # Find maximum among remaining points (matching numpy logic exactly)
        max_val = -np.inf
        max_idx = -1
        for k in range(n):
            if inds_left[k] and F[k] > max_val:
                max_val = F[k]
                max_idx = k

        # Record the F value using the selected index (matching external TDAvis)
        i = max_idx
        Fs[j] = F[i]
        inds[j] = i
        inds_left[i] = False

    return inds, Fs


@njit(fastmath=True)
def _build_distance_matrix_numba(X, inds):
    """Build final distance matrix efficiently with numba."""
    num_sample = len(inds)
    d = np.zeros((num_sample, num_sample))

    for j in range(num_sample):
        for k in range(num_sample):
            d[j, k] = X[inds[j], inds[k]]

    return d


@njit(fastmath=True)
def _smooth_knn_dist(distances, k, n_iter=64, local_connectivity=0.0, bandwidth=1.0):
    """
    Compute smoothed local distances for kNN graph with entropy balancing.

    Parameters:
        distances (ndarray): kNN distance matrix.
        k (int): Number of neighbors.
        n_iter (int): Number of binary search iterations.
        local_connectivity (float): Minimum local connectivity.
        bandwidth (float): Bandwidth parameter.

    Returns:
        sigmas (ndarray): Smoothed sigma values for each point.
        rhos (ndarray): Minimum distances (connectivity cutoff) for each point.
    """
    target = np.log2(k) * bandwidth
    #    target = np.log(k) * bandwidth
    #    target = k

    rho = np.zeros(distances.shape[0])
    result = np.zeros(distances.shape[0])

    mean_distances = np.mean(distances)

    for i in range(distances.shape[0]):
        lo = 0.0
        hi = np.inf
        mid = 1.0

        # TODO: This is very inefficient, but will do for now. FIXME
        ith_distances = distances[i]
        non_zero_dists = ith_distances[ith_distances > 0.0]
        if non_zero_dists.shape[0] >= local_connectivity:
            index = int(np.floor(local_connectivity))
            interpolation = local_connectivity - index
            if index > 0:
                rho[i] = non_zero_dists[index - 1]
                if interpolation > 1e-5:
                    rho[i] += interpolation * (non_zero_dists[index] - non_zero_dists[index - 1])
            else:
                rho[i] = interpolation * non_zero_dists[0]
        elif non_zero_dists.shape[0] > 0:
            rho[i] = np.max(non_zero_dists)

        for _ in range(n_iter):
            psum = 0.0
            for j in range(1, distances.shape[1]):
                d = distances[i, j] - rho[i]
                if d > 0:
                    psum += np.exp(-(d / mid))
                #                    psum += d / mid

                else:
                    psum += 1.0
            #                    psum += 0

            if np.fabs(psum - target) < 1e-5:
                break

            if psum > target:
                hi = mid
                mid = (lo + hi) / 2.0
            else:
                lo = mid
                if hi == np.inf:
                    mid *= 2
                else:
                    mid = (lo + hi) / 2.0
        result[i] = mid
        # TODO: This is very inefficient, but will do for now. FIXME
        if rho[i] > 0.0:
            mean_ith_distances = np.mean(ith_distances)
            if result[i] < 1e-3 * mean_ith_distances:
                result[i] = 1e-3 * mean_ith_distances
        else:
            if result[i] < 1e-3 * mean_distances:
                result[i] = 1e-3 * mean_distances

    return result, rho


@njit(parallel=True, fastmath=True)
def _compute_membership_strengths(knn_indices, knn_dists, sigmas, rhos):
    """
    Compute membership strength matrix from smoothed kNN graph.

    Parameters:
        knn_indices (ndarray): Indices of k-nearest neighbors.
        knn_dists (ndarray): Corresponding distances.
        sigmas (ndarray): Local bandwidths.
        rhos (ndarray): Minimum distance thresholds.

    Returns:
        rows (ndarray): Row indices for sparse matrix.
        cols (ndarray): Column indices for sparse matrix.
        vals (ndarray): Weight values for sparse matrix.
    """
    n_samples = knn_indices.shape[0]
    n_neighbors = knn_indices.shape[1]
    rows = np.zeros((n_samples * n_neighbors), dtype=np.int64)
    cols = np.zeros((n_samples * n_neighbors), dtype=np.int64)
    vals = np.zeros((n_samples * n_neighbors), dtype=np.float64)
    for i in range(n_samples):
        for j in range(n_neighbors):
            if knn_indices[i, j] == -1:
                continue  # We didn't get the full knn for i
            if knn_indices[i, j] == i:
                val = 0.0
            elif knn_dists[i, j] - rhos[i] <= 0.0:
                val = 1.0
            else:
                val = np.exp(-((knn_dists[i, j] - rhos[i]) / (sigmas[i])))
                # val = ((knn_dists[i, j] - rhos[i]) / (sigmas[i]))

            rows[i * n_neighbors + j] = i
            cols[i * n_neighbors + j] = knn_indices[i, j]
            vals[i * n_neighbors + j] = val

    return rows, cols, vals


def _second_build(data, indstemp, nbs=800, metric="cosine"):
    """
    Reconstruct distance matrix after denoising for persistent homology.

    Parameters:
        data (ndarray): PCA-reduced data matrix.
        indstemp (ndarray): Indices of sampled points.
        nbs (int): Number of neighbors in reconstructed graph.
        metric (str): Distance metric ('cosine', 'euclidean', etc).

    Returns:
        d (ndarray): Symmetric distance matrix used for persistent homology.
    """
    # Filter the data using the sampled point indices
    data = data[indstemp, :]

    # Compute the pairwise distance matrix
    X = squareform(pdist(data, metric))
    knn_indices = np.argsort(X)[:, :nbs]
    knn_dists = X[np.arange(X.shape[0])[:, None], knn_indices].copy()

    # Compute smoothed kernel widths
    sigmas, rhos = _smooth_knn_dist(knn_dists, nbs, local_connectivity=0)
    rows, cols, vals = _compute_membership_strengths(knn_indices, knn_dists, sigmas, rhos)

    # Construct a sparse graph
    result = coo_matrix((vals, (rows, cols)), shape=(X.shape[0], X.shape[0]))
    result.eliminate_zeros()
    transpose = result.transpose()
    prod_matrix = result.multiply(transpose)
    result = result + transpose - prod_matrix
    result.eliminate_zeros()

    # Build the final distance matrix
    d = result.toarray()
    # Match external TDAvis: direct negative log without epsilon handling
    # Temporarily suppress divide by zero warning to match external behavior
    with np.errstate(divide="ignore", invalid="ignore"):
        d = -np.log(d)
    np.fill_diagonal(d, 0)

    return d


def _run_shuffle_analysis(sspikes, num_shuffles=1000, num_cores=4, progress_bar=True, **kwargs):
    """Perform shuffle analysis with optimized computation."""
    return _run_shuffle_analysis_multiprocessing(
        sspikes, num_shuffles, num_cores, progress_bar, **kwargs
    )


def _run_shuffle_analysis_multiprocessing(
    sspikes, num_shuffles=1000, num_cores=4, progress_bar=True, **kwargs
):
    """Original multiprocessing implementation for fallback."""
    max_lifetimes = {0: [], 1: [], 2: []}

    # Estimate runtime with a test iteration
    logging.info("Running test iteration to estimate runtime...")

    _ = _process_single_shuffle((0, sspikes, kwargs))

    # Prepare task list
    tasks = [(i, sspikes, kwargs) for i in range(num_shuffles)]
    logging.info(
        f"Starting shuffle analysis with {num_shuffles} iterations using {num_cores} cores..."
    )

    # Use multiprocessing pool for parallel processing
    with mp.Pool(processes=num_cores) as pool:
        results = list(pool.imap(_process_single_shuffle, tasks))
        logging.info("Shuffle analysis completed")

    # Collect results
    for res in results:
        for dim, lifetime in res.items():
            max_lifetimes[dim].append(lifetime)

    return max_lifetimes


@njit(fastmath=True)
def _fast_pca_transform(data, components):
    """Fast PCA transformation using numba."""
    return np.dot(data, components.T)


def _process_single_shuffle(args):
    """Process a single shuffle task."""
    i, sspikes, kwargs = args
    try:
        shuffled_data = _shuffle_spike_trains(sspikes)
        persistence = _compute_persistence(shuffled_data, **kwargs)

        dim_max_lifetimes = {}
        for dim in [0, 1, 2]:
            if dim < len(persistence["dgms"]):
                # Filter out infinite values
                valid_bars = [bar for bar in persistence["dgms"][dim] if not np.isinf(bar[1])]
                if valid_bars:
                    lifetimes = [bar[1] - bar[0] for bar in valid_bars]
                    if lifetimes:
                        dim_max_lifetimes[dim] = max(lifetimes)
        return dim_max_lifetimes
    except Exception as e:
        print(f"Shuffle {i} failed: {str(e)}")
        return {}


def _shuffle_spike_trains(sspikes):
    """Perform random circular shift on spike trains."""
    shuffled = sspikes.copy()
    num_neurons = shuffled.shape[1]

    # Independent shift for each neuron
    for n in range(num_neurons):
        shift = np.random.randint(0, int(shuffled.shape[0] * 0.1))
        shuffled[:, n] = np.roll(shuffled[:, n], shift)

    return shuffled


def _plot_barcode(persistence):
    """
    Plot barcode diagram from persistent homology result.

    Parameters:
        persistence (dict): Persistent homology result with 'dgms' key.
    """
    cs = np.repeat([[0, 0.55, 0.2]], 3).reshape(3, 3).T  # RGB color for each dimension
    alpha = 1
    inf_delta = 0.1
    colormap = cs
    dgms = persistence["dgms"]
    maxdim = len(dgms) - 1
    dims = np.arange(maxdim + 1)
    labels = ["$H_0$", "$H_1$", "$H_2$"]

    # Determine axis range
    min_birth, max_death = 0, 0
    for dim in dims:
        persistence_dim = dgms[dim][~np.isinf(dgms[dim][:, 1]), :]
        if persistence_dim.size > 0:
            min_birth = min(min_birth, np.min(persistence_dim))
            max_death = max(max_death, np.max(persistence_dim))

    delta = (max_death - min_birth) * inf_delta
    infinity = max_death + delta
    axis_start = min_birth - delta

    # Create plot
    fig = plt.figure(figsize=(10, 6))
    gs = gridspec.GridSpec(len(dims), 1)

    for dim in dims:
        axes = plt.subplot(gs[dim])
        axes.axis("on")
        axes.set_yticks([])
        axes.set_ylabel(labels[dim], rotation=0, labelpad=20, fontsize=12)

        d = np.copy(dgms[dim])
        d[np.isinf(d[:, 1]), 1] = infinity
        dlife = d[:, 1] - d[:, 0]

        # Select top 30 bars by lifetime
        dinds = np.argsort(dlife)[-30:]
        if dim > 0:
            dinds = dinds[np.flip(np.argsort(d[dinds, 0]))]

        axes.barh(
            0.5 + np.arange(len(dinds)),
            dlife[dinds],
            height=0.8,
            left=d[dinds, 0],
            alpha=alpha,
            color=colormap[dim],
            linewidth=0,
        )

        axes.plot([0, 0], [0, len(dinds)], c="k", linestyle="-", lw=1)
        axes.plot([0, len(dinds)], [0, 0], c="k", linestyle="-", lw=1)
        axes.set_xlim([axis_start, infinity])

    plt.tight_layout()
    return fig


def _plot_barcode_with_shuffle(persistence, shuffle_max):
    """
    Plot barcode with shuffle region markers.
    """
    # Handle case where shuffle_max is None
    if shuffle_max is None:
        shuffle_max = {}

    cs = np.repeat([[0, 0.55, 0.2]], 3).reshape(3, 3).T
    alpha = 1
    inf_delta = 0.1
    colormap = cs
    maxdim = len(persistence["dgms"]) - 1
    dims = np.arange(maxdim + 1)

    min_birth, max_death = 0, 0
    for dim in dims:
        # Filter out infinite values
        valid_bars = [bar for bar in persistence["dgms"][dim] if not np.isinf(bar[1])]
        if valid_bars:
            min_birth = min(min_birth, np.min(valid_bars))
            max_death = max(max_death, np.max(valid_bars))

    # Handle case with no valid bars
    if max_death == 0 and min_birth == 0:
        min_birth = 0
        max_death = 1

    delta = (max_death - min_birth) * inf_delta
    infinity = max_death + delta

    # Create figure
    fig = plt.figure(figsize=(10, 8))
    gs = gridspec.GridSpec(len(dims), 1)

    # Get shuffle thresholds (99.9th percentile for each dimension)
    thresholds = {}
    for dim in dims:
        if dim in shuffle_max and shuffle_max[dim]:
            thresholds[dim] = np.percentile(shuffle_max[dim], 99.9)
        else:
            thresholds[dim] = 0

    for _, dim in enumerate(dims):
        axes = plt.subplot(gs[dim])
        axes.axis("off")

        # Add gray background to represent shuffle region
        if dim in thresholds:
            axes.axvspan(0, thresholds[dim], alpha=0.2, color="gray", zorder=-3)
            axes.axvline(x=thresholds[dim], color="gray", linestyle="--", alpha=0.7)

        # Filter out infinite values
        d = np.array([bar for bar in persistence["dgms"][dim] if not np.isinf(bar[1])])
        if len(d) == 0:
            d = np.zeros((0, 2))

        d = np.copy(d)
        d[np.isinf(d[:, 1]), 1] = infinity
        dlife = d[:, 1] - d[:, 0]

        # Select top 30 longest-lived bars
        if len(dlife) > 0:
            dinds = np.argsort(dlife)[-30:]
            if dim > 0:
                dinds = dinds[np.flip(np.argsort(d[dinds, 0]))]

            # Mark significant bars
            significant_bars = []
            for idx in dinds:
                if dlife[idx] > thresholds.get(dim, 0):
                    significant_bars.append(idx)

            # Draw bars
            for i, idx in enumerate(dinds):
                color = "red" if idx in significant_bars else colormap[dim]
                axes.barh(
                    0.5 + i,
                    dlife[idx],
                    height=0.8,
                    left=d[idx, 0],
                    alpha=alpha,
                    color=color,
                    linewidth=0,
                )

            indsall = len(dinds)
        else:
            indsall = 0

        axes.plot([0, 0], [0, indsall], c="k", linestyle="-", lw=1)
        axes.plot([0, indsall], [0, 0], c="k", linestyle="-", lw=1)
        axes.set_xlim([0, infinity])
        axes.set_title(f"$H_{dim}$", loc="left")

    plt.tight_layout()
    return fig


def decode_circular_coordinates(
    persistence_result: dict[str, Any],
    spike_data: dict[str, Any],
    real_ground: bool = True,
    real_of: bool = True,
    save_path: str | None = None,
) -> dict[str, Any]:
    """
    Decode circular coordinates (bump positions) from cohomology.

    Parameters:
        persistence_result : dict containing persistence analysis results with keys:
            - 'persistence': persistent homology result
            - 'indstemp': indices of sampled points
            - 'movetimes': selected time points
            - 'n_points': number of sampled points
        spike_data : dict, optional
            Spike data dictionary containing 'spike', 't', and optionally 'x', 'y'
        real_ground : bool
            Whether x, y, t ground truth exists
        real_of : bool
            Whether experiment was performed in open field
        save_path : str, optional
            Path to save decoding results. If None, saves to 'Results/spikes_decoding.npz'

    Returns:
        dict : Dictionary containing decoding results with keys:
            - 'coords': decoded coordinates for all timepoints
            - 'coordsbox': decoded coordinates for box timepoints
            - 'times': time indices for coords
            - 'times_box': time indices for coordsbox
            - 'centcosall': cosine centroids
            - 'centsinall': sine centroids
    """
    ph_classes = [0, 1]  # Decode the ith most persistent cohomology class
    num_circ = len(ph_classes)
    dec_tresh = 0.99
    coeff = 47

    # Extract persistence analysis results
    persistence = persistence_result["persistence"]
    indstemp = persistence_result["indstemp"]
    movetimes = persistence_result["movetimes"]
    n_points = persistence_result["n_points"]

    diagrams = persistence["dgms"]  # the multiset describing the lives of the persistence classes
    cocycles = persistence["cocycles"][1]  # the cocycle representatives for the 1-dim classes
    dists_land = persistence["dperm2all"]  # the pairwise distance between the points
    births1 = diagrams[1][:, 0]  # the time of birth for the 1-dim classes
    deaths1 = diagrams[1][:, 1]  # the time of death for the 1-dim classes
    deaths1[np.isinf(deaths1)] = 0
    lives1 = deaths1 - births1  # the lifetime for the 1-dim classes
    iMax = np.argsort(lives1)
    coords1 = np.zeros((num_circ, len(indstemp)))
    threshold = births1[iMax[-2]] + (deaths1[iMax[-2]] - births1[iMax[-2]]) * dec_tresh

    for c in ph_classes:
        cocycle = cocycles[iMax[-(c + 1)]]
        coords1[c, :], inds = _get_coords(cocycle, threshold, len(indstemp), dists_land, coeff)

    if real_ground:  # 用户所提供的数据是否有真实的xyt
        sspikes, xx, yy, tt = embed_spike_trains(
            spike_data, config=SpikeEmbeddingConfig(smooth=True, speed_filter=True)
        )
    else:
        sspikes = embed_spike_trains(
            spike_data, config=SpikeEmbeddingConfig(smooth=True, speed_filter=False)
        )

    num_neurons = sspikes.shape[1]
    centcosall = np.zeros((num_neurons, 2, n_points))
    centsinall = np.zeros((num_neurons, 2, n_points))
    dspk = preprocessing.scale(sspikes[movetimes[indstemp], :])

    for neurid in range(num_neurons):
        spktemp = dspk[:, neurid].copy()
        centcosall[neurid, :, :] = np.multiply(np.cos(coords1[:, :] * 2 * np.pi), spktemp)
        centsinall[neurid, :, :] = np.multiply(np.sin(coords1[:, :] * 2 * np.pi), spktemp)

    if real_ground:  # 用户所提供的数据是否有真实的xyt
        sspikes, xx, yy, tt = embed_spike_trains(
            spike_data, config=SpikeEmbeddingConfig(smooth=True, speed_filter=True)
        )
        spikes, __, __, __ = embed_spike_trains(
            spike_data, config=SpikeEmbeddingConfig(smooth=False, speed_filter=True)
        )
    else:
        sspikes = embed_spike_trains(
            spike_data, config=SpikeEmbeddingConfig(smooth=True, speed_filter=False)
        )
        spikes = embed_spike_trains(
            spike_data, config=SpikeEmbeddingConfig(smooth=False, speed_filter=False)
        )

    times = np.where(np.sum(spikes > 0, 1) >= 1)[0]
    dspk = preprocessing.scale(sspikes)
    sspikes = sspikes[times, :]
    dspk = dspk[times, :]

    a = np.zeros((len(sspikes[:, 0]), 2, num_neurons))
    for n in range(num_neurons):
        a[:, :, n] = np.multiply(dspk[:, n : n + 1], np.sum(centcosall[n, :, :], 1))

    c = np.zeros((len(sspikes[:, 0]), 2, num_neurons))
    for n in range(num_neurons):
        c[:, :, n] = np.multiply(dspk[:, n : n + 1], np.sum(centsinall[n, :, :], 1))

    mtot2 = np.sum(c, 2)
    mtot1 = np.sum(a, 2)
    coords = np.arctan2(mtot2, mtot1) % (2 * np.pi)

    if real_of:  # 用户的数据是否是来自真实的OF场地
        coordsbox = coords.copy()
        times_box = times.copy()
    else:
        sspikes, xx, yy, tt = embed_spike_trains(
            spike_data, config=SpikeEmbeddingConfig(smooth=True, speed_filter=True)
        )
        spikes, __, __, __ = embed_spike_trains(
            spike_data, config=SpikeEmbeddingConfig(smooth=False, speed_filter=True)
        )
        dspk = preprocessing.scale(sspikes)
        times_box = np.where(np.sum(spikes > 0, 1) >= 1)[0]
        dspk = dspk[times_box, :]

        a = np.zeros((len(times_box), 2, num_neurons))
        for n in range(num_neurons):
            a[:, :, n] = np.multiply(dspk[:, n : n + 1], np.sum(centcosall[n, :, :], 1))

        c = np.zeros((len(times_box), 2, num_neurons))
        for n in range(num_neurons):
            c[:, :, n] = np.multiply(dspk[:, n : n + 1], np.sum(centsinall[n, :, :], 1))

        mtot2 = np.sum(c, 2)
        mtot1 = np.sum(a, 2)
        coordsbox = np.arctan2(mtot2, mtot1) % (2 * np.pi)

    # Prepare results dictionary
    results = {
        "coords": coords,
        "coordsbox": coordsbox,
        "times": times,
        "times_box": times_box,
        "centcosall": centcosall,
        "centsinall": centsinall,
    }

    # Save results
    if save_path is None:
        os.makedirs("Results", exist_ok=True)
        save_path = "Results/spikes_decoding.npz"

    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    np.savez_compressed(save_path, **results)

    return results


def plot_3d_bump_on_torus(
    decoding_result: dict[str, Any] | str,
    spike_data: dict[str, Any],
    config: CANN2DPlotConfig | None = None,
    save_path: str | None = None,
    numangsint: int = 51,
    r1: float = 1.5,
    r2: float = 1.0,
    window_size: int = 300,
    frame_step: int = 5,
    n_frames: int = 20,
    fps: int = 5,
    show_progress: bool = True,
    show: bool = True,
    figsize: tuple[int, int] = (8, 8),
    **kwargs,
) -> animation.FuncAnimation:
    """
    Visualize the movement of the neural activity bump on a torus using matplotlib animation.

    This function follows the canns.analyzer.plotting patterns for animation generation
    with progress tracking and proper resource cleanup.

    Parameters:
        decoding_result : dict or str
            Dictionary containing decoding results with 'coordsbox' and 'times_box' keys,
            or path to .npz file containing these results
        spike_data : dict, optional
            Spike data dictionary containing spike information
        config : PlotConfig, optional
            Configuration object for unified plotting parameters
        **kwargs : backward compatibility parameters
        save_path : str, optional
            Path to save the animation (e.g., 'animation.gif' or 'animation.mp4')
        numangsint : int
            Grid resolution for the torus surface
        r1 : float
            Major radius of the torus
        r2 : float
            Minor radius of the torus
        window_size : int
            Time window (in number of time points) for each frame
        frame_step : int
            Step size to slide the time window between frames
        n_frames : int
            Total number of frames in the animation
        fps : int
            Frames per second for the output animation
        show_progress : bool
            Whether to show progress bar during generation
        show : bool
            Whether to display the animation
        figsize : tuple[int, int]
            Figure size for the animation

    Returns:
        matplotlib.animation.FuncAnimation : The animation object
    """
    # Handle backward compatibility and configuration
    if config is None:
        config = CANN2DPlotConfig.for_torus_animation(**kwargs)

    # Override config with any explicitly passed parameters
    for key, value in kwargs.items():
        if hasattr(config, key):
            setattr(config, key, value)

    # Extract configuration values
    save_path = config.save_path if config.save_path else save_path
    show = config.show
    figsize = config.figsize
    fps = config.fps
    show_progress = config.show_progress_bar
    numangsint = config.numangsint
    r1 = config.r1
    r2 = config.r2
    window_size = config.window_size
    frame_step = config.frame_step
    n_frames = config.n_frames

    # Load decoding results if path is provided
    if isinstance(decoding_result, str):
        f = np.load(decoding_result, allow_pickle=True)
        coords = f["coordsbox"]
        times = f["times_box"]
        f.close()
    else:
        coords = decoding_result["coordsbox"]
        times = decoding_result["times_box"]

    spk, *_ = embed_spike_trains(
        spike_data, config=SpikeEmbeddingConfig(smooth=False, speed_filter=True)
    )

    # Prepare animation data
    frame_data = []
    prev_m = None

    for frame_idx in tqdm(range(n_frames), desc="Processing frames"):
        start_idx = frame_idx * frame_step
        end_idx = start_idx + window_size
        if end_idx > np.max(times):
            break

        mask = (times >= start_idx) & (times < end_idx)
        coords_window = coords[mask]
        if len(coords_window) == 0:
            continue

        spk_window = spk[times[mask], :]
        activity = np.sum(spk_window, axis=1)

        m, x_edge, y_edge, _ = binned_statistic_2d(
            coords_window[:, 0],
            coords_window[:, 1],
            activity,
            statistic="sum",
            bins=np.linspace(0, 2 * np.pi, numangsint - 1),
        )
        m = np.nan_to_num(m)
        m = _smooth_tuning_map(m, numangsint - 1, sig=4.0, bClose=True)
        m = gaussian_filter(m, sigma=1.0)

        if prev_m is not None:
            m = 0.7 * prev_m + 0.3 * m
        prev_m = m

        # Store processed data for animation
        X, Y = np.meshgrid(x_edge, y_edge)
        X = (X + np.pi / 5) % (2 * np.pi)
        x = (r1 + r2 * np.cos(X)) * np.cos(Y)
        y = (r1 + r2 * np.cos(X)) * np.sin(Y)
        z = -r2 * np.sin(X)  # Flip torus surface orientation

        frame_data.append({"x": x, "y": y, "z": z, "m": m, "time": start_idx * frame_step})

    if not frame_data:
        raise ProcessingError("No valid frames generated for animation")

    # Create figure and animation following plotting package pattern
    fig = plt.figure(figsize=figsize)

    try:
        ax = fig.add_subplot(111, projection="3d")
        ax.set_zlim(-2, 2)
        ax.view_init(-125, 135)
        ax.axis("off")

        # Initialize with first frame
        first_frame = frame_data[0]
        surface = ax.plot_surface(
            first_frame["x"],
            first_frame["y"],
            first_frame["z"],
            facecolors=cm.viridis(first_frame["m"] / (np.max(first_frame["m"]) + 1e-9)),
            alpha=1,
            linewidth=0.1,
            antialiased=True,
            rstride=1,
            cstride=1,
            shade=False,
        )

        def animate(frame_idx):
            """Animation update function following plotting package pattern."""
            if frame_idx >= len(frame_data):
                return (surface,)

            frame = frame_data[frame_idx]

            # Clear and redraw surface
            ax.clear()
            ax.set_zlim(-2, 2)
            ax.view_init(-125, 135)
            ax.axis("off")

            new_surface = ax.plot_surface(
                frame["x"],
                frame["y"],
                frame["z"],
                facecolors=cm.viridis(frame["m"] / (np.max(frame["m"]) + 1e-9)),
                alpha=1,
                linewidth=0.1,
                antialiased=True,
                rstride=1,
                cstride=1,
                shade=False,
            )

            # Update time text
            time_text = ax.text2D(
                0.05,
                0.95,
                f"Frame: {frame_idx + 1}/{len(frame_data)}",
                transform=ax.transAxes,
                fontsize=12,
                bbox=dict(facecolor="white", alpha=0.7),
            )

            return new_surface, time_text

        # Create animation
        interval_ms = 1000 / fps
        ani = animation.FuncAnimation(
            fig, animate, frames=len(frame_data), interval=interval_ms, blit=False, repeat=True
        )

        # Save animation if path provided
        if save_path:
            if show_progress:
                pbar = tqdm(total=len(frame_data), desc=f"Saving animation to {save_path}")

                def progress_callback(current_frame, total_frames):
                    pbar.update(1)

                try:
                    writer = animation.PillowWriter(fps=fps)
                    ani.save(save_path, writer=writer, progress_callback=progress_callback)
                    pbar.close()
                    print(f"\nAnimation saved to: {save_path}")
                except Exception as e:
                    pbar.close()
                    print(f"\nError saving animation: {e}")
            else:
                try:
                    writer = animation.PillowWriter(fps=fps)
                    ani.save(save_path, writer=writer)
                    print(f"Animation saved to: {save_path}")
                except Exception as e:
                    print(f"Error saving animation: {e}")

        if show:
            plt.show()

        return ani

    except Exception as e:
        plt.close(fig)
        raise ProcessingError(f"Failed to create torus animation: {e}") from e


def _get_coords(cocycle, threshold, num_sampled, dists, coeff):
    """
    Reconstruct circular coordinates from cocycle information.

    Parameters:
        cocycle (ndarray): Persistent cocycle representative.
        threshold (float): Maximum allowable edge distance.
        num_sampled (int): Number of sampled points.
        dists (ndarray): Pairwise distance matrix.
        coeff (int): Finite field modulus for cohomology.

    Returns:
        f (ndarray): Circular coordinate values (in [0,1]).
        verts (ndarray): Indices of used vertices.
    """
    zint = np.where(coeff - cocycle[:, 2] < cocycle[:, 2])
    cocycle[zint, 2] = cocycle[zint, 2] - coeff
    d = np.zeros((num_sampled, num_sampled))
    d[np.tril_indices(num_sampled)] = np.nan
    d[cocycle[:, 1], cocycle[:, 0]] = cocycle[:, 2]
    d[dists > threshold] = np.nan
    d[dists == 0] = np.nan
    edges = np.where(~np.isnan(d))
    verts = np.array(np.unique(edges))
    num_edges = np.shape(edges)[1]
    num_verts = np.size(verts)
    values = d[edges]
    A = np.zeros((num_edges, num_verts), dtype=int)
    v1 = np.zeros((num_edges, 2), dtype=int)
    v2 = np.zeros((num_edges, 2), dtype=int)
    for i in range(num_edges):
        # Extract scalar indices from np.where results
        idx1 = np.where(verts == edges[0][i])[0]
        idx2 = np.where(verts == edges[1][i])[0]

        # Handle case where np.where returns multiple matches (shouldn't happen in valid data)
        if len(idx1) > 0:
            v1[i, :] = [i, idx1[0]]
        else:
            raise ValueError(f"No vertex found for edge {edges[0][i]}")

        if len(idx2) > 0:
            v2[i, :] = [i, idx2[0]]
        else:
            raise ValueError(f"No vertex found for edge {edges[1][i]}")

    A[v1[:, 0], v1[:, 1]] = -1
    A[v2[:, 0], v2[:, 1]] = 1

    L = np.ones((num_edges,))
    Aw = A * np.sqrt(L[:, np.newaxis])
    Bw = values * np.sqrt(L)
    f = lsmr(Aw, Bw)[0] % 1
    return f, verts


def _smooth_tuning_map(mtot, numangsint, sig, bClose=True):
    """
    Smooth activity map over circular topology (e.g., torus).

    Parameters:
        mtot (ndarray): Raw activity map matrix.
        numangsint (int): Grid resolution.
        sig (float): Smoothing kernel standard deviation.
        bClose (bool): Whether to assume circular boundary conditions.

    Returns:
        mtot_out (ndarray): Smoothed map matrix.
    """
    numangsint_1 = numangsint - 1
    mid = int((numangsint_1) / 2)
    indstemp1 = np.zeros((numangsint_1, numangsint_1), dtype=int)
    indstemp1[indstemp1 == 0] = np.arange((numangsint_1) ** 2)
    mid = int((numangsint_1) / 2)
    mtemp1_3 = mtot.copy()
    for i in range(numangsint_1):
        mtemp1_3[i, :] = np.roll(mtemp1_3[i, :], int(i / 2))
    mtot_out = np.zeros_like(mtot)
    mtemp1_4 = np.concatenate((mtemp1_3, mtemp1_3, mtemp1_3), 1)
    mtemp1_5 = np.zeros_like(mtemp1_4)
    mtemp1_5[:, :mid] = mtemp1_4[:, (numangsint_1) * 3 - mid :]
    mtemp1_5[:, mid:] = mtemp1_4[:, : (numangsint_1) * 3 - mid]
    if bClose:
        mtemp1_6 = _smooth_image(np.concatenate((mtemp1_5, mtemp1_4, mtemp1_5)), sigma=sig)
    else:
        mtemp1_6 = gaussian_filter(np.concatenate((mtemp1_5, mtemp1_4, mtemp1_5)), sigma=sig)
    for i in range(numangsint_1):
        mtot_out[i, :] = mtemp1_6[
            (numangsint_1) + i,
            (numangsint_1) + (int(i / 2) + 1) : (numangsint_1) * 2 + (int(i / 2) + 1),
        ]
    return mtot_out


def _smooth_image(img, sigma):
    """
    Smooth image using multivariate Gaussian kernel, handling missing (NaN) values.

    Parameters:
        img (ndarray): Input image matrix.
        sigma (float): Standard deviation of smoothing kernel.

    Returns:
        imgC (ndarray): Smoothed image with inpainting around NaNs.
    """
    filterSize = max(np.shape(img))
    grid = np.arange(-filterSize + 1, filterSize, 1)
    xx, yy = np.meshgrid(grid, grid)

    pos = np.dstack((xx, yy))

    var = multivariate_normal(mean=[0, 0], cov=[[sigma**2, 0], [0, sigma**2]])
    k = var.pdf(pos)
    k = k / np.sum(k)

    nans = np.isnan(img)
    imgA = img.copy()
    imgA[nans] = 0
    imgA = signal.convolve2d(imgA, k, mode="valid")
    imgD = img.copy()
    imgD[nans] = 0
    imgD[~nans] = 1
    radius = 1
    L = np.arange(-radius, radius + 1)
    X, Y = np.meshgrid(L, L)
    dk = np.array((X**2 + Y**2) <= radius**2, dtype=bool)
    imgE = np.zeros((filterSize + 2, filterSize + 2))
    imgE[1:-1, 1:-1] = imgD
    imgE = binary_closing(imgE, iterations=1, structure=dk)
    imgD = imgE[1:-1, 1:-1]

    imgB = np.divide(
        signal.convolve2d(imgD, k, mode="valid"),
        signal.convolve2d(np.ones(np.shape(imgD)), k, mode="valid"),
    )
    imgC = np.divide(imgA, imgB)
    imgC[imgD == 0] = -np.inf
    return imgC


if __name__ == "__main__":
    from canns.analyzer.experimental_data._datasets_utils import load_grid_data

    data = load_grid_data()

    spikes, *_ = embed_spike_trains(data)

    # import umap
    #
    # reducer = umap.UMAP(
    #     n_neighbors=15,
    #     min_dist=0.1,
    #     n_components=3,
    #     metric='euclidean',
    #     random_state=42
    # )
    #
    # reduce_func = reducer.fit_transform
    #
    # plot_projection(reduce_func=reduce_func, embed_data=spikes, show=True)
    results = tda_vis(embed_data=spikes, maxdim=1, do_shuffle=False, show=True)

    # results = tda_vis(embed_data=spikes, maxdim=1, do_shuffle=True, num_shuffles=10, show=True)
