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

write_files = False

n = 0
nn_total_levels = 0

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

    return i % round(nn_total_levels ** 0.5)

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

    return i // round(nn_total_levels ** 0.5)

def xy(i):
    return x(i), y(i)

def i(x, y):
    return x * round(nn_total_levels ** 0.5) + y

def nn(*args):
    if len(args) == 2:
        # legacy_mode
        cluster_no = args[0]
        node_name = args[1]

        if not nn_total_levels:
            return str(cluster_no) + '-' + node_name

        return str(cluster_no % nn_total_levels) + '-' + node_name
    else:
        # modern mode
        x = args[0]
        y = args[1]
        l = args[2]

        return str(i(x, y)) + '-' + l

def join_nodes(g, a, b, target_degree = 6):
    if a == b:
        return

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

        if g.degree[a] > target_degree or g.degree[b] > target_degree:
            g.remove_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 = ['ll', 'lu', 'lr', 'ld', 'rl', 'ru', 'rr', 'rd']
    edges = [('ll', 'lu'), ('lu', 'lr'), ('lr', 'ld'), ('ld', 'll'),
             ('rl', 'ru'), ('ru', 'rr'), ('rr', 'rd'), ('rd', 'rl'),
             ('lu', 'ru'), ('ld', 'rd'), ('ll', 'lr'), ('rl', 'rr')]

    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(l_n):
    class P():
        def __init__(self, x, y):
            self.x = x
            self.y = y

    g = build_cluster(0)

    global n, nn_total_levels
    n = l_n
    nn_total_levels = l_n ** 2

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

    for x in range(n):
        for y in range(n):
            # regular grid
            if x < n - 1:
                join_nodes(g, nn(x, y, 'lr'), nn(x + 1, y, 'rl'))
                join_nodes(g, nn(x, y, 'rr'), nn(x + 1, y, 'll'))

            if y < n - 1:
                join_nodes(g, nn(x, y, 'lu'), nn(x, y + 1, 'rd'))
                join_nodes(g, nn(x, y, 'ru'), nn(x, y + 1, 'ld'))

    # corner cross
    join_nodes(g, nn(0, 0, 'lu'), nn(n // 2, n // 2, 'lu'))
    join_nodes(g, nn(n // 2, n // 2, 'ld'), nn(n - 1, n - 1, 'ld'))

    join_nodes(g, nn(n - 1, 0, 'ru'), nn(n // 2, n // 2, 'ru'))
    join_nodes(g, nn(n // 2, n // 2, 'rd'), nn(0, n - 1, 'rd'))

    # grid fill
    i = 0
    while True:
        start = P(n // 2 ** (1 + i) - 1 + n % 2, 0)
        step = n // 2 ** (1 + i) - 1 + n % 2

        if step == 0:
            #print(f'Breaking at {i}')
            break

        while True:
        #for _ in range(2*i+1):
            y = start.y
            while y + step < n:
                #print('---', i, start.x, y, step, y + step, n)
                if step == 0:
                    exit(1)
                join_nodes(g, nn(start.x, y, 'll'), nn(start.x - step, y + step, 'll'))
                join_nodes(g, nn(start.x, y, 'rr'), nn(start.x + step, y + step, 'rr'))
                y += step * 2
                join_nodes(g, nn(start.x, y, 'lr'), nn(start.x - step, y - step, 'lr'))
                join_nodes(g, nn(start.x, y, 'rl'), nn(start.x + step, y - step, 'rl'))

            start.x += step * 2

            if start.x + step >= n:
                break

        i += 1
        #visualize(g)

    # inline skips
    for x in range(n):
        for y in range(n):
            join_nodes(g, nn(x, y, 'ld'), nn((x + n//2) % n, y, 'ld'))
            join_nodes(g, nn(x, y, 'rd'), nn((x + n//3) % n, y, 'rd'))

            join_nodes(g, nn(x, y, 'lu'), nn(x, (y + n//5) % n, 'lu'))
            join_nodes(g, nn(x, y, 'ru'), nn(x, (y + n//6) % n, 'ru'))

            join_nodes(g, nn(x, y, 'll'), nn((x + round(n//3)) % n, (y + round(n//3)) % n, 'll'))
            join_nodes(g, nn(x, y, 'rr'), nn((x + round(n//2)) % n, (y + round(n//2)) % n, 'rr'))

            join_nodes(g, nn(x, y, 'lr'), nn((x + round(n//5)) % n, (y - round(n//5)) % n, 'lr'))
            join_nodes(g, nn(x, y, 'rl'), nn((x + round(n//6)) % n, (y - round(n//6)) % n, 'rl'))

    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 visualize(g):
    def draw_edges_between(a, b, x_min = 0, x_max = False, y_min = 0, y_max = False, color = 'black', width = 1.0, debug = False, direct = False, fmode = True, fmask = [1, 1, 1, 1], fsplit = 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])]
        edgelist_filtered = []

        for i in edgelist:
            pa, pb = i

            constraints = [True] * 4

            if (x_max.__class__ != bool) and (abs(x(pa) - x(pb)) > x_max):
                constraints[1] = False

            if (y_max.__class__ != bool) and (abs(y(pa) - y(pb)) > y_max):
                constraints[3] = False

            if abs(x(pa) - x(pb)) < x_min:
                constraints[0] = False

            if abs(y(pa) - y(pb)) < y_min:
                constraints[2] = False

            #filt = [v for k, v in enumerate(constraints) if fmask[k]]
            if fsplit:
                if any([all(constraints[:2]), all(constraints[2:])]):
                    edgelist_filtered.append(i)
            else:
                if all(constraints):
                    edgelist_filtered.append(i)

            #if fmode:
            #    if all(filt):
            #        edgelist_filtered.append(i)
            #else:
            #    if any(filt):
            #        edgelist_filtered.append(i)

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

    pos = {
        'll': (-3, 0),
        'lu': (-2, 1),
        'lr': (-1, 0),
        'ld': (-2, -1),
        'rl': (1, 0),
        'ru': (2, 1),
        'rr': (3, 0),
        'rd': (2, -1)
    }

    exp_pos = {}

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

    nx.draw_networkx_nodes(g, exp_pos, node_size = 100)
    nx.draw_networkx_labels(g, exp_pos, font_size = 3)

    # cluster
    draw_edges_between('', '', x_max = 0, y_max = 0, color = '#888888', width = 0.5, direct = True)

    # regular grid
    draw_edges_between('', '', x_min = 1, x_max = 1, y_min = 1, y_max = 1, color = '#FF0000', width = 2.0, fsplit = True)

    # anything (debug)
    draw_edges_between('', '', x_min = 1, y_min = 1, color = '#00FF00', width = 2.0, direct = True)

    plt.show()

def main():
    global write_files

    write_files = False

    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() > 2000:
            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(9)
    visualize(g)

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