from jinja2 import Template
from collections import namedtuple
import subprocess
import cmath
import math
import random

shell_radius = 50
shell_1_neurons = 6
shell_2_neurons = 12
alpha= shell_radius
alpha += 2


Node = namedtuple('Node', ['name', 'x', 'y'])
Edge = namedtuple('Edge', ['start', 'end'])

def distance(n1, n2):
    return math.sqrt((n1.x -n2.x)**2+(n1.y -n2.y)**2)



def exp_distance(n1, n2):
    return math.exp(-distance(n1,n2)/alpha)


t = Template("""digraph {
 node [shape="point", width="0.3", color= "gray"]\
 {% for node in nodes %}
 {{ node.name }} [pos="{{ node.x }},{{ node.y }}"{% if node.name == 0 %} , color="black"{% endif %}];\
 {% endfor %}
 {% for edge in edges %}\
 {{ edge.start }} -> {{ edge.end }} {% if edge.start != 0 %}[color="gray"]{% else %}[penwidth="3"]{% endif %};
 {% endfor %}\
}""")


def shell(shell_nr, nr_neurons, offset_angle=0):
    for k in range(nr_neurons):
        pos = cmath.rect(shell_nr*shell_radius, k/nr_neurons*2*cmath.pi+offset_angle)
        yield pos.real, pos.imag
nodes = ( [Node(k+1, *pos) for k, pos in enumerate(shell(1,shell_1_neurons))]
         + [Node(k+1+shell_1_neurons, *pos) for k, pos in enumerate(shell(2,shell_2_neurons, offset_angle=.05))]
         + [Node(0,0,0)])
edges = [Edge(start.name, end.name) for start in nodes for end in nodes
         if random.random() < exp_distance(start,end)  and start != end]
s = t.render(nodes=nodes, edges=edges)


neato_file = '/home/emmanuel/tmp/exp_net.neato'
with open(neato_file, "w") as f:
    f.write(s)
print(s)
for extension in ['pdf', 'png']:
    res = subprocess.call(['neato', '-n', '-T' + extension, '-o',
                           '/home/emmanuel/tmp/exp_net.' + extension, neato_file])


