import pprint
import random

import matplotlib.pyplot as plt
import networkx as nx

MAT_ROWS = 7
MAT_COLS = 7
MAT_SIZE = MAT_ROWS * MAT_COLS
MAT_FOLD_SIZE = 3
HI = 3
LO = 2

assert MAT_ROWS == MAT_COLS

# Nodes matrix
#####

matrix = []

for i in range(MAT_ROWS):
    matrix.append([j + i * MAT_COLS for j in range(MAT_COLS)])

print("## Matrix\n")
pprint.pprint(matrix)
print("")

# LED index map matrix
#####

led_matrix = []

for i in range(MAT_ROWS):
    offset = max(led_matrix[-1]) + MAT_FOLD_SIZE + 1 if len(led_matrix) else 0
    row = [j + offset for j in range(MAT_COLS)]
    row = list(reversed(row)) if i % 2 == 1 else row
    led_matrix.append(row)

print("## LED index matrix\n")
pprint.pprint(led_matrix)
print("")

# Adjacency cost matrix
#####

adj_cost_mat = [
    [0 for j in range(MAT_SIZE)]
    for i in range(MAT_SIZE)
]

for i in range(MAT_ROWS):
    for j in range(MAT_COLS):
        adj_idx = i + j * 7

        cost_hi = [
            (i - 1, j - 1),
            (i - 1, j + 1),
            (i + 1, j - 1),
            (i + 1, j + 1)
        ]

        cost_lo = [
            (i - 1, j),
            (i, j - 1),
            (i, j),
            (i, j + 1),
            (i + 1, j)
        ]

        for tup in cost_hi + cost_lo:
            try:
                assert tup[0] >= 0 and tup[1] >= 0
                matrix[tup[0]][tup[1]]
                tup_idx = tup[0] + tup[1] * 7
                cost_val = HI if tup in cost_hi else LO
                adj_cost_mat[adj_idx][tup_idx] = cost_val
                adj_cost_mat[tup_idx][adj_idx] = cost_val
            except (IndexError, AssertionError):
                pass

print("## Adj. cost matrix\n")
for row in adj_cost_mat:
    print(row)
print("")

# Adjacency cost matrix (C literal)
#####

print("## Adj. cost matrix (C literal)\n")
print("{")

for i, adj_row in enumerate(adj_cost_mat):
    rstr = "  {}".format(str(adj_row).replace("[", "{").replace("]", "}"))
    print("{},".format(rstr) if i != len(adj_row) - 1 else "{}".format(rstr))

print("}")
print("")

# nx.Graph built from the previous data
#####

graph = nx.Graph()
graph.add_nodes_from(range(MAT_ROWS * MAT_COLS))

for i, adj_cost_row in enumerate(adj_cost_mat):
    for j, cost in enumerate(adj_cost_row):
        if cost > 0:
            graph.add_edge(i, j, weight=cost)

nx.draw_spectral(graph)
plt.show()

assert list(list(nx.connected_components(graph))[0]) == range(MAT_SIZE)

# Tests
#####

print("## Shortest paths\n")

test_paths = {
    (0, 3): [0, 1, 2, 3],
    (0, 21): [0, 7, 14, 21],
    (0, 24): [0, 8, 16, 24],
    (3, 6): [3, 4, 5, 6],
    (3, 21): [3, 9, 15, 21],
    (3, 24): [3, 10, 17, 24],
    (3, 27): [3, 11, 19, 27],
    (6, 24): [6, 12, 18, 24],
    (6, 27): [6, 13, 20, 27],
    (21, 24): [21, 22, 23, 24],
    (21, 42): [21, 28, 35, 42],
    (24, 27): [24, 25, 26, 27],
    (24, 42): [24, 30, 36, 42],
    (24, 45): [24, 31, 38, 45],
    (24, 48): [24, 32, 40, 48],
    (27, 45): [27, 33, 39, 45],
    (27, 48): [27, 34, 41, 48],
    (42, 45): [42, 43, 44, 45]
}

for (source, target), expected_path in test_paths.items():
    res_path = nx.dijkstra_path(graph, source=source, target=target)
    assert res_path == expected_path
    print("Shortest path {0:0=2d} -> {1:0=2d}: {2}".format(
        source, target,
        ["{0:0=2d}".format(item) for item in res_path]))

print("")

pprint.pprint([
    ["{0:0=2d}".format(item) for item in matrix[i]]
    for i in range(len(matrix))
])