import numpy as np
import matplotlib.pyplot as plt
from IPython import embed
from model_max import simulate, load_models
import matplotlib.gridspec as gridspec
from plot_eod_chirp import power_parameters
from scipy.ndimage import gaussian_filter
"""
Dependencies:
numpy 
matplotlib
numba (optional, speeds simulation up: pre-compiles functions to machine code)
"""


def main():
    # tiny example program:

    example_cell_idx = 11

    # load model parameter:
    parameters = load_models("models.csv")

    model_params = parameters[example_cell_idx]
    cell = model_params.pop('cell')
    eod_fr = model_params.pop('EODf')
    print("Example with cell:", cell)
    step = 20
    eod_fe = np.arange(0,eod_fr*5,step)
    # generate EOD-like stimulus with an amplitude step:
    deltat = model_params["deltat"]
    stimulus_length = 11.0  # in seconds
    time = np.arange(0, stimulus_length, deltat)
    time = np.arange(0, stimulus_length, deltat)
    # baseline EOD with amplitude 1:
    a_fr = 1  # amplitude fish reciever
    a_fe = 0.2  # amplitude fish emitter
    #results = [[]]*
    results = [[]] * 3
    counter = 0
    for e in range(len(eod_fe)):
        time_fish_r = time * 2 * np.pi * eod_fr
        eod_fish_r = a_fr * np.sin(time_fish_r)
        time_fish_e = time * 2 * np.pi * eod_fe[e]
        eod_fish_e = a_fe * np.sin(time_fish_e)
        stimulus = eod_fish_e+eod_fish_r

        # integrate the model:
        spikes = simulate(stimulus, **model_params)
        spikes_new = spikes[spikes > 1]
        sampling_rate = 1/deltat
        counter +=1
        if len(spikes_new)>0:
            spikes_mat = np.zeros(int(spikes_new[-1] * sampling_rate) + 2)
            spikes_idx = np.round((spikes_new) * sampling_rate)
            for spike in spikes_idx:
                spikes_mat[int(spike)] = 1
            spikes_mat = spikes_mat*sampling_rate
            window005 = 0.00005 * sampling_rate
            window05 = 0.0005 * sampling_rate
            window2 = 0.002 * sampling_rate
            smoothened_spikes_mat005 = gaussian_filter(spikes_mat, sigma=window005)
            smoothened_spikes_mat05 = gaussian_filter(spikes_mat, sigma=window05)
            smoothened_spikes_mat2 = gaussian_filter(spikes_mat, sigma=window2)
            nfft = 4096*2
            array = [spikes_mat, smoothened_spikes_mat05,smoothened_spikes_mat2]
            name = ['binary','05','2']

            for i in range(len(array)):
                results[i] = power_parameters(results[i], array[i], 1/deltat, nfft, name[i], eod_fr)
        else:
            print(counter)
    embed()
    ax = {}
    for i in range(len(results)):
        ax[i] = plt.subplot(2,3,i+1)
        plt.plot((eod_fe -eod_fr)/(eod_fr)+1,results[i]['f'],color = 'red')
        #plt.scatter((eod_fe - eod_fr) / (eod_fr) + 1, results[i]['f'],color = 'red')
        ax[0].set_ylabel('[Hz]')
        ax[i].set_ylim([0,eod_fr/2])
    for i in range(len(results)):
        ax[i+len(results)] = plt.subplot(2,3,i+len(results)+1)
        plt.plot((eod_fe -eod_fr)/(eod_fr)+1,results[i]['max'],color = 'steelblue')
        #plt.scatter((eod_fe - eod_fr) / (eod_fr) + 1, results[i]['max'],color = 'blue')
        ax[len(results)].set_ylabel('Modulation')
        ax[len(results)+i].set_xlabel('EOD multiples')
    plt.subplots_adjust(wspace = 0.3)
    plt.savefig('modell_single_cellmax.pdf')
    plt.savefig('../highbeats_pdf/cell_simulations/modell_single_cell'+cell+'max.pdf')
    plt.show()
    embed()
    # some analysis an dplotting:
    #embed()
    grid = gridspec.GridSpec(int(np.sqrt(len(parameters))), int(np.ceil(np.sqrt(len(parameters)))))
    parameters = load_models("models.csv")
    for i in range(4):
    #for i in range(len(parameters)):
        model_params = parameters[i]
        print(cell)
        cell = model_params.pop('cell')
        EODf = model_params.pop('EODf')


        # generate EOD-like stimulus
        deltat = model_params["deltat"]
        stimulus_length = 11.0  # in seconds
        time = np.arange(0, stimulus_length, deltat)
        # baseline EOD with amplitude 1:
        stimulus = np.sin(2 * np.pi * EODf * time)
        # das lasse ich eine sekunde integrieren dann weitere 10 sekunden integrieren und das nehmen
        spikes = simulate(stimulus, **model_params)

        # cut off first second of response
        new_spikes = spikes[spikes >1]

        freq,isis = calculate_isi_frequency(new_spikes, deltat)

        #embed()
        plt.subplot(grid[i])
        plt.title('B:'+np.mean(freq))
        plt.hist(isis, bins = 100, density = True)
    plt.savefig('isi_model.pdf')
    plt.savefig('../highbeats_pdf/isi_model.pdf')
    plt.show()

    freq_time = np.arange(spikes[0], spikes[-1], deltat)

    fig, axs = plt.subplots(2, 1, sharex="col")

    axs[0].plot(time, stimulus)
    axs[0].set_title("Stimulus")
    axs[0].set_ylabel("Amplitude in mV")

    axs[1].plot(freq_time, freq)
    axs[1].set_title("Model Frequency")
    axs[1].set_ylabel("Frequency in Hz")
    axs[1].set_xlabel("Time in s")
    plt.show()
    plt.close()


def calculate_isi_frequency(spikes, deltat):
    """
    calculates inter-spike interval frequency
    (wasn't tested a lot may give different length than time = np.arange(spikes[0], spikes[-1], deltat),
    or raise an index error for some inputs)

    :param spikes: spike time points
    :param deltat: integration time step of the model

    :return: the frequency trace:
                starts at the time of first spike
                ends at the time of the last spike.
    """

    isis = np.diff(spikes)
    freq_points = 1 / isis
    freq = np.zeros(int((spikes[-1] - spikes[0]) / deltat))

    current_idx = 0
    for i, isi in enumerate(isis):
        end_idx = int(current_idx + np.rint(isi / deltat))
        freq[current_idx:end_idx] = freq_points[i]
        current_idx = end_idx

    return freq,isis


if __name__ == '__main__':
    main()