# -*- coding: utf-8 -*-
"""
LFPs from a population of cells relying on MPI (Message Passing Interface)

Execution:

    <mpiexec> -n <processes> python example_mpi_2.py

Notes:
- on certain platforms and with mpirun, the --oversubscribe argument is needed
  to get more processes than the number of physical CPU cores.

Copyright (C) 2017 Computational Neuroscience Group, NMBU.

This program 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 3 of the License, or
(at your option) any later version.

This program 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.
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import PolyCollection, LineCollection
import os
from os.path import join
import sys
if sys.version < '3':
    from urllib2 import urlopen
else:
    from urllib.request import urlopen
import zipfile
import LFPy
import neuron
from mpi4py import MPI

#initialize the MPI interface
COMM = MPI.COMM_WORLD
SIZE = COMM.Get_size()
RANK = COMM.Get_rank()

#set the numpy random seeds
global_seed = 1234
np.random.seed(global_seed)

def stationary_poisson(nsyn,lambd,tstart,tstop):
    ''' Generates nsyn stationary possion processes with rate lambda between tstart and tstop'''
    interval_s = (tstop-tstart)*.001
    spiketimes = []
    for i in range(nsyn):
        spikecount = np.random.poisson(interval_s*lambd)
        spikevec = np.empty(spikecount)
        if spikecount==0:
            spiketimes.append(spikevec)
        else:
            spikevec = tstart + (tstop-tstart)*np.random.random(spikecount)
            spiketimes.append(np.sort(spikevec)) #sort them too!

    return spiketimes


#Fetch Mainen&Sejnowski 1996 model files
if not os.path.isfile(join('cells', 'cells', 'j4a.hoc')) and RANK==0:
    #get the model files:
    u = urlopen('http://senselab.med.yale.edu/ModelDB/eavBinDown.asp?o=2488&a=23&mime=application/zip')
    localFile = open('patdemo.zip', 'w')
    localFile.write(u.read())
    localFile.close()
    #unzip:
    myzip = zipfile.ZipFile('patdemo.zip', 'r')
    myzip.extractall('.')
    myzip.close()

#resync MPI threads
COMM.Barrier()

# Define cell parameters
cell_parameters = {          # various cell parameters,
    'morphology' : join('cells', 'cells', 'j4a.hoc'), # Mainen&Sejnowski, 1996
    'cm' : 1.0,         # membrane capacitance
    'Ra' : 150,         # axial resistance
    'v_init' : -65.,    # initial crossmembrane potential
    'passive' : True,   # turn on passive mechanism for all sections
    'passive_parameters' : {'g_pas' : 1./30000, 'e_pas' : -65}, # passive params
    'nsegs_method' : 'lambda_f',
    'lambda_f' : 100.,
    'dt' : 2.**-3,      # simulation time step size
    'tstart' :  0.,     # start time of simulation, recorders start at t=0
    'tstop' : 300.,     # stop simulation at 200 ms. These can be overridden
                        # by setting these arguments i cell.simulation()
}

# Define synapse parameters
synapse_parameters = {
    'idx' : 0, # to be set later
    'e' : 0.,                   # reversal potential
    'syntype' : 'ExpSyn',       # synapse type
    'tau' : 5.,                 # syn. time constant
    'weight' : .001,            # syn. weight
    'record_current' : True,
}

# Define electrode parameters
point_electrode_parameters = {
    'sigma' : 0.3,      # extracellular conductivity
    'x' : 0.,  # electrode requires 1d vector of positions
    'y' : 0.,
    'z' : 0.,
}

# number of units
n_cells = 6

# assign cell positions
x_cell_pos = np.linspace(-250., 250., n_cells)

# default rotation around x and y axis
xy_rotations = dict(x=4.99, y=-4.33)

# rotations around z-axis
if RANK == 0:
    z_rotation = COMM.bcast(np.random.permutation(np.arange(0., np.pi,
                                                            np.pi / n_cells)),
                            root=0)
else:
    z_rotation = COMM.bcast(None, root=0)


#synaptic spike times drawn on RANK 0 distributed to all processes
n_pre_syn = 1000
if RANK == 0:
    pre_syn_sptimes = COMM.bcast(stationary_poisson(nsyn=n_pre_syn, lambd=5.,
                                                    tstart=0, tstop=300),
                                 root=0)
else:
    pre_syn_sptimes = COMM.bcast(None, root=0)

# number of synapses on each cell
n_synapses = 100

# indices for presynaptic spike trains for each neuron also picked on RANK 0
# and scattered (for illustrating purposes, not efficiency)
if RANK == 0:
    # set up len SIZE nested list for spike train IDs.
    pre_syn_ids = [[]]*SIZE
    for cell_id in range(n_cells):
        pre_syn_ids[cell_id % SIZE] += [np.random.permutation(np.arange(
                                                    n_pre_syn))[0:n_synapses]]
else:
    pre_syn_ids = None
pre_syn_ids = COMM.scatter(pre_syn_ids, root=0)

# containers for per-cell LFP and summed LFPs
single_LFPs = []
summed_LFP = np.zeros(int(cell_parameters['tstop'] / cell_parameters['dt'] + 1))

# get state of random seed generator
state = np.random.get_state()

# iterate over cells in populations
for cell_id in range(n_cells):
    if cell_id % SIZE == RANK:
        # get set seed per cell in order to synapse locations
        np.random.seed(global_seed + cell_id)

        # Create cell
        cell = LFPy.Cell(**cell_parameters)

        #Have to position and rotate the cells!
        cell.set_rotation(z=z_rotation[cell_id], **xy_rotations)
        cell.set_pos(x=x_cell_pos[cell_id])

        for i_syn in range(n_synapses):
            syn_idx = cell.get_rand_idx_area_norm()
            synapse_parameters.update({'idx' : syn_idx})
            synapse = LFPy.Synapse(cell, **synapse_parameters)
            synapse.set_spike_times(pre_syn_sptimes[pre_syn_ids[
                                                        cell_id % SIZE][i_syn]])

        #run the cell simulation
        cell.simulate(rec_imem=True)

        #set up the extracellular device
        point_electrode = LFPy.RecExtElectrode(cell,
                                               **point_electrode_parameters)
        point_electrode.calc_lfp()

        # sum LFP on this RANK
        summed_LFP += point_electrode.LFP[0]

        # send LFP of this cell to RANK 0
        if RANK != 0:
            COMM.send(point_electrode.LFP[0], dest=0)
        else:
            single_LFPs += [point_electrode.LFP[0]]

    # collect single LFP contributions on RANK 0
    if RANK == 0:
        if cell_id % SIZE != RANK:
            single_LFPs += [COMM.recv(source=cell_id % SIZE)]

# we can also use MPI to sum arrays directly:
summed_LFP = COMM.reduce(summed_LFP)

# reset state of random number generator
np.random.set_state(state)

# plot output on RANK 0.
if RANK==0:
    #assign color to each unit
    color_vec = [plt.cm.rainbow(int(x*256./n_cells)) for x in range(n_cells)]

    #figure
    fig = plt.figure(figsize=(12, 8))

    # Morphologies axes:
    plt.axes([.175, .0, .65, 1], aspect='equal')
    plt.axis('off')

    for i_cell in range(n_cells):
        cell = LFPy.Cell(join('cells', 'cells', 'j4a.hoc'),
                         nsegs_method='lambda_f',
                         lambda_f=5)
        cell.set_rotation(z=z_rotation[i_cell], **xy_rotations)
        cell.set_pos(x=x_cell_pos[i_cell])

        zips = []
        for x, z in cell.get_idx_polygons():
            zips.append(list(zip(x, z)))
        linecol = LineCollection(zips,
                    edgecolor = 'none',
                    facecolor = color_vec[i_cell],
                    rasterized=False,
                    )

        ax = plt.gca()
        ax.add_collection(linecol)

    axis = ax.axis(ax.axis('equal'))
    ax.axis(np.array(axis) / 1.15)


    #adding a blue dot:
    ax.plot(point_electrode.x, point_electrode.z, 'o',
            markeredgecolor='none', markerfacecolor='b', markersize=3,
            zorder=10, clip_on=False)
    plt.annotate("Electrode",
            xy=(0., 0.), xycoords='data',
            xytext=(-100., 1000.),
            arrowprops=dict(arrowstyle='wedge',
                            shrinkA=1,
                            shrinkB=1,
                            #lw=0.5,
                            mutation_scale=20,
                            fc="0.6", ec="none",
                            edgecolor='k', facecolor='w'))

    plt.xlim([-700., 700.])

    ax.plot([100, 200], [-250, -250], 'k', lw=1, clip_on=False)
    ax.text(150, -300, r'100$\mu$m', va='center', ha='center')

    #presynaptic spike trains axes
    plt.axes([.05, .35, .25, .55])

    pop_sptimes = []
    for i_pre in range(n_pre_syn):
        sp = pre_syn_sptimes[i_pre]
        for i_sp in range(len(sp)):
            pop_sptimes.append(sp[i_sp])

    for i_pre in range(n_pre_syn):
        plt.scatter(pre_syn_sptimes[i_pre],
                    i_pre*np.ones(len(pre_syn_sptimes[i_pre])),
                    s=1, edgecolors='none', facecolors='k')

    plt.ylim([0,n_pre_syn])
    plt.xlim([0,cell_parameters['tstop']])
    plt.ylabel('train #', ha='left', labelpad=0)
    plt.title('Presynaptic spike times')

    ax = plt.gca()
    for loc, spine in ax.spines.items():
        if loc in ['right', 'top']:
            spine.set_color('none')
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')

    ax.set_xticklabels([])

    #spike rate axes
    plt.axes([.05,.12,.25,.2])

    binsize = 5
    bins=np.arange(0, cell_parameters['tstop']+1., binsize)
    count,b = np.histogram(pop_sptimes, bins=bins)
    rate = count*(1000./binsize)*(1./n_pre_syn)
    plt.plot(b[0:-1],rate,color='black',lw=1)

    plt.xlim([0,cell_parameters['tstop']])
    plt.ylim([0,10.])

    tvec = np.arange(point_electrode.LFP.shape[1])*cell.dt

    plt.xlabel('$t$ (ms)')
    plt.ylabel('Rate (spike/s)')

    ax = plt.gca()
    for loc, spine in ax.spines.items():
        if loc in ['right', 'top']:
            spine.set_color('none')
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')

    #single neuron EPs axes
    plt.axes([.7,.35,.25,.55])

    plt.title('Single neuron extracellular potentials')
    plt.axis('off')

    for cell_id in range(n_cells):
        plt.plot(tvec,
                        cell_id+2.e3*single_LFPs[cell_id],
                        color=color_vec[cell_id], lw=1,
                        )

    plt.ylim([-1,n_cells-.5])

    #Summed LFPs axes
    plt.axes([.7,.12,.25,.2])
    plt.plot(tvec, 1E3*summed_LFP, color='black', lw=1)
    plt.ylim([-5.e-1,5e-1])

    plt.title('Summed extracellular potentials')
    plt.xlabel(r'$t$ (ms)')
    plt.ylabel(r'$\mu$V',ha='left',rotation='horizontal')

    ax = plt.gca()
    for loc, spine in ax.spines.items():
        if loc in ['right', 'top']:
            spine.set_color('none')
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')


    fig.savefig('example_mpi_2.pdf', dpi=300)
    plt.show()