# -*- coding: utf-8 -*-
"""
Created on Sat May 29 21:10:20 2021

@author: nils
"""

import numpy as np
import os
import h5py
import json
from Utility import capacitance, stimulus_init, init_dict, NumpyEncoder

# model parameters
dt = 0.01
sec = 2
low = 0
high = 0.001
number_steps = 200
initial_period = 1000
num_gating = 10
num_current = 7
C, surf_area = capacitance(56.9, 1)
stim_time, I_in, stim_num, V_m = stimulus_init(low, high, number_steps, initial_period, dt, sec)
shift, scale, slope_shift, E, currents_included, b_param, g = init_dict(
    np.array(['m', 'h', 'n', 'q', 'r', 'p', 's', 'u', 's_mut', 'u_mut']))
tau_max_p = 502
V_init = -70
V_T = -57.9

# initialize arrays
current = np.zeros((num_current, stim_num))
gating_out = np.zeros((num_gating, stim_num))

# gating parameters
b_param = {}
b_param['m'] = np.array([-34.33054521, -8.21450277, 1.42295686])
b_param['h'] = np.array([-34.51951036, 4.04059373, 1., 0.])
b_param['n'] = np.array([-63.76096946, -13.83488194, 7.35347425])
b_param['q'] = np.array([-39.03684525, -5.57756176, 2.25190197])
b_param['r'] = np.array([-57.37, 20.98, 1.])
b_param['p'] = np.array([-45., -9.9998807337, 1.])
b_param['s'] = np.array([-14.16, -10.15, 1.])
b_param['u'] = np.array([-31., 5.256, 1., 0.245])
b_param['s_mut'] = np.array([-14.16, -10.15, 1.])
b_param['u_mut'] = np.array([-31., 5.256, 1., 0.245])

mut_act_Vhalf_wt = -30.01851851851851
mut_act_k_wt = -7.7333333333333325
s_diff_Vhalf = mut_act_Vhalf_wt - b_param['s'][0]
s_diff_k = mut_act_k_wt - b_param['s'][1]
b_param['s'][1] = b_param['s'][1] + s_diff_k
b_param['u'][1] = b_param['u'][1] + s_diff_k
b_param['s'][0]  = b_param['s'][0] + s_diff_Vhalf
b_param['u'][0]  = b_param['u'][0] + s_diff_Vhalf
b_param['s_mut'][1] = b_param['s_mut'][1] + s_diff_k
b_param['u_mut'][1] = b_param['u_mut'][1] + s_diff_k
b_param['s_mut'][0]  = b_param['s_mut'][0] + s_diff_Vhalf
b_param['u_mut'][0]  = b_param['u_mut'][0] + s_diff_Vhalf

# reversal potentials
E["Na"] = 50.
E["K"] = -90.
E["Ca"] = 120.
E["Leak"] = -70.4

# model currents
currents_included["Na"] = True
currents_included["Kd"] = True
currents_included["Kv"] = True
currents_included["Kv_mut"] = True
currents_included["L"] = False
currents_included["M"] = True
currents_included["Leak"] = True

# model conductances
Kv_ratio = 0.1
g["Na"] = 58. * surf_area
g["Kd"] = 3.9 * (1 - Kv_ratio) * surf_area
g["M"] = 0.075 * surf_area
if currents_included["Kv_mut"] == True:
    g["Kv"] = 3.9 * Kv_ratio / 2 * surf_area
else:
    g["Kv"] = 3.9 * Kv_ratio / 2 * surf_area * 2
g["Kv_mut"] = 3.9 * Kv_ratio / 2 * surf_area
g["L"] = 0. * surf_area
g["Leak"] = 0.038 * surf_area

# save folder
folder = './KCNA1_mutations/FS'
if not os.path.isdir(folder):
    os.makedirs(folder)

# mutation properties
mutations = json.load(open("mutations_effects_dict.json"))


# prominence = 50
# # min_spike_height = 0
# desired_AUC_width = high/5
#
# Parallel(n_jobs=8, verbose=9)(
#     delayed(Pospischil_multi)(V_init, V_T, g, E, I_in, dt, currents_included, stim_time, stim_num, C, tau_max_p,
#                               shift, scale, b_param, slope_shift, gating_out, current,  prominence,
#                               desired_AUC_width, mutations, mut, folder, high,low, number_steps, initial_period, sec, save_gating=True)
#     for mut in list(mutations.keys()))

#%% Get pd Dataframes for certain variables
import pandas as pd
AUC = pd.DataFrame(columns=mutations.keys())
AUC_rel = pd.DataFrame(columns=mutations.keys())
rheobase = pd.DataFrame(columns=mutations.keys())
rheobase_fit = pd.DataFrame(columns=mutations.keys())
rheobase_null_fit = pd.DataFrame(columns=mutations.keys())

for mut in list(mutations.keys()):
    fname = os.path.join(folder, "{}.hdf5".format(mut.replace(" ", "_")))
    f = h5py.File(fname, "r")
    AUC['{}'.format(mut.replace(" ", "_"))] = f['analysis']['AUC']
    AUC_rel['{}'.format(mut.replace(" ", "_"))] = f['analysis']['AUC_rel']
    rheobase['{}'.format(mut.replace(" ", "_"))] = f['analysis']['rheobase']
    rheobase_fit['{}'.format(mut.replace(" ", "_"))] = f['analysis']['rheobase_fit']
    rheobase_null_fit['{}'.format(mut.replace(" ", "_"))] = f['analysis']['rheobase_null_fit']

top_dir = '../KCNA1_mut'
model_name = 'FS'
save_folder = os.path.join(top_dir, 'mut_summary_df')
if not os.path.isdir(save_folder):
    os.makedirs(save_folder)
AUC.to_json(os.path.join(save_folder, '{}_AUC.json'.format(model_name)))
AUC_rel.to_json(os.path.join(save_folder, '{}_AUC_rel.json'.format(model_name)))
rheobase.to_json(os.path.join(save_folder, '{}_rheobase.json'.format(model_name)))
rheobase_fit.to_json(os.path.join(save_folder, '{}_rheobase_fit.json'.format(model_name)))
rheobase_null_fit.to_json(os.path.join(save_folder, '{}_rheobase_null_fit.json'.format(model_name)))