Plots¶

def plot_circle(
    results_df: pd.DataFrame,
    angles: list[float] | np.ndarray,
    *,
    profile_indices: list[int] | None = None,
    amax: float | None = None,
    angle_labels: list[str] | None = None,
    colors: str | list[str] | None = "Set2",
    fontsize: float = 12,
    drop_lowfit: bool = False,
    figsize: tuple[float, float] = (8, 8),
    title: str | None = None,
) -> Figure:
    """Plot SSM profiles on a circumplex circle.

    Creates a circular plot showing amplitude and displacement of SSM profiles,
    with arc bars representing confidence intervals. Automatically handles both
    single and multiple profiles.

    Parameters
    ----------
    results_df
        DataFrame with SSM results. Must contain columns:
        - Label: str, profile name
        - e_est, x_est, y_est, a_est, d_est: float, parameter estimates
        - e_lci, x_lci, ..., d_uci: float, confidence intervals
        - fit_est: float, model fit (0-1)
    angles
        Angular positions of scales in degrees
        (e.g., [90, 135, 180, 225, 270, 315, 360, 45]).
    profile_indices
        Which rows of results_df to plot. If None, plots all profiles.
    amax
        Maximum amplitude for scaling. If None, auto-computed using pretty_max().
    angle_labels
        Labels for each angle. If None, shows degree symbols (e.g., "90°").
        Pass empty strings to hide labels.
    colors
        Colors for profiles. Can be:
        - Seaborn palette name: "Set2", "husl", "deep", etc.
        - List of color specifications: ['red', 'blue'] or ['#FF0000', '#0000FF']
        - None: single blue color with no legend
    fontsize
        Base font size in points.
    drop_lowfit
        If True, omit profiles with fit < 0.70. If False, show with dashed borders.
    figsize
        Figure size in inches (width, height).
    title
        Title for the plot. If None, no title is added.

    Returns
    -------
    :
        Matplotlib Figure object.

    Examples
    --------
    Plot all profiles from an SSM analysis:

    >>> from circumplex import ssm_analyze
    >>> results = ssm_analyze(data, scales=list(range(8)))
    >>> fig = plot_circle(results.results, results.details.angles)
    >>> fig.savefig('profiles.png')

    Plot specific profiles with custom styling:

    >>> fig = plot_circle(
    ...     results.results,
    ...     results.details.angles,
    ...     profile_indices=[0, 1],
    ...     colors="husl",
    ...     fontsize=14,
    ...     figsize=(10, 10),
    ... )

    Use custom colors:

    >>> fig = plot_circle(
    ...     results.results,
    ...     results.details.angles,
    ...     colors=['red', 'blue', 'green'],
    ... )

    See Also
    --------
    [`plot_curve`](../plots/#circumplex.visualization.plots.plot_curve) :
        Plot SSM fitted curves with observed scores
    [`plot_contrast`](../plots/#circumplex.visualization.plots.plot_contrast) :
        Plot SSM parameter contrasts

    """
    # Validate and prepare data
    _validate_results_df(results_df)
    plot_df = _prepare_plot_data(results_df, profile_indices, drop_lowfit=drop_lowfit)

    # Determine amax
    if amax is None:
        if "a_uci" in plot_df.columns:
            amax = pretty_max(plot_df["a_uci"].values)
        else:
            amax = pretty_max(plot_df["a_est"].values)

    # Create figure and draw base
    fig, ax = plt.subplots(figsize=figsize)
    _draw_circle_base(ax, angles, amax, fontsize, angle_labels)

    # Scale factor: radius 5 corresponds to amax
    scale_factor = 5.0 / amax

    # Set up colors and legend
    n_profiles = len(plot_df)
    color_list, show_legend = _setup_colors(n_profiles, colors)

    # Plot each profile
    for i, (_idx, row) in enumerate(plot_df.iterrows()):
        color = color_list[i]
        label = row["Label"]

        # Scale parameters to circle coordinates
        x_plot = row["x_est"] * scale_factor
        y_plot = row["y_est"] * scale_factor

        # Draw confidence interval arc if available
        ci_cols = ["a_lci", "a_uci", "d_lci", "d_uci"]
        has_ci = all(col in row and not pd.isna(row[col]) for col in ci_cols)

        if has_ci:
            a_lci_plot = row["a_lci"] * scale_factor
            a_uci_plot = row["a_uci"] * scale_factor
            d_lci = row["d_lci"]
            d_uci = row["d_uci"]

            # Handle displacement wrapping (CI crosses 0/360)
            if d_uci < d_lci:
                d_uci += 360

            # Draw arc bar (wedge)
            wedge = mpatches.Wedge(
                center=(0, 0),
                r=a_uci_plot,
                theta1=d_lci,
                theta2=d_uci,
                width=a_uci_plot - a_lci_plot,
                facecolor=color,
                edgecolor=color,
                alpha=0.4,
                linestyle=row["linestyle"],
                linewidth=1.0,
            )
            ax.add_patch(wedge)

        # Draw point at (x, y)
        ax.scatter(
            x_plot,
            y_plot,
            s=100,
            facecolor=color,
            edgecolor="black",
            linewidth=1.0,
            zorder=10,
            label=label if show_legend else None,
        )

    # Add legend if multiple profiles
    if show_legend:
        ax.legend(
            title="Profile",
            fontsize=fontsize * 0.9,
            title_fontsize=fontsize,
            loc="upper left",
            bbox_to_anchor=(1.02, 1),
            frameon=True,
        )

    # Add title if provided
    if title is not None:
        fig.suptitle(title, fontsize=fontsize * 1.2, fontweight="bold")

    plt.tight_layout()

    return fig

def plot_curve(
    results_df: pd.DataFrame,
    scores_df: pd.DataFrame,
    angles: list[float] | np.ndarray,
    *,
    profile_indices: list[int] | None = None,
    angle_labels: list[str] | None = None,
    c_scores: str = "red",
    c_fit: str = "black",
    base_size: float = 11,
    drop_lowfit: bool = False,
    figsize: tuple[float, float] | None = None,
    incl_pred: bool = True,
    incl_fit: bool = False,
    incl_disp: bool = False,
    incl_amp: bool = False,
    incl_elev: bool = False,
) -> Figure:
    """Plot SSM fitted curves with observed scores.

    Creates a faceted plot showing the fitted cosine curve overlaid on the
    observed circumplex scale scores. Each profile is shown in a separate subplot.

    Parameters
    ----------
    results_df
        DataFrame with SSM results. Must contain columns:
        - Label: str, profile name
        - e_est, a_est, d_est: float, parameter estimates
        - fit_est: float, model fit (0-1)
    scores_df
        DataFrame with observed circumplex scores. Must have columns:
        - Label: str, profile name (matching results_df)
        - Scale columns: float, one column per circumplex scale
    angles
        Angular positions of scales in degrees, matching score column order.
    profile_indices
        Which rows to plot. If None, plots all profiles.
    angle_labels
        Labels for each angle on x-axis. If None, shows degree symbols (e.g., "90°").
    c_scores
        Color for observed scores (points and lines).
    c_fit
        Color for fitted curve.
    base_size
        Base font size in points for labels and text.
    drop_lowfit
        If True, omit profiles with fit < 0.70. If False, show with dashed curves.
    figsize
        Figure size in inches (width, height). If None, auto-computed based on
        number of profiles.

    Returns
    -------
    :
        Matplotlib Figure object.

    Examples
    --------
    Plot curves from an SSM analysis:

    >>> from circumplex import ssm_analyze
    >>> results = ssm_analyze(data, scales=list(range(8)))
    >>> fig = plot_curve(results.results, results.scores, results.details.angles)
    >>> fig.savefig('curves.png')

    Use custom angle labels:

    >>> fig = plot_curve(
    ...     results.results,
    ...     results.scores,
    ...     results.details.angles,
    ...     angle_labels=['PA', 'BC', 'DE', 'FG', 'HI', 'JK', 'LM', 'NO'],
    ... )

    See Also
    --------
    [`plot_circle`](../plots/#circumplex.visualization.plots.plot_circle) :
        Plot SSM profiles on a circumplex circle
    [`plot_contrast`](../plots/#circumplex.visualization.plots.plot_contrast) :
        Plot SSM parameter contrasts

    """
    # Validate results
    _validate_results_df(results_df)

    # Prepare plot data (filter profiles, handle low fit)
    plot_results = _prepare_plot_data(
        results_df, profile_indices, drop_lowfit=drop_lowfit
    )

    # Filter scores to match plot_results
    plot_scores = scores_df[scores_df["Label"].isin(plot_results["Label"])].copy()

    n_profiles = len(plot_results)
    if n_profiles == 0:
        msg = "No profiles to plot"
        raise ValueError(msg)

    # Determine subplot layout
    ncols = min(3, n_profiles)
    nrows = (n_profiles + ncols - 1) // ncols

    # Auto-compute figure size if not provided
    if figsize is None:
        figsize = (ncols * 6, nrows * 4)

    # Set up colors
    # color_list, _ = _setup_colors(n_profiles, colors)  # noqa: ERA001

    # Prepare angle labels for x-axis
    if angle_labels is None:
        xlabel = "Angle"
        tick_labels = [f"{int(a)}°" for a in angles]
    else:
        xlabel = "Scale"
        if len(angle_labels) != len(angles):
            msg = f"angle_labels must have same length as angles ({len(angles)})"
            raise ValueError(msg)
        tick_labels = angle_labels

    # Create figure with subplots
    fig, axes = plt.subplots(nrows, ncols, figsize=figsize, squeeze=False)
    # Convert ndarray of Axes to a plain list with proper typing for type checkers
    axes_flat: list[plt.Axes] = [cast("plt.Axes", a) for a in axes.flatten()]

    # Hide unused subplots
    for idx in range(n_profiles, len(axes_flat)):
        axes_flat[idx].set_visible(False)

    # Get scale column names (everything except Label, Model, Fit)
    info_cols = {"Label", "Model", "Fit"}
    scale_cols = [col for col in plot_scores.columns if col not in info_cols]

    if len(scale_cols) != len(angles):
        msg = (
            f"Number of scale columns ({len(scale_cols)}) must match "
            f"angles ({len(angles)})"
        )
        raise ValueError(msg)

    # Prepare angles array for sorting
    angles_array = np.array(angles)

    # Plot each profile
    for idx, (_ridx, result_row) in enumerate(plot_results.iterrows()):
        ax = axes_flat[idx]
        label = result_row["Label"]

        # Get scores for this profile
        score_row = plot_scores[plot_scores["Label"] == label].iloc[0]
        observed_scores = score_row[scale_cols].to_numpy().astype(float)

        # Sort angles and scores together for proper line connection
        sorted_indices = np.argsort(angles_array)
        angles_sorted = angles_array[sorted_indices]
        scores_sorted = observed_scores[sorted_indices]

        # Plot using helper function
        _plot_single_curve(
            ax,
            angles_array,
            angles_sorted,
            observed_scores,
            scores_sorted,
            result_row,
            c_scores,
            c_fit,
            tick_labels,
            xlabel,
            base_size,
            incl_pred=incl_pred,
            incl_fit=incl_fit,
            incl_disp=incl_disp,
            incl_amp=incl_amp,
            incl_elev=incl_elev,
        )

    plt.tight_layout()

    return fig

def plot_contrast(
    results_df: pd.DataFrame,
    *,
    drop_xy: bool = False,
    sig_color: str = "#fc8d62",
    ns_color: str = "white",
    linewidth: float = 1.25,
    fontsize: float = 12,
    figsize: tuple[float, float] | None = None,
) -> Figure:
    """Plot SSM parameter contrasts with confidence intervals.

    Creates a faceted plot showing the difference between two profiles for each
    SSM parameter (elevation, x-value, y-value, amplitude, displacement). Points
    are colored based on statistical significance (whether CI includes zero).

    This function requires results from a contrast analysis (e.g., comparing two
    groups or measures).

    Parameters
    ----------
    results_df
        DataFrame with SSM contrast results. Must contain the contrast row
        (typically the last row) with columns:
        - Label: str, contrast label (e.g., "Group 1 - Group 2")
        - e_est, x_est, y_est, a_est, d_est: float, parameter differences
        - e_lci, x_lci, ..., d_uci: float, confidence intervals
    drop_xy
        Whether to omit x-value and y-value parameters from the plot. This can
        simplify the plot when only interested in elevation, amplitude, and
        displacement (default = False).
    sig_color
        Color for significant contrasts (CI excludes zero).
    ns_color
        Color for non-significant contrasts (CI includes zero).
    linewidth
        Width of error bars and point outlines in points.
    fontsize
        Base font size in points for labels and text.
    figsize
        Figure size in inches (width, height). If None, uses (10, 4) for all
        parameters or (7, 4) if drop_xy=True.

    Returns
    -------
    :
        Matplotlib Figure object.

    Examples
    --------
    Plot contrasts from an SSM analysis:

    >>> from circumplex import ssm_analyze
    >>> results = ssm_analyze(
    ...     data, scales=list(range(8)),
    ...     grouping='condition', contrast=True
    ... )
    >>> fig = plot_contrast(results.results)
    >>> fig.savefig('contrasts.png')

    Drop x and y parameters for simpler plot:

    >>> fig = plot_contrast(results.results, drop_xy=True)

    Use custom colors:

    >>> fig = plot_contrast(
    ...     results.results,
    ...     sig_color='red',
    ...     ns_color='lightgray',
    ... )

    See Also
    --------
    [`plot_circle`](../plots/#circumplex.visualization.plots.plot_circle) :
        Plot SSM profiles on a circumplex circle
    [`plot_curve`](../plots/#circumplex.visualization.plots.plot_curve) :
        Plot SSM fitted curves with observed scores

    """
    # Build data
    contrast_row = _get_contrast_row(results_df)
    plot_data = _build_contrast_plot_data(contrast_row, drop_xy=drop_xy)

    n_params = len(plot_data)

    # Set figure size
    if figsize is None:
        figsize = (7, 4) if drop_xy else (10, 4)

    # Create figure with subplots
    fig, axes = plt.subplots(
        1, n_params, figsize=figsize, sharey=False, constrained_layout=False
    )

    # Normalize axes into a list of Axes for consistent typing/iteration
    if n_params == 1:
        axes_list: list[plt.Axes] = [cast("plt.Axes", axes)]
    else:
        axes_list = [cast("plt.Axes", a) for a in cast("np.ndarray", axes)]

    # Plot each parameter
    for i, (ax, data) in enumerate(zip(axes_list, plot_data, strict=False)):
        _draw_contrast_subplot(
            ax=ax,
            data=data,
            sig_color=sig_color,
            ns_color=ns_color,
            linewidth=linewidth,
            fontsize=fontsize,
            is_leftmost=(i == 0),
        )

    # Add legend with better positioning
    _add_contrast_legend(fig, sig_color, ns_color, fontsize)

    # Add contrast label as suptitle
    contrast_label = contrast_row["Label"]
    fig.suptitle(
        f"Contrast: {contrast_label}",
        fontsize=fontsize * 1.15,
        fontweight="bold",
        y=0.88,
    )

    plt.tight_layout(rect=[0, 0, 1, 0.86])

    return fig