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


def test_lifac():
    model = LIFACModel()
    stimulus = StepStimulus(0.5, 1, 15)

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

        v, spiketimes = model(stimulus, 2)

        plt.plot(np.arange(0, 2, step_size/1000), v, 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.show()
    plt.close()


def test_plot_inverses(ficurve):
    var = ficurve.boltzmann_fit_vars

    fig, ax1 = plt.subplots(1, 1, figsize=(4.5, 4.5), tight_layout=True)
    start = min(ficurve.stimulus_value)
    end = max(ficurve.stimulus_value)
    x_values = np.arange(start, end, (end-start)/5000)
    ax1.plot(x_values, [fu.full_boltzmann(x, var[0], var[1], var[2], var[3]) for x in x_values], label="fit")
    ax1.set_ylabel('freq')
    ax1.set_xlabel('stimulus')

    start = var[1]
    end = var[0]
    x_values = np.arange(start, end, (end - start) / 50)
    ax1.plot([fu.inverse_full_boltzmann(x, var[0], var[1], var[2], var[3]) for x in x_values], x_values,
             '.', c="red", label='inverse')
    plt.legend()
    plt.show()