import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.backends.backend_pdf import PdfPages
import numpy as np
from sim_vs_net import sim_vs_net

run_time_vars = {'g_gap':1000, # nS
                     'g_inh_scale':0.06, # Factor which multiplies g_gap to give g_inh
                     'g_L_dend':180, # nS
                     'g_L_axon':30, # nS
                     'g_axon_dend':110, #nS
                     'g_exc_in':10000, # nS
                     'g_inh_in':15000, # nS
                     'E_E':60, # mV
                     'E_I':-40, # mV,
                     'tau_m':1.4, # ms
                     'noise_std':6, #10**2.5, # pA
                     'tau_hp':50, # ms
                     'tau_lp':20, # ms
                     'dt_sim':0.01, # ms
                     'dt_im':1, # ms
                     'rot_angle':90,
                     'deg_per_ms': 0.125,
                     'T':100,
                     } 
  
run_time_vars['image'] = {'type':'stripe','x_pixels':360,'y_pixels':180,'strip_rad':5,'freq_s':8}
T = run_time_vars['T']
t = np.linspace(0,T,T/run_time_vars['dt_sim'],endpoint=False)

nRotAngles = 361 #for rapid debugging and/or visualization use 3 or 37 instead of 361
rot_angles = np.linspace(-180,180,nRotAngles)
v_r_ss_ang = np.zeros([nRotAngles,10])
v_l_ss_ang = np.zeros([nRotAngles,10])

for cAng_ind in range(rot_angles.size):
    run_time_vars['rot_angle'] = rot_angles[cAng_ind]
    (V_L,V_R) = sim_vs_net(run_time_vars)
    
    #uncomment to see how individual responses look like over time
    #fig1 = plt.figure()
    #ax1 = fig1.add_subplot(121)
    #ax2 = fig1.add_subplot(122)
    #for i in range(10):
    #    ax1.plot(t,V_L[i,:].T,label='VS'+str(i+1))
    #    ax2.plot(t,V_R[i,:].T)
    #
    #ax1.legend()
    #ax1.set_title('Left side responses - rot: ' + str(rot_angles[cAng_ind]))
    #ax2.set_title('Right side responses')
    #disp_lim = np.max([np.max(np.abs(V_L)),np.max(np.abs(V_R))])*1.1
    #ax1.set_ylim([-disp_lim,disp_lim])
    #ax2.set_ylim([-disp_lim,disp_lim])
    #plt.show()

    #average over the last 10 ms of stimulation to get steady state responses
    v_r_ss_ang[cAng_ind,:] = np.mean(V_R[:,10000-1000:10000],axis=1)
    v_l_ss_ang[cAng_ind,:] = np.mean(V_L[:,10000-1000:10000],axis=1)

#plot the tuning curves for all VS model cells
fig2 = plt.figure()
ax1f2 = fig2.add_subplot(111, projection='3d')
for vsInd in range(10):
    ax1f2.plot(rot_angles,vsInd*np.ones(nRotAngles),v_r_ss_ang[:,vsInd])
    ax1f2.plot(rot_angles,-vsInd*np.ones(nRotAngles),v_l_ss_ang[:,vsInd])
plt.show() 

#uncomment to save as pdf
#with PdfPages('vs_stripe_tuning.pdf') as pdf:
#    pdf.savefig(fig2)