#! /usr/bin/env python
#
# layers.py
#
# This file is part of NEST.
#
# Copyright (C) 2004 The NEST Initiative
#
# NEST is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# NEST is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with NEST.  If not, see <http://www.gnu.org/licenses/>.
# Run as python layers.py > layers.log

import matplotlib.pyplot as plt
import nest
import numpy as np

def beautify_layer(l, fig=plt.gcf(), xlabel=None, ylabel=None,
                   xlim=None, ylim=None, xticks=None, yticks=None, dx=0, dy=0):
    """Assume either x and ylims/ticks given or none"""

    top = nest.GetStatus(l)[0]['topology']
    ctr = top['center']
    ext = top['extent']

    if xticks == None:
        if top.has_key('rows'):
            dx = float(ext[0]) / top['columns']
            dy = float(ext[1]) / top['rows']
            xticks = ctr[0]-ext[0]/2.+dx/2. + dx*np.arange(top['columns'])
            yticks = ctr[1]-ext[1]/2.+dy/2. + dy*np.arange(top['rows'])

    if xlim == None:
        xlim = [ctr[0]-ext[0]/2.-dx/2., ctr[0]+ext[0]/2.+dx/2.] # extra space so extent is visible
        ylim = [ctr[1]-ext[1]/2.-dy/2., ctr[1]+ext[1]/2.+dy/2.]
    else:
        ext = [xlim[1]-xlim[0], ylim[1]-ylim[0]]
        
    ax = fig.gca()
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    ax.set_aspect('equal', 'box')
    ax.set_xticks(xticks)
    ax.set_yticks(yticks)
    ax.grid(True)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    return

# --------------------------------------------------

nest.ResetKernel()

#{ layer1 #}
import nest.topology as tp
l = tp.CreateLayer({'rows'    : 5,
                    'columns' : 5,
                    'elements': 'iaf_neuron'})
#{ end #}

fig = tp.PlotLayer(l, nodesize=50)
beautify_layer(l, fig, xlabel='x-axis (columns)', ylabel='y-axis (rows)')
ax = fig.gca()
tx=[]
for r in range(5):
    tx.append(ax.text( 0.65, 0.4-r*0.2, str(r),
                       horizontalalignment='center',verticalalignment='center'))
    tx.append(ax.text(-0.4+r*0.2, 0.65, str(r),
                      horizontalalignment='center',verticalalignment='center'))

# For bbox_extra_artists, see http://old.nabble.com/bbox_inches%3D'tight'-issue-with-text-outside-axes-td28767991.html
plt.savefig('../user_manual_figures/layer1.png', bbox_inches='tight', bbox_extra_artists=tx)

print "#{ layer1s.log #}"
#{ layer1s #}
print nest.GetStatus(l)[0]['topology']
#{ end #}
print "#{ end.log #}"


print "#{ layer1p.log #}"
#{ layer1p #}
nest.PrintNetwork(depth=2)
#{ end #}
print "#{ end.log #}"

# --------------------------------------------------

nest.ResetKernel()

#{ layer2 #}
l = tp.CreateLayer({'rows'    : 5,
                    'columns' : 5,
                    'extent'  : [2.0, 0.5],
                    'elements': 'iaf_neuron'})
#{ end #}

fig = tp.PlotLayer(l, nodesize=50)
beautify_layer(l, fig, xlabel='x-axis (columns)', ylabel='y-axis (rows)')
ax = fig.gca()
tx = []

for r in range(5):
    tx.append(fig.gca().text( 1.25, 0.2-r*0.1, str(r),
                              horizontalalignment='center',verticalalignment='center'))
    tx.append(fig.gca().text(-0.8+r*0.4, 0.35, str(r),
                             horizontalalignment='center',verticalalignment='center'))

# for bbox_extra_artists, see http://old.nabble.com/bbox_inches%3D'tight'-issue-with-text-outside-axes-td28767991.html
plt.savefig('../user_manual_figures/layer2.png', bbox_inches='tight', bbox_extra_artists=tx)

# --------------------------------------------------

nest.ResetKernel()

#{ layer3 #}
l1 = tp.CreateLayer({'rows': 5, 'columns' : 5, 'elements': 'iaf_neuron'})
l2 = tp.CreateLayer({'rows': 5, 'columns' : 5, 'elements': 'iaf_neuron',
                     'center': [-1.,1.]})
l3 = tp.CreateLayer({'rows': 5, 'columns' : 5, 'elements': 'iaf_neuron',
                     'center': [1.5,0.5]})
#{ end #}

fig = tp.PlotLayer(l1, nodesize=50)
tp.PlotLayer(l2, nodesize=50, nodecolor='g', fig=fig)
tp.PlotLayer(l3, nodesize=50, nodecolor='r', fig=fig)
beautify_layer(l1, fig, xlabel='x-axis (columns)', ylabel='y-axis (rows)',
               xlim=[-1.6, 2.1], ylim=[-0.6,1.6],
               xticks=np.arange(-1.4,2.05,0.2),
               yticks=np.arange(-0.4,1.45,0.2))

plt.savefig('../user_manual_figures/layer3.png', bbox_inches='tight')

# --------------------------------------------------

nest.ResetKernel()

#{ layer3a #}
nc, nr = 5, 3
d = 0.1
l = tp.CreateLayer({'columns': nc, 'rows': nr, 'elements': 'iaf_neuron',
                    'extent': [nc*d, nr*d], 'center': [nc*d/2., 0.]})

#{ end #}

fig = tp.PlotLayer(l, nodesize=100)
plt.plot(0,0,'x',markersize=20, c='k', mew=3)
plt.plot(nc*d/2,0,'o',markersize=20, c='k', mew=3,mfc='none',
         zorder=100)
beautify_layer(l, fig, xlabel='x-axis (columns)', ylabel='y-axis (rows)',
               xticks=np.arange(0.,0.501,0.05),
               yticks=np.arange(-0.15,0.151,0.05),
               xlim=[-0.05, 0.55], ylim=[-0.2,0.2])

plt.savefig('../user_manual_figures/layer3a.png', bbox_inches='tight')


# --------------------------------------------------

nest.ResetKernel()

#{ layer4 #}
import numpy as np
pos = [[np.random.uniform(-0.5,0.5),np.random.uniform(-0.5,0.5)]
       for j in xrange(50)]
l = tp.CreateLayer({'positions': pos,
                    'elements': 'iaf_neuron'})
#{ end #}

fig = tp.PlotLayer(l, nodesize=50)
beautify_layer(l, fig, xlabel='x-axis (columns)', ylabel='y-axis (rows)',
               xlim=[-0.55,0.55],ylim=[-0.55,0.55],
               xticks=[-0.5,0.,0.5],yticks=[-0.5,0.,0.5])

plt.savefig('../user_manual_figures/layer4.png', bbox_inches='tight')

# --------------------------------------------------

nest.ResetKernel()

#{ layer4_3d #}
import numpy as np
pos = [[np.random.uniform(-0.5,0.5),np.random.uniform(-0.5,0.5),
        np.random.uniform(-0.5,0.5)] for j in xrange(200)]
l = tp.CreateLayer({'positions': pos,
                    'elements': 'iaf_neuron'})
#{ end #}

fig = tp.PlotLayer(l, nodesize=50)

plt.savefig('../user_manual_figures/layer4_3d.png', bbox_inches='tight')

# --------------------------------------------------

nest.ResetKernel()

#{ player #}
lp = tp.CreateLayer({'rows': 1, 'columns': 5, 'extent': [5., 1.],
                     'elements': 'iaf_neuron',
                     'edge_wrap': True})
#{ end #}

# fake plot with layer on line and circle
import matplotlib.cm as cm
clist = [(0,0,1),(0.35,0,1),(0.6,0,1),(0.8,0,1),(1.0,0,1)]
fig = plt.figure()
ax1 = fig.add_subplot(221)
ax1.plot([0.5, 5.5], [0,0], 'k-', lw=2)
ax1.scatter(range(1,6), [0]*5, s=200, c=clist) #=range(1,6), cmap=cm.jet)
#c=np.array([(w,w,1) for w in np.arange(0,0.95,0.18)]))
ax1.set_xlim([0,6])
ax1.set_ylim([-0.5, 1.25])
ax1.set_aspect('equal', 'box')
ax1.set_xticks([])
ax1.set_yticks([])
for j in range(1,6):
    ax1.text(j, 0.5, str('(%d,0)'%(j-3)),
             horizontalalignment='center',verticalalignment='bottom')

ax1a = fig.add_subplot(223)
ax1a.plot([0.5, 5.5], [0,0], 'k-', lw=2)
ax1a.scatter(range(1,6), [0]*5, s=200, c=[clist[0],clist[1],clist[2],clist[2],clist[1]])
#c=np.array([(w,w,1) for w in np.arange(0,0.95,0.18)]))
ax1a.set_xlim([0,6])
ax1a.set_ylim([-0.5, 1.25])
ax1a.set_aspect('equal', 'box')
ax1a.set_xticks([])
ax1a.set_yticks([])
for j in range(1,6):
    ax1a.text(j, 0.5, str('(%d,0)'%(j-3)),
              horizontalalignment='center',verticalalignment='bottom')


ax2 = fig.add_subplot(122)
phic = np.arange(0., 2*np.pi+0.5, 0.1)
r = 5. / (2*np.pi)
ax2.plot(r*np.cos(phic), r*np.sin(phic), 'k-', lw=2)
phin = np.arange(0., 4.1, 1.) * 2*np.pi/5
ax2.scatter(r*np.sin(phin), r*np.cos(phin), s=200, c=[clist[0],clist[1],clist[2],clist[2],clist[1]])
            #[1,2,3,3,2], cmap=cm.jet)
#            c=np.array([(w,w,1) for w in [0., 0.18, 0.36, 0.36, 0.18]]))
ax2.set_xlim([-1.3,1.3])
ax2.set_ylim([-1.2,1.2])
ax2.set_aspect('equal', 'box')
ax2.set_xticks([])
ax2.set_yticks([])
for j in range(5):
    ax2.text(1.4*r*np.sin(phin[j]), 1.4*r*np.cos(phin[j]), str('(%d,0)' % (j+1-3)),
             horizontalalignment='center',verticalalignment='center')



plt.savefig('../user_manual_figures/player.png', bbox_inches='tight')


# --------------------------------------------------

nest.ResetKernel()

#{ layer6 #}
l = tp.CreateLayer({'rows': 1, 'columns': 2,
                    'elements': ['iaf_cond_alpha', 'poisson_generator']})
#{ end #}

print "#{ layer6 #}"
nest.PrintNetwork(depth=3)
print "#{ end #}"

# --------------------------------------------------

nest.ResetKernel()

#{ layer7 #}
l = tp.CreateLayer({'rows': 1, 'columns': 2,
                    'elements': ['iaf_cond_alpha', 10, 'poisson_generator',
                                 'noise_generator', 2]})
#{ end #}

print "#{ layer7 #}"
nest.PrintNetwork(depth=3)
print "#{ end #}"

# --------------------------------------------------

nest.ResetKernel()

#{ layer10 #}
for lyr in ['L23', 'L4', 'L56']:
    nest.CopyModel('iaf_neuron', lyr+'pyr')
    nest.CopyModel('iaf_neuron', lyr+'in', {'V_th': -52.})
l = tp.CreateLayer({'rows': 20, 'columns' : 20, 'extent': [0.5, 0.5],
                    'elements': ['L23pyr', 3, 'L23in',
                                 'L4pyr',  3, 'L4in',
                                 'L56pyr', 3, 'L56in']})
#{ end #}

# --------------------------------------------------

nest.ResetKernel()

#{ vislayer #}
l = tp.CreateLayer({'rows': 21, 'columns' : 21,
                    'elements': 'iaf_neuron'})
conndict = {'connection_type': 'divergent',
            'mask'  : {'circular': {'radius': 0.4}},
            'kernel': {'gaussian': {'p_center': 1.0, 'sigma': 0.15}}}
tp.ConnectLayers(l, l, conndict)
fig = tp.PlotLayer(l, nodesize=80)
ctr = tp.FindCenterElement(l)
tp.PlotTargets(ctr, l, fig=fig,
               mask=conndict['mask'], kernel=conndict['kernel'],
               src_size=250, tgt_color='red', tgt_size=20,
               kernel_color='green')
#{ end #}
plt.savefig('../user_manual_figures/vislayer.png', bbox_inches='tight')