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]

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

    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 = []

    data = sorted(np.load('%s files.npy' %ident), 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()

        sinv, sinc = curve_fit(sin_response, time, jm - cutf, [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)
        if f not in amfreq:
            amfreq.append(f)

        # root mean square
        RMS = np.sqrt(np.mean(((jm - cutf) - sin_response(cutt, sinv[0], sinv[1], sinv[2]))**2))

        thresh = A / np.sqrt(2)

        # 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 = []
            meanrms = []
            meanthresh = []
            for x in idxlist:
                meanf.append(gain[x])
                meanp.append(phaseshift[x])
                meanrms.append(RMS)
                meanthresh.append(thresh)
            meanedf = np.mean(meanf)
            meanedp = np.mean(meanp)
            meanedrms = np.mean(meanrms)
            meanedthresh = np.mean(meanthresh)
            mgain.append(meanedf)
            mphaseshift.append(meanedp)
            rootmeansquare.append(meanedrms)
            threshold.append(meanedthresh)
            currf = d[1]    # set back for next loop
            idxlist = [i]
    meanf = []
    meanp = []
    meanrms = []
    meanthresh = []
    for y in idxlist:
        meanf.append(gain[y])
        meanp.append(phaseshift[y])
        meanrms.append(RMS)
        meanthresh.append(thresh)
    meanedf = np.mean(meanf)
    meanedp = np.mean(meanp)
    meanedrms = np.mean(meanrms)
    meanedthresh = np.mean(meanthresh)
    mgain.append(meanedf)
    mphaseshift.append(meanedp)
    rootmeansquare.append(meanedrms)
    threshold.append(meanedthresh)

    # 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()
    ax0 = fig.add_subplot(2, 1, 1)
    ax0.plot(amfreq, mgain(RMS<threshold), 'o')
    #ax0.plot(amf, predict)
    ax0.set_yscale('log')
    ax0.set_xscale('log')
    ax0.set_title('%s' % data[0][0])
    ax0.set_ylabel('gain [Hz/(mV/cm)]')
    ax0.set_xlabel('envelope_frequency [Hz]')
    #plt.savefig('%s gain' % data[0][0])

    ax1 = fig.add_subplot(2, 1, 2, sharex = ax0)
    ax1.plot(amfreq, threshold, 'o-', label = 'threshold', color = 'b')
    ax1.set_xscale('log')
    ax1.plot(amfreq, rootmeansquare, 'o-', label = 'RMS', color = 'orange')
    ax1.set_xscale('log')
    ax1.set_xlabel('envelope_frequency [Hz]')
    ax1.set_ylabel('RMS')
    plt.legend()
    pylab.show()

embed()

# zu eigenmannia: jeden fisch mit amplituden von max und min von modulationstiefe und evtl 1 oder 2 dazwischen
                # und dann für die am frequenzen von apteronotus für 15Hz delta f messen

# mit zu hohem RMS rauskicken: gain/rms < ... (?)
# gain kurven als array abspeichern
# daten von natalie zu eigenmannia mit + / - delta f anschauen ob unterschiede
# unterschiedliche nffts auf anderem rechner laufen lassen evtl um unterschiede zu sehen?


# long term: extra datei mit script drin um fertige daten darzustellen, den code hier als datenverarbeitung allein verwenden
#            darstellung: specgram --> rausgezogene jarspur darüber --> filterung --> fit und daten zusammen dargestellt, das ganze für verschiedene frequenzen
#            liste mit eigenschaften der fische (dominanz/größe) und messvariablen (temp/conductivity) machen um diese plotten zu können

# phase in degree