# -*- coding: utf-8 -*-

"""

Example plot for LFPy: Single-synapse contribution to the potential through head

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 LFPy

import numpy as np

import matplotlib.pyplot as plt

from os.path import join

# four-sphere properties

radii = [79000., 80000., 85000., 90000.]

radii_name = ["Brain", "CSF", "Skull", "Scalp"]

sigmas = [0.3, 1.5, 0.015, 0.3]

rad_tol = 1e-2

somapos = np.array([0., 0., 78750.])

xlim = [-1000, 1000]

ylim = [radii[0]-500, radii[1] + 500]

# cell with simplified morphology

cellParameters = {

    'morphology' : join("morphologies", 'example_morphology.hoc'),

    'tstart' : 0.,

    'tstop' : 5.,

    'dt' : 2**-4,

    'v_init' : -65,

    'cm' : 1.,

    'Ra' : 150.,

    'passive' : True,

    'passive_parameters' : {'g_pas' : 1./3E4, 'e_pas' : -65.},

    'pt3d' : True,

}

cell = LFPy.Cell(**cellParameters)

synidx = cell.get_closest_idx(z=200)

synapseParameters = {

    'idx' : synidx,

    'e' : 0,                                # reversal potential

    'syntype' : 'Exp2Syn',                   # synapse type

    'tau1' : 0.1,                              # syn. rise time constant

    'tau2' : 1.,                              # syn. decay time constant

    'weight' : 0.002,                        # syn. weight

    'record_current' : True                 # syn. current record

}

synapse = LFPy.Synapse(cell, **synapseParameters)

synapse.set_spike_times(np.array([1.]))

cell.set_pos(z=somapos[2])

cell.simulate(rec_imem=True, rec_vmem=True, rec_current_dipole_moment=True)

# Setting up recording positions

elec_z = np.array([radii[0],

                  (radii[0] + radii[1]) / 2,

                   radii[1],

                   radii[1] + 500])

elec_x = np.zeros(elec_z.shape)

elec_y = np.zeros(elec_x.shape)

eeg_coords = np.array([elec_x.flatten(), elec_y.flatten(), elec_z.flatten()]).T

MD_4s = LFPy.FourSphereVolumeConductor(radii, sigmas, eeg_coords)

phi = MD_4s.calc_potential_from_multi_dipoles(cell) * 1e6  # from mV to nV

# Plotting results

plt.close('all')

fig = plt.figure(figsize=[8, 10])

fig.subplots_adjust(left=0.02, hspace=0.5, right=0.98, top=0.95)

ax = fig.add_subplot(111, aspect=1, frameon=False, xlim=xlim, ylim=ylim,

                     xticks=[], yticks=[])

# Plotting cell

for sec in LFPy.cell.neuron.h.allsec():

    idx = cell.get_idx(sec.name())

    l_cell, = ax.plot(np.r_[cell.xstart[idx], cell.xend[idx][-1]],

            np.r_[cell.zstart[idx], cell.zend[idx][-1]],

            color='r', zorder=3, lw=np.sqrt(np.average(cell.diam[idx])))

l_syn, = ax.plot(cell.xmid[synapseParameters['idx']],

                 cell.zmid[synapseParameters['idx']],

                 '*', c='orange', zorder=5)

# Plotting the layers of the head model

max_angle = np.abs(np.rad2deg(np.arcsin(xlim[0] / ylim[0])))

plot_angle = np.linspace(-max_angle, max_angle, 100)

for b_idx in range(len(radii)):

    x_ = radii[b_idx] * np.sin(np.deg2rad(plot_angle))

    z_ = radii[b_idx] * np.cos(np.deg2rad(plot_angle))

    l_curved, = ax.plot(x_, z_, ':', c="gray")

ax.text(xlim[0], radii[0] - 70, "Brain", va="top", ha="left", color="k")

ax.text(xlim[0], radii[0] + 70, "CSF", va="top", ha="left", color="k")

ax.text(xlim[0], radii[1] + 70, "Skull", va="top", ha="left", color="k")

ax.plot([-700, -200], [radii[0]-300, radii[0]-300], 'k', lw=2)

ax.text(-550, radii[0]-290, "500 $\mu$m")

# Plotting the potential at different positions

y_norm = 200

x_norm = 500

elec_clr = lambda elec_idx: plt.cm.viridis(elec_idx / (len(elec_z) - 1))

for elec_idx in range(len(elec_z)):

    l_elec, = ax.plot(elec_x[elec_idx], elec_z[elec_idx],

                      'o', c=elec_clr(elec_idx), clip_on=False)

    y_ = phi[elec_idx, :]

    norm_const = np.max(np.abs(y_))

    y_ = y_ * y_norm / norm_const + elec_z[elec_idx]

    x_ = cell.tvec / cell.tvec[-1] * x_norm + elec_x[elec_idx]

    ax.plot(x_, y_, c=elec_clr(elec_idx), clip_on=False)

    ax.plot([x_[-1] + 200, x_[-1] + 200],

            [elec_z[elec_idx], elec_z[elec_idx] - y_norm], c='k')

    scale_text = "{:2.1f} nV".format(norm_const)

    ax.text(x_[-1] + 250, elec_z[elec_idx] - y_norm / 2, scale_text)

ax.plot([x_[0], x_[-1]],

        [elec_z[1] - y_norm*1.1, elec_z[1] - y_norm*1.1],

        c='k', clip_on=False, lw=2)

ax.text((x_[0] + x_[-1])/2, elec_z[1] - y_norm * 1.2,

        "{} ms".format(cell.tvec[-1]),

        va="top", ha="center", clip_on=False)

ax_inset = fig.add_axes([0.7, 0.05, 0.25, 0.15], title="Synaptic current",

                        ylabel="nA", xlabel="Time (ms)")

l_isyn, = ax_inset.plot(cell.tvec, synapse.i, c='orange')

ax.legend([l_cell, l_curved, l_syn, l_isyn, l_elec],

          ["Cell", "Four-sphere boundary", "Synapse",

           "Synaptic current", "Electrode"],

          ncol=1, loc=(0.01, -0.1), frameon=False)

plt.savefig(join('example_potential_through_head.pdf'))

plt.show()