#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Example plot for LFPy: Single-synapse contribution to the EEG

Execution:

    python example_EEG.py
    
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 matplotlib import colorbar

def plot_cell_to_ax(cell, ax, synidxs):
    for idx in range(cell.totnsegs):
        if idx == 0:
            ax.plot(cell.xmid[idx], cell.zmid[idx], 'ko')
        else:
            ax.plot([cell.xstart[idx], cell.xend[idx]],
                    [cell.zstart[idx], cell.zend[idx]], c='k')

    for synidx in synidxs:
        l, = ax.plot(cell.xmid[synidx], cell.zmid[synidx], '*',
                c="r", ms=10)
    ax.legend([l], ["Synapse"], frameon=False, bbox_to_anchor=[1, -0.1])

def plot_EEG_sphere(fig, eeg, x_eeg, y_eeg, z_eeg):
    from mpl_toolkits.mplot3d import Axes3D
    ax = fig.add_subplot(322, projection='3d',
                         title="Max EEG potential\nat 4-sphere surface")
    vmax = 6
    vmin = -vmax
    clr = lambda phi: plt.cm.PRGn((phi - vmin) / (vmax - vmin))
    clrs = clr(eeg)
    surf = ax.plot_surface(x_eeg.reshape(num_theta, num_phi),
                           y_eeg.reshape(num_theta, num_phi),
                           z_eeg.reshape(num_theta, num_phi),
                           rstride=1, cstride=1, facecolors=clrs,
                           linewidth=0, antialiased=False)

    #ax.set_aspect('equal') # Raises NotImplementedError in MPL currently
    ax.axis('off')
    ax.set_xlim3d(-65000, 65000)
    ax.set_ylim3d(-65000, 65000)
    ax.set_zlim3d(-65000, 65000)
    ax.view_init(10, 0)

    # colorbar
    cax = fig.add_axes([0.65, 0.75, 0.25, 0.01])
    m = plt.cm.ScalarMappable(cmap=plt.cm.PRGn)
    ticks = np.linspace(vmin, vmax, 5) # global normalization
    m.set_array(ticks)
    cbar = fig.colorbar(m, cax=cax,
                        extend='both', orientation='horizontal')
    cbar.outline.set_visible(False)
    cbar.set_ticks(ticks)
    cbar.set_label(r'$\phi$ (pV)', labelpad=1.)

if __name__ == '__main__':

    # four_sphere properties
    radii = [79000., 80000., 85000., 90000.]
    sigmas = [0.3, 1.5, 0.015, 0.3]
    rad_tol = 1e-2

    # simulate cell
    syn_loc = (0, 0, 1000)

    cell_params = {'morphology': 'morphologies/L5_Mainen96_LFPy.hoc',
                   'cm' : 1.0,                 # membrane capacitance
                   'Ra' : 150,                 # axial resistance
                   'tstart': 0.,
                   'passive' : True,           # switch on passive mechs
                   'nsegs_method' : 'lambda_f',# method for setting number of segments,
                   'lambda_f' : 100,           # segments are isopotential at this frequency
                   'passive_parameters' : {'g_pas' : 1./30000, 'e_pas' : -70}, # passive params
                   'tstop': 40
                   }

    synapse_params = {'e': 0.,  # reversal potential
                      'syntype': 'ExpSyn',  # exponential synapse
                      'tau': 5.,  # synapse time constant
                      'weight': 0.001,  # 0.001, # synapse weight
                      'record_current': True  # record synapse current
                     }
    # create cell with parameters in dictionary
    cell = LFPy.Cell(**cell_params)
    pos = syn_loc
    synapse_params['idx'] = cell.get_closest_idx(x=pos[0], y=pos[1], z=pos[2])
    synapse = LFPy.Synapse(cell, **synapse_params)
    synapse.set_spike_times(np.array([5.]))

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

    # compute dipole
    P = cell.current_dipole_moment

    somapos = np.array([0., 0., 77500])
    r_soma_syns = [cell.get_intersegment_vector(idx0=0, idx1=i) for i in cell.synidx]
    r_mid = np.average(r_soma_syns, axis=0)
    r_mid = somapos + r_mid/2.

    eeg_coords_top = np.array([[0., 0., radii[3] - rad_tol]])
    four_sphere_top = LFPy.FourSphereVolumeConductor(radii, sigmas, eeg_coords_top)
    pot_db_4s_top = four_sphere_top.calc_potential(P, r_mid)
    eeg_top = np.array(pot_db_4s_top) * 1e9

    #measurement points
    # for nice plot use theta_step = 1 and phi_step = 1. NB: Long computation time.
    theta_step = 5
    phi_step = 5
    theta, phi_angle = np.mgrid[0.:180.:theta_step, 0.:360.+phi_step:phi_step]

    num_theta = theta.shape[0]
    num_phi = theta.shape[1]
    theta = theta.flatten()
    phi_angle = phi_angle.flatten()

    theta_r = np.deg2rad(theta)
    phi_angle_r = np.deg2rad(phi_angle)

    x_eeg = (radii[3] - rad_tol) * np.sin(theta_r) * np.cos(phi_angle_r)
    y_eeg = (radii[3] - rad_tol) * np.sin(theta_r) * np.sin(phi_angle_r)
    z_eeg = (radii[3] - rad_tol) * np.cos(theta_r)
    eeg_coords = np.vstack((x_eeg, y_eeg, z_eeg)).T

    # potential in 4S with db
    time_max = np.argmax(np.linalg.norm(P, axis=1))
    p = P[time_max, None]

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

    pot_db_4s = four_sphere.calc_potential(p, r_mid)
    eeg = pot_db_4s.reshape(num_theta, num_phi)*1e9# from mV to pV

    P_mag = np.sqrt(P[:, 0]**2 + P[:, 1]**2 + P[:, 2]**2)

    plt.close('all')
    fig = plt.figure(figsize=[6, 8])
    fig.subplots_adjust(left=0.15, hspace=0.5, right=0.98, top=0.95)

    morph_ax = fig.add_subplot(321, aspect=1, frameon=False,
                               title="Cell and synapse",
                               xticks=[], yticks=[])
    plot_cell_to_ax(cell, morph_ax, cell.synidx)

    plot_EEG_sphere(fig, eeg, x_eeg, y_eeg, z_eeg)

    ax_p = fig.add_subplot(312, title="Current dipole moment", ylabel="nA$\cdot\mu$m")
    ax_p.plot(cell.tvec, P[:, 0], label="P$_x$")
    ax_p.plot(cell.tvec, P[:, 1], label="P$_y$")
    ax_p.plot(cell.tvec, P[:, 2], label="P$_z$")
    ax_p.axvline(cell.tvec[time_max], c='gray', ls='--')
    ax_p.legend(frameon=False, ncol=4, bbox_to_anchor=[1, -0.1])

    ax_eeg = fig.add_subplot(313, title="EEG at top", xlabel="Time (ms)",
                             ylabel='pV')
    ax_eeg.plot(cell.tvec, eeg_top[0, :], 'k')

    ax_eeg.axvline(cell.tvec[time_max], c='gray', ls='--')

    # plt.savefig('example_EEG.pdf')
    plt.show()
