import sys
sys.path.append('../../../')
import easygraph as eg

__all__ = [
    'effective_size',
    'efficiency',
    'constraint',
    'hierarchy'
]


def mutual_weight(G, u, v, weight=None):
    try:
        a_uv = G[u][v].get(weight, 1)
    except KeyError:
        a_uv = 0
    try:
        a_vu = G[v][u].get(weight, 1)
    except KeyError:
        a_vu = 0
    return a_uv + a_vu

sum_nmw_rec = {}
max_nmw_rec = {}

def normalized_mutual_weight(G, u, v, norm=sum, weight=None):
    if norm == sum:
        try:
            # res = sum_nmw_rec[(u, v)]
            # print('yes')
            return sum_nmw_rec[(u, v)]
        except KeyError:
            scale = norm(mutual_weight(G, u, w, weight)
                    for w in G.all_neighbors(u))
            nmw = 0 if scale == 0 else mutual_weight(G, u, v, weight) / scale
            sum_nmw_rec[(u,v)] = nmw
            return nmw
    elif norm == max:
        try:
            return max_nmw_rec[(u, v)]
        except KeyError:
            scale = norm(mutual_weight(G, u, w, weight)
                    for w in G.all_neighbors(u))
            nmw = 0 if scale == 0 else mutual_weight(G, u, v, weight) / scale
            max_nmw_rec[(u,v)] = nmw
            return nmw


def effective_size(G, nodes=None, weight=None):
    sum_nmw_rec.clear()
    max_nmw_rec.clear()
    def redundancy(G, u, v, weight=None):
        nmw = normalized_mutual_weight
        r = sum(nmw(G, u, w, weight=weight) * nmw(G, v, w, norm=max, weight=weight)
                for w in set(G.all_neighbors(u)))
        return 1 - r
    effective_size = {}
    if nodes is None:
        nodes = G
    # Use Borgatti's simplified formula for unweighted and undirected graphs
    if not G.is_directed() and weight is None:
        for v in nodes:
            # Effective size is not defined for isolated nodes
            if len(G[v]) == 0:
                effective_size[v] = float('nan')
                continue
            E = G.ego_subgraph(v)
            effective_size[v] = len(E) - 1 - (2 * E.size()) / (len(E) - 1)
    else:
        for v in nodes:
            # Effective size is not defined for isolated nodes
            if len(G[v]) == 0:
                effective_size[v] = float('nan')
                continue
            effective_size[v] = sum(redundancy(G, v, u, weight)
                                    for u in set(G.all_neighbors(v)))
    return effective_size


def efficiency(G, nodes=None, weight=None):
    e_size = effective_size(G=G, nodes=nodes, weight=weight)
    degree = G.degree(weihgt=weight)
    efficiency = {n: v / degree[n] for n, v in e_size.items()}


def constraint(G, nodes=None, weight=None, n_workers=None):
    sum_nmw_rec.clear()
    max_nmw_rec.clear()
    local_constraint_rec.clear()
    if nodes is None:
        nodes = G.nodes
    constraint = {}
    def compute_constraint_of_v(v):
        neighbors_of_v = set(G.all_neighbors(v))
        if len(neighbors_of_v) == 0:
            constraint_of_v = float('nan')
        else:
            constraint_of_v = sum(local_constraint(G, v, n, weight)
                                for n in neighbors_of_v)
        return v, constraint_of_v

    if n_workers is None:
        constraint_results = []
        for v in nodes:
            constraint_results.append(compute_constraint_of_v(v))
    else:
        from joblib import Parallel, delayed
        constraint_results = Parallel(n_jobs=n_workers)(delayed(compute_constraint_of_v)(v) for v in nodes)
    for v, constraint_of_v in constraint_results:
        constraint[v] = constraint_of_v
    return constraint

local_constraint_rec = {}
def local_constraint(G, u, v, weight=None):
    try:
        return local_constraint_rec[(u, v)]
    except KeyError:
        nmw = normalized_mutual_weight
        direct = nmw(G, u, v, weight=weight)
        indirect = sum(nmw(G, u, w, weight=weight) * nmw(G, w, v, weight=weight)
                    for w in set(G.all_neighbors(u)))
        result = (direct + indirect) ** 2
        local_constraint_rec[(u, v)] = result
        return result

def hierarchy(G,nodes=None,weight=None):
    if nodes is None:
        nodes = G.nodes
    hierarchy = {}
    con=constraint(G)
    for v in G.nodes:
        E = G.ego_subgraph(v)
        n=len(E)-1
        C=0
        c={}
        neighbors_of_v = set(G.all_neighbors(v))
        for w in neighbors_of_v:
            C+=local_constraint(G,v,w)
            c[w]=local_constraint(G,v,w)
        if n>1:
            hierarchy[v]=sum(c[w]/C*n*math.log(c[w]/C*n)/(n*math.log(n)) for w in neighbors_of_v)
    return hierarchy

