"""
This module contains visualizations functions for gene annotation.
"""
import textwrap
import matplotlib.pyplot as plt

from pandas import DataFrame

def annotate_gene_groups_with_ai_biological_process(
    plot_obj: dict,
    df: DataFrame,
    groupby: str,
    adt_key: tuple[str,...] | None = None
) -> None:
    """
    Annotate gene groups on a heatmap with their biological process labels.

    This function adds annotations below the x-axis of a heatmap generated by
    :func:`sc.pl.rank_genes_groups_heatmap`, labeling gene groups with their associated
    biological process descriptions from :func:`adt.ai_annotate_biological_process`.

    Parameters
    ----------
    plot_obj
        Output from `sc.pl.rank_genes_groups_heatmap`, containing matplotlib axes,
        specifically with the key 'heatmap_ax' representing the heatmap axis.

    df
        Output from :func:`adt.ai_annotate_biological_process`, containing gene group information
        with the following columns:
            - ``'groupby'``: Identifier for each gene group.
            - ``'top_10_genes'``: List of gene names in the group.
            - ``'ai_biological_process'``: Biological process label for the group.

    group_by
        The grouping column name.

    adt_key
        Used by :func:`adata_dict_fapply` and :func:`adata_dict_fapply_return` when passing this function.

    Returns
    -------
    ``None``

    See Also
    ----------
    :func:`sc.pl.rank_genes_groups_heatmap` : To generate the ``plot_obj``.
    :func:`adt.ai_annotate_biological_process` : To annotate the gene sets.
    """

    if adt_key:
        print(adt_key)

    heatmap_ax = plot_obj['heatmap_ax']

    # Get the genes and their positions from the x-axis
    xticks = heatmap_ax.get_xticks()
    xlabels = [tick.get_text() for tick in heatmap_ax.get_xticklabels()]
    gene_positions = {gene: xtick for gene, xtick in zip(xlabels, xticks)}

    group_positions = {}
    for idx, row in df.iterrows():
        group = row[groupby]
        genes_in_group = row['top_10_genes']
        # Filter genes that are actually in the plot
        positions = [gene_positions[gene] for gene in genes_in_group if gene in gene_positions]
        if positions:
            start_pos = min(positions)
            end_pos = max(positions)
            center_pos = (start_pos + end_pos) / 2.0
            group_positions[group] = {'start': start_pos, 'end': end_pos, 'center': center_pos}

    # Adjust the figure to make space for annotations
    fig = heatmap_ax.figure
    bottom_margin = 0.18  # Adjust this value if more space is needed
    fig.subplots_adjust(bottom=bottom_margin)

    # Annotate each group with the ai_biological_process label
    for group, pos_dict in group_positions.items():
        process_label = df.loc[df[groupby] == group, 'ai_biological_process'].values[0]
        # Wrap the text to fit within the group width
        # Convert group width from data to display coordinates
        inv = heatmap_ax.transData.inverted()
        start_display = heatmap_ax.transData.transform((pos_dict['start'], 0))
        end_display = heatmap_ax.transData.transform((pos_dict['end'], 0))
        group_display_width = end_display[0] - start_display[0]
        # Convert display width to figure coordinates
        start_fig = fig.transFigure.inverted().transform(start_display)
        end_fig = fig.transFigure.inverted().transform(end_display)
        group_fig_width = end_fig[0] - start_fig[0]
        # Estimate character width in figure coordinates
        char_width = 0.006  # Adjust this value based on font size and figure size
        max_chars = int((group_fig_width - 0.15) / char_width)  # Subtract padding
        max_chars = max(max_chars, 0.5)
        wrapped_label = textwrap.fill(process_label, width=max_chars)

        # Get the center position in figure coordinates
        center_display = heatmap_ax.transData.transform((pos_dict['center'], 0))
        center_fig = fig.transFigure.inverted().transform(center_display)

        # Place the text below the x-axis labels
        y_pos = 0.02  # Position in figure coordinates, adjust as needed
        fig.text(center_fig[0], y_pos, wrapped_label, ha='center', va='top', fontsize=10, fontweight='bold')

    # Redraw the plot to show the annotations
    fig.canvas.draw()

    plt.figure(plot_obj['gene_groups_ax'].figure)  # Select the figure associated with the heatmap axis
    plt.show()
