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.SinusoidalStepStimulus import SinusoidalStepStimulus
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 = 1
    duration = 3
    total_time = duration + 2*start
    step_size = model.get_parameters()["step_size"]
    time = np.arange(0, total_time, step_size)
    stimulus = SinusoidalStepStimulus(700, 0.2, start, duration)

    model.simulate(stimulus, total_time)

    fig, axes = plt.subplots(nrows=4, 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()

    input_voltage = model.input_voltage
    axes[1].plot(time, input_voltage)
    axes[1].set_title("Stimulus after dendtritic filter")
    axes[1].set_ylabel("stimulus strength")

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

    spiketimes = model.get_spiketimes()
    t, f = hf.calculate_time_and_frequency_trace(spiketimes, step_size)
    axes[3].plot(t, f)
    axes[3].set_title("ISI frequency trace")
    axes[3].set_ylabel("Frequency")

    plt.show()

    # 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_time_and_frequency_trace(model.get_spiketimes(), 0.02)
    #
    # 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()