import matplotlib.pyplot as plt
import numpy as np
from my_util import helperFunctions as hf, functions as fu
from models.FirerateModel import FirerateModel
from models.LIFACnoise import LifacNoiseModel
from stimuli.StepStimulus import StepStimulus
from stimuli.SinusoidalStepStimulus import SinusoidalStepStimulus


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()