from stimuli.SinusAmplitudeModulation import SinusAmplitudeModulationStimulus as SAM
from models.LIFACnoise import LifacNoiseModel
import numpy as np
import matplotlib.pyplot as plt
import helperFunctions as hF
from CellData import CellData


def main():
    # 2012-12-13_ao fit and eod frequency:
    parameters = {'mem_tau': 0.0133705462739553, 'tau_a': 0.06682759542588587, 'input_scaling': 60.766243690761144,
                  'v_base': 0, 'step_size': 5e-05, 'dend_tau': 0.0008667253013050408, 'v_zero': 0, 'v_offset': -6.25,
                  'noise_strength': 0.03337309379328535, 'a_zero': 2, 'threshold': 1, 'delta_a': 0.0726267312975076}
    eod_freq = 658
    cell_data = CellData("./data/2012-12-13-ao-invivo-1/")
    model = LifacNoiseModel(parameters)
    mean_duration = np.mean(cell_data.get_sam_durations())
    contrasts = cell_data.get_sam_contrasts()
    spiketimes = cell_data.get_sam_spiketimes()
    for i, m_freq in enumerate(cell_data.get_sam_delta_frequencies()):

        stimulus = SAM(eod_freq, contrasts[i], m_freq)
        prob_desnity_function_model = generate_pdf(model, stimulus, sim_length=mean_duration)
        for spikes in spiketimes[i]:
            prob_density_cell = spiketimes_calculate_pdf(spikes, cell_data.get_sampling_interval())

            plt.plot(prob_density_cell)
            plt.plot(prob_desnity_function_model)
            plt.show()
            plt.close()


    # __init__(carrier_frequency, contrast, modulation_frequency, start_time=0, duration=np.inf, amplitude=1)
    mod_freqs = np.arange(-60, eod_freq*4 + 61, 10)
    sigma_of_pdfs = []
    for m_freq in mod_freqs:
        print(m_freq, "max: {:.2f}".format(mod_freqs[-1]))
        stimulus = SAM(eod_freq, 0.2, m_freq)

        prob_density_function = generate_pdf(model, stimulus)
        buffer = 0.25
        buffer_idx = int(buffer / model.get_parameters()["step_size"])

        sigma_of_pdfs.append(np.std(prob_density_function[buffer_idx:-buffer_idx]))

    normed_mod_freqs = (mod_freqs + eod_freq) / eod_freq
    plt.plot(normed_mod_freqs, sigma_of_pdfs)
    plt.savefig("./figures/sam/test.png")
    plt.close()

    pass


def generate_pdf(model, stimulus, trials=4, sim_length=3, kernel_width=0.005):

    trials_rate_list = []
    step_size = model.get_parameters()["step_size"]
    for _ in range(trials):
        v1, spikes = model.simulate(stimulus, total_time_s=sim_length)

        binary = np.zeros(int(sim_length/step_size))
        spikes = [int(s / step_size) for s in spikes]
        for s_idx in spikes:
            binary[s_idx] = 1

        kernel = gaussian_kernel(kernel_width, step_size)
        rate = np.convolve(binary, kernel, mode='same')
        trials_rate_list.append(rate)

    times = [np.arange(0, sim_length, step_size) for _ in range(trials)]
    t, mean_rate = hF.calculate_mean_of_frequency_traces(times, trials_rate_list, step_size)

    return mean_rate


def spiketimes_calculate_pdf(spikes, step_size, kernel_width=0.005):
    length = int(spikes[-1] / step_size)+1
    binary = np.zeros(length)
    spikes = [int(s / step_size) for s in spikes]
    for s_idx in spikes:
        binary[s_idx] = 1

    kernel = gaussian_kernel(kernel_width, step_size)
    rate = np.convolve(binary, kernel, mode='same')

    return rate


def gaussian_kernel(sigma, dt):
    x = np.arange(-4. * sigma, 4. * sigma, dt)
    y = np.exp(-0.5 * (x / sigma) ** 2) / np.sqrt(2. * np.pi) / sigma
    return y


if __name__ == '__main__':
    main()