import matplotlib.pyplot as plt
import numpy as np
import helperFunctions as hf
from models.FirerateModel import FirerateModel
from models.LIFACnoise import LifacNoiseModel
from stimuli.StepStimulus import StepStimulus
from stimuli.SinusAmplitudeModulation import SinusAmplitudeModulationStimulus
import functions as fu


def main():
    # test_firerate_model()
    # test_stepsize_influence()
    test_lifac_noise()


def test_stepsize_influence():
    # model = LIFACModel()
    model = FirerateModel()
    stimulus = StepStimulus(0.5, 1, 15)

    for step_size in [0.001, 0.1]:
        model.set_variable("step_size", step_size)

        model.simulate(stimulus, 2)

        # y = model.get_voltage_trace()
        y = model.get_frequency()

        plt.plot(np.arange(0, 2, step_size/1000), y, label="step_size:" + str(step_size))

    plt.xlabel("time in seconds")
    plt.ylabel("Voltage")
    plt.title("Voltage in the LIFAC-Model with different step sizes")
    plt.legend()
    plt.show()
    plt.close()


def test_firerate_model():
    model = FirerateModel()
    duration = 0.2
    stimulus = StepStimulus(duration/2, duration, 10)

    model.simulate(stimulus, duration*2)
    plt.plot(np.arange(0, duration*2, model.get_sampling_interval()/1000), model.get_frequency())
    plt.show()


def test_lifac_noise():

    model = LifacNoiseModel()

    start = 0.2
    duration = 30
    total_time = duration + 2*start
    step_size = model.get_parameters()["step_size"]/1000
    time = np.arange(0, total_time, step_size)
    stimulus = SinusAmplitudeModulationStimulus(700, 0.2, 10, 20, start, duration)

    model.simulate(stimulus, total_time)

    fig, axes = plt.subplots(nrows=3, sharex="col")
    sparse_time = np.arange(0, total_time, 1/5000)
    axes[0].plot(sparse_time, [stimulus.value_at_time_in_s(x) for x in sparse_time], label="given stimulus")
    axes[0].plot(sparse_time, [fu.rectify(stimulus.value_at_time_in_s(x)) for x in sparse_time], label="seen stimulus")
    axes[0].set_title("Stimulus")
    axes[0].set_ylabel("stimulus strength")
    axes[0].legend()

    axes[1].plot(time, model.get_voltage_trace())
    axes[1].set_title("Voltage trace")
    axes[1].set_ylabel("voltage")

    t, f = hf.calculate_isi_frequency(model.get_spiketimes(), 0, step_size)
    axes[2].plot(t, f)
    axes[2].set_title("ISI frequency trace")
    axes[2].set_ylabel("Frequency")

    spiketimes_small_step = model.get_spiketimes()

    model.set_variable("step_size", 0.02)
    model.simulate(stimulus, total_time)
    print(model.get_adaption_trace()[int(0.1/(0.01/1000))])
    step_size = model.get_parameters()["step_size"] / 1000
    time = np.arange(0, total_time, step_size)
    t, f = hf.calculate_isi_frequency(model.get_spiketimes(), 0, step_size)

    axes[1].plot(time, model.get_voltage_trace())
    axes[2].plot(t, f)

    spiketimes_big_step = model.get_spiketimes()

    print("CV:")
    print("small step:", hf.calculate_coefficient_of_variation(spiketimes_small_step))
    print("big   step:", hf.calculate_coefficient_of_variation(spiketimes_big_step))

    plt.show()
    plt.close()

    max_lag = 5
    x = np.arange(1, max_lag+1, 1)
    serial_cor_small = hf.calculate_serial_correlation(spiketimes_small_step, max_lag)
    serial_cor_big = hf.calculate_serial_correlation(spiketimes_big_step, max_lag)

    print(serial_cor_small)
    print(serial_cor_big)
    plt.plot(x, serial_cor_small, 'o', label='small step',)
    plt.plot(x, serial_cor_big, 'o', label='big step')
    plt.ylim(-1, 1)
    plt.legend()
    plt.show()
    plt.close()

    bins = np.arange(0, max(np.diff(spiketimes_small_step)), 0.0001)
    plt.hist(np.diff(spiketimes_small_step), bins=bins, alpha=0.5)
    plt.hist(np.diff(spiketimes_big_step), bins=bins, alpha=0.5)
    plt.show()


if __name__ == '__main__':
    main()