Visualisation¶

def hex_heatmap(
    df: pd.Series | pd.DataFrame,
    dataset: pd.DataFrame,
    style: dict | None = None,
    sizing: dict | None = None,
    dpi: int = 72,
    custom_colorscale: list | None = None,
) -> go.Figure:
    """Generate an interactive hexagonal heatmap.

    Visualizes data on a hexagonal grid layout. Single-column data produces
    a static heatmap; multi-column DataFrames produce an animated heatmap
    with a slider to navigate through timepoints or conditions.

    Args:
        df: Data to visualize. Single Series for static plot, or DataFrame
            with multiple columns for animated plot (one frame per column).
            Must have 'p' and 'q' columns representing hex grid coordinates.
        dataset: Reference dataset defining the full hex grid background.
            Used to render empty hexagons for spatial context.
        style: Styling options dict with keys:
            - "font_type": Font family (default: "arial")
            - "markerlinecolor": Marker line color
            - "linecolor": Axis/line color (default: "black")
            - "papercolor": Background color (default: "rgba(255,255,255,255)")
        sizing: Size configuration dict with keys:
            - "fig_width", "fig_height": Figure dimensions in mm
            - "markersize": Hexagon marker size (default: 16)
            - "cbar_thickness", "cbar_len": Colorbar dimensions
        dpi: Dots per inch for pixel calculations (default: 72).
        custom_colorscale: Custom Plotly colorscale. Default is white-to-blue.

    Returns:
        Plotly Figure with hexagonal heatmap. Static for Series input,
        animated with slider for DataFrame input.

    Raises:
        ValueError: If `df` is not a Series or DataFrame.

    Example:
        >>> # Static heatmap
        >>> fig = hex_heatmap(data_series, background_dataset)
        >>> fig.show()
        >>>
        >>> # Animated heatmap with timepoints
        >>> fig = hex_heatmap(timepoint_df, background_dataset)
        >>> fig.write_html("animated_hexmap.html")
    """

    def bg_hex():
        goscatter = go.Scatter(
            x=background_hex["x"],
            y=background_hex["y"],
            mode="markers",
            marker_symbol=symbol_number,
            marker={
                "size": sizing["markersize"],
                "color": "white",
                "line": {
                    "width": sizing["markerlinewidth"],
                    "color": "lightgrey",
                },
            },
            showlegend=False,
        )
        return goscatter

    def data_hex(aseries):
        goscatter = go.Scatter(
            x=x_vals,
            y=y_vals,
            mode="markers",
            marker_symbol=symbol_number,
            marker={
                "cmin": global_min,
                "cmax": global_max,
                "size": sizing["markersize"],
                "color": aseries.values,
                "line": {
                    "width": sizing["markerlinewidth"],
                    "color": "lightgrey",
                },
                "colorbar": {
                    "orientation": "v",
                    "outlinecolor": style["linecolor"],
                    "outlinewidth": sizing["axislinewidth"],
                    "thickness": sizing["cbar_thickness"],
                    "len": sizing["cbar_len"],
                    "tickmode": "array",
                    "ticklen": sizing["ticklen"],
                    "tickwidth": sizing["tickwidth"],
                    "tickcolor": style["linecolor"],
                    "tickfont": {
                        "size": fsize_ticks_px,
                        "family": style["font_type"],
                        "color": style["linecolor"],
                    },
                    "tickformat": ".5f",
                    "title": {
                        "font": {
                            "family": style["font_type"],
                            "size": fsize_title_px,
                            "color": style["linecolor"],
                        },
                        "side": "right",
                    },
                },
                "colorscale": custom_colorscale,
            },
            showlegend=False,
        )
        return goscatter

    default_style = {
        "font_type": "arial",
        "markerlinecolor": "rgba(0,0,0,0)",
        "linecolor": "black",
        "papercolor": "rgba(255,255,255,255)",
    }

    markersize = 16

    default_sizing = {
        "fig_width": 260,
        "fig_height": 220,
        "fig_margin": 0,
        "fsize_ticks_pt": 20,
        "fsize_title_pt": 20,
        "markersize": markersize,
        "ticklen": 15,
        "tickwidth": 5,
        "axislinewidth": 3,
        "markerlinewidth": 0.9,
        "cbar_thickness": 20,
        "cbar_len": 0.75,
    }

    if style is not None:
        default_style.update(style)
    style = default_style

    if sizing is not None:
        default_sizing.update(sizing)
    sizing = default_sizing

    POINTS_PER_INCH = 72
    MM_PER_INCH = 25.4

    pixelsperinch = dpi
    pixelspermm = pixelsperinch / MM_PER_INCH

    if custom_colorscale is None:
        custom_colorscale = [[0, "rgb(255, 255, 255)"], [1, "rgb(0, 20, 200)"]]

    area_width = (sizing["fig_width"] - sizing["fig_margin"]) * pixelspermm
    area_height = (sizing["fig_height"] - sizing["fig_margin"]) * pixelspermm

    fsize_ticks_px = sizing["fsize_ticks_pt"] * (1 / POINTS_PER_INCH) * pixelsperinch
    fsize_title_px = sizing["fsize_title_pt"] * (1 / POINTS_PER_INCH) * pixelsperinch

    global_min = min(0, df.values.min())
    global_max = df.values.max()

    symbol_number = 15

    background_hex = dataset
    background_hex = background_hex.drop_duplicates(subset=["p", "q"])[
        ["p", "q"]
    ].astype(float)
    x, y = hex_to_pixel(background_hex.p, background_hex.q, mode="flat")
    background_hex["x"], background_hex["y"] = x, y

    fig = go.Figure()
    fig.update_layout(
        autosize=False,
        height=area_height,
        width=area_width,
        margin={"l": 0, "r": 0, "b": 0, "t": 0, "pad": 0},
        paper_bgcolor=style["papercolor"],
        plot_bgcolor=style["papercolor"],
    )
    fig.update_xaxes(
        showgrid=False, showticklabels=False, showline=False, visible=False
    )
    fig.update_yaxes(
        showgrid=False,
        showticklabels=False,
        showline=False,
        visible=False,
        scaleanchor="x",
        scaleratio=1,
    )

    df["x"], df["y"] = hex_to_pixel(df.p, df.q, mode="flat")
    x_vals, y_vals = df.x, df.y
    df = df.drop(columns=["p", "q", "x", "y"])

    if len(df.columns) == 1:
        if isinstance(df, pd.DataFrame):
            df = df.iloc[:, 0]
        fig.add_trace(bg_hex())
        fig.add_trace(data_hex(df))

    elif isinstance(df, pd.DataFrame):
        slider_height = 100
        area_height += slider_height

        frames = []
        slider_steps = []

        fig.update_layout(
            autosize=False,
            height=area_height,
            width=area_width,
            margin={
                "l": 0,
                "r": 0,
                "b": slider_height,
                "t": 0,
                "pad": 0,
            },
            paper_bgcolor=style["papercolor"],
            plot_bgcolor=style["papercolor"],
            sliders=[
                {
                    "active": 0,
                    "currentvalue": {
                        "font": {"size": 16},
                        "visible": True,
                        "xanchor": "right",
                    },
                    "pad": {"b": 10, "t": 0},
                    "len": 0.9,
                    "x": 0.1,
                    "y": 0,
                    "steps": [],
                }
            ],
        )

        for i, col_name in enumerate(df.columns):
            series = df[col_name]
            frame_data = [
                bg_hex(),
                data_hex(series),
            ]

            frames.append(go.Frame(data=frame_data, name=str(i)))

            slider_steps.append(
                {
                    "args": [
                        [str(i)],
                        {"frame": {"duration": 0, "redraw": True}, "mode": "immediate"},
                    ],
                    "label": col_name,
                    "method": "animate",
                }
            )

            if i == 0:
                fig.add_traces(frame_data)

        fig.layout.sliders[0].steps = slider_steps
        fig.frames = frames

        fig.update_xaxes(
            showgrid=False, showticklabels=False, showline=False, visible=False
        )
        fig.update_yaxes(
            showgrid=False, showticklabels=False, showline=False, visible=False
        )

    else:
        raise ValueError("df must be a pd.Series or pd.DataFrame")

    return fig

def plot_avalanche_analysis(
    spikes: np.ndarray | torch.Tensor,
    bin_size: int = 1,
    dt: float = 1.0,  # Added for consistency
) -> tuple[Figure, dict]:
    """Plot avalanche size and duration distributions to analyze criticality.

    Creates a 3-panel figure showing:
    1. Avalanche size distribution P(S) with power-law fit
    2. Avalanche duration distribution P(T) with power-law fit
    3. Average size vs duration scaling relation `<S>(T)`

    Criticality is indicated by power-law distributions and specific
    scaling exponents (tau, alpha, gamma).

    Args:
        spikes: Spike matrix with shape (time, neurons).
        bin_size: Bin size for avalanche detection in timesteps.
        dt: Timestep in ms (unused, kept for interface consistency).

    Returns:
        Tuple of (figure, results) where results contains fitted exponents
        (tau, alpha, gamma), CCC (criticality consistency check), and
        power-law fit objects.

    Example:
        >>> fig, results = plot_avalanche_analysis(spikes, bin_size=5)
        >>> print(f"Tau: {results['tau']:.2f}, CCC: {results['CCC']:.2f}")
    """
    from ..analysis.dynamic_tools.criticality import compute_avalanche_statistics

    spikes = _to_numpy(spikes)
    results = compute_avalanche_statistics(spikes, bin_size=bin_size)

    fig = plt.figure(figsize=(15, 4))

    # 1. Size Distribution P(S)
    ax1 = fig.add_subplot(1, 3, 1)
    if results["fit_S"]:
        results["fit_S"].plot_pdf(color="b", linewidth=2, ax=ax1, label="Data")
        results["fit_S"].power_law.plot_pdf(
            color="b", linestyle="--", ax=ax1, label=f"Fit (tau={results['tau']:.2f})"
        )
    ax1.set_xlabel("Avalanche Size (S)")
    ax1.set_ylabel("P(S)")
    ax1.set_title("Size Distribution")
    if ax1.get_legend_handles_labels()[0]:
        ax1.legend()

    # 2. Duration Distribution P(T)
    ax2 = fig.add_subplot(1, 3, 2)
    if results["fit_T"]:
        results["fit_T"].plot_pdf(color="r", linewidth=2, ax=ax2, label="Data")
        results["fit_T"].power_law.plot_pdf(
            color="r",
            linestyle="--",
            ax=ax2,
            label=f"Fit (alpha={results['alpha']:.2f})",
        )
    ax2.set_xlabel("Avalanche Duration (T)")
    ax2.set_ylabel("P(T)")
    ax2.set_title("Duration Distribution")
    if ax2.get_legend_handles_labels()[0]:
        ax2.legend()

    # 3. Average Size vs Duration <S>(T)
    ax3 = fig.add_subplot(1, 3, 3)
    if (
        "avg_size_by_duration" in results
        and results["avg_size_by_duration"] is not None
    ):
        durations, mean_sizes = results["avg_size_by_duration"]
        ax3.loglog(durations, mean_sizes, "ko", markersize=4, label="Data")

        # Plot fit
        if not np.isnan(results["gamma"]):
            if "gamma_stats" in results and "popt" in results["gamma_stats"]:
                popt = results["gamma_stats"]["popt"]
                a, gamma = popt
                x_fit = np.logspace(
                    np.log10(durations.min()), np.log10(durations.max()), 100
                )
                y_fit = a * np.power(x_fit, gamma)
                ax3.loglog(
                    x_fit,
                    y_fit,
                    "g--",
                    label=f"Fit (gamma={results['gamma']:.2f})",
                )

    # Annotate CCC
    if not np.isnan(results["CCC"]):
        txt = f"CCC = {results['CCC']:.2f}\nPred gamma = {results['gamma_pred']:.2f}"
        ax3.text(
            0.05,
            0.95,
            txt,
            transform=ax3.transAxes,
            verticalalignment="top",
            bbox=dict(boxstyle="round", facecolor="white", alpha=0.8),
        )

    ax3.set_xlabel("Duration (T)")
    ax3.set_ylabel("Average Size <S>")
    ax3.set_title("Scaling Relation")
    if ax3.get_legend_handles_labels()[0]:
        ax3.legend()

    plt.tight_layout()
    return fig, results

def plot_dfa_analysis(
    # Dataclass interface
    data: DynamicsData | None = None,
    config: DFAConfig | None = None,
    format: DynamicsPlotFormat | None = None,
    # Plain args interface
    spikes: np.ndarray | torch.Tensor | None = None,
    dt: float = 1.0,
    min_window: int = 4,
    max_window: int | None = None,
    bin_size: int = 1,
    mode: Literal["individual", "grouped", "distribution"] = "individual",
    neurons_df: pd.DataFrame | None = None,
    **kwargs,
) -> Figure:
    """Plot DFA (Detrended Fluctuation Analysis) results.

    DFA quantifies long-range temporal correlations in spike trains.
    The scaling exponent (alpha) indicates:
    - alpha ≈ 0.5: Uncorrelated (random) activity
    - alpha > 0.5: Long-range positive correlations
    - alpha < 0.5: Long-range anti-correlations

    Supports both dataclass and plain argument interfaces.

    Args:
        data: DynamicsData container with spikes.
        config: DFAConfig with window and bin settings.
        format: DynamicsPlotFormat (mode affects plot style).
        spikes: Spike trains (time, neurons). Required if `data` not provided.
        dt: Timestep in milliseconds (for label consistency).
        min_window: Minimum window size for DFA in timesteps.
        max_window: Maximum window size. Auto-calculated if None.
        bin_size: Bin size for spike binning in timesteps.
        mode: Visualization mode (affects annotation style).
        neurons_df: Neuron metadata for potential grouping.
        **kwargs: Additional arguments.

    Returns:
        Figure with DFA summary and interpretation guide.

    Raises:
        ValueError: If spikes are not provided.

    Example:
        >>> fig = plot_dfa_analysis(spikes, bin_size=10)
    """
    # Resolve dataclass vs plain args
    if data is not None:
        spikes = data.spikes if spikes is None else spikes
        dt = data.dt if dt == 1.0 else dt
        neurons_df = data.neurons_df if neurons_df is None else neurons_df

    if config is not None:
        max_window = config.max_window if max_window is None else max_window
        bin_size = config.bin_size

    # Validate
    if spikes is None:
        raise ValueError("spikes is required")

    spikes = _to_numpy(spikes)

    # Compute DFA
    from ..analysis.dynamic_tools.criticality import calculate_dfa

    alpha = calculate_dfa(spikes, bin_size=bin_size)

    # Create simple plot showing the result
    fig, ax = plt.subplots(1, 1, figsize=(8, 6))
    ax.text(
        0.5,
        0.5,
        f"DFA Exponent (α): {alpha:.3f}",
        ha="center",
        va="center",
        fontsize=16,
        transform=ax.transAxes,
    )
    ax.text(
        0.5,
        0.4,
        "α ≈ 0.5: uncorrelated\n"
        "α > 0.5: long-range correlations\n"
        "α < 0.5: anti-correlations",
        ha="center",
        va="center",
        fontsize=10,
        transform=ax.transAxes,
    )
    ax.set_title("Detrended Fluctuation Analysis")
    ax.axis("off")
    return fig

def plot_eigenvalue_spectrum(
    weight_matrix: np.ndarray | torch.Tensor, ax: Axes | None = None
) -> tuple[Figure, Axes, dict]:
    """Plot the eigenvalue spectrum of a weight matrix.

    Visualizes eigenvalues in the complex plane with the spectral radius
    indicated by a dashed circle. Outliers (eigenvalues outside the bulk)
    are highlighted in red.

    Args:
        weight_matrix: Square connectivity matrix (N, N).
        ax: Existing axes to plot on. Creates new figure if None.

    Returns:
        Tuple of (figure, axes, results) where results contains:
        - "eigenvalues": Complex array of all eigenvalues
        - "spectral_radius": Radius of spectral bulk
        - "outliers": Array of outlier eigenvalues

    Example:
        >>> fig, ax, results = plot_eigenvalue_spectrum(W)
        >>> print(f"Spectral radius: {results['spectral_radius']:.2f}")
    """
    from ..analysis.dynamic_tools.attractor_dynamics import (
        calculate_structural_eigenvalue_outliers,
    )

    W = _to_numpy(weight_matrix)
    results = calculate_structural_eigenvalue_outliers(W)

    if ax is None:
        fig, ax = plt.subplots(figsize=(6, 6))
    else:
        fig = ax.get_figure()

    evals = results["eigenvalues"]
    r_spec = results["spectral_radius"]

    # Draw unit circle / spectral radius
    circle = plt.Circle(
        (0, 0),
        r_spec,
        color="black",
        fill=False,
        linestyle="--",
        alpha=0.5,
        label=f"R={r_spec:.2f}",
    )
    ax.add_artist(circle)

    # Scatter eigenvalues
    ax.scatter(evals.real, evals.imag, s=10, alpha=0.6, c="gray", edgecolors="none")

    # Highlight outliers
    outliers = results["outliers"]
    if len(outliers) > 0:
        ax.scatter(
            outliers.real,
            outliers.imag,
            s=30,
            c="red",
            label=f"Outliers ({len(outliers)})",
        )

    ax.set_aspect("equal")
    ax.set_xlabel(r"Re($\lambda$)")
    ax.set_ylabel(r"Im($\lambda$)")
    ax.set_title("Eigenvalue Spectrum")
    ax.grid(True, linestyle=":", alpha=0.3)
    ax.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
    plt.tight_layout()

    return fig, ax, results

def plot_firing_rate_distribution(
    spikes: np.ndarray | torch.Tensor,
    dt: float = 1.0,
    ax: Axes | None = None,
) -> tuple[Figure, dict]:
    """Plot the distribution of firing rates across neurons.

    Computes per-neuron firing rates and displays as a histogram with
    mean indicator.

    Args:
        spikes: Spike matrix with shape (time, neurons).
        dt: Timestep in milliseconds for rate calculation.
        ax: Existing axes to plot on. Creates new figure if None.

    Returns:
        Tuple of (figure, stats) where stats contains:
        - "rates": Array of firing rates per neuron (Hz)
        - "mean", "std", "median": Summary statistics

    Example:
        >>> fig, stats = plot_firing_rate_distribution(spikes, dt=1.0)
        >>> print(f"Mean rate: {stats['mean']:.1f} Hz")
    """
    from ..analysis.dynamic_tools.micro_scale import calculate_fr_distribution

    spikes = _to_numpy(spikes)
    stats = calculate_fr_distribution(spikes, dt=dt)

    if ax is None:
        fig, ax = plt.subplots(figsize=(6, 4))
    else:
        fig = ax.get_figure()

    rates = stats["rates"]
    ax.hist(rates, bins=30, color="skyblue", edgecolor="black", alpha=0.7)
    ax.axvline(
        stats["mean"],
        color="red",
        linestyle="--",
        label=f"Mean={stats['mean']:.1f} Hz",
    )
    ax.set_xlabel("Firing Rate (Hz)")
    ax.set_ylabel("Count")
    ax.set_title("Rate Distribution")
    ax.legend()
    ax.grid(alpha=0.3)

    if ax is None:
        plt.tight_layout()

    return fig, stats

def plot_gain_stability(data: tuple) -> tuple[Figure, Axes]:
    """Plot gain stability analysis results.

    Visualizes the relationship between network gain (g) and stability
    metrics (e.g., maximum Lyapunov exponent or spectral abscissa).
    A linear fit indicates consistent scaling behavior.

    Args:
        data: Tuple of (slope, intercept, g_values, lambda_values) where:
            - slope, intercept: Linear fit parameters
            - g_values: Array of gain values tested
            - lambda_values: Corresponding stability metrics

    Returns:
        Tuple of (figure, axes) with scatter plot and fit line.

    Example:
        >>> data = (slope, intercept, g_vals, lyap_vals)
        >>> fig, ax = plot_gain_stability(data)
    """
    slope, intercept, g_values, lambda_values = data

    fig, ax = plt.subplots(figsize=(6, 4))

    # Plot scatter of metrics
    ax.scatter(g_values, lambda_values, label="Data", color="blue", alpha=0.6)

    # Plot fit line
    x_range = np.linspace(min(g_values), max(g_values), 100)
    y_fit = slope * x_range + intercept
    ax.plot(x_range, y_fit, "r--", label=f"Fit: y={slope:.2f}x+{intercept:.2f}")

    ax.set_xlabel("Gain (g)")
    ax.set_ylabel("Lyapunov / Eigenvalue metric")
    ax.set_title("Gain Stability Analysis")
    ax.legend()
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    return fig, ax

def plot_group_box(
    values: TensorLike,
    neurons_df: pd.DataFrame,
    group_by: str,
    **kwargs,
) -> tuple[Figure, Axes]:
    """Convenience wrapper for `plot_group_distribution(..., kind='box')`."""
    return _plot_group_distribution_kind(
        values=values,
        neurons_df=neurons_df,
        group_by=group_by,
        kind="box",
        **kwargs,
    )

def plot_group_distribution(
    values: TensorLike,
    neurons_df: pd.DataFrame,
    group_by: str,
    *,
    kind: GroupPlotKind = "violin",
    simple_id_col: str = "simple_id",
    value_name: str = "value",
    group_order: Sequence | None = None,
    dropna: bool = True,
    ax: Axes | None = None,
    figsize: tuple[float, float] = (9.0, 5.0),
    title: str | None = None,
    showfliers: bool = False,
    linewidth: float = 1.5,
    alpha: float = 0.8,
) -> tuple[Figure, Axes]:
    """Plot grouped value distributions as violin, box, or ECDF.

    Groups per-neuron values by a categorical column in `neurons_df` and
    visualizes the distribution using the specified plot type.

    Args:
        values: Per-neuron values with shape (neurons,) or (time, neurons).
            If 2D, values are flattened across time for each neuron.
        neurons_df: DataFrame with neuron metadata. Must contain `simple_id_col`
            and `group_by` columns.
        group_by: Column name in `neurons_df` to group by (e.g., "cell_type").
        kind: Plot type - "violin", "box", or "ecdf".
        simple_id_col: Column name for neuron identifiers in `neurons_df`.
        value_name: Label for the y-axis / value dimension.
        group_order: Explicit order for groups. If None, uses natural sort.
        dropna: Whether to drop NaN values before plotting.
        ax: Existing axes to plot on. If None, creates new figure.
        figsize: Figure size (width, height) in inches.
        title: Plot title. If None, auto-generated from `kind` and `group_by`.
        showfliers: For box plots, whether to show outlier points.
        linewidth: Line width for ECDF curves.
        alpha: Opacity for violin/box fill.

    Returns:
        Tuple of (figure, axes) containing the plot.

    Raises:
        ValueError: If `kind` is not one of "violin", "box", "ecdf".

    Example:
        >>> fig, ax = plot_group_distribution(
        ...     firing_rates, neurons_df, group_by="cell_type", kind="violin"
        ... )
    """
    fig, ax = _resolve_figure_ax(ax=ax, figsize=figsize)

    if kind in {"violin", "box"}:
        grouped = group_values(
            values,
            neurons_df,
            group_by,
            simple_id_col=simple_id_col,
            value_name=value_name,
            group_order=group_order,
            dropna=dropna,
        )
        order = list(grouped.keys())
        grouped_arrays = [grouped[group] for group in order]

        if kind == "violin":
            _plot_violin(ax, grouped_arrays, order, alpha=alpha)
        else:
            _plot_box(
                ax,
                grouped_arrays,
                order,
                showfliers=showfliers,
                alpha=alpha,
            )

        ax.set_xlabel(group_by)
        ax.set_ylabel(value_name)
    elif kind == "ecdf":
        ecdf_by_group = group_ecdf(
            values,
            neurons_df,
            group_by,
            simple_id_col=simple_id_col,
            value_name=value_name,
            group_order=group_order,
            dropna=dropna,
        )
        _plot_ecdf(
            ax,
            ecdf_by_group,
            group_by=group_by,
            value_name=value_name,
            linewidth=linewidth,
        )
    else:
        raise ValueError(f"Unsupported kind `{kind}`.")

    if title is None:
        title = f"{kind.upper()} grouped by {group_by}"
    ax.set_title(title)
    ax.grid(alpha=0.2)

    return fig, ax