# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02_metrics.ipynb (unless otherwise specified).

__all__ = ['sparse_norm', 'sparse_2norm', 'sparse_maxnorm', 'laplacian_distance', 'LineDistances', 'average_gmdegree',
           'edge_degree_gm']

# Cell
from nbdev.showdoc import *
from .graphtools import laplacian
from .data import get_benchmark
from functools import lru_cache
from pathlib import Path
from hashlib import md5
import networkx as nx
import numpy as np
import scipy.sparse as sp
import pickle
import os

# Cell
def sparse_norm(A, ord=2):
    """Like scipy.sparse.lingalg.norm but with the 2-norm and max norm implemented.

    If `ord=2` or `ord='max'` a grapht implementation is used, otherwise scipy.sparse.lingalg.norm is used.
    """
    if not sp.issparse(A):
        raise TypeError('input must be sparse')
    if ord == 2:
        return sparse_2norm(A)
    elif ord == 'max':
        return sparse_maxnorm(A)
    else:
        return sp.linalg.norm(A, ord=ord)

def sparse_2norm(A):
    """Returns the matrix 2-norm of a sparse matrix `A`."""
    if not sp.issparse(A):
        raise TypeError('input must be sparse')
    return sp.linalg.svds(A, k=1, which='LM', return_singular_vectors=False)[0]

def sparse_maxnorm(A):
    """Returns the max |A_ij| for a sparse matrix `A`."""
    if not sp.issparse(A):
        raise TypeError('input must be sparse')
    return max(-A.min(), A.max())

# Cell
def laplacian_distance(G, Gp, setdiag=False):
    """Calculates $|| \mathcal{L}(G) -  \mathcal{L}(G_p) ||$ using the matrix 2-norm."""
    L = laplacian(G, setdiag)
    Lp = laplacian(Gp, setdiag)
    E = Lp - L
    return sparse_2norm(E)

# Cell
class LineDistances():
    """
    An object which computes the distances of edges in the graphs line graph.

    If precompute is True a distance matrix is computed for all pairs of edges or loaded if a file exists at precompute_dir.
    """

    def __init__(self, G, precompute=False, precompute_dir='/tmp', verbose=False):
        """G is a networkx graph."""
        self.G = G
        self.line_graph = nx.line_graph(G)
        self.line_graph_nodes = list(self.line_graph.nodes())
        self.precompute = precompute
        self.verbose = verbose

        if self.precompute:
            self.precompute_dir = precompute_dir
            graph_hash = self.hash_graph(G)
            self.fname = os.path.join(self.precompute_dir, f'grapht_{graph_hash}.npy')
            if os.path.isfile(os.path.isfile(self.fname)):
                self.load_precompute()
            else:
                self.precompute_and_save()


    def precompute_and_save(self):
        """Compute all path lengths and save to disk."""
        L = nx.to_scipy_sparse_matrix(self.line_graph)
        if self.verbose:
            print(f'Precomputing all paths (n = {L.shape[0]})')
            print('Destination', self.fname, flush=True)
        self.all_path_lengths = sp.csgraph.dijkstra(L, directed=False, unweighted=True)
        np.save(self.fname, self.all_path_lengths)


    def load_precompute(self):
        """Load the precompute path lengths matrix."""
        if self.verbose:
            print('Loading ', self.fname)
        self.all_path_lengths = np.load(self.fname)


    def hash_graph(self, G):
        """Return a string hash of a networkx graph."""
        pickle_str = pickle.dumps(G)
        return md5(pickle_str).hexdigest()


    def __call__(self, edge1, edge2):
        """Calculates the linegraph distance between `edge1` and `edge2`."""
        edge1, edge2 = self.sort_edge(edge1), self.sort_edge(edge2)
        if self.precompute:
            i, j = self.edge_index(edge1), self.edge_index(edge2)
            return self.all_path_lengths[i, j]
        else:
            return nx.shortest_path_length(self.line_graph, edge1, edge2)

    def sort_edge(self, edge):
        """Makes sure edges are of the form (u, v) where u <= v."""
        if edge[0] <= edge[1]:
            return edge
        else:
            return (edge[1], edge[0])

    @lru_cache(maxsize=None)
    def edge_index(self, edge):
        """Returns the index of the matrix which corresponds to `edge`."""
        return self.line_graph_nodes.index(edge)


    def pairwise_distances(self, edges):
        """Calculates the linegraph distance between all pairs of edges in `edges`."""
        distances = []
        for i in range(len(edges)):
            for j in range(i+1, len(edges)):
                distances.append(self(edges[i], edges[j]))
        return np.array(distances)

# Cell
def average_gmdegree(G, edges):
    """The average edge degree geometric mean over all edges in `edges`."""
    return np.mean([edge_degree_gm(G, edge) for edge in edges])

def edge_degree_gm(G, edge):
    """For an edge (u, v) with degree du, dv this function returns the geometric mean of du and dv."""
    return np.sqrt(G.degree(edge[0]) * G.degree(edge[1]))