#!/usr/bin/env python

# -*- coding: utf-8 -*-



__author__ = 'Stefano Vissicchio <stefano.vissicchio@uclouvain.be>'



#    Copyright (C) 2013-2017 by 

#    Stefano Vissicchio <stefano.vissicchio@uclouvain.be>

#    All rights reserved.

#    BSD license.



import unittest, os

import heapq





#####################

# Compute next hops #

#####################



def computeNextHops(graph,destinations,sources=None):

    nh = dict()

    if sources is None:

        sources=graph.nodes()



    # compute best paths

    best_paths = computeBestPaths(graph,destinations,sources=None)

    

    # extract and store next-hops

    for dest in destinations:

        nh[dest] = dict()

        for k in sources:

            #print best_paths[dest][k]

            nh[dest][k] = set([])

            for path in best_paths[dest][k]:

                penultimate = path.split()

                if len(penultimate) >= 2:

                    penultimate = penultimate[1:2]

                nh[dest][k].add(penultimate.pop())

    

    return nh



def computeBestPaths(graph,destinations,sources=None):

    if sources is None:

        sources=graph.nodes()

    bestpaths = dict()

    for source in sources:

        sps = single_source_dijkstra_ecmp_paths(graph,source)

        for dest in destinations:

            if dest not in bestpaths.keys():

                bestpaths[dest] = dict()

            try:

                bestpaths[dest][source] = sps[dest]

            except Exception:

                bestpaths[dest][source] = []

    return bestpaths





#########################################################################

# Dijkstra algorithm, modified to keep track of all ecmp shortest paths #

#########################################################################



def single_source_dijkstra_ecmp_paths(G,source, weight = 'weight'):

    (length,path)=single_source_dijkstra(G,source, weight = weight)

    return path



def single_source_dijkstra(G,source,target=None,cutoff=None,weight='weight'):

    if source==target: return (0, [source])

    dist = {}  # dictionary of final distances

    paths = {source:[source]}  # dictionary of paths

    seen = {source:0}

    fringe=[] # use heapq with (distance,label) tuples

    heapq.heappush(fringe,(0,source))

    while fringe:

        (d,v)=heapq.heappop(fringe)

        if v in dist: continue # already searched this node.

        dist[v] = d

        if v == target: break

        if G.is_multigraph():

            edata=[]

            for w,keydata in list(G[v].items()):

                minweight=min((dd.get(weight,1)

                               for k,dd in keydata.items()))

                edata.append((w,{weight:minweight}))

        else:

            try:

                edata=iter(G[v].items())

            except Exception:

                return (sys.maxint,[""])

        for w,edgedata in edata:

            vw_dist = dist[v] + edgedata.get(weight,1)

            if cutoff is not None:

                if vw_dist>cutoff:

                    continue

            if w in dist:

                if vw_dist < dist[w]:

                    raise ValueError("Contradictory paths found: negative weights?")

            elif w not in seen or vw_dist < seen[w]:

                seen[w] = vw_dist

                heapq.heappush(fringe,(vw_dist,w))

                paths[w] = []

                for i in range(0,len(paths[v])):

                    paths[w].append(str(paths[v][i]) + " " + str(w))

            elif vw_dist == seen[w]:

                for i in range(0,len(paths[v])):

                    paths[w].append(str(paths[v][i]) + " " + str(w))

    return (dist,paths)





def single_destination_dijkstra_ecmp_paths(G, dest, weight = 'weight'):

    (length,paths)=single_destination_dijkstra(G, dest, weight = weight)

    for n in paths:

        paths[n] = set(paths[n])

    return paths



def single_destination_dijkstra(G,source,target=None,cutoff=None,weight='weight'):

    if source==target: return (0, [source])

    dist = {}  # dictionary of final distances

    paths = {source:[source]}  # dictionary of paths

    seen = {source:0}

    fringe=[] # use heapq with (distance,label) tuples

    heapq.heappush(fringe,(0,source))

    while fringe:

        (d,v)=heapq.heappop(fringe)

        if v in dist: continue # already searched this node.

        dist[v] = d

        if v == target: break

        if G.is_multigraph():

            edata=[]

            for w,keydata in list(G[v].items()):

                minweight=min((dd.get(weight,1)

                               for k,dd in keydata.items()))

                edata.append((w,{weight:minweight}))

        else:

            try:

                edata=iter(G[v].items())

            except Exception:

                return (sys.maxint,[""])

        for w,edgedata in edata:

            vw_dist = dist[v] + edgedata.get(weight,1)

            if cutoff is not None:

                if vw_dist>cutoff:

                    continue

            if w in dist:

                if vw_dist < dist[w]:

                    raise ValueError("Contradictory paths found: negative weights?")

            elif w not in seen or vw_dist < seen[w]:

                seen[w] = vw_dist

                heapq.heappush(fringe,(vw_dist,w))

                paths[w] = []

                for i in range(0,len(paths[v])):

                    paths[w].append(str(w) + " " + str(paths[v][i]))

            elif vw_dist == seen[w]:

                for i in range(0,len(paths[v])):

                    paths[w].append(str(w) + " " + str(paths[v][i]))

    return (dist,paths)





###########################

# Other support functions #

###########################



def flatten_list(l):

    return [item for sublist in l for item in sublist]