# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/Core (graphtools).ipynb.

# %% auto 0
__all__ = ['DiffusionCurvature', 'get_graph_type', 'plot_manifold_curvature']

# %% ../nbs/Core (graphtools).ipynb 9
from fastcore.all import *
import graphtools
import numpy as np
import graphtools
import scipy.sparse
import math
import scipy
from scipy.stats import entropy
from .distances import *
import skdim
class DiffusionCurvature():
    def __init__(
    self,
    t:int, # Number of diffusion steps to use when measuring curvature. TODO: Heuristics
    distance_type = "PHATE",
    dimest = None, # Dimensionality estimator to use. If None, defaults to KNN with default params
    use_entropy:bool = False, # If true, uses KL Divergence instead of Wasserstein Distances. Faster, seems empirically as good, but less proven.
    **kwargs
    ):
        store_attr()
        if self.dimest is None:
            self.dimest = skdim.id.KNN()

# %% ../nbs/Core (graphtools).ipynb 12
@patch
def power_diffusion_matrix(self:DiffusionCurvature, G:graphtools.api.Graph, t=None):
        # Raises the diffusion matrix of G to t and stores the result in G.Pt
        if not t:
            t = self.t
        if type(G.P) == np.ndarray:
            G.Pt = np.linalg.matrix_power(G.P, t)
        elif type(G.P) == scipy.sparse.csr_matrix:
            if 'n_landmark' in G.get_params().keys():
                # the graph was built with landmarking we'll power the diffusion matrix that way
                G.Pt = scipy.sparse.csr_matrix(
                    G.transitions @ np.linalg.matrix_power(G.landmark_op,t) @ G.transitions.T
                )
            else:
                G.Pt = G.P ** t
        else:
            raise NotImplementedError("Unknown Matrix Type")
        return G

# %% ../nbs/Core (graphtools).ipynb 18
@patch
def distances(self:DiffusionCurvature, G):
        if self.distance_type == "PHATE":
            G = phate_distances(G)
        else:
           raise NotImplementedError(f"No distance with name {self.distance_type}")
        return G

# %% ../nbs/Core (graphtools).ipynb 23
@patch
def wasserstein_spread_of_diffusion(self:DiffusionCurvature, G:graphtools.graphs.DataGraph, idx = None):
        """
        Returns how "spread out" each diffusion is, with wasserstein distance"
        Presumes that the manifold distances have been separately calculated
        If idx is passed, only computes wsd at that index
        """
        assert G.Pt is not None
        assert G.D is not None
        # G.P is our distribution on each node. We compute the sum of the distances each bit of "diffusion mass" has to be carried 
        # back to the center
        # Convert to numpy array;
        if type(G.Pt) == np.ndarray:
            Pt = G.Pt
        else:
            Pt = G.Pt.toarray()
        if idx is None:
            return np.sum(G.D * Pt, axis=1)
        else:
            return np.sum(G.D[idx]* Pt[idx])

# %% ../nbs/Core (graphtools).ipynb 27
def get_graph_type(G):
    tt = str(type(G)).split('.')[-1][:3].lower()
    # TODO: Support for MNN graphs
    if tt == 'mnn': raise NotImplementedError("MNN Graphs require manual indices. These are not yet supported by our graph flattener.")
    if tt in ['knn']: return tt
    else: return 'exact'

# %% ../nbs/Core (graphtools).ipynb 32
@patch
def flattened_facsimile_of_graph(self:DiffusionCurvature, G:graphtools.graphs.DataGraph, dimension):
        """
        Constructs a flat graph, hewn from uniform random noise of the supplied dimension.
        Calculates the powered diffusion matrix on this graph.
        """
        noise = np.concatenate([np.zeros((1,dimension)), 2*np.random.rand(G.K.shape[0]-1,dimension)-1])
        # Build a graph out of the noise, with all of the same kernel settings as our first graph
        # NOTE: The graph building settings must be scale invariant! 
        params = G.get_params()
        needed_keys = ['data', 'n_pca', 'knn', 'decay', 'bandwidth', 'bandwidth_scale', 
                    'knn_max', 'anisotropy', 'beta', 'adaptive_k', 'n_landmark', 'n_landmark', 
                    'n_svd', 'n_jobs']
        found_keys = {} # TODO: Likely a more elegant way to do this
        for nk in needed_keys:
            if nk in params.keys():
                found_keys[nk] = params[nk]
            else:
                found_keys[nk] = None
        
        G_flat = graphtools.Graph(
            data = noise,
            n_pca = found_keys['n_pca'],
            knn = found_keys['knn'],
            decay = found_keys['decay'],
            bandwidth = found_keys['bandwidth'],
            bandwidth_scale = found_keys['bandwidth_scale'],
            knn_max = found_keys['knn_max'],
            anisotropy = found_keys['anisotropy'],
            beta = found_keys['beta'],
            adaptive_k = found_keys['adaptive_k'],
            n_landmark = found_keys['n_landmark'],
            n_svd = found_keys['n_svd'],
            n_jobs = found_keys['n_jobs'],
            graphtype = get_graph_type(G),
        )
        G_flat = self.power_diffusion_matrix(G_flat, t = self.t)
        # TODO: This is a tad wasteful. We only need the center diffusion powered.
        return G_flat

# %% ../nbs/Core (graphtools).ipynb 35
@patch
def entropy_of_diffusion(self:DiffusionCurvature, G:graphtools.graphs.DataGraph, idx=None):
        """
        Returns the pointwise entropy of diffusion from the powered diffusion matrix in the inpiut 
        """
        assert G.Pt is not None
        # TODO: Entropy may be compatible with scipy.sparse matrices
        if type(G.Pt) == np.ndarray:
            Pt = G.Pt
        else:
            Pt = G.Pt.toarray()
        if idx is None:
            return entropy(Pt, axis=1)
        else:
            return entropy(Pt[idx])

# %% ../nbs/Core (graphtools).ipynb 38
@patch
def curvature(self:DiffusionCurvature, 
    G:graphtools.graphs.DataGraph, # A graphtools graph to compute the curvature of
    t=None,  # The number of steps within the random walks. Corresponds to how local/global the curvature estimate is.
    dimension = None,  # If supplied, the intrinsic dimension of the manifold, either a list of the intrinsic dimension per point or an int of the global intrinsic dimension. If not supplied, estimates the local dimension of each point on the manifold using the estimator passed to DiffusionCurvature.
    ):
    """
    Computes diffusion curvature of input graph. Stores it in G.ks
    """
    if t is None:
        t = self.t
    G = self.power_diffusion_matrix(G,t)
    G = self.distances(G)
    spreads_on_manifold = self.wasserstein_spread_of_diffusion(G) if not self.use_entropy else self.entropy_of_diffusion(G)
    # Create flattened version of manifold and compute stuff
    if dimension is None: # The dimension wasn't supplied; we'll estimate it pointwise
        print("estimating local dimension of each point... may take a while")
        ldims = self.dimest.fit_pw(G.data,
                            n_neighbors = 100,
                            n_jobs = 1)
        dims_per_point = np.round(ldims.dimension_pw_).astype(int)
    else: # the dimension *was* supplied, but it may be either a single global dimension or a local dimension for each point
        if isinstance(dimension, int):
            dims_per_point = np.ones(G.P.shape[0], dtype=int)*dimension
        else:
            dims_per_point = dimension
    # TODO: Currently this requires underlying points.
    unique_dims = set(dims_per_point)
    unique_flats = {}
    for d in unique_dims:
        G_flat = self.flattened_facsimile_of_graph(G, dimension=d)
        G_flat = self.power_diffusion_matrix(G_flat,t)
        G_flat = self.distances(G_flat)
        unique_flats[d] = self.wasserstein_spread_of_diffusion(G_flat, idx=0) if not self.use_entropy else self.entropy_of_diffusion(G, idx=0)
    divided_pts = np.array([spreads_on_manifold[idx]/unique_flats[localdim] for idx, localdim in enumerate(dims_per_point)])
    G.ks = 1 - divided_pts
    return G

# %% ../nbs/Core (graphtools).ipynb 45
from .kernels import plot_3d
def plot_manifold_curvature(G, title = None):
    X = G.data
    ks = G.ks
    plot_3d(X, ks, colorbar=True, title=title)
