import matplotlib.pyplot as plt
import numpy as np
import pylab
from IPython import embed

def avgNestedLists(nested_vals):
    """
    Averages a 2-D array and returns a 1-D array of all of the columns
    averaged together, regardless of their dimensions.
    """
    output = []
    maximum = 0
    for lst in nested_vals:
        if len(lst) > maximum:
            maximum = len(lst)
    for index in range(maximum): # Go through each index of longest list
        temp = []
        for lst in nested_vals: # Go through each list
            if index < len(lst): # If not an index error
                temp.append(lst[index])
        output.append(np.nanmean(temp))
    return output

identifier = ['2018lepto4',
              '2018lepto1',
              '2018lepto5',
              '2018lepto76',
              '2018lepto98',
              '2019lepto03',
              '2019lepto24',
              '2019lepto27',
              '2019lepto30',
              '2020lepto04',
              '2020lepto06',
              '2020lepto16',
              '2020lepto19',
              '2020lepto20'
              ]

amf = [0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1]

all = []

for ident in identifier:
    data = np.load('gain_%s.npy' %ident)
    all.append(data)

av = avgNestedLists(all)

fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(amf, av, 'o')
ax.set_xscale('log')
ax.set_yscale('log')
ax.set_title('gaincurve_average_allfish')
ax.set_ylabel('gain [Hz/(mV/cm)]')
ax.set_xlabel('envelope_frequency [Hz]')
plt.show()
embed()

'''len_arr = []
for a in all:
    len_arr.append(len(a))
max_a = np.max(len_arr)

arr = np.ma.empty((1,len(all),max_a))
arr.mask = True

for x, a in enumerate(all):
    arr[:a.shape[0],x] = arr[0][x]
embed()

print(arr.mean(axis = 2))
embed()'''