from pathlib import Path
from typing import Optional, Union, Tuple, List
import matplotlib.pyplot as plt
from anndata import AnnData
import scanpy as sc
from typing import Optional

def plot_spatial(  
    adata: AnnData,
    color: str = 'DePass',  
    save_path: Optional[Union[str, Path]] = None,
    save_name: str = 'spatial_plot',
    title: Optional[str] = None,
    s: int = 35,
    figsize: Tuple[float, float] = (3, 3),
    dpi: int = 300,
    format: str = "png",
    frameon: bool = True,
    adjust_margins: bool = True,
    legend_loc: Optional[str] = 'right margin',
    colorbar_loc: Optional[str] = 'right', 
    show: bool = False,
    **kwargs
) -> None:
    r"""
    Plot spatial data with specified parameters and save the figure.

    Parameters:
        adata (AnnData): Input AnnData object containing spatial data.
        color (str): Column name in adata.obs for coloring the plot (default: 'DePass').
        save_path (Optional[Union[str, Path]]): Path to save the figure (default: None).
        save_name (str): Name of the saved figure (default: 'spatial_plot').
        title (Optional[str]): Title of the plot (default: None).
        s (int): Size of the markers (default: 35).
        figsize (Tuple[float, float]): Figure size (default: (3, 3)).
        dpi (int): DPI of the saved figure (default: 300).
        format (str): Format of the saved figure (default: "png").
        frameon (bool): Whether to show the frame (default: True).
        adjust_margins (bool): Whether to adjust margins (default: True).
        legend_loc (Optional[str]): Location of the legend (default: 'right margin').
        colorbar_loc (Optional[str]): Colorbar position (default: 'right'). Set to None to disable colorbar.
        show (bool): Whether to show the plot (default: False).
        **kwargs: Additional keyword arguments for sc.pl.embedding.

    Returns:
        None
    """

    if not isinstance(adata, AnnData):
        raise TypeError("Expected AnnData object, got {}".format(type(adata)))
    
    if not save_name.strip():
        raise ValueError("save_name must contain non-whitespace characters")
    
    file_format = format.lower().lstrip('.')
    allowed_formats = {'png', 'pdf', 'svg', 'tiff', 'jpg', 'jpeg'}
    if file_format not in allowed_formats:
        raise ValueError(f"Invalid format: {format}. Choose from {allowed_formats}")

    if save_path is not None:
        save_path = Path(save_path)
        save_path.mkdir(parents=True, exist_ok=True)

    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(111)

    if title is None:
        title = color if isinstance(color, str) else ', '.join(color)

    try:
        sc.pl.embedding(
            adata,
            basis='spatial',
            color=color,
            title=title,
            s=s,
            ax=ax,
            show=False,
            frameon=frameon,
            legend_loc=legend_loc,
            colorbar_loc=colorbar_loc, 
            **kwargs
        )
    except KeyError as e:
        raise ValueError(f"Missing required data: {e}") from None

    if save_path is not None:
        output_path = save_path / f"{save_name}.{file_format}"
        try:
            plt.gca().set_rasterized(True)  
            fig.savefig(
                output_path,
                dpi=dpi,
                bbox_inches='tight' if adjust_margins else None,
                pad_inches=0.1 if adjust_margins else 0.5
            )
        except Exception as e:
            raise IOError(f"Failed to save figure: {e}") from None

    if show:
        plt.show()
    
    plt.close(fig)

import pandas as pd
import numpy as np

def getLogFC(
    target_genes: list,
    target_groups: list,
    logfoldchanges: dict,
    gene_names: dict
) -> pd.DataFrame:
    r"""
    Get log fold changes for target genes and groups.

    Parameters:
        target_genes (list): List of target genes.
        target_groups (list): List of target groups.
        logfoldchanges (dict): Dictionary of log fold changes.
        gene_names (dict): Dictionary of gene names.

    Returns:
        pd.DataFrame: DataFrame containing log fold changes for target genes and groups.
    """
    
    if len(target_genes) != len(target_groups):
        raise ValueError("Lengths of `target_genes` and `target_groups` must match.")

    results = []
    for gene, group in zip(target_genes, target_groups):
        # Extract gene names for the current group
        group_genes = gene_names[group]
        # Find index of the target gene in the group's gene list
        gene_idx = np.where(group_genes == gene)[0]
        
        if len(gene_idx) == 0:
            print(f"Warning: Gene '{gene}' not found in group '{group}'.")
            results.append((gene, group, None))  # Store None if gene is missing
        else:
            # Get the first matching LogFC value (assumes unique gene names per group)
            logfc = logfoldchanges[group][gene_idx[0]]
            results.append((gene, group, logfc))
    
    # Convert results to a structured DataFrame
    results_df = pd.DataFrame(results, columns=["Gene", "Group", "LogFC"])
    return results_df


def rank_genes_groups(
    adata,
    groupby: str = "DePass",       
    method: str = "wilcoxon",        
    n_genes: int = 10,             
    standard_scale: str = "var",    
    dpi: int = 300,                
    show: bool = True, 
    save_path: Optional[str] = None,                
    figname: str = 'rank_genes_dotplot',  
    figsize: Tuple[float, float] = (6, 3),
) -> None:
    r"""
    Rank genes by groups and optionally plot the results.

    Parameters:
        adata (AnnData): Input AnnData object.
        groupby (str): Column name in adata.obs for grouping (default: "DePass").
        method (str): Method for ranking genes (default: "wilcoxon").
        n_genes (int): Number of top genes to show (default: 10).
        standard_scale (str): Scaling method (default: "var").
        dpi (int): DPI of the saved figure (default: 300).
        show (bool): Whether to show the results (default: True).
        save_path (Optional[str]): Path to save the figure (default: None).
        figname (str): Name of the saved figure (default: 'rank_genes_dotplot').

    Returns:
        None
    """
    
    from sklearn.preprocessing import MinMaxScaler
    import matplotlib.pyplot as plt
     
    if save_path is not None:
        os.makedirs(save_path, exist_ok=True)

    # Normalize data to [0,1] range for comparative analysis
    scaler = MinMaxScaler()
    adata.obs[groupby] = adata.obs[groupby].astype('str').astype('category')  # Ensure categorical type
    adata.X = scaler.fit_transform(adata.X) 
    
    # Perform differential expression analysis
    sc.tl.rank_genes_groups(adata, groupby=groupby, method=method, use_raw=False)

    if show or save_path is not None:
        sc.pl.rank_genes_groups_dotplot(
            adata,
            groupby=groupby,
            standard_scale=standard_scale,  
            n_genes=n_genes,
            show=False ,
            dendrogram=False,
            figsize=figsize,
        )
        if save_path is not None:
           plt.savefig(
               os.path.join(save_path, figname+".png"),
               dpi=dpi,
               bbox_inches="tight"  
           )

        if show: plt.show()
        plt.close()

 
from typing import Optional
import numpy as np
import matplotlib.pyplot as plt
import scanpy as sc
from sklearn.preprocessing import MinMaxScaler
import os
import warnings

def get_logfc(
    target_gene: str,
    target_group: str,
    logfoldchanges: np.ndarray,  # Structured array from rank_genes_groups
    gene_names: np.ndarray,      # Structured array from rank_genes_groups
) -> float:
   
    # Validate group existence
    if target_group not in gene_names.dtype.names:
        available_groups = list(gene_names.dtype.names)
        raise KeyError(f"Group '{target_group}' not found. Available groups: {available_groups}")
    
    # Locate gene index
    group_genes = gene_names[target_group]
    gene_idx = np.flatnonzero(group_genes == target_gene)
    
    if not gene_idx.size:
        raise KeyError(f"Gene '{target_gene}' not found in group '{target_group}'")
        
    return float(logfoldchanges[target_group][gene_idx[0]])

def plot_marker_comparison(
    adata1: sc.AnnData,
    adata2: sc.AnnData,
    target_gene: str,
    save_path: Optional[str] = None,
    save_name: str = "gene_comparison",
    show: bool = False,
    s: int = 80,
    cmap: str = "turbo",
    dpi: int = 300,
    colorbar_loc: Optional[str] = None, 
    figsize: tuple = (7, 3),
    frameon=False,
) -> None:
    r"""
    Plot marker comparison between two datasets.

    Parameters:
        adata1 (AnnData): First AnnData object.
        adata2 (AnnData): Second AnnData object.
        target_gene (str): Target gene name.
        save_path (Optional[str]): Path to save the figure (default: None).
        save_name (str): Name of the saved figure (default: "gene_comparison").
        show (bool): Whether to show the plot (default: False).
        s (int): Size of the markers (default: 80).
        cmap (str): Colormap (default: "viridis").
        dpi (int): DPI of the saved figure (default: 300).
        colorbar_loc (Optional[str]): Colorbar position. Set to None to disable colorbar.
        figsize (tuple): Figure size (default: (7, 3)).
        frameon (bool): Whether to show the frame (default: False).

    Returns:
        None
    """
    
    for adata, name in [(adata1, 'adata1'), (adata2, 'adata2')]:
        if 'spatial' not in adata.obsm:
            raise KeyError(f"Missing spatial coordinates in {name}.obsm['spatial']")
    
    if save_path is not None:
        os.makedirs(save_path, exist_ok=True)
    
    vis_params = {
        'basis': 'spatial',
        'color': f'{target_gene}_expr',
        's': s,
        'frameon':   frameon,
        'colorbar_loc': colorbar_loc,
        'cmap': cmap,
    }

    def _scaler_data(adata: sc.AnnData) -> None:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            expr = adata[:, target_gene].X.toarray()
        adata.obs[f'{target_gene}_expr'] = MinMaxScaler().fit_transform(expr)

    def _create_plot(adata: sc.AnnData, ax: plt.Axes,name: str) -> None:
        sc.pl.embedding(
            adata,
            title=f"{name + target_gene}",
            ax=ax,
            show=False,
            **vis_params
        )

    
    _scaler_data(adata1)
    _scaler_data(adata2)
    
    fig, axes = plt.subplots(1, 2, figsize=figsize)
    _create_plot(adata1, axes[0],'Raw - ')
    _create_plot(adata2, axes[1],'Enhanced - ')
    
    if save_path is not None:
        try:
            plt.gca().set_rasterized(True)  
            fig.savefig(
            os.path.join(save_path, f"{save_name}_combined.png"),
            dpi=dpi,
            bbox_inches="tight"
               )
        except Exception as e:
            raise IOError(f"Failed to save figure: {e}") from None

    if show:
        plt.show()
    plt.close(fig)


def plot_marker_comparison_with_logFC(
    adata1: sc.AnnData,
    adata2: sc.AnnData,
    target_gene: str,
    target_group: str,
    save_path: Optional[str] = None,
    save_name: str = "gene_comparison",
    show: bool = False,
    s: int = 80,
    cmap: str = "turbo",
    dpi: int = 300,
    colorbar_loc: Optional[str] = None, 
    figsize: tuple = (7, 3),
    frameon=False,
) -> None:
    r"""
    Visualize marker comparison between two datasets with log fold changes.

    Parameters:
        adata1 (AnnData): First AnnData object.
        adata2 (AnnData): Second AnnData object.
        target_gene (str): Target gene name.
        target_group (str): Target group name.
        save_path (Optional[str]): Path to save the figure (default: None).
        save_name (str): Name of the saved figure (default: "gene_comparison").
        show (bool): Whether to show the plot (default: False).
        s (int): Size of the markers (default: 80).
        cmap (str): Colormap (default: "viridis").
        dpi (int): DPI of the saved figure (default: 300).
        colorbar_loc (Optional[str]): Colorbar position. Set to None to disable colorbar.
        figsize (tuple): Figure size (default: (7, 3)).
        frameon (bool): Whether to show the frame (default: False).

    Returns:
        None
    """
    
    for adata, name in [(adata1, 'adata1'), (adata2, 'adata2')]:
        if 'rank_genes_groups' not in adata.uns:
            raise KeyError(f"Missing DEG results in {name}. Run sc.tl.rank_genes_groups first.")
        if 'spatial' not in adata.obsm:
            raise KeyError(f"Missing spatial coordinates in {name}.obsm['spatial']")
   
    if save_path is not None:
        os.makedirs(save_path, exist_ok=True)
    
    logfc1 = get_logfc(target_gene, target_group,
                               adata1.uns['rank_genes_groups']['logfoldchanges'],
                               adata1.uns['rank_genes_groups']['names'])
    
    logfc2 = get_logfc(target_gene, target_group,
                               adata2.uns['rank_genes_groups']['logfoldchanges'],
                               adata2.uns['rank_genes_groups']['names'])
    
    if save_path is not None:
        os.makedirs(save_path, exist_ok=True)
    
    # Shared visualization parameters
    vis_params = {
        'basis': 'spatial',
        'color': f'{target_gene}_expr',
        's': s,
        'frameon': frameon,
        'colorbar_loc': colorbar_loc,
        'cmap': cmap,
    }

    def _scaler_data(adata: sc.AnnData) -> None:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            expr = adata[:, target_gene].X.toarray()
        adata.obs[f'{target_gene}_expr'] = MinMaxScaler().fit_transform(expr)

    def _create_plot(adata: sc.AnnData, logfc: float, ax: plt.Axes, name:str) -> None:
        sc.pl.embedding(
            adata,
            title=f"{name+target_gene}\n(LogFC={logfc:.3f})",
            ax=ax,
            show=False,
            **vis_params
        )

    for adata in [adata1, adata2]:
        _scaler_data(adata)
    
    fig, axes = plt.subplots(1, 2, figsize=figsize)
    _create_plot(adata1, logfc1, axes[0],'Raw - ')
    _create_plot(adata2, logfc2, axes[1],'Enhanced - ')
    
    if save_path is not None:
        plt.gca().set_rasterized(True)  # 关键：启用栅格化
        fig.savefig(
            os.path.join(save_path, f"{save_name}_combined_logFC.png"),
            dpi=dpi,
            bbox_inches="tight"
        )
    if show:
        plt.show()
    plt.close(fig)

from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import os
import matplotlib.cm as cm
from matplotlib.colors import to_rgb
import matplotlib.patches as patches

# cluster visualization
def cluster_and_visualize_superpixel(
    final_embeddings,
    data_dict,
    n_clusters,
    mode="joint",  # 'joint' or 'independent'
    defined_labels=None,
    vis_basis="spatial",
    random_state=0,
    colormap=None,
    swap_xy=False,
    invert_x=False,
    invert_y=False,
    offset=False,
    save_path=None,
    dpi=300,
    remove_title = False,
    remove_legend = False,
    remove_spine = False,
    figscale = 35
):
    import numpy as np
    from sklearn.cluster import KMeans
    import os
    import numpy as np


    adata_list = []
    embeddings = []
    coords_all = []
    section_names = []

    for section, embedding in final_embeddings.items():
        idx = int(section[1:]) - 1
        for modality, adata_list_per_mod in data_dict.items():
            if idx < len(adata_list_per_mod) and adata_list_per_mod[idx] is not None:
                adata = adata_list_per_mod[idx]
                adata_list.append(adata)
                embeddings.append(embedding)
                coords = adata.obsm[vis_basis].copy()
                if swap_xy:
                    coords = coords[:, [1, 0]]
                coords = coords.astype(int)
                if offset:
                    offset_value = coords.min(axis=0)     
                    coords -= offset_value               
                coords_all.append(coords)
                section_names.append(section)
                break

    cluster_labels = {}

    if mode == "joint":
        print("Perform joint clustering...")
        combined_embedding = np.vstack(embeddings)
        kmeans = KMeans(n_clusters=n_clusters, random_state=random_state)
        all_clusters = kmeans.fit_predict(combined_embedding)
        start = 0
        for section, emb in zip(section_names, embeddings):
            end = start + emb.shape[0]
            cluster_labels[section] = all_clusters[start:end]
            start = end
    elif mode == "independent":
        print("Perform independent clustering...")
        for section, emb in zip(section_names, embeddings):
            kmeans = KMeans(n_clusters=n_clusters, random_state=random_state)
            cluster_labels[section] = kmeans.fit_predict(emb)
    elif mode == 'defined':
        if defined_labels is None:
            raise ValueError("If mode='defined', you must provide `defined_labels`.")
        cluster_labels = defined_labels
    else:
        raise ValueError("mode must be 'joint' or 'independent'")

    for section, coords, labels in zip(section_names, coords_all, cluster_labels.values()):
        max_y, max_x = coords.max(axis=0) + 1
        image = np.full((max_y, max_x), fill_value=-1, dtype=int)
        for (y, x), label in zip(coords, labels):
            image[y, x] = label
        if invert_x:
            image = image[:, ::-1]
        if invert_y:
            image = image[::-1, :]
        section_save_path = None
        if save_path:
            base, ext = os.path.splitext(save_path)
            section_save_path = f"{base}_section_{section}{ext or '.png'}"
        
        plot_histology_clusters(
            he_clusters_image=image,
            num_he_clusters=n_clusters,
            section_title=f"Section {section} ({mode})",
            colormap=colormap,
            save_path=section_save_path,
            dpi=dpi,
            figscale = figscale,
            remove_title = remove_title,
            remove_legend = remove_legend,
            remove_spine=remove_legend, 
        )

    return cluster_labels



def plot_histology_clusters(he_clusters_image,
                            num_he_clusters,
                            section_title=None,
                            colormap=None,
                            save_path=None,
                            figscale = 35,
                            remove_title = False,
                            remove_legend = False,
                            remove_spine=False, 
                            dpi=300):


    if colormap is None:
        color_list = [[255,127,14],[44,160,44],[214,39,40],[148,103,189],
                      [140,86,75],[227,119,194],[127,127,127],[188,189,34],
                      [23,190,207],[174,199,232],[255,187,120],[152,223,138],
                      [255,152,150],[197,176,213],[196,156,148],[247,182,210],
                      [199,199,199],[219,219,141],[158,218,229],[16,60,90],
                      [128,64,7],[22,80,22],[107,20,20],[74,52,94],[70,43,38],
                      [114,60,97],[64,64,64],[94,94,17],[12,95,104],[0,0,0]]

    elif isinstance(colormap, list):
        color_list = colormap

    else:
        cmap = cm.get_cmap(colormap)
        color_list = [ [int(255 * c) for c in to_rgb(cmap(i))] for i in range(len(cmap.colors)) ]

    image_rgb = 255 * np.ones([he_clusters_image.shape[0], he_clusters_image.shape[1], 3])
    for cluster in range(num_he_clusters):
        image_rgb[he_clusters_image == cluster] = color_list[cluster]
    image_rgb = np.array(image_rgb, dtype='uint8')

    plt.figure(figsize=(he_clusters_image.shape[1] // figscale, he_clusters_image.shape[0] // figscale))
    if remove_title:
        plt.title("")
    else:
        title = section_title if section_title else "Histology Clusters"
        plt.title(title, fontsize=18)
    plt.imshow(image_rgb, interpolation='none')
    # plt.show()
    ax = plt.gca()
    ax.set_xticks([])
    ax.set_yticks([])

    if remove_spine:
        for spine in ax.spines.values():
            spine.set_visible(False)

    if not remove_legend:
        legend_elements = [patches.Patch(facecolor=np.array(color_list[i]) / 255,
                                         label=f'Cluster {i}')
                           for i in range(num_he_clusters)]
        plt.legend(handles=legend_elements,
                   bbox_to_anchor=(1.05, 1),
                   loc='upper left',
                   borderaxespad=0.,
                   fontsize=12)

    if save_path is not None:
        plt.savefig(save_path, dpi=dpi, bbox_inches="tight")
        print(f"Saved: {save_path}")

    plt.show()
    # plt.close()

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.cm as cm
from matplotlib.colors import to_rgb
import matplotlib.patches as patches
import os


def plot_superpixel(
    adata,
    label_key='label',  
    vis_basis='spatial',  
    colormap=None,
    save_path=None,
    save_name='visualization',
    title=None,
    figscale=100,
    format='png',
    show=True,
    remove_title=False,
    remove_legend=False,
    remove_spine=False,
    dpi=300,
    random_state=2024,  
    swap_xy=False,  
    invert_x=False,  
    invert_y=False  
):
    """
    Visualize clusters using labels and spatial coordinates from adata.

    Parameters:
        adata: AnnData object containing cell data and cluster labels.
        label_key: Key for cluster labels in adata.obs.
        vis_basis: Key for coordinates in adata.obsm.
        colormap: Custom color mapping.
        save_path: Directory path to save the image.
        save_name: File name to save the image.
        title: Title of the plot.
        figscale: Figure size scaling factor.
        format: Format to save the image (e.g., 'pdf', 'png').
        show: Whether to display the plot.
        remove_title: Whether to remove the title.
        remove_legend: Whether to remove the legend.
        remove_spine: Whether to remove the border.
        dpi: Image resolution.
        random_state: Random seed for reproducibility.
        swap_xy: Whether to swap x and y coordinates.
        invert_x: Whether to invert x-axis.
        invert_y: Whether to invert y-axis.
    """
   
    np.random.seed(random_state)

   
    labels = adata.obs[label_key].values.astype(int)
    coords = adata.obsm[vis_basis].copy().astype(int)
    

    if swap_xy:
        coords = coords[:, [1, 0]]

    unique_labels = np.unique(labels)
    num_clusters = len(unique_labels)

    label_to_index = {label: idx for idx, label in enumerate(unique_labels)}
    zero_based_labels = np.array([label_to_index[label] for label in labels])


    if colormap is None:
        color_list = [
        [60, 142, 204],    
        [187, 187, 187],   
        [246, 216, 208],   
        [254, 238, 237],   
        [215, 102, 102],   
        [177, 157, 177],   
        [60, 162, 254],    
        [151, 215, 243],   
        [208, 163, 239],   
        [246, 216, 212],   
        [255, 247, 180],   
        [241, 91, 108],    
        [60, 188, 60],     
        [104, 220, 104],   
        [247, 172, 188],   
        [222, 171, 138],   
        [255, 188, 188],   
        [199, 133, 89],    
        [60, 251, 255],    
        [195, 236, 255],   
        [204, 238, 204],   
        [254, 220, 189],   
        [239, 91, 156],    
        [176,224,230],    
        [220,220,220]
]
    elif isinstance(colormap, list):
        color_list = colormap
    else:
        cmap = cm.get_cmap(colormap)
        color_list = [[int(255 * c) for c in to_rgb(cmap(i))] for i in range(len(cmap.colors))]

    if len(color_list) < num_clusters:
        raise ValueError("Color list is not long enough to cover all clusters.")
    
    

    max_y, max_x = coords.max(axis=0) + 1
    image = np.full((max_y, max_x), fill_value=-1, dtype=int)
    for (y, x), label in zip(coords, zero_based_labels):
        if 0 <= x < max_x and 0 <= y < max_y:
            image[y, x] = label

    if invert_x:
        image = image[:, ::-1]
    if invert_y:
        image = image[::-1, :]


    image_rgb = np.ones([image.shape[0], image.shape[1], 3])
    for cluster in range(num_clusters):
        image_rgb[image == cluster] = np.array(color_list[cluster]) / 255.0
    
    # image_rgb = 255 * np.ones([image.shape[0], image.shape[1], 3])
    # for cluster in range(num_clusters):
    #     image_rgb[image == cluster] = color_list[cluster]
    # image_rgb = np.array(image_rgb, dtype='uint8')
    
    plt.figure(figsize=(image.shape[1] // figscale, image.shape[0] // figscale))
    plt.imshow(image_rgb, interpolation='none')

    if remove_title or title is None:
        plt.title("")
    else:
        plt.title(title, fontsize=18)

    ax = plt.gca()
    ax.set_xticks([])
    ax.set_yticks([])

    if remove_spine:
        for spine in ax.spines.values():
            spine.set_visible(False)

    if not remove_legend:
        legend_elements = [patches.Patch(facecolor=np.array(color_list[i]) / 255,
                                          label=f'Cluster {unique_labels[i]}') for i in range(num_clusters)]
        plt.legend(handles=legend_elements,
                   bbox_to_anchor=(1.05, 1),
                   loc='upper left',
                   borderaxespad=0.,
                   fontsize=12)

    if save_path:
        os.makedirs(save_path, exist_ok=True)
        file_path = os.path.join(save_path, f"{save_name}.{format}")
        plt.savefig(file_path, dpi=dpi, bbox_inches="tight")
        print(f"Image saved to: {file_path}")


    if show:
        plt.show()
    else:
        plt.close()


# marker visualization
def prepare_image(adata, molecule_name, basis, swap_xy, invert_x, invert_y, offset, scale):
    coords = adata.obsm[basis].copy()
    if swap_xy:
        coords = coords[:, [1, 0]]
    coords = coords.astype(int)
    if offset:
        offset_value = coords.min(axis=0)
        coords -= offset_value 

    values = adata[:, molecule_name].X
  

    if hasattr(values, "toarray"):
        values = values.toarray().flatten()
    else:
        values = np.array(values).flatten()

    if scale:
        from sklearn.preprocessing import MinMaxScaler
        scaler = MinMaxScaler()
        values = values.reshape(-1, 1)  
        values = scaler.fit_transform(values)
        values = values.flatten()  
  

    max_y, max_x = coords.max(axis=0) + 1
    image = np.full((max_y, max_x), np.nan, dtype=float)
    for (y, x), val in zip(coords, values):
        image[y, x] = val

    if invert_x:
        image = image[:, ::-1]
    if invert_y:
        image = image[::-1, :]

    return image

def plot_marker_comparison_superpixel(
    molecule_name: str,
    adata1,
    adata2,
    section1_label: str = 'Section 1',
    section2_label: str = 'Section 2',
    basis: str = 'spatial',
    colormap: str = "viridis",
    plot_style: str = "original",
    scale: bool = True,
    swap_xy: bool = False,
    invert_x: bool = False,
    invert_y: bool = False,
    offset: bool = False,
    figscale: int = 35,
    dpi: int = 300,
    remove_title: bool = False,     
    remove_spine: bool = False,    
    remove_legend: bool = False,      
    save_path: str = None,
    format: str = 'pdf'
):


    img1 = prepare_image(adata1, molecule_name, basis, swap_xy, invert_x, invert_y, offset, scale)
    img2 = prepare_image(adata2, molecule_name, basis, swap_xy, invert_x, invert_y, offset, scale)


    figsize1 = (img1.shape[1] / figscale, img1.shape[0] / figscale)
    figsize2 = (img2.shape[1] / figscale, img2.shape[0] / figscale)
    figsize = (figsize1[0] + figsize2[0], max(figsize1[1], figsize2[1]))

    fig, axes = plt.subplots(1, 2, figsize=figsize)
    for ax, img, title in zip(axes, [img1, img2], [section1_label, section2_label]):
        im = ax.imshow(img, cmap=colormap, interpolation='none')
        if not remove_title:
            ax.set_title(f"{title} - {molecule_name}", fontsize=16)
        else:
            ax.set_title("")
        ax.set_xticks([])
        ax.set_yticks([])
        if remove_spine:
            for spine in ax.spines.values():
                spine.set_visible(False)
        if plot_style == "equal":
            ax.set_aspect("equal")

        if not remove_legend:
            cbar = fig.colorbar(im, ax=ax, shrink=0.7, pad=0.02)  

    if save_path:
        os.makedirs(save_path, exist_ok=True)
        file_path = os.path.join(save_path, f"{molecule_name+'_combined'}.{format}")
        plt.savefig(file_path, dpi=dpi, bbox_inches="tight")
        print(f"Saving marker comparison to: {save_path}")

    plt.show()
    plt.close()


import pandas as pd
from typing import List, Dict, Union
from pathlib import Path
from anndata import AnnData


def get_logfc_df(
    adata_list: List[AnnData],
    adata_names: List[str],
    target_genes: List[str],
    target_groups: List[str],
    save_path: Union[str, Path] = "results",
    save_name: str = "logfc_comparison"
) -> pd.DataFrame:
    r"""
    Get log fold change DataFrame for multiple datasets.

    Parameters:
        adata_list (List[AnnData]): List of AnnData objects.
        adata_names (List[str]): List of dataset names.
        target_genes (List[str]): List of target genes.
        target_groups (List[str]): List of target groups.
        save_path (Union[str, Path]): Path to save the DataFrame (default: "results").
        save_name (str): Name of the saved file (default: "logfc_comparison").

    Returns:
        pd.DataFrame: DataFrame containing log fold changes.
    """
    
    if len(target_genes) != len(target_groups):
        raise ValueError("Length of target_genes and target_groups must match!")
    
    if len(adata_list) != len(adata_names):
        raise ValueError("Length of adata_list and adata_names must match!")
    
    results = {
        'Gene': target_genes,
        'Group': target_groups
    }
    
    # Extract logFC values for each dataset
    for adata, name in zip(adata_list, adata_names):
        logfc_values = []
        for gene, group in zip(target_genes, target_groups):
            try:
                logfc = get_logfc(
                    target_gene=gene,
                    target_group=group,
                    logfoldchanges=adata.uns['rank_genes_groups']['logfoldchanges'],
                    gene_names=adata.uns['rank_genes_groups']['names']
                )
                logfc_values.append(logfc)
            except KeyError as e:
                print(f"[Warning] Failed to retrieve {gene}@{group} in dataset {name}: {str(e)}")
                logfc_values.append(None)
        
        results[f'logFC_{name}'] = logfc_values
    
    # Convert to DataFrame
    logfc_df = pd.DataFrame(results)
    
    # Melt the DataFrame for visualization
    value_vars = [col for col in logfc_df.columns if col.startswith('logFC_')]
    long_df = pd.melt(
        logfc_df,
        id_vars=['Gene', 'Group'],
        value_vars=value_vars,
        var_name='type',
        value_name='logFC'
    )
    
    long_df['type'] = long_df['type'].str.replace('logFC_', '')
    
    output_path = Path(save_path)
    output_path.mkdir(parents=True, exist_ok=True)
    long_df.to_csv(output_path / f"{save_name}.csv", sep='\t', index=False)
    
    return long_df

import pandas as pd
from anndata import AnnData
from typing import Optional, Union

def get_top_degs_df(
    adata: AnnData,
    n_top_genes: int = 20,
    groupby: Optional[str] = None,
    key: str = 'rank_genes_groups'
) -> pd.DataFrame:
    r"""
    Extract top N differentially expressed genes (DEGs) and their statistics from Scanpy's rank_genes_groups results.
    
    Parameters
    ----------
    adata : AnnData
        AnnData object containing rank_genes_groups results.
    n_top_genes : int, optional
        Number of top DEGs to extract per group (default: 20).
    groupby : str, optional
        Column name in adata.obs used for grouping. If None, automatically reads from rank_genes_groups.
    key : str, optional
        Key name for rank_genes_groups in adata.uns (default: 'rank_genes_groups').
    
    Returns
    -------
    pd.DataFrame
        DataFrame containing the following columns:
        - Group: Group name
        - Gene: Gene symbol
        - LogFC: Log fold change
        - PValue: Raw p-value
        - AdjPValue: Adjusted p-value (e.g., FDR)
    """

    if key not in adata.uns:
        raise KeyError(f"'{key}' not found in adata.uns. Run sc.tl.rank_genes_groups first.")
    
    rank_data = adata.uns[key]
    required_fields = ['names', 'logfoldchanges', 'pvals', 'pvals_adj']
    for field in required_fields:
        if field not in rank_data:
            raise ValueError(f"Missing required field '{field}' in rank_genes_groups data.")

    if groupby is None:
        groupby = rank_data['params']['groupby'] if 'params' in rank_data else None
        if groupby is None:
            raise ValueError("Please specify them manually through the 'groupby' parameter.")


    groups = rank_data['names'].dtype.names
    gene_names = rank_data['names']
    logfcs = rank_data['logfoldchanges']
    pvals = rank_data['pvals']
    pvals_adj = rank_data['pvals_adj']

    top_genes = []
    for group in groups:
        genes = gene_names[group][:n_top_genes]
        valid_idx = ~pd.isnull(genes)
        genes = genes[valid_idx]
        
        group_logfcs = logfcs[group][:n_top_genes][valid_idx]
        group_pvals = pvals[group][:n_top_genes][valid_idx]
        group_padjs = pvals_adj[group][:n_top_genes][valid_idx]
        
        for gene, lfc, pval, padj in zip(genes, group_logfcs, group_pvals, group_padjs):
            top_genes.append({
                'Group': group,
                'Gene': gene,
                'LogFC': lfc,
                'PValue': pval,
                'AdjPValue': padj
            })

    return pd.DataFrame(top_genes)

from typing import Optional, Dict
import os
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from anndata import AnnData

def plot_modality_weights(
    adata: AnnData,
    modality_names: tuple = ("RNA", "Protein"),
    cluster_column: str = "DePass",
    save_path: Optional[str] = None,
    save_name: str = "modality_weights",
    show: bool = True,
    figsize: tuple = (5, 3),
    palette: Dict[str, str] = None,
    **kwargs
) -> plt.Axes:
    r"""
    Plot modality weights for clusters with legend outside right.
    """
 
    if 'alpha' not in adata.obsm:
        raise KeyError("Missing modality weights in adata.obsm['alpha']")
    if cluster_column not in adata.obs:
        raise KeyError(f"Cluster column '{cluster_column}' not found in adata.obs")
    if adata.obsm['alpha'].shape[1] != 2:
        raise ValueError(f"Expected 2 columns in adata.obsm['alpha'], found {adata.obsm['alpha'].shape[1]}")

    default_palette = {modality_names[0]: "#CAB82E", modality_names[1]: "#9368A6"}
    palette = palette or default_palette

    plot_df = pd.DataFrame({
        modality_names[0]: adata.obsm['alpha'][:, 0],
        modality_names[1]: adata.obsm['alpha'][:, 1],
        'Cluster': adata.obs[cluster_column].astype(str)
    })
    melted_df = plot_df.melt(
        id_vars='Cluster',
        value_vars=modality_names,
        var_name='Modality',
        value_name='Weight'
    )
    clusters = sorted(plot_df['Cluster'].unique(), key=lambda x: int(x))
    melted_df['Cluster'] = pd.Categorical(melted_df['Cluster'], categories=clusters, ordered=True)

    plt.figure(figsize=figsize)
    ax = sns.violinplot(
        data=melted_df,
        x='Cluster',
        y='Weight',
        hue='Modality',
        inner="quart",
        linewidth=0.6,
        palette=palette,
        **kwargs
    )

    ax.set_title(f"{modality_names[0]} vs {modality_names[1]}", pad=15)
    ax.set_xlabel("Cluster", labelpad=10)
    ax.set_ylabel("Attention Weight", labelpad=10)

    ax.legend(
        loc='center left',
        bbox_to_anchor=(1.05, 0.5),
        frameon=True,
        title='Modality'
    )
    
    plt.subplots_adjust(right=0.75)
    plt.tight_layout()

    if save_path:
        os.makedirs(save_path, exist_ok=True)
        plt.savefig(
            os.path.join(save_path, f"{save_name}.png"),
            dpi=300,
            bbox_inches="tight",
            transparent=True
        )
    if show:
        plt.show()
    plt.close()

    return ax


def Generate_masked_data(adata_omics1,adata_omics2,mask_ratio):
    r"""
    Generate masked data.

    Parameters:
        mask_ratio (float): Masking ratio.

    Returns:
        tuple: Masked data.
    """
    import numpy as np
    from scipy import sparse
    import scanpy as sc

    adata_omics1.var_names_make_unique()
    adata_omics2.var_names_make_unique()

    def mask(adata, ratio):
        X = adata.X.copy()  

        if sparse.issparse(X):
            nonzero_positions = X.nonzero()
        else:
            nonzero_positions = np.nonzero(X)

        num_to_mask = int(ratio * len(nonzero_positions[0]))
        mask_indices = np.random.choice(len(nonzero_positions[0]), num_to_mask, replace=False)
        
        if sparse.issparse(X):
            X[(nonzero_positions[0][mask_indices], nonzero_positions[1][mask_indices])] = 0
        else:
            X[nonzero_positions[0][mask_indices], nonzero_positions[1][mask_indices]] = 0

        adata.X = X
        return adata

    sc.pp.filter_genes(adata_omics1, min_cells=adata_omics1.shape[0] * 0.01)
    sc.pp.highly_variable_genes(adata_omics1, flavor="seurat_v3", n_top_genes=1000)
    adata_omics1 = adata_omics1[:, adata_omics1.var['highly_variable']]
    adata_omics1 = mask(adata_omics1, mask_ratio)
    sc.pp.normalize_total(adata_omics1, target_sum=1e4)
    sc.pp.log1p(adata_omics1)
    sc.pp.scale(adata_omics1)

    from utils import clr_normalize_each_cell
    adata_omics2 = mask(adata_omics2, mask_ratio)
    adata_omics2 = clr_normalize_each_cell(adata_omics2)
    sc.pp.scale(adata_omics2)
    return adata_omics1.X, adata_omics2.X


import scanpy as sc
import pandas as pd
import numpy as np
from scipy.stats import pearsonr
import matplotlib.pyplot as plt
import seaborn as sns

def calculate_correlation(adata1, adata2, gene_adt_mapping):
    r"""
    Calculate correlations between genes and ADTs.

    Parameters:
        adata1 (AnnData): First AnnData object (genes).
        adata2 (AnnData): Second AnnData object (ADTs).
        gene_adt_mapping (dict): Mapping from genes to ADTs.

    Returns:
        pd.DataFrame: DataFrame containing correlation results.
    """
    results = []
    for adt, genes in gene_adt_mapping.items():
        if adt not in adata2.var_names:
            print(f"Warning: {adt} not found in ADT data.")
            continue
        
        for gene in genes:
            if gene not in adata1.var_names:
                print(f"Warning: {gene} not found in gene data.")
                continue
            
            # Extract expression values
            gene_expression = adata1[:, gene].X.flatten()  # Gene expression
            adt_expression = adata2[:, adt].X.flatten()    # ADT expression
            
            # Calculate Pearson correlation coefficient
            correlation, p_value = pearsonr(gene_expression, adt_expression)
            
            results.append({
                "ADT": adt,
                "Gene": gene,
                "Gene_ADT": f"{gene}_{adt}",  
                "Correlation": correlation,
                "P_value": p_value
            })
    
    return pd.DataFrame(results)