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

"""

Created on Wed Oct 29 13:46:55 2014

@author: Lindsay

This module takes stress files into a generator function to generate groups of

generator currents, and then call the LIF Model to output spikes.

"""

import numpy as np

import matplotlib.pyplot as plt

from model_constants import (LIF_RESOLUTION, DURATION, MC_GROUPS)

from gen_function import prepare_stress, stress_to_current

from lif_model import get_spikes

# %% Convert fine stress to current in all groups of Merkel Cells

def sed_to_group_current(fine_time, fine_stress, tau1, tau2, k1, mode, groups):

return group_current

# %% Transfer spike times to spike trace

def spike_time_to_trace(spike_time):

return spike_trace

# %% Main function

if __name__ == '__main__':

# Select stimulation method among 1, 2, and 3, and comment out the others.

# %% 1. For force control input

# rough_time = np.genfromtxt('./fcon_disp6_time.csv', delimiter=',')

# rough_stress = np.genfromtxt('./fcon_disp6_stress.csv', delimiter=',')

# fine_time, fine_stress = prepare_stress(rough_time, rough_stress)

# %% 2. For disp control input

fine_stress = np.genfromtxt('./Shawn model data/dcon_disp3_stress.csv',

delimiter=',')

fine_time = np.genfromtxt('./Shawn model data/dcon_disp3_time.csv',

delimiter=',')

# %% 3. For sinusoid input

# freq = 1.

# dur = 5.0001

# fine_time = np.arange(0,dur,LIF_RESOLUTION/1000)

# sin_stress = 3000*np.sin(2*np.pi*(fine_time)*freq)

# fine_stress = sin_stress.copy()

# fine_stress[np.nonzero(fine_stress<0)[0]]=0

# %% Generator function decay parameters

tau1_s = 0.003

tau1_m = 0.008

tau1_l = 0.013

tau2_s = 0.05

tau2_m = 0.2

tau2_l = 0.35

k1_s = 0.1

k1_m = 0.5

k1_l = 0.9

tau1_ap1 = 0.5

tau2_ap1 = 1

k1_ap1 = 0.1

tau2_ap2 = 1

k1_ap2 = 0.1

tau1_ap3 = 0.5

tau2_ap3 = 1

k1_ap3 = 0.5

# %% Choose 1 stress-current based on decay parameters

# # The standard generator current

gen_current = sed_to_group_current(fine_time, fine_stress, tau1_m, tau2_m, k1_m, 'gen', MC_GROUPS)

# The standard neurite current

# gen_current = sed_to_group_current(fine_time, fine_stress, tau1_m, tau2_m, k1_m, 'nr', MC_GROUPS)

# # Change tau1

# gen_current = sed_to_group_current(fine_time, fine_stress, tau1_s, tau2_m, k1_m, 'gen', MC_GROUPS)

# gen_current = sed_to_group_current(fine_time, fine_stress, tau1_l, tau2_m, k1_m, 'gen', MC_GROUPS)

# # Change tau2

# gen_current = sed_to_group_current(fine_time, fine_stress, tau1_m, tau2_s, k1_m, 'gen', MC_GROUPS)

# gen_current = sed_to_group_current(fine_time, fine_stress, tau1_m, tau2_l, k1_m, 'gen', MC_GROUPS)

# # Change k1

# gen_current = sed_to_group_current(fine_time, fine_stress, tau1_m, tau2_m, k1_s, 'gen', MC_GROUPS)

# gen_current = sed_to_group_current(fine_time, fine_stress, tau1_m, tau2_m, k1_l, 'gen', MC_GROUPS)

# New approach 1: modify parameters

# gen_current = sed_to_group_current(fine_time, fine_stress, tau1_ap1, tau2_ap1, k1_ap1, 'mc', MC_GROUPS)

# New approach 2: add a mechanism

# gen_current = sed_to_group_current(fine_time, fine_stress, tau1_m, tau2_ap2, k1_ap2, 'mc', MC_GROUPS)

# New approach 3: add K to neurite mechanism

# gen_current = sed_to_group_current(fine_time, fine_stress, tau1_ap3, tau2_ap3, k1_ap3, 'mc', MC_GROUPS)

# %% Current to spike calculation

spike_time = get_spikes(gen_current)

spike_time = np.array(spike_time)

spike_trace = spike_time_to_trace(spike_time)

inst_fr = np.r_[0, 1000/np.diff(spike_time)]

inst_fr_index = spike_trace[:, 1].nonzero()[0]

inst_fr_time = inst_fr_index * LIF_RESOLUTION

# %% Stress to spike calculations for looped sinusoid stimulations

# freq_set = np.array([5, 10, 20])

# stress_range = np.arange(0, 1500, 50)

# stress_len = stress_range.shape[0]

# spike_array = []

# for freq in freq_set:

# for stress_val in stress_range:

# print(freq, ', ', stress_val)

# sin_stress = stress_val*np.sin(2*np.pi*(fine_time)*freq)

# fine_stress = sin_stress.copy()

# fine_stress[np.nonzero(fine_stress<0)[0]]=0

# gen_current = sed_to_group_current(fine_time, fine_stress, tau1_m, tau2_m, k1_m, 'gen', MC_GROUPS)

# spike_list = get_spikes(gen_current)

# spike_time = spike_list[0]

# spike_time = np.array(spike_time)

# spike_per_cyc = spike_time.shape[0] / (dur*freq)

# spike_array.append(spike_per_cyc)

# spike_array = np.array(spike_array)

# spike_per_cyc_5hz = spike_array[0:stress_len]

# spike_per_cyc_10hz = spike_array[stress_len:2*stress_len]

# spike_per_cyc_20hz = spike_array[2*stress_len:3*stress_len]

# fig, axs = plt.subplots()

# axs.plot(stress_range, spike_per_cyc_5hz, color='0')

# axs.plot(stress_range, spike_per_cyc_10hz, color='0.4')

# axs.plot(stress_range, spike_per_cyc_20hz, color='0.8')

# np.savetxt("sin_stress.csv", stress_range, delimiter=",")

# np.savetxt("sin_spike_per_cyc_5hz.csv", spike_per_cyc_5hz, delimiter=",")

# np.savetxt("sin_spike_per_cyc_10hz.csv", spike_per_cyc_10hz, delimiter=",")

# np.savetxt("sin_spike_per_cyc_20hz.csv", spike_per_cyc_20hz, delimiter=",")

# %% Firing rate calculations for disp and force control

# Get ramp-up fr

max_index = np.argmax(fine_stress)

ramp_index = spike_trace[0:max_index, 1].nonzero()[0]

ramp_spike_time = spike_trace[ramp_index, 0]

ramp_fr = np.mean(1000/np.diff(ramp_spike_time))

# Get early-hold fr

early_index = spike_trace[max_index:max_index+500, 1].nonzero()[0]

early_spike_time = spike_trace[early_index, 0]

early_fr = np.mean(1000/np.diff(early_spike_time))

# Get late-hold fr

late_index = spike_trace[2000/LIF_RESOLUTION:4500/LIF_RESOLUTION, 1]. \

nonzero()[0]

late_spike_time = spike_trace[late_index, 0]

late_fr = np.mean(1000/np.diff(late_spike_time))

print('ramp fr = ', ramp_fr)

print('early hold fr = ', early_fr)

print('late hold fr = ', late_fr)

# %% Plot

fig, axs = plt.subplots()

stress_plot = axs.plot(fine_time, fine_stress / 600, 'b',

label='scaled stress')

current_plot = axs.plot(fine_time[:50000], gen_current[:50000, 0]*1e8, 'k',

label='scaled current')

instfr_plot = axs.plot(inst_fr_time/1000, inst_fr*1, 'g.',

label='instant fr')

spike_plot = axs.plot(spike_trace[:, 0]/1000, spike_trace[:, 1],

label='spikes', color='r')

# %% Save files

# np.savetxt("ap4_disp3_gen_spike_time.csv", spike_trace[:, 0]/1000,

# delimiter=",")

# np.savetxt("ap4_disp3_gen_spike_spike.csv", spike_trace[:, 1],

# delimiter=",")

# np.savetxt("ap4_disp3_gen_inst_fr_time.csv", inst_fr_time/1000,

# delimiter=",")

# np.savetxt("ap4_disp3_gen_inst_fr_fr.csv", inst_fr, delimiter=",")

# np.savetxt("ap4_disp3_gen_spike.csv", spike_time/1000, delimiter=",")

# np.savetxt("ap4_disp3_gen_current.csv", gen_current/8, delimiter=",")

# np.savetxt("disp2_stress.csv", fine_stress, delimiter=",")

# np.savetxt("disp2_time.csv", fine_time, delimiter=",")