# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/40_diff.ipynb.

# %% auto 0
__all__ = ['Diffusion', 'DiffusionDistance', 'DiffusionAffinity', 'DiffusionMap', 'PHATEDistance']

# %% ../nbs/40_diff.ipynb 6
#| export


# %% ../nbs/40_diff.ipynb 8
from typing import Self

# %% ../nbs/40_diff.ipynb 10
#| export


# %% ../nbs/40_diff.ipynb 12
#| export


# %% ../nbs/40_diff.ipynb 14
try: import numpy as np
except ImportError: ...

try: import pandas as pd
except ImportError: ...

# %% ../nbs/40_diff.ipynb 15
try: import sklearn
except: ...

try: from sklearn.metrics.pairwise import pairwise_distances
except ImportError: ...

# %% ../nbs/40_diff.ipynb 16
try: from scipy.sparse import csr_matrix
except ImportError: ...

try: from scipy.sparse.linalg import eigs
except ImportError: ...

try: from scipy.spatial.distance import pdist
except: ...

# %% ../nbs/40_diff.ipynb 17
try: import phate, graphtools as gt
except ImportError: ...

# %% ../nbs/40_diff.ipynb 19
#| export


# %% ../nbs/40_diff.ipynb 21
#| export


# %% ../nbs/40_diff.ipynb 23
from nlit import (L1, L2, DEVICE, ANISOTROPY)
from quac import deviceq
from chck import isnone, notnone
from asto import asdev, asarr
from etrc import Framework

# %% ../nbs/40_diff.ipynb 25
from .kwds import DiffusionKeywords

# %% ../nbs/40_diff.ipynb 28
class Diffusion:   
    def __init__(
        self: Self, 
        knn: int = 5,
        tmax: int = 5,
        anisotropy: int = 1,
        tdiff: int = 1,
        topeig: int = 100,
        log: bool = False,
        normalize: bool = False,
        symmetrize: bool = False,
        verbose: bool = False,
        nemb: int = 10,
        njobs: int = -1,
        device: deviceq = None,
    ) -> None:
        self.knn = knn               # 5 number of neighbors for the KNN in the alpha decay kernel construction, same default as in PHATE
        self.tmax = tmax             # 2 ^ tmax is the max scale of the Diffusion kernel
        self.anisotropy = anisotropy # Anisotropy (int): the alpha in Coifman Lafon 2006, 1: double normalization 0: usual random walk
        self.tdiff = tdiff           # t_diff (int) the power of the diffusion affinity matrix
        self.topeig = topeig         # in the the top k eigenvalues to consider in the spectral decomposition
        self.log = log
        self.normalize = normalize
        self.symmetrize = symmetrize
        self.K = None  # kernel
        self.P = None  # diffusion_operator
        self.pi = None # stationary_distribution
        self.G = None  # diffusion_distance
        self.A = None  # diffusion_affinity
        self.verbose = verbose
        self.nemb = nemb # the dimension of the emb space
        self.njobs = njobs
        self.diff_pot = None # diffusion_potential
        self.graph = None
        self.device = asdev(device)

    def to(self: Self, device: deviceq) -> Self:
        if notnone(device): self.device = asdev(device)
        return self
    
    def mount(self: Self, X):
        X = self._swap_framework(X, Framework.torch)
        return X
        
    def calculate_stationary_distribution(self, K):
        pi = np.sum(self.K, axis=1)
        pi /= np.sum(pi)
        return asarr(pi)

    @property
    def stationary_distribution(self): 
        try:
            return self._stationary_distribution
        except AttributeError:
            self._stationary_distribution = self.calculate_stationary_distribution(self.K)
            return self._stationary_distribution
                
    def _eigs_to_use(self, X = None):
        if isnone(X): X = asarr(self.A)
        n_rows = X.shape[0]
        return min(self.topeig, n_rows - 2)
    
    def get_P(self):
        P = self.P.copy()
        if self.log:
            P = csr_matrix((np.log(P.data), P.indices, P.indptr), shape=P.shape)
        return P
    
    def _pairwise_distances(self, X, metric: str = L1):
        dist = pairwise_distances(X ,X, metric=metric, n_jobs=self.njobs)
        if self.log:
            np.fill_diagonal(dist, 1)
            dist = (-1) * np.log(dist)
        return asarr(dist)

    def l1_pairwise_distances(self, X):
        return self._pairwise_distances(X, L1)
    
    def l2_pairwise_distances(self, X):
        return self._pairwise_distances(X, L2)
        
    def diffuse_G(self, G, t, dist):
        return G + 2 ** ( -t / 2.0) * dist
    
    def diffuse_A(self, A, t = None):
        if isnone(t): t = self.tdiff
        if t == 1: return asarr(A)
        
        # NOTE: originally used X to check shape      
        # k = self._eigs_to_use(X)
        k = self._eigs_to_use(A)
        w, v = eigs(A, k=k, which='LR')
        W = np.diag(w)
        A = np.real(v @ (W ** self.t_diff) @ v.T)
        return asarr(A)
    
    def normalize_G(self, G):
        return (G - np.min(G)) / (np.max(G) - np.min(G))
    
    def symmetrize_G(self, G):
        return (G + np.transpose(G)) / 0.5

    def compute_diffusion_distance(self):
        P = self.get_P()
        G = self.l1_pairwise_distances(P)
        for t in range(1, self.tmax):
            P = P @ P
            dist = self.l1_pairwise_distances(P)
            G = self.diffuse_G(G, t, dist)
            
        dist = self.l1_pairwise_distances(self.pi)
        G = self.diffuse_G(G, (self.tmax + 1), dist)
        if self.normalize: G = self.normalize_G(G)
        
        self.G = asarr(G)
        return self.G
    
    def make_graph(self, X):
        self.graph = gt.Graph(X, knn=self.knn, anisotropy=self.anisotropy)
        return self.graph

    def _swap_framework(self: Self, X, desired: Framework = Framework.numpy):
        self.device = asdev(getattr(X, DEVICE, None))
        return Framework.swap(X, desired, self.device)
            
    def fit(self, X):
        X = self._swap_framework(X, Framework.numpy)
        graph = self.make_graph(X)
        self.K = graph.K
        self.P = graph.diff_op
        self.pi = self.stationary_distribution
        G = self.compute_diffusion_distance()       
        if self.symmetrize:
            G = self.symmetrize_G(G)
        self.G = G
        return self.mount(self.G)

# %% ../nbs/40_diff.ipynb 30
class DiffusionDistance(Diffusion):    
    def __init__(self: Self, *args, **kwargs: DiffusionKeywords) -> None:  
        super(DiffusionDistance, self).__init__(*args, **kwargs)
    
    def fit(self, X):
        return super(DiffusionDistance, self).fit(X)

# %% ../nbs/40_diff.ipynb 32
class DiffusionAffinity(Diffusion):
    def __init__(self: Self, anisotropy: int = 0, *args, **kwargs: DiffusionKeywords) -> None:
        kwargs[ANISOTROPY] = anisotropy
        super(DiffusionAffinity, self).__init__(*args, **kwargs)        
                
    def fit(self: Self, X):
        '''return the l2 between the row of the affinity matrix A^t'''
        X = self._swap_framework(X, Framework.numpy)
        graph = self.make_graph(X)
        self.A = graph.diff_aff
        A = self.diffuse_A(A, self.tdiff)
        G = self.l2_pairwise_distances(self.A)
        self.G = G
        return self.mount(self.G)

# %% ../nbs/40_diff.ipynb 34
class DiffusionMap(Diffusion):
    def __init__(self: Self, anisotropy: int = 0, *args, **kwargs: DiffusionKeywords) -> None:
        kwargs[ANISOTROPY] = anisotropy
        super(DiffusionMap, self).__init__(*args, **kwargs)
        
    def fit(self, X):
        '''return the pairwise dist. in the diffusion map embedding space'''
        X = self._swap_framework(X, Framework.numpy)
        graph = self.make_graph(X)
        self.A = graph.diff_aff
        k = self._eigs_to_use(X)
        
        w, v = eigs(self.A, k=k, which='LR')
        w, v = np.real(w), np.real(v)
        v = v / v[:, 0, None]
        self.emb = np.vstack([
            (w[i] ** self.tdiff) * v[:, i]
            for i in range(1, self.nemb + 1)
        ])

        self.G = pdist(self.emb.T)
        return self.mount(self.G)

# %% ../nbs/40_diff.ipynb 36
class PHATEDistance(Diffusion):
    def __init__(self: Self, anisotropy: int = 0, *args, **kwargs: DiffusionKeywords) -> None:
        kwargs[ANISOTROPY] = anisotropy
        super(PHATEDistance, self).__init__(*args, **kwargs)

    def fit(self, X):
        '''return PHATE distance the L2 between Potential of Heat-diffusion'''
        X = self._swap_framework(X, Framework.numpy)
        graph = phate.PHATE(knn=self.knn, verbose=self.verbose, n_landmark=X.shape[0]).fit(X)
        self.diff_pot = graph.diff_potential 
        self.G = self.l2_pairwise_distances(self.diff_pot)
        return self.mount(self.G)
    
