import matplotlib.pyplot as plt
import networkx as nx
import pandas as pd

write_files = False

nn_total_levels = 0

def su(s):
    return 2 ** (s + 1) - 1

def lvl(i):
    # backwards compatibility
    if i.__class__ == str:
        i = int(i.split('-')[0])

    y = 0
    while i - su(y) >= 0:
        y += 1

    return y

def ps(i):
    # backwards compatibility
    if i.__class__ == str:
        i = int(i.split('-')[0])

    y = lvl(i)
    return i - su(y-1)

def neigh(i, o):
    y = lvl(i)
    a = 2 ** y
    i = i - su(y - 1)
    return su(y - 1) + (i + o) % a

def nn(cluster_no, node_name):
    if not nn_total_levels:
        return str(cluster_no) + '-' + node_name
    return str(cluster_no % nn_total_levels) + '-' + node_name

def join_nodes(g, a, b, target_degree = 6):
    if not (g.degree[a] == target_degree or g.degree[b] == target_degree):
        g.add_edge(a, b)

def start_file(p):
    if not write_files:
        return

    index = ['i', 'N', 'E', 'S', 'D', 'D_avg', 'T', 'C']
    open(p, 'w').write(','.join(index) + '\n')

def write_file(p, data):
    if not write_files:
        return

    index = ['i', 'N', 'E', 'S', 'D', 'D_avg', 'T', 'C']
    open(p, 'a').write(','.join([str(data[i]) for i in index]) + '\n')

def build_cluster(n):
    g = nx.Graph()
    nodes = ['ls', 'lu', 'lm', 'ld', 'm', 'rs', 'ru', 'rm', 'rd']
    edges = [('ls', 'lu'), ('lu',  'm'), ( 'm', 'ld'), ('ld', 'ls'),
             ('ls', 'lm'), ('lu', 'lm'), ( 'm', 'lm'), ('ld', 'lm'),
             ('rs', 'ru'), ('ru',  'm'), ( 'm', 'rd'), ('rd', 'rs'),
             ('rs', 'rm'), ('ru', 'rm'), ( 'm', 'rm'), ('rd', 'rm'),
             ('lu', 'ru'), ('ld', 'rd')]

    g.add_nodes_from([nn(n, x) for x in nodes])
    g.add_edges_from([tuple(nn(n, y) for y in x) for x in edges])

    return g

def total_clusters(s):
    return 2 ** s - 1

def build_network(n):
    throw_dist = 4

    g = build_cluster(0)

    global nn_total_levels
    nn_total_levels = total_clusters(n)

    for i in range(1, nn_total_levels):
        c = build_cluster(i)
        g = nx.union(g, c)

    for l_lvl in range(n):
        for c in range(2 ** l_lvl):
            i = total_clusters(l_lvl) + c

            # child nodes
            cl = total_clusters(l_lvl + 1) + 2 * c
            cr = total_clusters(l_lvl + 1) + 2 * c + 1

            # child nodes ^ 2 crossed
            ccl = total_clusters(l_lvl + 2) + 4 * c + 1
            ccr = total_clusters(l_lvl + 2) + 4 * c + 2

            # child nodes ^ 3 straight
            cccl = total_clusters(l_lvl + 3) + 8 * c
            cccr = total_clusters(l_lvl + 3) + 8 * c + 7

            # child nodes ^ 4 crossed
            ccccl = total_clusters(l_lvl + 4) + 16 * c + 5
            ccccr = total_clusters(l_lvl + 4) + 16 * c + 13

            # child nodes ^ 5 crimped
            cccccl = total_clusters(l_lvl + 4) + 32 * c + 9
            cccccr = total_clusters(l_lvl + 4) + 32 * c + 25

            # regular set 1
            if l_lvl < n - 1:
                join_nodes(g, nn(i, 'lm'), nn(cl, 'rm'))
                join_nodes(g, nn(i, 'rm'), nn(cr, 'lm'))

            # irregular: grid ???
            if l_lvl < n - 2:
                join_nodes(g, nn(i, 'lu'), nn(ccl, 'rd'))
                join_nodes(g, nn(i, 'ru'), nn(ccr, 'ld'))

            # irregular: grid ?????
            if l_lvl < n - 3:
                join_nodes(g, nn(i, 'ld'), nn(cccl, 'ru'))
                join_nodes(g, nn(i, 'rd'), nn(cccr, 'lu'))

            # irregulars: inline skips
            #if lvl(i) == lvl(i + l_lvl) and l_lvl:
            join_nodes(g, nn(i, 'ru'), nn(neigh(i, +l_lvl), 'rd'))
            #if lvl(i) == lvl(i - l_lvl) and l_lvl:
            join_nodes(g, nn(i, 'lu'), nn(neigh(i, -l_lvl * 2), 'ld'))

            join_nodes(g, nn(i, 'rs'), nn(neigh(i, +l_lvl * 3), 'ls'))
            join_nodes(g, nn(i, 'ls'), nn(neigh(i, -l_lvl * 5), 'rs'))

            # irregular: inline ring
            join_nodes(g, nn(i, 'rs'), nn(neigh(i, +1), 'ls'))
            join_nodes(g, nn(i, 'ls'), nn(neigh(i, -1), 'rs'))

            #if l_lvl >= 4:
                #join_nodes(g, nn(i, 'rs'), nn(neigh(i, +12), 'ls'), 6)
                #join_nodes(g, nn(i, 'ls'), nn(neigh(i, -12), 'rs'), 6)

    return g

def analyze(G: nx.Graph) -> dict:
    N = G.number_of_nodes()
    E = G.number_of_edges()
    degs = G.degree()
    s_vals = [s[1] for s in degs]
    S = max(s_vals)
    D = nx.diameter(G)
    D_average = nx.average_shortest_path_length(G)
    Q = N*D_average/E
    C = N*D*S
    return {
            'N': N,
            'E': E,
            'S': S,
            'D': D,
            'D_avg': round(D_average, 4),
            'T': round(Q, 4),
            'C': C,
    }

def get_cluster_base_coords(cluster_i):
    if cluster_i == 0:
        return 0, 0

    y = lvl(cluster_i)

    s = 2 ** (2-y)
    r = 0

    for i in range(y):
        r += 2 ** -i

    x = r + (cluster_i - su(y) + 1) * s

    return x, float(y)

def visualize(g):
    def draw_edges_between(a, b, lvl_min = 0, lvl_max = False, ps_min = 0, ps_max = False, color = 'black', width = 1.0, debug = False, direct = False):
        edgelist = [e for e in g.edges() if all([True if (a in i) or (b in i) else False for i in e])]

        if lvl_max.__class__ == bool:
            edgelist = [(pa, pb) for pa, pb in edgelist if not (abs(lvl(pa) - lvl(pb)) < lvl_min or nn_total_levels - abs(lvl(pa) - lvl(pb)) < lvl_min)]
        else:
            edgelist = [(pa, pb) for pa, pb in edgelist if not ((abs(lvl(pa) - lvl(pb)) < lvl_min or nn_total_levels - abs(lvl(pa) - lvl(pb)) < lvl_min) or (abs(lvl(pa) - lvl(pb)) > lvl_max and nn_total_levels - abs(lvl(pa) - lvl(pb)) > lvl_max))]

        if ps_max.__class__ == bool:
            edgelist = [(pa, pb) for pa, pb in edgelist if not (abs(ps(pa) - ps(pb)) < ps_min or nn_total_levels - abs(ps(pa) - ps(pb)) < ps_min)]
        else:
            edgelist = [(pa, pb) for pa, pb in edgelist if not ((abs(ps(pa) - ps(pb)) < ps_min or nn_total_levels - abs(ps(pa) - ps(pb)) < ps_min) or (abs(ps(pa) - ps(pb)) > ps_max and nn_total_levels - abs(ps(pa) - ps(pb)) > ps_max))]

        if direct:
            nx.draw_networkx_edges(g, exp_pos, edgelist = edgelist, width = width, edge_color = color)
        else:
            nx.draw_networkx_edges(g, exp_pos, edgelist = edgelist, connectionstyle='arc3,rad=0.2', arrows=True, width = width, edge_color = color)

    pos = {
        'ls': (-2, 0),
        'lu': (-1, 1),
        'ld': (-1, -1),
        'lm': (-1, 0),
        'm' : (0, 0),
        'rm': (1, 0),
        'ru': (1, 1),
        'rd': (1, -1),
        'rs': (2, 0)
    }

    exp_pos = {}

    for i in range(nn_total_levels):
        for j, v in pos.items():
            from math import log
            bx, by = get_cluster_base_coords(i)
            scaler = (log(nn_total_levels, 2) + 1) * 2
            exp_pos[nn(i, j)] = (v[0] + bx * scaler, v[1] + by * 3.5)

    nx.draw_networkx_nodes(g, exp_pos, node_size = 600)
    nx.draw_networkx_labels(g, exp_pos, font_size = 8)

    # cluster
    draw_edges_between('', '', lvl_max = 0, ps_max = 0, color = '#888888', width = 0.5, direct = True)

    # basic tree
    draw_edges_between('', '', lvl_min = 1, lvl_max = 1, color = '#FF0000', width = 2.0)

    # inline ring links
    draw_edges_between('ls', 'rs', lvl_max = 0, ps_min = 1, color = '#00FF00', width = 2.0)
    draw_edges_between('lu', 'ld', lvl_max = 0, ps_min = 1, color = '#0000FF', width = 2.0)
    draw_edges_between('ru', 'rd', lvl_max = 0, ps_min = 1, color = '#AA00AA', width = 2.0)

    # grid links
    draw_edges_between('', '', lvl_min = 2, color = '#AAAA00', width = 2.0)

    plt.show()

def main():
    global write_files

    write_files = True

    i = 0

    start_file('data.csv')
    d = None

    while True:
        i += 1
        g = build_network(i)

        r = analyze(g)
        r.update({"i": i})

        print(r)
        write_file('data.csv', r)

        if g.number_of_nodes() > 3000:
            break

def test():
    g = build_network(15)
    print(analyze(g))
    print([i for i, v in g.degree if v > 6])
    print(g.adj['14-ld'])
    nx.draw(g, with_labels = True)
    plt.show()

def test2():
    g = build_network(144)
    print(analyze(g))
    print(set([v for i, v in g.degree()]))

def test3():
    g = build_network(3)
    visualize(g)

if __name__=="__main__":
    #test3()
    main()