Api pl

def violin(
    adata: AnnData,
    keys: Union[str, Sequence[str]],
    groupby: Optional[str] = None,
    *,
    log: bool = False,
    use_raw: Optional[bool] = None,
    stripplot: bool = True,
    jitter: Union[float, bool] = True,
    size: int = 1,
    layer: Optional[str] = None,
    density_norm: DensityNorm = "width",
    order: Optional[Sequence[str]] = None,
    multi_panel: Optional[bool] = None,
    xlabel: str = "",
    ylabel: Optional[Union[str, Sequence[str]]] = None,
    rotation: Optional[float] = None,
    show: Optional[bool] = None,
    save: Optional[Union[bool, str]] = None,
    ax: Optional[Axes] = None,
    # Enhanced features
    enhanced_style: bool = True,
    show_means: bool = False,
    show_boxplot: bool = False,
    jitter_method: str = 'uniform',  # 'uniform', 't_dist'
    jitter_alpha: float = 0.4,
    violin_alpha: float = 0.8,
    background_color: str = 'white',
    spine_color: str = '#b4aea9',
    grid_lines: bool = True,
    statistical_tests: bool = False,
    custom_colors: Optional[Sequence[str]] = None,
    figsize: Optional[tuple] = None,
    fontsize=13,
    ticks_fontsize=None,
    **kwds
) -> Union[Axes, None]:
    r"""
    Enhanced violin plot compatible with scanpy's interface.

    This function provides all the functionality of scanpy's violin plot
    with additional customization options for enhanced visualization,
    implemented using pure matplotlib.

    Arguments:
        adata: AnnData. Annotated data matrix.
        keys: str | Sequence[str]. Keys for accessing variables of `.var_names` or fields of `.obs`.
        groupby: str | None. The key of the observation grouping to consider. (None)
        log: bool. Plot on logarithmic axis. (False)
        use_raw: bool | None. Whether to use `raw` attribute of `adata`. Defaults to `True` if `.raw` is present. (None)
        stripplot: bool. Add a stripplot on top of the violin plot. (True)
        jitter: float | bool. Add jitter to the stripplot (only when stripplot is True). (True)
        size: int. Size of the jitter points. (1)
        layer: str | None. Name of the AnnData object layer that wants to be plotted. (None)
        density_norm: str. The method used to scale the width of each violin. If 'width' (the default), each violin will have the same width. If 'area', each violin will have the same area. If 'count', a violin's width corresponds to the number of observations. ('width')
        order: Sequence[str] | None. Order in which to show the categories. (None)
        multi_panel: bool | None. Display keys in multiple panels also when `groupby is not None`. (None)
        xlabel: str. Label of the x axis. ('')
        ylabel: str | Sequence[str] | None. Label of the y axis. (None)
        rotation: float | None. Rotation of xtick labels. (None)
        show: bool | None. Whether to show the plot. (None)
        save: bool | str | None. Path to save the figure. (None)
        ax: Axes | None. A matplotlib axes object. (None)
        enhanced_style: bool. Whether to apply enhanced styling. (True)
        show_means: bool. Whether to show mean values with annotations. (False)
        show_boxplot: bool. Whether to overlay box plots on violins. (False)
        jitter_method: str. Method for jittering: 'uniform' or 't_dist'. ('uniform')
        jitter_alpha: float. Transparency of jittered points. (0.4)
        violin_alpha: float. Transparency of violin plots. (0.8)
        background_color: str. Background color of the plot. ('white')
        spine_color: str. Color of plot spines. ('#b4aea9')
        grid_lines: bool. Whether to show horizontal grid lines. (True)
        statistical_tests: bool. Whether to perform and display statistical tests. (False)
        custom_colors: Sequence[str] | None. Custom colors for groups. (None)
        figsize: tuple | None. Figure size (width, height). (None)
        fontsize: int. Font size for labels and ticks. (13)
        ticks_fontsize: int | None. Font size for axis ticks. If None, uses fontsize-1. (None)
        **kwds: Additional keyword arguments passed to violinplot.

    Returns:
        ax: matplotlib.axes.Axes | None. A matplotlib axes object if `ax` is `None` else `None`.
    """

    # Handle AnnData availability
    if not ANNDATA_AVAILABLE:
        raise ImportError("AnnData is required for this function. Install with: pip install anndata")

    # Ensure keys is a list
    if isinstance(keys, str):
        keys = [keys]
    keys = list(OrderedDict.fromkeys(keys))  # remove duplicates, preserving order

    # Handle ylabel
    if isinstance(ylabel, (str, type(None))):
        ylabel = [ylabel] * (1 if groupby is None else len(keys))
    ylabel=keys

    # Validate ylabel length
    if groupby is None:
        if len(ylabel) != 1:
            raise ValueError(f"Expected number of y-labels to be `1`, found `{len(ylabel)}`.")
    elif len(ylabel) != len(keys):
        raise ValueError(f"Expected number of y-labels to be `{len(keys)}`, found `{len(ylabel)}`.")

    # Extract data from AnnData object
    obs_df = _extract_data_from_adata(adata, keys, groupby, layer, use_raw)

    # Colorful data analysis and parameter optimization suggestions
    print(f"{Colors.HEADER}{Colors.BOLD}🎻 Violin Plot Analysis:{Colors.ENDC}")
    print(f"   {Colors.CYAN}Total cells: {Colors.BOLD}{len(obs_df)}{Colors.ENDC}")
    print(f"   {Colors.BLUE}Variables to plot: {Colors.BOLD}{len(keys)} {keys}{Colors.ENDC}")

    if groupby is not None:
        group_counts = obs_df[groupby].value_counts()
        print(f"   {Colors.GREEN}Groupby variable: '{Colors.BOLD}{groupby}{Colors.ENDC}{Colors.GREEN}' with {Colors.BOLD}{len(group_counts)}{Colors.ENDC}{Colors.GREEN} groups{Colors.ENDC}")

        # Show group distribution
        for group, count in group_counts.head(10).items():  # Show top 10 groups
            if count < 10:
                color = Colors.WARNING
            elif count < 50:
                color = Colors.BLUE
            else:
                color = Colors.GREEN
            print(f"     {color}• {group}: {Colors.BOLD}{count}{Colors.ENDC}{color} cells{Colors.ENDC}")

        if len(group_counts) > 10:
            print(f"     {Colors.CYAN}... and {Colors.BOLD}{len(group_counts) - 10}{Colors.ENDC}{Colors.CYAN} more groups{Colors.ENDC}")

        # Check for imbalanced groups
        min_count = group_counts.min()
        max_count = group_counts.max()
        if max_count / min_count > 10:
            print(f"   {Colors.WARNING}⚠️  Imbalanced groups detected: {Colors.BOLD}{min_count}-{max_count}{Colors.ENDC}{Colors.WARNING} cells per group{Colors.ENDC}")
    else:
        print(f"   {Colors.BLUE}Groupby: {Colors.BOLD}None{Colors.ENDC}{Colors.BLUE} (comparing variables){Colors.ENDC}")

    # Analyze data distribution for each variable
    print(f"\n{Colors.HEADER}{Colors.BOLD}📊 Data Distribution Analysis:{Colors.ENDC}")
    for key in keys:
        if key in obs_df.columns:
            data_vals = obs_df[key].dropna()
            if len(data_vals) > 0:
                data_range = data_vals.max() - data_vals.min()
                zero_fraction = (data_vals == 0).sum() / len(data_vals)

                # Determine if data might be log-transformed already
                log_suggestion = ""
                if data_vals.min() >= 0 and data_vals.max() > 100:
                    log_suggestion = f" {Colors.BLUE}(consider log=True){Colors.ENDC}"
                elif data_vals.min() < 0:
                    log_suggestion = f" {Colors.WARNING}(negative values detected){Colors.ENDC}"

                print(f"   {Colors.BLUE}'{key}': range {Colors.BOLD}{data_vals.min():.2f}-{data_vals.max():.2f}{Colors.ENDC}{Colors.BLUE}, {Colors.BOLD}{zero_fraction*100:.1f}%{Colors.ENDC}{Colors.BLUE} zeros{log_suggestion}")

    # Display current function parameters
    print(f"\n{Colors.HEADER}{Colors.BOLD}⚙️  Current Function Parameters:{Colors.ENDC}")
    print(f"   {Colors.BLUE}Plot style: enhanced_style={Colors.BOLD}{enhanced_style}{Colors.ENDC}{Colors.BLUE}, stripplot={Colors.BOLD}{stripplot}{Colors.ENDC}{Colors.BLUE}, jitter={Colors.BOLD}{jitter}{Colors.ENDC}")
    print(f"   {Colors.BLUE}Additional features: show_means={Colors.BOLD}{show_means}{Colors.ENDC}{Colors.BLUE}, show_boxplot={Colors.BOLD}{show_boxplot}{Colors.ENDC}{Colors.BLUE}, statistical_tests={Colors.BOLD}{statistical_tests}{Colors.ENDC}")
    print(f"   {Colors.BLUE}Figure settings: figsize={Colors.BOLD}{figsize}{Colors.ENDC}{Colors.BLUE}, fontsize={Colors.BOLD}{fontsize}{Colors.ENDC}{Colors.BLUE}, violin_alpha={Colors.BOLD}{violin_alpha}{Colors.ENDC}")
    if custom_colors is not None:
        print(f"   {Colors.BLUE}Colors: {Colors.BOLD}{len(custom_colors)} custom colors specified{Colors.ENDC}")
    else:
        print(f"   {Colors.BLUE}Colors: {Colors.BOLD}Default palette{Colors.ENDC}")

    # Parameter optimization suggestions
    print(f"\n{Colors.HEADER}{Colors.BOLD}💡 Parameter Optimization Suggestions:{Colors.ENDC}")
    suggestions = []

    # Check for too many groups
    if groupby is not None:
        n_groups = len(obs_df[groupby].unique())
        if n_groups > 8:
            suggestions.append(f"   {Colors.WARNING}▶ Many groups detected ({n_groups}):{Colors.ENDC}")
            suggestions.append(f"     {Colors.CYAN}Current: figsize={Colors.BOLD}{figsize}{Colors.ENDC}")
            suggested_width = max(8, n_groups * 1.2)
            suggestions.append(f"     {Colors.GREEN}Suggested: {Colors.BOLD}figsize=({suggested_width}, {figsize[1] if figsize else 6}){Colors.ENDC}")

            if rotation is None:
                suggestions.append(f"     {Colors.GREEN}Consider adding: {Colors.BOLD}rotation=45{Colors.ENDC} for better label readability")

        # Check for small sample sizes
        if groupby is not None:
            group_counts = obs_df[groupby].value_counts()
            if group_counts.min() < 10:
                suggestions.append(f"   {Colors.WARNING}▶ Small sample sizes detected (min: {group_counts.min()}):{Colors.ENDC}")
                suggestions.append(f"     {Colors.GREEN}Suggested: {Colors.BOLD}stripplot=True, jitter=0.3{Colors.ENDC} to show individual points")

    # Check data distribution and log scale
    for key in keys:
        if key in obs_df.columns:
            data_vals = obs_df[key].dropna()
            if len(data_vals) > 0 and data_vals.min() >= 0 and data_vals.max() / data_vals.min() > 100:
                suggestions.append(f"   {Colors.WARNING}▶ Wide data range for '{key}' ({data_vals.max()/data_vals.min():.1f}x):{Colors.ENDC}")
                suggestions.append(f"     {Colors.CYAN}Current: log={Colors.BOLD}{log}{Colors.ENDC}")
                suggestions.append(f"     {Colors.GREEN}Suggested: {Colors.BOLD}log=True{Colors.ENDC} for better visualization")
                break

    # Check figure size vs number of variables
    if len(keys) > 3 and (figsize is None or figsize[0] < len(keys) * 2):
        suggestions.append(f"   {Colors.WARNING}▶ Multiple variables with small figure:{Colors.ENDC}")
        suggested_width = max(len(keys) * 2, 8)
        suggestions.append(f"     {Colors.CYAN}Current: figsize={Colors.BOLD}{figsize}{Colors.ENDC}")
        suggestions.append(f"     {Colors.GREEN}Suggested: {Colors.BOLD}figsize=({suggested_width}, 6){Colors.ENDC} or {Colors.BOLD}multi_panel=True{Colors.ENDC}")

    # Font size optimization
    if groupby is not None:
        n_groups = len(obs_df[groupby].unique())
        max_label_length = max(len(str(x)) for x in obs_df[groupby].unique())
        if max_label_length > 8 and fontsize > 12:
            suggestions.append(f"   {Colors.WARNING}▶ Long group labels detected:{Colors.ENDC}")
            suggestions.append(f"     {Colors.CYAN}Current: fontsize={Colors.BOLD}{fontsize}{Colors.ENDC}")
            suggestions.append(f"     {Colors.GREEN}Suggested: {Colors.BOLD}fontsize=10, rotation=45{Colors.ENDC}")

    # Enhanced features suggestions
    if not show_means and not show_boxplot and not statistical_tests:
        suggestions.append(f"   {Colors.BLUE}▶ Consider enhancing your plot:{Colors.ENDC}")
        suggestions.append(f"     {Colors.GREEN}Options: {Colors.BOLD}show_means=True{Colors.ENDC}{Colors.GREEN}, {Colors.BOLD}show_boxplot=True{Colors.ENDC}{Colors.GREEN}, or {Colors.BOLD}statistical_tests=True{Colors.ENDC}")

    if suggestions:
        for suggestion in suggestions:
            print(suggestion)

        print(f"\n   {Colors.BOLD}📋 Copy-paste ready function call:{Colors.ENDC}")

        # Generate optimized function call
        optimized_params = []

        # Core parameters
        if isinstance(keys, list) and len(keys) == 1:
            optimized_params.append(f"adata, keys='{keys[0]}'")
        else:
            optimized_params.append(f"adata, keys={keys}")

        if groupby is not None:
            optimized_params.append(f"groupby='{groupby}'")

        # Add optimized parameters based on suggestions
        if groupby is not None:
            n_groups = len(obs_df[groupby].unique())
            if n_groups > 8:
                suggested_width = max(8, n_groups * 1.2)
                optimized_params.append(f"figsize=({suggested_width}, 6)")

            if rotation is None and n_groups > 6:
                optimized_params.append("rotation=45")

            group_counts = obs_df[groupby].value_counts()
            if group_counts.min() < 10:
                optimized_params.append("stripplot=True")
                optimized_params.append("jitter=0.3")

        # Check for log scale suggestion
        for key in keys:
            if key in obs_df.columns:
                data_vals = obs_df[key].dropna()
                if len(data_vals) > 0 and data_vals.min() >= 0 and data_vals.max() / data_vals.min() > 100:
                    optimized_params.append("log=True")
                    break

        # Multi-panel suggestion
        if len(keys) > 3 and (figsize is None or figsize[0] < len(keys) * 2):
            optimized_params.append("multi_panel=True")

        # Enhanced features
        if not show_means and not show_boxplot:
            optimized_params.append("show_means=True")

        optimized_call = f"   {Colors.GREEN}ov.pl.violin({', '.join(optimized_params)}){Colors.ENDC}"
        print(optimized_call)
    else:
        print(f"   {Colors.GREEN}✅ Current parameters are optimal for your data!{Colors.ENDC}")

    print(f"{Colors.CYAN}{'─' * 60}{Colors.ENDC}")

    # Prepare data for plotting
    if groupby is None:
        obs_tidy = pd.melt(obs_df, value_vars=keys)
        x_col = "variable"
        y_cols = ["value"]
        group_categories = keys
    else:
        obs_tidy = obs_df
        x_col = groupby
        y_cols = keys
        obs_df[groupby] = obs_df[groupby].astype('category')
        group_categories = obs_df[groupby].cat.categories if order is None else order

    # Set up colors
    colors = _setup_colors(custom_colors, group_categories, adata, groupby)

    # Handle multi-panel case
    if multi_panel and groupby is None and len(y_cols) == 1:
        return _create_multi_panel_plot(
            obs_tidy, keys, y_cols[0], colors, log, stripplot, jitter, 
            size, density_norm, enhanced_style, **kwds
        )

    # Create single or multiple axis plots
    if ax is None and figsize is not None:
        fig, ax = plt.subplots(figsize=figsize)
    elif ax is None:
        fig, ax = plt.subplots(figsize=(10, 6))

    # Apply enhanced styling
    if enhanced_style:
        _apply_enhanced_styling(ax, background_color, spine_color, grid_lines)

    # Create plots for each y variable
    for i, (y_col, y_label) in enumerate(zip(y_cols, ylabel)):
        if len(y_cols) > 1:
            # Create subplots for multiple keys
            if i == 0:
                fig, axes = plt.subplots(1, len(y_cols), figsize=figsize or (5*len(y_cols), 6))
                if len(y_cols) == 1:
                    axes = [axes]
            current_ax = axes[i] if len(y_cols) > 1 else ax
        else:
            current_ax = ax

        # Prepare data for current y variable
        plot_data = _prepare_plot_data(obs_tidy, x_col, y_col, group_categories, order)

        # Create violin plots
        _create_violin_plots(
            current_ax, plot_data, group_categories, colors, density_norm, 
            violin_alpha, enhanced_style, **kwds
        )

        # Add box plots if requested
        if show_boxplot:
            _add_box_plots(current_ax, plot_data, group_categories)

        # Add strip plot (jittered points)
        if stripplot:
            _add_strip_plot(
                current_ax, plot_data, group_categories, jitter, jitter_method, 
                size, jitter_alpha, colors
            )

        # Add mean annotations
        if show_means:
            _add_mean_annotations(current_ax, plot_data, group_categories)

        # Add statistical tests
        if statistical_tests:
            _add_statistical_tests(current_ax, plot_data, group_categories)

        # Customize axis
        _customize_axis(
            current_ax, group_categories, xlabel, y_label, groupby, 
            rotation, log, order
        )
        current_ax.spines['top'].set_visible(False)
        current_ax.spines['right'].set_visible(False)
        current_ax.spines['bottom'].set_visible(True)
        current_ax.spines['left'].set_visible(True)
        current_ax.spines['left'].set_position(('outward', 10))
        current_ax.spines['bottom'].set_position(('outward', 10))
        if ticks_fontsize==None:
            ticks_fontsize=fontsize-1

        current_ax.set_xticklabels(current_ax.get_xticklabels(),fontsize=ticks_fontsize,rotation=rotation)
        current_ax.set_yticklabels(current_ax.get_yticklabels(),fontsize=ticks_fontsize)
        current_ax.set_xlabel(groupby,fontsize=fontsize)
        current_ax.set_ylabel(y_label,fontsize=fontsize)
        #print(y_label)

    #plt.tight_layout()

        current_ax.grid(False)



    # Save figure if requested
    if save:
        save_path = save if isinstance(save, str) else "violin_plot.pdf"
        plt.savefig(save_path, dpi=300, bbox_inches='tight')

    # Show figure if requested
    if show is True or (show is None and ax is None):
        plt.show()
        return None

    return ax

def dotplot(
    adata: AnnData,
    var_names: Union[_VarNames, Mapping[str, _VarNames]],
    groupby: Union[str, Sequence[str]],
    *,
    use_raw: Optional[bool] = None,
    log: bool = False,
    num_categories: int = 7,
    categories_order: Optional[Sequence[str]] = None,
    expression_cutoff: float = 0.0,
    mean_only_expressed: bool = False,
    standard_scale: Optional[Literal['var', 'group']] = None,
    title: Optional[str] = None,
    colorbar_title: Optional[str] = 'Mean expression\nin group',
    size_title: Optional[str] = 'Fraction of cells\nin group (%)',
    figsize: Optional[Tuple[float, float]] = None,
    dendrogram: Union[bool, str] = False,
    gene_symbols: Optional[str] = None,
    var_group_positions: Optional[Sequence[Tuple[int, int]]] = None,
    var_group_labels: Optional[Sequence[str]] = None,
    var_group_rotation: Optional[float] = None,
    layer: Optional[str] = None,
    swap_axes: Optional[bool] = False,
    dot_color_df: Optional[pd.DataFrame] = None,
    show: Optional[bool] = None,
    save: Optional[Union[str, bool]] = None,
    ax: Optional[_AxesSubplot] = None,
    return_fig: Optional[bool] = False,
    vmin: Optional[float] = None,
    vmax: Optional[float] = None,
    vcenter: Optional[float] = None,
    norm: Optional[Normalize] = None,
    cmap: Union[Colormap, str, None] = 'Reds',
    dot_max: Optional[float] = None,
    dot_min: Optional[float] = None,
    smallest_dot: float = 0.0,
    fontsize: int = 12,
    preserve_dict_order: bool = False,
    **kwds,
) -> Optional[Union[Dict, 'DotPlot']]:
    r"""
    Make a dot plot of the expression values of `var_names`.

    For each var_name and each `groupby` category a dot is plotted.
    Each dot represents two values: mean expression within each category
    (visualized by color) and fraction of cells expressing the `var_name` in the
    category (visualized by the size of the dot).

    Arguments:
        adata: AnnData
            Annotated data matrix.
        var_names: str or list of str or dict
            Variables to plot.
        groupby: str or list of str
            The key of the observation grouping to consider.
        use_raw: bool, optional (default=None)
            Use `raw` attribute of `adata` if present.
        log: bool, optional (default=False)
            Whether to log-transform the data.
        num_categories: int, optional (default=7)
            Number of categories to show.
        categories_order: list of str, optional (default=None)
            Order of categories to display.
        expression_cutoff: float, optional (default=0.0)
            Expression cutoff for calculating fraction of expressing cells.
        mean_only_expressed: bool, optional (default=False)
            Whether to calculate mean only for expressing cells.
        standard_scale: {'var', 'group'} or None, optional (default=None)
            Whether to standardize data.
        title: str, optional (default=None)
            Title for the plot.
        colorbar_title: str, optional (default='Mean expression\nin group')
            Title for the color bar.
        size_title: str, optional (default='Fraction of cells\nin group (%)')
            Title for the size legend.
        figsize: tuple, optional (default=None)
            Figure size (width, height) in inches. If provided, the plot dimensions will be scaled accordingly.
        dendrogram: bool or str, optional (default=False)
            Whether to add dendrogram to the plot.
        gene_symbols: str, optional (default=None)
            Key for gene symbols in `adata.var`.
        var_group_positions: list of tuples, optional (default=None)
            Positions for variable groups.
        var_group_labels: list of str, optional (default=None)
            Labels for variable groups.
        var_group_rotation: float, optional (default=None)
            Rotation angle for variable group labels.
        layer: str, optional (default=None)
            Layer to use for expression data.
        swap_axes: bool, optional (default=False)
            Whether to swap x and y axes.
        dot_color_df: pandas.DataFrame, optional (default=None)
            DataFrame for dot colors.
        show: bool, optional (default=None)
            Whether to show the plot.
        save: str or bool, optional (default=None)
            Whether to save the plot.
        ax: matplotlib.axes.Axes, optional (default=None)
            Axes object to plot on.
        return_fig: bool, optional (default=False)
            Whether to return the figure object.
        vmin: float, optional (default=None)
            Minimum value for color scaling.
        vmax: float, optional (default=None)
            Maximum value for color scaling.
        vcenter: float, optional (default=None)
            Center value for diverging colormap.
        norm: matplotlib.colors.Normalize, optional (default=None)
            Normalization object for colors.
        cmap: str or matplotlib.colors.Colormap, optional (default='Reds')
            Colormap for the plot.
        dot_max: float, optional (default=None)
            Maximum dot size.
        dot_min: float, optional (default=None)
            Minimum dot size.
        smallest_dot: float, optional (default=0.0)
            Size of the smallest dot.
        fontsize: int, optional (default=12)
            Font size for labels and legends. Titles will be one point larger.
        preserve_dict_order: bool, optional (default=False)
            When var_names is a dictionary, whether to preserve the original dictionary order.
            If True, genes will be ordered according to the dictionary's insertion order.
            If False (default), genes will be ordered according to cell type categories.

    Returns:
        If `return_fig` is True, returns the figure object.
        If `show` is False, returns axes dictionary.
    """
    # Convert var_names to list if string
    original_var_names_dict = None
    if isinstance(var_names, str):
        var_names = [var_names]
    elif isinstance(var_names, Mapping):
        # Save original dictionary reference for color bar ordering
        if preserve_dict_order:
            original_var_names_dict = var_names

        # Get gene groups
        gene_groups = []
        var_names_list = []

        if preserve_dict_order:
            # Preserve the original dictionary order
            for group, genes in var_names.items():
                if isinstance(genes, str):
                    genes = [genes]
                var_names_list.extend(genes)
                gene_groups.extend([group] * len(genes))
        else:
            # Get cell type order (original behavior)
            if categories_order is not None:
                group_order = categories_order
            elif pd.api.types.is_categorical_dtype(adata.obs[groupby]):
                group_order = list(adata.obs[groupby].cat.categories)
            else:
                group_order = list(adata.obs[groupby].unique())

            # Order gene groups according to cell types
            for group in group_order:
                if group in var_names:
                    genes = var_names[group]
                    if isinstance(genes, str):
                        genes = [genes]
                    var_names_list.extend(genes)
                    gene_groups.extend([group] * len(genes))

            # Add any remaining groups that weren't in the cell types
            for group, genes in var_names.items():
                if group not in group_order:
                    if isinstance(genes, str):
                        genes = [genes]
                    var_names_list.extend(genes)
                    gene_groups.extend([group] * len(genes))

        var_names = var_names_list

    # Get expression matrix
    if use_raw and adata.raw is not None:
        matrix = adata.raw.X
        var_names_idx = [adata.raw.var_names.get_loc(name) for name in var_names]
    else:
        matrix = adata.X if layer is None else adata.layers[layer]
        var_names_idx = [adata.var_names.get_loc(name) for name in var_names]

    # Determine category order
    if categories_order is not None:
        cats = categories_order
    else:
        # Use the categorical order from adata if available
        if pd.api.types.is_categorical_dtype(adata.obs[groupby]):
            cats = adata.obs[groupby].cat.categories
        else:
            # If not categorical, get unique values
            cats = adata.obs[groupby].unique()

    # Get aggregated data with specified order
    agg = adata.obs[groupby].value_counts().reindex(cats)
    cell_counts = agg.to_numpy()

    # Get colors for cell types if available
    cell_colors = None
    color_dict = None
    try:
        color_key = f"{groupby}_colors"
        if color_key in adata.uns:
            colors = adata.uns[color_key]
            # Create color dictionary mapping cell types to colors
            if pd.api.types.is_categorical_dtype(adata.obs[groupby]):
                # Use categorical order for colors
                color_dict = dict(zip(adata.obs[groupby].cat.categories, colors))
            else:
                # Use unique order for colors
                unique_cats = adata.obs[groupby].unique()
                color_dict = dict(zip(unique_cats, colors[:len(unique_cats)]))

            # Get colors for the actual categories in the plot
            cell_colors = [color_dict.get(cat, '#CCCCCC') for cat in agg.index]
    except (KeyError, IndexError):
        cell_colors = None
        color_dict = None

    # Calculate mean expression and fraction of expressing cells
    means = np.zeros((len(agg), len(var_names)))
    fractions = np.zeros_like(means)

    for i, group in enumerate(agg.index):
        mask = adata.obs[groupby] == group
        group_matrix = matrix[mask][:, var_names_idx]

        # Calculate mean expression
        if mean_only_expressed:
            expressed = group_matrix > expression_cutoff
            means[i] = np.array([
                group_matrix[:, j][expressed[:, j]].mean() if expressed[:, j].any() else 0
                for j in range(group_matrix.shape[1])
            ])
        else:
            means[i] = np.mean(group_matrix, axis=0)

        # Calculate fraction of expressing cells
        fractions[i] = np.mean(group_matrix > expression_cutoff, axis=0)

    # Scale if requested
    if standard_scale == 'group':
        means = (means - means.min(axis=1, keepdims=True)) / (means.max(axis=1, keepdims=True) - means.min(axis=1, keepdims=True))
    elif standard_scale == 'var':
        means = (means - means.min(axis=0)) / (means.max(axis=0) - means.min(axis=0))

    # Handle dot size limits
    if dot_max is not None:
        fractions = np.minimum(fractions, dot_max)
    if dot_min is not None:
        fractions = np.maximum(fractions, dot_min)

    # Scale dot sizes to account for smallest_dot
    if smallest_dot > 0:
        fractions = smallest_dot + (1 - smallest_dot) * fractions

    # Create the plot
    h, w = means.shape

    # Calculate dimensions based on figsize if provided
    if figsize is not None:
        # Use figsize to determine height and width
        # Adjust for the number of rows and columns to maintain aspect ratio
        base_height = figsize[1] * 0.7  # Use 70% of figsize height for main plot
        base_width = figsize[0] * 0.7   # Use 70% of figsize width for main plot

        # Scale based on data dimensions
        height = base_height * (h / max(h, w))
        width = base_width * (w / max(h, w))
    else:
        # Default behavior
        height = h / 3
        width = w / 3


    # Create SizedHeatmap
    m = ma.SizedHeatmap(
        size=fractions,
        color=means,
        cluster_data=fractions if dendrogram else None,
        height=height,
        width=width,
        edgecolor="lightgray",
        cmap=cmap,
        vmin=vmin,
        vmax=vmax,
        #norm=norm,
        size_legend_kws=dict(
            colors="#c2c2c2",
            title=size_title,
            labels=[f"{int(x*100)}%" for x in [0.2, 0.4, 0.6, 0.8, 1.0]],
            show_at=[0.2, 0.4, 0.6, 0.8, 1.0],
            fontsize=fontsize,
            ncol=3,
            title_fontproperties={"size": fontsize + 1, "weight": 100}
        ),
        color_legend_kws=dict(
            title=colorbar_title,
            fontsize=fontsize,
            orientation="horizontal",
            title_fontproperties={"size": fontsize + 1, "weight": 100}
        ),
    )

    # Add labels
    m.add_top(mp.Labels(var_names, fontsize=fontsize), pad=0.1)

    # Group genes if var_names was a dictionary
    if 'gene_groups' in locals():
        # Get colors for gene groups
        try:
            # Use the same color_dict that was created for cell types
            if color_dict is not None:
                # Get unique groups and check which ones are not in color_dict
                unique_groups = list(dict.fromkeys(gene_groups))
                missing_groups = [g for g in unique_groups if g not in color_dict]

                # If there are missing groups, add colors from palette
                if missing_groups:
                    if len(palette_28) >= len(color_dict) + len(missing_groups):
                        extra_colors = palette_28[len(color_dict):len(color_dict) + len(missing_groups)]
                    else:
                        extra_colors = palette_56[len(color_dict):len(color_dict) + len(missing_groups)]
                    color_dict.update(dict(zip(missing_groups, extra_colors)))
            else:
                # If no colors found in uns, use default palette
                unique_groups = list(dict.fromkeys(gene_groups))
                if len(unique_groups) <= 28:
                    palette = palette_28
                else:
                    palette = palette_56
                color_dict = dict(zip(unique_groups, palette[:len(unique_groups)]))
        except (KeyError, AttributeError):
            # If colors not found in uns, use default palette
            unique_groups = list(dict.fromkeys(gene_groups))
            if len(unique_groups) <= 28:
                palette = palette_28
            else:
                palette = palette_56
            color_dict = dict(zip(unique_groups, palette[:len(unique_groups)]))

        # Add color bars with matching order
        # Add group labels
        # Only show used colors in legend and increase group spacing
        if preserve_dict_order and original_var_names_dict is not None:
            # When preserving dict order, use the original dictionary key order
            used_groups = list(original_var_names_dict.keys())
        else:
            # Use the order as they appear in gene_groups
            used_groups = list(dict.fromkeys(gene_groups))

        used_color_dict = {k: color_dict[k] for k in used_groups}
        m.add_top(
            mp.Colors(gene_groups, palette=used_color_dict),
            pad=0.1,
            size=0.15,
        )
        # Add group labels with increased spacing
        m.group_cols(gene_groups,order=used_groups)

    # Add cell type colors if available
    if color_dict is not None:
        # Add color bar using the properly created color_dict
        m.add_left(
            mp.Colors(agg.index, palette=color_dict),
            size=0.15,
            pad=0.1,
            legend=False,
        )

    # Add cell type labels
    m.add_left(mp.Labels(agg.index, align="right", fontsize=fontsize), pad=0.1)

    # Add cell counts
    m.add_right(
        mp.Numbers(
            cell_counts,
            color="#EEB76B",
            label="Count",
            label_props={'size': fontsize},
            props={'size': fontsize},
            show_value=False
        ),
        size=0.5,
        pad=0.1,
    )

    # Add dendrogram if requested
    if dendrogram:
        m.add_dendrogram("right", pad=0.1)

    # Add legends
    m.add_legends(box_padding=2)

    # Render the plot
    fig = m.render()

    if return_fig:
        return fig
    elif not show:
        return m
    return None

def rank_genes_groups_dotplot(
    adata: AnnData,
    plot_type: str = "dotplot",
    *,
    groups: Optional[Union[str, Sequence[str]]] = None,
    n_genes: Optional[int] = None,
    groupby: Optional[str] = None,
    values_to_plot: Optional[str] = None,
    var_names: Optional[Union[Sequence[str], Mapping[str, Sequence[str]]]] = None,
    min_logfoldchange: Optional[float] = None,
    key: Optional[str] = None,
    show: Optional[bool] = None,
    save: Optional[bool] = None,
    return_fig: bool = False,
    gene_symbols: Optional[str] = None,
    **kwds: Any,
) -> Optional[Union[Dict, Any]]:
    """
    Create a dot plot from rank_genes_groups results.

    Parameters
    ----------
    adata : AnnData
        Annotated data matrix.
    plot_type : str
        Currently only 'dotplot' is supported.
    groups : str or list of str, optional
        Groups to include in the plot.
    n_genes : int, optional
        Number of genes to include in the plot.
    groupby : str, optional
        Key in `adata.obs` to group by.
    values_to_plot : str, optional
        Key in rank_genes_groups results to plot (e.g. 'logfoldchanges', 'scores').
    var_names : str or list of str or dict, optional
        Variables to include in the plot. Can be:
        - A list of gene names: ['gene1', 'gene2', ...]
        - A dictionary mapping group names to gene lists: {'group1': ['gene1', 'gene2'], 'group2': ['gene3', 'gene4']}
        When a dictionary is provided, genes will be grouped and labeled accordingly in the plot.
    min_logfoldchange : float, optional
        Minimum log fold change to include in the plot.
    key : str, optional
        Key in `adata.uns` to use for rank_genes_groups results.
    show : bool, optional
        Whether to show the plot.
    save : bool, optional
        Whether to save the plot.
    return_fig : bool
        Whether to return the figure object.
    gene_symbols : str, optional
        Key for gene symbols in `adata.var`.
    **kwds : dict
        Additional keyword arguments to pass to dotplot.

    Returns
    -------
    If `return_fig` is True, returns the figure object.
    If `show` is False, returns axes dictionary.

    Examples
    --------
    >>> # Basic usage with top genes
    >>> sc.pl.rank_genes_groups_dotplot(adata, n_genes=5)

    >>> # Using logfoldchanges for coloring
    >>> sc.pl.rank_genes_groups_dotplot(adata, n_genes=5, values_to_plot='logfoldchanges')

    >>> # Grouping genes manually
    >>> gene_groups = {
    ...     'Group1': ['gene1', 'gene2'],
    ...     'Group2': ['gene3', 'gene4']
    ... }
    >>> sc.pl.rank_genes_groups_dotplot(adata, var_names=gene_groups)
    """
    if plot_type != "dotplot":
        raise ValueError("Only 'dotplot' is currently supported")

    if var_names is not None and n_genes is not None:
        msg = (
            "The arguments n_genes and var_names are mutually exclusive. Please "
            "select only one."
        )
        raise ValueError(msg)

    if key is None:
        key = "rank_genes_groups"

    if groupby is None:
        groupby = str(adata.uns[key]["params"]["groupby"])
    group_names = adata.uns[key]["names"].dtype.names if groups is None else groups

    if var_names is not None:
        if isinstance(var_names, Mapping):
            # get a single list of all gene names in the dictionary
            var_names_list = functools.reduce(
                operator.iadd, [list(x) for x in var_names.values()], []
            )
        elif isinstance(var_names, str):
            var_names_list = [var_names]
        else:
            var_names_list = var_names
    else:
        # set n_genes = 10 as default when none of the options is given
        if n_genes is None:
            n_genes = 10

        # dict in which each group is the key and the n_genes are the values
        var_names = {}
        var_names_list = []
        for group in group_names:
            df = rank_genes_groups_df(
                adata,
                group,
                key=key,
                gene_symbols=gene_symbols,
                log2fc_min=min_logfoldchange,
            )

            if gene_symbols is not None:
                df["names"] = df[gene_symbols]

            genes_list = df.names[df.names.notnull()].tolist()

            if len(genes_list) == 0:
                print(f"Warning: No genes found for group {group}")
                continue
            genes_list = genes_list[n_genes:] if n_genes < 0 else genes_list[:n_genes]
            var_names[group] = genes_list
            var_names_list.extend(genes_list)

    # by default add dendrogram to plots
    kwds.setdefault("dendrogram", True)

    # Get values to plot if specified
    title = None
    values_df = None
    if values_to_plot is not None:
        values_df = _get_values_to_plot(
            adata,
            values_to_plot,
            var_names_list,
            key=key,
            gene_symbols=gene_symbols,
        )
        title = values_to_plot
        if values_to_plot == "logfoldchanges":
            title = "log fold change"
        else:
            title = values_to_plot.replace("_", " ").replace("pvals", "p-value")

    # Create the plot
    _pl = dotplot(
        adata,
        var_names,
        groupby,
        dot_color_df=values_df,
        return_fig=True,
        gene_symbols=gene_symbols,
        preserve_dict_order=True,
        **kwds,
    )

    if title is not None and "colorbar_title" not in kwds:
        _pl.legend(colorbar_title=title)

    if return_fig:
        return _pl
    elif not show:
        return _pl
    return None