from __future__ import annotations

import pandas as pd
from scipy.cluster import hierarchy as sch


def cluster(
    hotspot_df: pd.DataFrame,
    AI_series: pd.Series,
    n_svgs: int = 1000,
    n_svg_clusters: int = 8,
) -> pd.Series:
    """
    Cluster SVGs using hotspot matrix.

    Parameters
    ==========
    hotspot_df : pd.DataFrame
        A hotspot DataFrame generated by svgbit.

    AI_series : pd.Series
        A Series for AI value.

    n_svgs : int, default 1000
        Number of SVGs to find clusters.

    n_svg_clusters : int, default 8
        Number of SVG clusters to find.

    Returns
    =======
    gene_result : pd.Series
        A Series of clustering result.

    """
    selected_genes = AI_series.sort_values(ascending=False)[:n_svgs].index
    hotspot_set = hotspot_df[selected_genes]
    gene_distmat = sch.distance.pdist(hotspot_set.T, metric="jaccard")
    Z_gene = sch.linkage(gene_distmat, method="ward")
    gene_result = pd.Series(
        sch.fcluster(Z_gene, t=n_svg_clusters, criterion="maxclust"),
        index=selected_genes,
    ).sort_values()
    return gene_result
