from stimuli.SinusoidalStepStimulus import SinusoidalStepStimulus
from stimuli.SinusAmplitudeModulation import SinusAmplitudeModulationStimulus
from parser.CellData import CellData
import numpy as np
import matplotlib.pyplot as plt
import Figure_constants as consts


def main():
    # stimuli_protocol_examples()
    am_generation()
    pass


def am_generation():
    cell = "data/final/2013-04-17-ac-invivo-1"
    cell_data = CellData(cell)
    fig, axes = plt.subplots(3, 1, sharey=True, sharex=True, figsize=consts.FIG_SIZE_MEDIUM)

    start = 0
    end = 0.05

    time_start = -0.025
    time_end = 0.075
    step_size = 0.00005

    # frequency = 55
    contrast = 0.3

    eod = cell_data.get_base_traces(cell_data.EOD)[0]
    step = cell_data.get_sampling_interval()

    values = eod[:int(round((time_end-time_start)/step))]

    # base_stim = SinusoidalStepStimulus(frequency, 0, start, end - start)
    #
    # values = base_stim.as_array(time_start, time_end - time_start, step_size)
    time = np.arange(time_start, time_end, step_size) * 1000

    axes[0].set_title("Fish EOD")
    axes[0].plot(time, values)

    am = np.zeros(len(values))
    am[int((start-time_start)/step_size):int((end-time_start)/step_size)] = contrast

    axes[1].set_title("Amplitude Modulation")
    axes[1].plot(time, values*am)
    axes[1].plot(time, am)
    axes[1].set_ylabel("Voltage [mV]")

    axes[2].set_title("Resulting stimulus")
    axes[2].plot(time, (values*am) + values)
    axes[2].set_xlabel("Time [ms]")
    axes[2].set_xlim((time[0], time[-1]))
    axes[2].set_xticks([-25, 0, 25, 50, 75])
    plt.tight_layout()

    consts.set_figure_labels(xoffset=-2.5, yoffset=1.5)
    fig.label_axes()

    plt.savefig("thesis/figures/amGeneration.pdf")
    plt.close()


def stimuli_protocol_examples():

    fig, axes = plt.subplots(3, 1, sharey=True, sharex=True, figsize=consts.FIG_SIZE_MEDIUM)

    start = 0
    end = 1

    time_start = -0.2
    time_end = 1.2
    step_size = 0.00005

    frequency = 55
    contrast = 0.5

    base_stim = SinusoidalStepStimulus(frequency, 0, start, end - start)

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

    axes[0].set_title("Baseline Stimulus")
    axes[0].plot(time, values)
    # axes[1].set_ylabel("Voltage [mV]")

    # Sinusoidal STEP stimulus
    # 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)

    axes[1].set_title("Step Stimulus")
    axes[1].plot(time, values)
    axes[1].set_ylabel("Voltage [mV]")

    # SAM Stimulus:
    mod_freq = 10
    step_stim = SinusAmplitudeModulationStimulus(frequency, 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
    axes[2].plot(beat_time, beat_values)

    axes[2].set_title("SAM Stimulus")
    axes[2].set_xlabel("Time [s]")
    # axes[2].set_ylabel("Voltage [mV]")

    axes[2].set_xlim((time_start, time_end))
    plt.tight_layout()

    consts.set_figure_labels(xoffset=-2.5)
    fig.label_axes()

    plt.savefig("thesis/figures/stimuliExamples.pdf")
    plt.close()


if __name__ == '__main__':
    main()