import unittest
import numpy as np
import helperFunctions as hF
import matplotlib.pyplot as plt
from models.LIFACnoise import LifacNoiseModel
from stimuli.SinusAmplitudeModulation import SinusAmplitudeModulationStimulus


class HelperFunctionsTester(unittest.TestCase):

    # TODO specify exact parameters for all test (multiple sets?)
    def setUp(self) -> None:
        self.model = LifacNoiseModel()
        self.model.set_variable("noise_strength", 0)

        self.step_sinus = SinusAmplitudeModulationStimulus(700, 0.5, 10, start_time=0.5, duration=0.5)
        self.base_sinus = SinusAmplitudeModulationStimulus(700, 0, 0)
        self.none_stimulus = SinusAmplitudeModulationStimulus(700, 0, 0, amplitude=0)
        self.simulation_time = 1.5
        self.voltage_precision = 0.000001

    def test_simulate_fast_step_sinus(self):
        v1, spikes = self.model.simulate(self.step_sinus, self.simulation_time)
        v_fast, spikes_fast = self.model.simulate_fast(self.step_sinus, self.simulation_time)
        test = [abs(v1[i] - v_fast[i]) > self.voltage_precision for i in range(len(v1))]

        self.assertTrue(sum(test) == 0, msg="The voltage traces between the fast and slow simulation aren't equal diff: {:}".format(sum(test)))
        self.assertTrue(np.array_equal(spikes, spikes_fast), msg="The spike times between the fast and slow simulation aren't equal")

    def test_simulate_fast_base_sinus(self):
        v1, spikes = self.model.simulate(self.base_sinus, self.simulation_time)
        v_fast, spikes_fast = self.model.simulate_fast(self.base_sinus, self.simulation_time)

        test = [abs(v1[i] - v_fast[i]) > self.voltage_precision for i in range(len(v1))]

        self.assertTrue(sum(test) == 0, msg="The voltage traces between the fast and slow simulation aren't equal diff: {:}".format(sum(test)))
        self.assertTrue(np.array_equal(spikes, spikes_fast), msg="The spike times between the fast and slow simulation aren't equal")

    def test_simulate_fast_no_stimulus(self):
        v1, spikes = self.model.simulate(self.none_stimulus, self.simulation_time)
        v_fast, spikes_fast = self.model.simulate_fast(self.none_stimulus, self.simulation_time)

        test = [abs(v1[i] - v_fast[i]) > self.voltage_precision for i in range(len(v1))]

        self.assertTrue(sum(test) == 0, msg="The voltage traces between the fast and slow simulation aren't equal diff: {:}".format(sum(test)))
        self.assertTrue(np.array_equal(spikes, spikes_fast), msg="The spike times between the fast and slow simulation aren't equal")

    def test_find_v_offset(self):
        v_offsets = [50, 100, 150, 250, 1000]
        threshold = 0.01
        test_model = self.model.get_model_copy()
        stimulus = SinusAmplitudeModulationStimulus(700, 0, 0, amplitude=0)  # no stimulus
        for offset in v_offsets:
            test_model.set_variable("v_offset", offset)

            _, spikes = test_model.simulate_fast(stimulus, 5)
            goal_freq = hF.mean_freq_of_spiketimes_after_time_x(spikes, 0.0005, 1)

            if goal_freq <= threshold:
                print("test Offset ({:.1f}) generates a too low frequency: {:.2f}".format(offset, goal_freq))
                continue

            measured_offset = self.model.find_v_offset(goal_freq, stimulus, threshold)
            # print(offset, measured_offset)
            self.assertTrue(abs(offset - measured_offset) < 1)

    def test_fin_v_offset_threshold_limit(self):
        for threshold in [0, -0.0001, -2, -500]:
            self.assertRaises(ValueError, self.model.find_v_offset, 700, self.base_sinus, threshold)