# %%
import numpy as np
import pickle

with open('tuning_3.4.pickle', 'rb') as f:
    data = pickle.load(f)

#%%
data['neuron1'].shape

import matplotlib.pyplot as plt


# plt.plot(data['stim'],data['neuron1'].mean(axis=0))
# plt.show()


# plt.plot(data['stim'],data['neuron2'].mean(axis=0))
# plt.show()


# plt.plot(data['stim'],data['neuron3'].mean(axis=0))
# plt.show()


# plt.plot(data['stim'],data['neuron4'].mean(axis=0))
# plt.show()

import numpy as np

def fano(neuron):
  d = data[neuron]
  f = (100/10) * d.std(axis=0) / d.mean(axis=0)
  return f

# calc r_max per neuron
r_max={}
for i in range(1,5):
  n = f'neuron{i}'
  print(f'fano of {n}:', np.nanmean(fano(n)))
  # r_max_index = data[n].mean(axis=0).argmax()
  # r_max[n] = data['stim'][r_max_index]
  r_max[n] = data[n].mean(axis=0).max()

print('r_max: ', r_max)

#%%

# part 2 (last question)
with open('pop_coding_3.4.pickle', 'rb') as f:
    pdata = pickle.load(f)

r0 = 0
v_total = [0,0]
for n in range(1,5):
  v = pdata[f'c{n}'] * np.mean(pdata[f'r{n}']) / r_max[f'neuron{n}']
  v_total += v

# map to the expected coord (clockwise, y-axis is zero degree)
v_deg = np.rad2deg(np.arctan(v_total)[1])
np.round(-(v_deg - 90) % 360)