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

# %% auto 0
__all__ = ['ManifoldGraph', 'power_diffusion_matrix', 'phate_distances', 'wasserstein_spread_of_diffusion',
           'entropy_of_diffusion', 'flattened_facsimile_of_graph', 'diffusion_curvature', 'diffusion_entropy_curvature',
           'plot_manifold_curvature']

# %% ../nbs/Core (manifoldgraph).ipynb 3
from fastcore.all import *
import numpy as np
from .kernels import *
class ManifoldGraph():
    """
    Class for Manifold Graphs. Stores the basic graph information used for geometric analysis: 
    A (the affinity/adjacency matrix), P (the diffusion matrix), P^t (the powered diffusion matrix), and D (the manifold distance matrix)
    Takes raw points, X. Can also optionally pass in any of the above to save the trouble of recomputing it.
    """
    def __init__(self, 
    X=None, # raw points
    dimension=None, # dimension of manifold. 
    kernel_type:str = "adaptive", # "fixed" or "adaptive"
    anisotropic_density_normalization:float=1,  # normalize out this amount of the density. 0 is none. 1 is all.
    num_neighbors = 10, # for adaptive kernel, constructs graph to connect this number of neighbors
    t = 10, # steps of diffusion
    A = None, 
    P = None, 
    Pt = None, 
    D = None, 
    *kw_args
    ):
        """
        We compute the affinity matrix and diffusion matrix automatically
        """
        self.node_attributes = {}
        store_attr()
        if self.A is None:
            self.A = gaussian_kernel(self.X, kernel_type=self.kernel_type, k = self.num_neighbors, anisotropic_density_normalization=self.anisotropic_density_normalization)
            np.fill_diagonal(self.A,0)
        if self.P is None:
            self.P = compute_anisotropic_diffusion_matrix_from_graph(self.A, self.anisotropic_density_normalization)        
    def num_nodes(self):
        return len(self.A)

# %% ../nbs/Core (manifoldgraph).ipynb 4
import numpy as np
def power_diffusion_matrix(G:ManifoldGraph, t:int=None):
    # Raises the diffusion matrix to t
    if not t:
        t = G.t
    G.Pt = np.linalg.matrix_power(G.P, t)
    return G


# %% ../nbs/Core (manifoldgraph).ipynb 6
from sklearn.metrics import pairwise_distances
import numpy as np
@Transform
def phate_distances(G:ManifoldGraph):
    assert G.Pt is not None
    log_Pts = -np.log(G.Pt + 1e-6)
    D = pairwise_distances(log_Pts)
    G.D = D
    return G

# %% ../nbs/Core (manifoldgraph).ipynb 8
def wasserstein_spread_of_diffusion(G:ManifoldGraph, 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
    if idx is None:
        return np.sum(G.D * G.Pt, axis=1)
    else:
        return np.sum(G.D[idx]* G.Pt[idx])

# %% ../nbs/Core (manifoldgraph).ipynb 9
from scipy.stats import entropy
def entropy_of_diffusion(G:ManifoldGraph, idx=None):
    assert G.Pt is not None
    if idx is None:
        return entropy(G.Pt, axis=1)
    else:
        return entropy(G.Pt[idx])

# %% ../nbs/Core (manifoldgraph).ipynb 10
def flattened_facsimile_of_graph(G:ManifoldGraph, dimension=None):
    if not dimension:
        dimension=G.dimension
    noise = np.concatenate([np.zeros((1,dimension)), 2*np.random.rand(G.num_nodes()-1,dimension)-1])
    # compute the adjacency matrix and P
    G_flat = ManifoldGraph(noise, kernel_type=G.kernel_type, num_neighbors=G.num_neighbors, anisotropic_density_normalization=G.anisotropic_density_normalization)
    G_flat = power_diffusion_matrix(G_flat, t = G.t)
    # TODO: This is a tad wasteful. We only need the center diffusion powered.
    return G_flat

# %% ../nbs/Core (manifoldgraph).ipynb 11
@Transform
def diffusion_curvature(G:ManifoldGraph, t=None,):
    if t is None:
        t = G.t
    if G.Pt is None or t != G.t:
        G = power_diffusion_matrix(G,t)
    if G.D is None:
        G = phate_distances(G)
    spreads_on_manifold = wasserstein_spread_of_diffusion(G)
    # Create flattened version of manifold and compute stuff
    G_flat = flattened_facsimile_of_graph(G)
    G_flat = power_diffusion_matrix(G_flat)
    G_flat = phate_distances(G_flat)
    spreads_on_flat = wasserstein_spread_of_diffusion(G_flat, idx=0)
    G.ks = 1 - (spreads_on_manifold/(spreads_on_flat))
    return G

# %% ../nbs/Core (manifoldgraph).ipynb 12
@Transform
def diffusion_entropy_curvature(G:ManifoldGraph, t=None,):
    if t is None:
        t = G.t
    if G.Pt is None:
        G = power_diffusion_matrix(G,t)
    spreads_on_manifold = entropy_of_diffusion(G)
    # Create flattened version of manifold and compute stuff
    G_flat = flattened_facsimile_of_graph(G)
    G_flat = power_diffusion_matrix(G_flat)
    spreads_on_flat = entropy_of_diffusion(G_flat, idx=0)
    print(spreads_on_flat)
    G.ks = spreads_on_flat - spreads_on_manifold
    return G

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