### Extract data from pickle file from simulation and calculate
###  muscle force.

import pickle
import math
import numpy as np

# Open and load pickle file data
with open('model output_traces.pkl', 'rb') as f:
	data = pickle.load(f)

# To check keys (names of data) in pickle file
#keys = data.keys()
#print(keys)

# Extract data recorded using traces
tracesData = data['tracesData']		#data is a dict; tracesData is a list

# Create time vector from simulation
cell_data = [tracesData[0]]
cell_data_0 = cell_data[0]
time = cell_data_0['t']

# Extract cli and mgi for each cell from the data
N = 310     # number of cells in network
cli = np.zeros([N,len(time)-1])
mgi = np.zeros([N,len(time)-1])

for i in range(0,N):
    cell_data_cli = [tracesData[(2*i)]]
    cell_data_mgi = [tracesData[((2*i)+1)]]
    cell_data_0_cli = cell_data_cli[0]
    cell_data_0_mgi = cell_data_mgi[0]
    cli[i,:] = cell_data_0_cli['cell_'+str(i)+'_cli_musc']
    mgi[i,:] = cell_data_0_mgi['cell_'+str(i)+'_mgi_musc']

# Calculate force (written from module3.mod [Kim2017])
a0 = 2.35
b0 = 24.35
c0 = -7.4
d0 = 30.0
p0 = 300
g1 = -8
g2 = 21.4
Kse = 0.4
xm_init = -8.0
xce_init = -8
dt_NEURON = 0.025

xm = xm_init
A = 0.0

# Scaling when average force on neuron 180 (determined using peak_twitch_force_calculation.m)
force_scaling = [0.0677518020847825,
0.0687657943485401,
0.0697949623017898,
0.0708395330681100,
0.0718997371702737,
0.0729758085811217,
0.0740679847751973,
0.0751765067811535,
0.0763016192349452,
0.0774435704338166,
0.0786026123910976,
0.0797790008918194,
0.0809729955491629,
0.0821848598617518,
0.0834148612718035,
0.0846632712241496,
0.0859303652261406,
0.0872164229084462,
0.0885217280867662,
0.0898465688244652,
0.0911912374961439,
0.0925560308521622,
0.0939412500841284,
0.0953472008913680,
0.0967741935483871,
0.0982225429733464,
0.0996925687975591,
0.101184595436029,
0.102698952159046,
0.104235973164849,
0.105795997653380,
0.107379369901144,
0.108986439337182,
0.110617560620187,
0.112273093716771,
0.113953403980910,
0.115658862234565,
0.117389844849524,
0.119146733830457,
0.120929916899225,
0.122739787580439,
0.124576745288310,
0.126441195414794,
0.128333549419052,
0.130254224918261,
0.132203645779770,
0.134182242214644,
0.136190450872607,
0.138228714938406,
0.140297484229610,
0.142397215295884,
0.144528371519743,
0.146691423218813,
0.148886847749622,
0.151115129612951,
0.153376760560753,
0.155672239704676,
0.158002073626213,
0.160366776488495,
0.162766870149762,
0.165202884278528,
0.167675356470474,
0.170184832367088,
0.172731865776076,
0.175317018793588,
0.177940861928260,
0.180603974227115,
0.183306943403360,
0.186050365966076,
0.188834847351867,
0.191661002058469,
0.194529453780360,
0.197440835546405,
0.200395789859552,
0.203394968838628,
0.206439034362251,
0.209528658214897,
0.212664522235157,
0.215847318466205,
0.219077749308527,
0.222356527674929,
0.225684377147866,
0.229062032139132,
0.232490238051928,
0.235969751445370,
0.239501340201445,
0.243085783694475,
0.246723872963116,
0.250416410884924,
0.254164212353548,
0.257968104458559,
0.261828926667980,
0.265747531013545,
0.269724782278733,
0.273761558189610,
0.277858749608535,
0.282017260730762,
0.286238009283978,
0.290521926730840,
0.294869958474533,
0.299283064067408,
0.303762217422743,
0.308308407029672,
0.312922636171331,
0.317605923146272,
0.322359301493182,
0.327183820218977,
0.332080544030302,
0.337050553568495,
0.342094945648075,
0.347214833498787,
0.352411347011285,
0.357685632986479,
0.363038855388618,
0.368472195602165,
0.373986852692509,
0.379584043670584,
0.385265003761449,
0.391030986676882,
0.396883264892061,
0.402823129926379,
0.408851892628468,
0.414970883465481,
0.421181452816717,
0.427484971271620,
0.433882829932260,
0.440376440720325,
0.446967236688715,
0.453656672337797,
0.460446223936394,
0.467337389847579,
0.474331690859338,
0.481430670520198,
0.488635895479857,
0.495948955834929,
0.503371465479854,
0.510905062463063,
0.518551409348478,
0.526312193582409,
0.534189127865958,
0.542183950532988,
0.550298425933747,
0.558534344824238,
0.566893524761415,
0.575377810504296,
0.583989074421071,
0.592729216902314,
0.601600166780371,
0.610603881755026,
0.619742348825543,
0.629017584729169,
0.638431636386198,
0.647986581351701,
0.657684528274016,
0.667527617360096,
0.677518020847825,
0.687657943485401,
0.697949623017898,
0.708395330681100,
0.718997371702737,
0.729758085811217,
0.740679847751973,
0.751765067811535,
0.763016192349452,
0.774435704338166,
0.786026123910976,
0.797790008918194,
0.809729955491629,
0.821848598617518,
0.834148612718035,
0.846632712241497,
0.859303652261406,
0.872164229084462,
0.885217280867662,
0.898465688244653,
0.911912374961439,
0.925560308521622,
0.939412500841285,
0.953472008913680,
0.967741935483871,
0.982225429733464,
0.996925687975591,
1.01184595436029,
1.02698952159046,
1.04235973164849,
1.05795997653380,
1.07379369901144,
1.08986439337182,
1.10617560620187,
1.12273093716771,
1.13953403980910,
1.15658862234565,
1.17389844849524,
1.19146733830457,
1.20929916899225,
1.22739787580439,
1.24576745288310,
1.26441195414794,
1.28333549419052,
1.30254224918261,
1.32203645779770,
1.34182242214644,
1.36190450872607,
1.38228714938406,
1.40297484229610,
1.42397215295884,
1.44528371519743,
1.46691423218813,
1.48886847749622,
1.51115129612951,
1.53376760560753,
1.55672239704676,
1.58002073626213,
1.60366776488495,
1.62766870149762,
1.65202884278528,
1.67675356470474,
1.70184832367088,
1.72731865776076,
1.75317018793588,
1.77940861928260,
1.80603974227115,
1.83306943403360,
1.86050365966076,
1.88834847351867,
1.91661002058469,
1.94529453780360,
1.97440835546405,
2.00395789859552,
2.03394968838628,
2.06439034362251,
2.09528658214897,
2.12664522235157,
2.15847318466205,
2.19077749308527,
2.22356527674929,
2.25684377147866,
2.29062032139132,
2.32490238051928,
2.35969751445371,
2.39501340201445,
2.43085783694475,
2.46723872963115,
2.50416410884924,
2.54164212353548,
2.57968104458559,
2.61828926667980,
2.65747531013545,
2.69724782278733,
2.73761558189610,
2.77858749608535,
2.82017260730762,
2.86238009283978,
2.90521926730840,
2.94869958474533,
2.99283064067408,
3.03762217422743,
3.08308407029672,
3.12922636171331,
3.17605923146272,
3.22359301493182,
3.27183820218977,
3.32080544030302,
3.37050553568496,
3.42094945648075,
3.47214833498787,
3.52411347011285,
3.57685632986479,
3.63038855388619,
3.68472195602165,
3.73986852692509,
3.79584043670584,
3.85265003761449,
3.91030986676882,
3.96883264892061,
4.02823129926380,
4.08851892628468,
4.14970883465482,
4.21181452816717,
4.27484971271620,
4.33882829932260,
4.40376440720325,
4.46967236688715,
4.53656672337797,
4.60446223936395,
4.67337389847579,
4.74331690859339,
4.81430670520199,
4.88635895479857,
4.95948955834929,
5.03371465479854,
5.10905062463063,
5.18551409348478,
5.26312193582409,
5.34189127865959,
5.42183950532989,
5.50298425933747,
5.58534344824238,
5.66893524761415,
5.75377810504296,
5.83989074421071,
5.92729216902314,
6.01600166780371,
6.10603881755026,
6.19742348825543,
6.29017584729169,
6.38431636386197,
6.47986581351701,
6.57684528274016,
6.67527617360096]

force = np.zeros([N,len(time)-1])

for j in range(0,N):
    xce = xce_init
    F_old = 0.00001

    for k in range(0,len(time)-1):

        xm = cli[j,k]
        A = mgi[j,k]

        g_xm = math.exp(-((xm-g1)/g2)**2)

        #Fc = p0*g_xm*A
        Fc = force_scaling[j]*g_xm*A

        if F_old <= Fc:
            dxdt = (0.001*(-b0*(Fc-F_old)))/(F_old+(a0*(Fc/p0)))
        else:
            gain_length = (-d0*(Fc-F_old))/((2*Fc)-F_old+(c0*Fc/p0))

            if gain_length <= 0:
                dxdt = 100
            else:
                dxdt = 0.001*gain_length

        xce = xce - (dt_NEURON*(-dxdt))

        d_xm = xm - xm_init
        d_xce = xce - xce_init
        d_se = d_xm - d_xce

        if d_se <= 0:
            xse = 0
        else:
            xse = d_se

        #F = p0*Kse*xse
        F = force_scaling[j]*Kse*xse
        F_old = F

        force[j,k] = F

# Linear Summation of Twitch Forces for Muscle Force
muscle_force = np.zeros(len(time)-1)

for m in range(1,len(time)-1):
    muscle_force[m] = sum(force[:,m])


# Output forces to file for use in fortran/Abaqus
outF = open("force_output_results.txt","a+")
for line in muscle_force:
	outF.write(str(line))
	outF.write("\n")
outF.close()

# Create second file that is a replica of the first for antagonist muscle
outF2 = open("force_output_results_copy.txt","a+")
for line in muscle_force:
	outF2.write(str(line))
	outF2.write("\n")
outF2.close()