import matplotlib.pyplot as plt
import numpy as np
import pylab
from IPython import embed
from scipy.optimize import curve_fit
from jar_functions import gain_curve_fit
from jar_functions import avgNestedLists
import matplotlib as mpl
from matplotlib import cm

identifier_uniform = ['2018lepto1',
              #'2018lepto4',
              '2018lepto5',
              '2018lepto76',
              '2018lepto98',
              #'2019lepto03',
              '2019lepto24',
              #'2019lepto27',
              #'2019lepto30',
              #'2020lepto04',
              '2020lepto06',
              #'2020lepto16',
              #'2020lepto19',
              #'2020lepto20'
              ]
identifier = ['2018lepto1',
              '2018lepto4',
              '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]
#colors = ['dimgray', 'dimgrey', 'gray', 'grey', 'darkgray', 'darkgrey', 'silver', 'lightgray', 'lightgrey', 'gainsboro', 'whitesmoke']
colorss = ['g', 'b', 'r', 'y', 'c', 'm', 'k']
all = []
new_all = []
for ident in identifier:
    data = np.load('gain_%s.npy' %ident)
    all.append(data)
for ident in identifier_uniform:
    data = np.load('gain_%s.npy' % ident)
    new_all.append(data)

av = avgNestedLists(all)
new_av = avgNestedLists(new_all)

fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(amf, av, 'o', color = 'darkorange', label = 'normal')
ax.plot(amf, new_av, 'o', label = 'uniform')

tau = []
f_c = []
fit = []
fit_amf = []
for ID in identifier:
    print(ID)
    amf = np.load('amf_%s.npy' %ID)
    gain = np.load('gain_%s.npy' %ID)

    sinv, sinc = curve_fit(gain_curve_fit, amf, gain)
    #print('tau:', sinv[0])
    tau.append(sinv[0])
    f_cutoff = abs(1 / (2*np.pi*sinv[0]))
    print('f_cutoff:', f_cutoff)
    f_c.append(f_cutoff)
    fit.append(gain_curve_fit(amf, *sinv))
    fit_amf.append(amf)

# uniformed: 2018lepto5, 2018lepto1, 2018lepto76,  2018lepto98, 2020lepto06, 2019lepto24, 2020lepto4

tau_uniform = []
f_c_uniform = []
fit_uniform = []
fit_amf_uniform = []
for ID in identifier_uniform:
    print(ID)
    amf = np.load('amf_%s.npy' %ID)
    gain = np.load('gain_%s.npy' %ID)

    sinv, sinc = curve_fit(gain_curve_fit, amf, gain)
    print('tau:', sinv[0])
    tau_uniform.append(sinv[0])
    f_cutoff = abs(1 / (2*np.pi*sinv[0]))
    print('f_cutoff:', f_cutoff)
    print('alpha:', sinv[1])
    f_c_uniform.append(f_cutoff)
    fit_uniform.append(gain_curve_fit(amf, *sinv))
    fit_amf_uniform.append(amf)

colors = plt.cm.flag(np.linspace(0,1,len(fit_uniform)))

#for ff ,f in enumerate(fit_uniform):
#    ax.plot(fit_amf_uniform[ff], fit_uniform[ff],color =  colorss[ff])
#    ax.axvline(x=f_c_uniform[ff], ymin=0, ymax=5, ls = '-', alpha = 0.5, color= colorss[ff])#colors_uniform[ff])

ax.set_xscale('log')
ax.set_yscale('log')
ax.set_title('gain average all fish uniform')
ax.set_ylabel('gain [Hz/(mV/cm)]')
ax.set_xlabel('envelope_frequency [Hz]')
ax.set_ylim(0.0008, )
ax.plot(f_c, np.full(len(identifier), 0.0015), 'o', alpha = 0.5, c = 'darkorange', label = 'normal cutoff frequencies')
ax.plot(f_c_uniform, np.full(len(identifier_uniform), 0.002), 'o', alpha = 0.5, c = 'C0', label = 'uniform cutoff frequencies')
ax.legend(loc = 'center left')

plt.show()
embed()