import torch
from matplotlib import pyplot as plt
import numpy as np

def block_mask(sparsity, blocks = None, size = (64, 64), block_size = 8):
    assert((size[0] % block_size) == 0)
    assert((size[1] % block_size) == 0)
    m = torch.zeros(((size[0] // block_size), size[1] // block_size))
    
    if blocks != None:
        for b in blocks:
            m[b[0],b[1]] = 1
    else:
        m = torch.rand(size[0] // block_size, size[1] // block_size)
        m = (m > sparsity) 
        
    m = np.repeat(np.repeat(m, repeats=block_size, axis=0), repeats=block_size, axis = 1)
    return m
    
    
blocks = []

def add_layer(blocks, offset=(0,0), size=(1,1)):
    for i in range(size[0]):
        for j in range(size[1]):
            blocks.append((i+offset[0], j + offset[1]))

layer_sizes = [8,16,32,8]
offset = np.array([0, layer_sizes[0]])
for i in range(len(layer_sizes) - 1):
    size = layer_sizes[i:i+2]
    add_layer(blocks, offset=offset, size=size)
    offset += size

block_size = 1
full_size = sum(layer_sizes) * block_size
bm = block_mask(0.8, blocks, size = (full_size, full_size), block_size = block_size)
conventional_multilayer_matrix = torch.randn(full_size,full_size).abs()
conventional_multilayer_matrix *= bm
plt.imshow(conventional_multilayer_matrix)
def forward_layers(input_state, matrix, iterations, last_layer_size, name):
    state = input_state
    fig, axes = plt.subplots(iterations + 2, 1, figsize=(5,10))

    axes[0].plot(state)
    axes[0].set_ylabel("input")

    for l in range(iterations):    
        state = state.matmul(matrix)
        state = torch.nn.functional.relu(state)
        axes[l+1].plot(state)
        axes[l+1].set_ylabel(name % (l+1))
        
    axes[-1].plot(state[-last_layer_size:])
    axes[-1].set_ylabel(f"final output")

    return state
  
def forward(input, matrix, iterations, last_layer_size, name):
    input_state = torch.zeros(full_size)
    input_state[:input.shape[0]] = input
    output_state = forward_layers(input_state, matrix, iterations, last_layer_size, name)
    output = output_state[-layer_sizes[0] * block_size:]
    return output


# Fill the first dimensions with the input
input = torch.randn(layer_sizes[0] * block_size)

output = forward(input,
                 conventional_multilayer_matrix,
                 len(layer_sizes) - 1, 
                 layer_sizes[-1] * block_size,
                 "layer %d output")

bm2 = block_mask(0.95, blocks = None, size = (full_size, full_size), block_size = block_size)
print(bm2.shape)
arbitrary_matrix = torch.randn((full_size,full_size)).abs()
arbitrary_matrix *= bm2


plt.imshow(arbitrary_matrix)iterations = 1
output = forward(input, arbitrary_matrix, iterations, layer_sizes[0] * block_size, "state %d")iterations = 4
output = forward(input, arbitrary_matrix, iterations, layer_sizes[0] * block_size, "state %d")