from scipy import signal
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 sin_response
from jar_functions import mean_noise_cut

def take_second(elem):      # function for taking the names out of files
    return elem[1]

predict = []

rootmeansquare = []
threshold = []

gain = []
mgain = []
phaseshift = []
mphaseshift = []
amfreq = []
amf = [0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1]

currf = None
idxlist = []

identifier = '2020lepto06'

data = sorted(np.load('%s files.npy' %identifier), key = take_second)      # list with filenames in it

for i, d in enumerate(data):
    dd = list(d)
    jar = np.load('%s.npy' %dd)     # load data for every file name
    jm = jar - np.mean(jar)         # low-pass filtering by subtracting mean
    print(dd)

    time = np.load('%s time.npy' %dd)       # time file
    dt = time[1] - time[0]

    n = int(1/float(d[1])/dt)
    cutf = mean_noise_cut(jm, time, n = n)
    cutt = time
    plt.plot(time, jm-cutf, label='cut amfreq')
    plt.plot(time, jm, label='spec')
    #plt.legend()
    #plt.show()

    b, a = signal.butter(4, (float(d[1])) / (0.5/dt), 'high', analog=True)     # high pass filtering so our fit gets a bit better
    y = signal.filtfilt(b, a, jm)
    #plt.plot(time, y)
    #plt.plot(time, jm)

    sinv, sinc = curve_fit(sin_response, time, y, [float(d[1]), 2, 0.5])        # fitting
    print('frequency, phaseshift, amplitude:', sinv)
    p = np.sqrt(sinv[1]**2)
    A = np.sqrt(sinv[2] ** 2)
    f = float(d[1])
    phaseshift.append(p)
    gain.append(A)
    amfreq.append(f)
    '''
    # root mean square
    RMS = np.sqrt(np.mean((cutf_arr - sin_response(cutt_arr, sinv[0], sinv[1], sinv[2]))**2))
    rootmeansquare.append(RMS)
    thresh = A / np.sqrt(2)
    threshold.append(thresh)
    '''
    plt.plot(time, sin_response(time, *sinv), label='fit: f=%f, p=%.2f, A=%.2f' % tuple(sinv))
    plt.legend()
    plt.show()
    # mean over same amfreqs for phase and gain
    if currf is None or currf == d[1]:
        currf = d[1]
        idxlist.append(i)

    else:  # currf != f
        meanf = []  # lists to make mean of
        meanp = []
        for x in idxlist:
            meanf.append(gain[x])
            meanp.append(phaseshift[x])
        meanedf = np.mean(meanf)
        meanedp = np.mean(meanp)
        mgain.append(meanedf)
        mphaseshift.append(meanedp)
        currf = d[1]    # set back for next loop
        idxlist = [i]
meanf = []
meanp = []
for y in idxlist:
    meanf.append(gain[y])
    meanp.append(phaseshift[y])
meanedf = np.mean(meanf)
meanedp = np.mean(meanp)
mgain.append(meanedf)
mphaseshift.append(meanedp)

# predict of gain
for f in amf:
    G = np.max(mgain) / np.sqrt(1 + (2*((np.pi*f*3.14)**2)))
    predict.append(G)

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.plot(amf, mgain, 'o')
#ax.plot(amf, predict)
ax.set_yscale('log')
ax.set_xscale('log')
ax.set_title('%s' % data[0][0])
ax.set_ylabel('gain [Hz/(mV/cm)]')
ax.set_xlabel('envelope_frequency [Hz]')
#plt.savefig('%s gain' % data[0][0])
pylab.show()

plt.plot(threshold, label = 'threshold')
plt.plot(rootmeansquare, label = 'RMS')
plt.legend()
plt.show()

embed()

# phase in degree
# Q10 / conductivity
# AM-frequency / envelope-frequency scale title?

# einfach passendes n verwenden um AM-beats rauszufiltern? ...Menge Datenpunkte zu gering