from stimuli.SinusoidalStepStimulus import SinusoidalStepStimulus
from stimuli.SinusAmplitudeModulation import SinusAmplitudeModulationStimulus
import numpy as np
import matplotlib.pyplot as plt


def main():
    plot_step_stimulus()
    plot_sam_stimulus()
    pass


def plot_step_stimulus():
    start = 0
    end = 1

    time_start = -0.2
    time_end = 1.2
    step_size = 0.00005

    frequency = 20
    contrast = 0.5
    # frequency, contrast, start_time=0, duration=np.inf, amplitude=1
    step_stim= SinusoidalStepStimulus(frequency, contrast, start, end-start)

    values = step_stim.as_array(time_start, time_end-time_start, step_size)
    time = np.arange(time_start, time_end, step_size)

    plt.plot(time, values)
    plt.xlabel("Time [s]")
    plt.ylabel("Voltage [mV]")
    plt.savefig("thesis/figures/sin_step_stim_example.pdf")
    plt.close()


def plot_sam_stimulus():
    start = 0
    end = 1

    time_start = -0.2
    time_end = 1.2
    step_size = 0.00005

    contrast = 0.5
    mod_freq = 10
    carrier_freq = 53
    # carrier_frequency, contrast, modulation_frequency, start_time=0, duration=np.inf, amplitude=1
    step_stim = SinusAmplitudeModulationStimulus(carrier_freq, contrast, mod_freq, start, end - start)

    values = step_stim.as_array(time_start, time_end - time_start, step_size)
    time = np.arange(time_start, time_end, step_size)
    plt.plot(time, values)

    beat_time = np.arange(start, end, step_size)
    beat_values = np.sin(beat_time*2*np.pi*mod_freq) * contrast + 1
    plt.plot(beat_time, beat_values)

    plt.xlabel("Time [s]")
    plt.ylabel("Voltage [mV]")
    # plt.show()
    plt.savefig("thesis/figures/sam_stim_example.pdf")
    plt.close()


if __name__ == '__main__':
    main()