import matplotlib.pyplot as plt
import numpy as np
from models.FirerateModel import FirerateModel
from models.LIFAC import LIFACModel
from stimuli.StepStimulus import StepStimulus


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


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


if __name__ == '__main__':
    main()