Api plot

def palette()->list:
    r"""Returns the default OmicVerse color palette.

    Returns:
        List of hex color codes for plotting
    """ 
    return sc_color

def red_palette()->list:
    r"""Returns a red-themed color palette.

    Returns:
        List of red-themed hex color codes
    """ 
    return red_color

def green_palette()->list:
    r"""Returns a green-themed color palette.

    Returns:
        List of green-themed hex color codes
    """ 
    return green_color

def orange_palette()->list:
    r"""Returns an orange-themed color palette.

    Returns:
        List of orange-themed hex color codes
    """ 
    return orange_color

def blue_palette()->list:
    r"""Returns a blue-themed color palette.

    Returns:
        List of blue-themed hex color codes
    """ 
    return blue_color

def plot_cellproportion(adata:anndata.AnnData,celltype_clusters:str,visual_clusters:str,
                       visual_li=None,visual_name:str='',figsize:tuple=(4,6),
                       ticks_fontsize:int=12,labels_fontsize:int=12,
                       legend:bool=False):
    r"""Plot stacked bar chart showing cell type proportions across groups.

    Arguments:
        adata: AnnData object
        celltype_clusters: Column name for cell types
        visual_clusters: Column name for grouping variable
        visual_li: List of groups to plot (None)
        visual_name: Label for x-axis ('')
        figsize: Figure size ((4,6))
        ticks_fontsize: Font size for tick labels (12)
        labels_fontsize: Font size for axis labels (12)
        legend: Whether to show legend (False)

    Returns:
        Tuple of (figure, axes) objects
    """

    b=pd.DataFrame(columns=['cell_type','value','Week'])

    if visual_li==None:
        adata.obs[visual_clusters]=adata.obs[visual_clusters].astype('category')
        visual_li=adata.obs[visual_clusters].cat.categories

    for i in visual_li:
        b1=pd.DataFrame()
        test=adata.obs.loc[adata.obs[visual_clusters]==i,celltype_clusters].value_counts()
        b1['cell_type']=test.index
        b1['value']=test.values/test.sum()
        b1['Week']=i.replace('Retinoblastoma_','')
        b=pd.concat([b,b1])

    plt_data2=adata.obs[celltype_clusters].value_counts()
    plot_data2_color_dict=dict(zip(adata.obs[celltype_clusters].cat.categories,adata.uns['{}_colors'.format(celltype_clusters)]))
    plt_data3=adata.obs[visual_clusters].value_counts()
    plot_data3_color_dict=dict(zip([i.replace('Retinoblastoma_','') for i in adata.obs[visual_clusters].cat.categories],adata.uns['{}_colors'.format(visual_clusters)]))
    b['cell_type_color'] = b['cell_type'].map(plot_data2_color_dict)
    b['stage_color']=b['Week'].map(plot_data3_color_dict)

    fig, ax = plt.subplots(figsize=figsize)
    #用ax控制图片
    #sns.set_theme(style="whitegrid")
    #sns.set_theme(style="ticks")
    n=0
    all_celltype=adata.obs[celltype_clusters].cat.categories
    for i in all_celltype:
        if n==0:
            test1=b[b['cell_type']==i]
            ax.bar(x=test1['Week'],height=test1['value'],width=0.8,color=list(set(test1['cell_type_color']))[0], label=i)
            bottoms=test1['value'].values
        else:
            test2=b[b['cell_type']==i]
            ax.bar(x=test2['Week'],height=test2['value'],bottom=bottoms,width=0.8,color=list(set(test2['cell_type_color']))[0], label=i)
            test1=test2
            bottoms+=test1['value'].values
        n+=1
    if legend!=False:
        plt.legend(bbox_to_anchor=(1.05, -0.05), loc=3, borderaxespad=0,fontsize=10)

    plt.grid(False)

    plt.grid(False)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['bottom'].set_visible(True)
    ax.spines['left'].set_visible(True)

    # 设置左边和下边的坐标刻度为透明色
    #ax.yaxis.tick_left()
    #ax.xaxis.tick_bottom()
    #ax.xaxis.set_tick_params(color='none')
    #ax.yaxis.set_tick_params(color='none')

    # 设置左边和下边的坐标轴线为独立的线段
    ax.spines['left'].set_position(('outward', 10))
    ax.spines['bottom'].set_position(('outward', 10))

    plt.xticks(fontsize=ticks_fontsize,rotation=90)
    plt.yticks(fontsize=ticks_fontsize)
    plt.xlabel(visual_name,fontsize=labels_fontsize)
    plt.ylabel('Cells per Stage',fontsize=labels_fontsize)
    fig.tight_layout()
    return fig,ax

def plot_embedding_celltype(adata:anndata.AnnData,figsize:tuple=(6,4),basis:str='umap',
                            celltype_key:str='major_celltype',title:str=None,
                            celltype_range:tuple=(2,9),
                            embedding_range:tuple=(3,10),
                            xlim:int=-1000)->tuple:
    r"""Create combined embedding plot with cell type legend and counts.

    Arguments:
        adata: AnnData object
        figsize: Figure size ((6,4))
        basis: Embedding basis name ('umap')
        celltype_key: Column name for cell types ('major_celltype')
        title: Plot title (None)
        celltype_range: Grid range for cell type panel ((2,9))
        embedding_range: Grid range for embedding panel ((3,10))
        xlim: X-axis limit for counts (-1000)

    Returns:
        Tuple of (figure, [embedding_axis, celltype_axis])
    """

    adata.obs[celltype_key]=adata.obs[celltype_key].astype('category')
    cell_num_pd=pd.DataFrame(adata.obs[celltype_key].value_counts())
    if '{}_colors'.format(celltype_key) in adata.uns.keys():
        cell_color_dict=dict(zip(adata.obs[celltype_key].cat.categories.tolist(),
                        adata.uns['{}_colors'.format(celltype_key)]))
    else:
        if len(adata.obs[celltype_key].cat.categories)>28:
            cell_color_dict=dict(zip(adata.obs[celltype_key].cat.categories,sc.pl.palettes.default_102))
        else:
            cell_color_dict=dict(zip(adata.obs[celltype_key].cat.categories,sc.pl.palettes.zeileis_28))

    if figsize==None:
        if len(adata.obs[celltype_key].cat.categories)<10:
            fig = plt.figure(figsize=(6,4))
        else:
            print('The number of cell types is too large, please set the figsize parameter')
            return
    else:
        fig = plt.figure(figsize=figsize)
    grid = plt.GridSpec(10, 10)
    ax1 = fig.add_subplot(grid[:, embedding_range[0]:embedding_range[1]])       # 占据第一行的所有列
    ax2 = fig.add_subplot(grid[celltype_range[0]:celltype_range[1], :2]) 
    # 定义子图的大小和位置
         # 占据第二行的前两列
    #ax3 = fig.add_subplot(grid[1:, 2])      # 占据第二行及以后的最后一列
    #ax4 = fig.add_subplot(grid[2, 0])       # 占据最后一行的第一列
    #ax5 = fig.add_subplot(grid[2, 1])       # 占据最后一行的第二列

    sc.pl.embedding(
        adata,
        basis=basis,
        color=[celltype_key],
        title='',
        frameon=False,
        #wspace=0.65,
        ncols=3,
        ax=ax1,
        legend_loc=False,
        show=False
    )

    for idx,cell in zip(range(cell_num_pd.shape[0]),
                        adata.obs[celltype_key].cat.categories):
        ax2.scatter(100,
                cell,c=cell_color_dict[cell],s=50)
        ax2.plot((100,cell_num_pd.loc[cell,celltype_key]),(idx,idx),
                c=cell_color_dict[cell],lw=4)
        ax2.text(100,idx+0.2,
                cell+'('+str("{:,}".format(cell_num_pd.loc[cell,celltype_key]))+')',fontsize=11)
    ax2.set_xlim(xlim,cell_num_pd.iloc[1].values[0]) 
    ax2.text(xlim,idx+1,title,fontsize=12)
    ax2.grid(False)
    ax2.spines['top'].set_visible(False)
    ax2.spines['right'].set_visible(False)
    ax2.spines['bottom'].set_visible(False)
    ax2.spines['left'].set_visible(False)
    ax2.axis('off')

    return fig,[ax1,ax2]

def plot_embedding(adata:anndata.AnnData,basis:str,color:str,color_dict=None,
                   figsize:tuple=(4,4),**kwargs):
    r"""Create embedding plot with customizable colors.

    Arguments:
        adata: AnnData object
        basis: Embedding basis name
        color: Column name for coloring
        color_dict: Custom color mapping (None)
        figsize: Figure size ((4,4))
        **kwargs: Additional parameters for sc.pl.embedding

    Returns:
        Tuple of (figure, axes) objects
    """
    if type(color)!=str:
        print("Only one color could be input, don't input list")
        return
    fig,ax=plt.subplots(1,1,figsize=figsize)
    adata.obs[color]=adata.obs[color].astype('category')

    if '{}_colors'.format(color) in adata.uns.keys():
        print('{}_colors'.format(color))
        type_color_all=dict(zip(adata.obs[color].cat.categories,adata.uns['{}_colors'.format(color)]))
    else:
        if len(adata.obs[color].cat.categories)>28:
            type_color_all=dict(zip(adata.obs[color].cat.categories,sc.pl.palettes.default_102))
        else:
            type_color_all=dict(zip(adata.obs[color].cat.categories,sc.pl.palettes.zeileis_28))
    if color_dict is not None:
        for color_key in color_dict.keys():
            type_color_all[color_key]=color_dict[color_key]

    adata.uns['{}_colors'.format(color)]=np.array([i for i in type_color_all.values()])
    sc.pl.embedding(adata,basis=basis,
                    color=color,ax=ax,**kwargs)
    return fig,ax