#!/usr/bin/env python

#

# NN Plot

# Tim O'Shea (c) 2016

#

#

#

import matplotlib.pyplot as plt

import matplotlib.patches as patches

import matplotlib.lines as lines

# load the config def

import sys,json

fn = sys.argv[1]

print fn

cfg = json.loads(open(fn).read())

io = cfg["io"]

layers = cfg["layers"]

# ...

maxy = max(map(lambda x: x[3], layers))

print maxy

fig = plt.figure()

ax = fig.add_subplot(111)

for i,l in enumerate(layers):

ll=0.25 + i

w = 0.5

h = layers[i][3]

lr=(maxy-h)/2.0

print (ll,lr,w,h)

ax.add_patch(

patches.Rectangle( (ll, lr), w, h, fill=False)

)

ax.annotate(l[0], xy=(ll+w/2.0, lr+h/2.0),

ha='center', va='center',

rotation=90)

if(i < len(layers)-1):

h_n = layers[i+1][3]

lr_n=(maxy-h_n)/2.0

i_x = ll+w

o_x = ll+1

for j in range(0,layers[i][2]):

for k in range(0,layers[i+1][2]):

j_rel = (j*1.0/(layers[i][2]-1))

k_rel = (k*1.0/(layers[i+1][2]-1))

if(l[1]==None or abs(j_rel - k_rel)<l[1]):

i_ht = lr+h*j_rel

o_ht = lr_n+h_n*k_rel

ax.add_line(lines.Line2D([i_x,o_x], [i_ht,o_ht], color='black', linestyle='solid'))

if(i == 0):

for j in range(0,l[2]):

j_rel = (j*1.0/(l[2]-1))

ax.add_patch( patches.Circle([i_x-w,lr+h*j_rel], 0.025, color='blue', ec="none") )

if(i == len(layers)-2):

for j in range(0,layers[i+1][2]):

j_rel = (j*1.0/(layers[i+1][2]-1))

ax.add_patch( patches.Circle([o_x+w,lr_n+h_n*j_rel], 0.025, color='blue', ec="none") )

ax.annotate("Inputs (2x88)", xy=(0,maxy/2.0),

ha='center',va='center',

rotation=90)

ax.annotate("Outputs (2x88)", xy=(len(layers),maxy/2.0),

ha='center',va='center',

rotation=90)

ymarg = 0.05

xmarg = 0.05

plt.ylim(-ymarg, maxy+0.5+ymarg)

plt.xlim(-xmarg,len(layers)+xmarg)

plt.axis('off')

plt.savefig('test.png', bbix_inches='tight')

plt.show()