import torch
import pandas as pd
import numpy as np
import torch.nn.functional as F
from sklearn.covariance import GraphicalLasso

def gen_similarity_graph(X:pd.DataFrame, k:int = 15, metric:str = "cosine", mutual:bool = False, per_node:bool = True, self_loops:bool = True) -> pd.DataFrame:
    """
    Build a normalized knn similarity graph from feature vectors. Computes pairwise cosine or ecledian disntace,then sparsifies via knn or global a threshold.
    Optionally prunes to mutual neighbors and/or adds self-loops.

    Args:
        - X: pandas.DataFrame of shape (N, D) (rows = nodes, cols = features)
        - k (int): Number of neighbors to keep per node.
        - metric (str): "cosine" or "euclidean" (uses gaussian kernel on distances).
        - mutual (bool): If True, retain only mutual edges (i->j and j->i).
        - per_node (bool): If True, use per-node top_k; else global cutoff.
        - self_loops (bool): If True, add self-loop weight of 1.

    Returns:

        - DataFrame of shape (N, N) the normalized adjacency matrix
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if isinstance(X, pd.DataFrame):
        nodes = X.index
        x_torch = torch.tensor(X.values, dtype=torch.float32, device=device)
    else:
        raise TypeError("X must be a pandas.DataFrame")
    
    N = x_torch.size(0)

    # full similarity matrix
    if metric == "cosine":
        X_normalized = F.normalize(x_torch, p=2, dim=1)
        S = torch.mm(X_normalized, X_normalized.t())
    else:
        D2 = torch.cdist(x_torch, x_torch).pow(2)
        median_d2 = D2.median()
        S = torch.exp(-D2 / (median_d2 + 1e-8))

    # building the knn graph or global threshold mask
    if per_node:
        _, index = torch.topk(S, k=k+1, dim=1) 
        mask = torch.zeros(N, N, dtype=torch.bool, device=device)
        for i in range(N):
            for j in index[i, 1:k+1]:
                mask[i, j] = True
    else:
        flat = S.reshape(-1)
        threshold = torch.kthvalue(flat, k * N).values
        mask = S >= threshold
        mask.fill_diagonal_(False)

    # mutual pruning option
    if mutual:
        mask = torch.logical_and(mask, mask.t())

    # mask and add self loops if set to true
    A = S * mask.float()
    if self_loops:
        A = A + torch.eye(N, device=device, dtype=x_torch.dtype)

    # row normalize
    A = F.normalize(A, p=1, dim=1)

    A_numpy = A.cpu().numpy()
    final_graph =pd.DataFrame(A_numpy, index=nodes, columns=nodes)
    
    return final_graph


def gen_correlation_graph(X: pd.DataFrame, k: int = 15,method: str = 'pearson', mutual: bool = False, per_node: bool = True,threshold: float = None, self_loops:bool = True) -> pd.DataFrame:
    """
    Build a graph based on pairwise Pearson or Spearman correlations.

    Args:

        - X pd.dataframe (N, D) data tensor where rows are nodes.
        - k (int): Number of neighbors to keep per node if per_node is True.
        - method (str): 'pearson' or 'spearman'.
        - mutual (bool): If True, only mutual knn edges.
        - per_node (bool): If True, use per node topk selection, else global threshold.
        - threshold (float): Correlation cutoff when per_node is False.
        - self_loops (bool): If True, adds weight 1 to diagonal.

    Returns:

        - pandas.DataFrame. Normalized adjacency matrix (N x N) of the sparse correlation graph.

    Note:

        - Correlation is very expensive to compute, so this function is not recommended for large datasets.

    """

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if isinstance(X, pd.DataFrame):
        nodes = X.index
        x_torch = torch.tensor(X.values, dtype=torch.float32, device=device)
    else:
        raise TypeError("X must be a pandas.DataFrame")
    
    N = x_torch.size(0)

    # rank transform for Spearman
    if method == 'spearman':
        x_ranked = x_torch.argsort(dim=1).argsort(dim=1).float()
        x_correlation = x_ranked - x_ranked.mean(dim=1, keepdim=True)
    else:
        x_correlation = x_torch - x_torch.mean(dim=1, keepdim=True)

    # correlation with clamping and normalization
    num = torch.mm(x_correlation, x_correlation.t())
    sum_sq = (x_correlation**2).sum(dim=1, keepdim=True)
    denom = torch.sqrt(torch.mm(sum_sq, sum_sq.t())).clamp(min=1e-8)
    C = num / denom
    S = C.abs()

    # build mask
    if per_node:
        _, index = torch.topk(S, k=k+1, dim=1)
        mask = torch.zeros(N, N, dtype=torch.bool, device=device)
        for i in range(N):
            for j in index[i, 1:k+1]:
                mask[i, j] = True
    else:
        if threshold is not None:
            mask = S >= threshold
            mask.fill_diagonal_(False)
        else:
            flat = S.reshape(-1)
            thresh = torch.kthvalue(flat, k * N).values
            mask = S >= thresh
            mask.fill_diagonal_(False)

    if mutual:
        mask = torch.logical_and(mask, mask.t())

    W = S * mask.float()
    if self_loops:
        W.fill_diagonal_(1.0)

    W = F.normalize(W, p=1, dim=1)
    final_graph = pd.DataFrame(W.cpu().numpy(), index=nodes, columns=nodes)

    return final_graph


def gen_threshold_graph(X:pd.DataFrame, b: float = 6.0,k: int = 15, mutual: bool = False, self_loops: bool = True) -> pd.DataFrame:
    """
    Generate a soft threshold co-xpression network this is very similar to how WGCNA works

    Args:

        - X (pd.DataFrame): Data matrix where rows are nodes and columns are features.
        - b (float): Thresholding exponent applied to absolute correlations.
        - k (int): Number of neighbors to keep per node.
        - mutual (bool): If True, only mutual knn edges.
        - self_loops (bool): If True, adds weight 1 to diagonal.

    Returns:

        - pandas.DataFrame: Normalized adjacency matrix (N x N) of the soft-thresholded graph.
    """

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if isinstance(X, pd.DataFrame):
        nodes = X.index
        x_torch = torch.tensor(X.values, dtype=torch.float32, device=device)
    else:
        raise TypeError("X must be a pandas.DataFrame")
    
    N = x_torch.size(0)

    # pearson correlation matrix

    Xc = x_torch - x_torch.mean(dim=1, keepdim=True)
    num = torch.mm(Xc, Xc.t())
    sum_sq = (Xc**2).sum(dim=1, keepdim=True)
    denom = torch.sqrt(torch.mm(sum_sq, sum_sq.t())).clamp(min=1e-8)
    C = num / denom

    # threshold
    S = C.abs().pow(b)

    # mask building
    _, index = torch.topk(S, k=k+1, dim=1)
    mask = torch.zeros(N, N, dtype=torch.bool, device=device)

    for i in range(N):
        for j in index[i, 1:k+1]:
            mask[i, j] = True
    if mutual:
        mask = torch.logical_and(mask, mask.t())

    # apply mask (and self-loops if set)
    W = S * mask.float()

    if self_loops:
        W.fill_diagonal_(1.0)

    W = F.normalize(W, p=1, dim=1)
    final_graph = pd.DataFrame(W.cpu().numpy(), index=nodes, columns=nodes) 

    return final_graph

def gen_gaussian_knn_graph(X: pd.DataFrame,k: int = 15,sigma: float = None,mutual: bool = False,self_loops: bool = True) -> pd.DataFrame:
    """
    Build a normalized knn similarity graph from feature vectors. Computes pairwise cosine or Euclidean similarities, sparsifies via k-nearest neighbors or
    a global threshold. Optionally prunes to mutual neighbors and/or adds self-loops.

    Args:

        - X (pd.DataFrame): Feature matrix where rows are nodes and columns are features.
        - k (int): Number of neighbors to keep per node.
        - metric (str): 'cosine' or 'euclidean', uses Gaussian kernel for distances.
        - mutual (bool): If True, only mutual knn edges.
        - per_node (bool): If True, use per-node topk selection; else global threshold.
        - self_loops (bool): If True, adds weight 1 to diagonal.

    Returns:

        - pandas.DataFrame: Normalized adjacency matrix (N x N) of the sparse similarity graph.
    """

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if isinstance(X, pd.DataFrame):
        nodes = X.index
        X_torch = torch.tensor(X.values, dtype=torch.float32, device=device)
    else:
        raise TypeError("X must be a pandas.DataFrame")
    
    N = X_torch.size(0)

    D2 = torch.cdist(X_torch, X_torch).pow(2)
    if sigma is None:
        sigma = D2.median().item()

    S = torch.exp(-D2 / (2 * sigma))

    _, index = torch.topk(S, k=k+1, dim=1)
    mask = torch.zeros(N, N, dtype=torch.bool, device=device)

    for i in range(N):
        for j in index[i, 1:k+1]:
            mask[i, j] = True

    if mutual:
        mask = torch.logical_and(mask, mask.t())

    W = S * mask.float()

    if self_loops:
        W.fill_diagonal_(1.0)

    W = F.normalize(W, p=1, dim=1)
    final_graph = pd.DataFrame(W.cpu().numpy(), index=nodes, columns=nodes)

    return final_graph


def gen_lasso_graph(X: pd.DataFrame, alpha: float = 0.01, self_loops: bool = True) -> pd.DataFrame:
    """
    Infer a sparse network via Graphical Lasso.

    Args:

        - X (pd.DataFrame): Data matrix where rows are nodes and columns are features.
        - alpha (float): Regularization parameter for Graphical Lasso.
        - self_loops (bool): If True, adds weight 1 to diagonal.

    Returns:

        - pandas.DataFrame: Normalized adjacency matrix (N x N) of the inferred network.

    """
    if isinstance(X, pd.DataFrame):
        nodes = X.index
        x_numpy = X.values
    else:
        raise TypeError("X must be a pandas.DataFrame")
    
    model = GraphicalLasso(alpha=alpha, max_iter=200)
    model.fit(x_numpy)

    P = torch.from_numpy(model.precision_).abs()
    W = (P + P.t()) / 2

    if self_loops:
        W.fill_diagonal_(W.diagonal() + 1.0)

    W = F.normalize(W, p=1, dim=1)
    final_graph = pd.DataFrame(W.cpu().numpy(), index=nodes, columns=nodes)

    return final_graph

def gen_mst_graph(X: pd.DataFrame, self_loops: bool = True) -> pd.DataFrame:
    """
    Compute the minimum spanning tree (MST) on Euclidean distances.

    Args:

        - X (pd.DataFrame): Feature matrix where rows are nodes and columns are features.
        - self_loops (bool): If True, adds weight 1 to diagonal.

    Returns:

        - pandas.DataFrame: Normalized adjacency matrix (N x N) of the MST graph.
    """

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if isinstance(X, pd.DataFrame):
        nodes = X.index
        X_torch = torch.tensor(X.values, dtype=torch.float32, device=device)
    else:
        raise TypeError("X must be a pandas.DataFrame")
    
    N = X_torch.size(0)
    D = torch.cdist(X_torch, X_torch)

    visited = torch.zeros(N, dtype=torch.bool, device=device)
    visited[0] = True

    min_cost = D[0].clone()
    closest = torch.zeros(N, dtype=torch.long, device=device)
    mst = torch.zeros((N, N), dtype=X_torch.dtype, device=device)

    for tmp in range(N - 1):
        cost_masked = min_cost.clone()
        cost_masked[visited] = float('inf')

        j = torch.argmin(cost_masked)
        i = closest[j].item()
        w = D[i, j]

        mst[i, j] = w
        mst[j, i] = w

        visited[j] = True
        distance_j = D[j]

        update = torch.logical_not(visited) & (distance_j < min_cost)
        min_cost = torch.where(update, distance_j, min_cost)
        closest = torch.where(update, j, closest)

    W = mst
    if self_loops:
        W.fill_diagonal_(1.0)

    W = F.normalize(W, p=1, dim=1)
    final_graph = pd.DataFrame(W.cpu().numpy(), index=nodes, columns=nodes)

    return final_graph

def gen_snn_graph(X: pd.DataFrame,k: int = 15,mutual: bool = False, self_loops: bool = True) -> pd.DataFrame:
    """
    Build a shared nearest neighbor (SNN) graph.

    Args:

        - X (pd.DataFrame): Feature matrix where rows are nodes and columns are features.
        - k (int): Number of neighbors to keep per node.
        - mutual (bool): If True, only mutual knn edges.
        - self_loops (bool): If True, adds weight 1 to diagonal.

    Returns:

        - pandas.DataFrame: Normalized adjacency matrix (N x N) of the SNN graph.
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if isinstance(X, pd.DataFrame):
        nodes = X.index
        X_torch = torch.tensor(X.values, dtype=torch.float32, device=device)
    else:
        raise TypeError("X must be a pandas.DataFrame")
    
    N = X_torch.size(0)
    S = torch.mm(X_torch, X_torch.t())
    _, index = torch.topk(S, k=k+1, dim=1)

    neighbors = []
    for i in range(N):
        neighbors.append(set(index[i, 1:k+1].tolist()))

    W = torch.zeros((N, N), dtype=X_torch.dtype, device=device)

    for i in range(N):
        for j in index[i, 1:k+1]:
            j = j.item()
            shared = len(neighbors
            [i].intersection(neighbors[j]))
            W[i, j] = shared

    if mutual:
        Wij = W.clone()
        W = torch.min(Wij, Wij.t())

    if self_loops:
        W.fill_diagonal_(1.0)

    W = F.normalize(W, p=1, dim=1)
    final_graph = pd.DataFrame(W.cpu().numpy(), index=nodes, columns=nodes)

    return final_graph
