from stimuli.AbstractStimulus import AbstractStimulus
from models.AbstractModel import AbstractModel
import numpy as np


class LIFACModel(AbstractModel):
    # all times in milliseconds
    # possible mem_res: 100 * 1000000

    DEFAULT_VALUES = {"mem_tau": 20,
                      "v_base": 0,
                      "v_zero": 0,
                      "threshold": 10,
                      "step_size": 0.01,
                      "delta_a": 1,
                      "tau_a": 20}

    # membrane time constant tau = mem_cap*mem_res
    def __init__(self, params: dict = None):
        super().__init__(params)

        self.voltage_trace = []
        self.adaption_trace = []
        self.spiketimes = []
        self.stimulus = None
        # self.frequency_trace = []

    def simulate(self, stimulus: AbstractStimulus, total_time_s):
        self.stimulus = stimulus
        output_voltage = []
        adaption = []
        spiketimes = []
        current_v = self.parameters["v_zero"]
        current_a = 0

        for time_point in np.arange(0, total_time_s*1000, self.parameters["step_size"]):
            v_next = self._calculate_voltage_step(current_v, stimulus.value_at_time_in_ms(time_point) - current_a)
            a_next = self._calculate_adaption_step(current_a)

            if v_next > self.parameters["threshold"]:
                v_next = self.parameters["v_base"]
                spiketimes.append(time_point/1000)
                a_next += self.parameters["delta_a"] / (self.parameters["tau_a"] / 1000)

            output_voltage.append(v_next)
            adaption.append(a_next)

            current_v = v_next
            current_a = a_next

        self.voltage_trace = output_voltage
        self.adaption_trace = adaption
        self.spiketimes = spiketimes

        return output_voltage, spiketimes

    def _calculate_voltage_step(self, current_v, input_v):
        v_base = self.parameters["v_base"]
        step_size = self.parameters["step_size"]
        # mem_res = self.parameters["mem_res"]
        mem_tau = self.parameters["mem_tau"]
        return current_v + (step_size * (v_base - current_v + input_v)) / mem_tau

    def _calculate_adaption_step(self, current_a):
        step_size = self.parameters["step_size"]
        return current_a + (step_size * (-current_a)) / self.parameters["tau_a"]

    def min_stimulus_strength_to_spike(self):
        return self.parameters["threshold"] - self.parameters["v_base"]

    def get_sampling_interval(self):
        return self.parameters["step_size"]

    def get_frequency(self):
        # TODO also change simulates_frequency() if any calculation is added!
        raise NotImplementedError("No calculation implemented yet for the frequency.")

    def get_spiketimes(self):
        return self.spiketimes

    def get_voltage_trace(self):
        return self.voltage_trace

    def get_adaption_trace(self):
        return self.adaption_trace

    def simulates_frequency(self) -> bool:
        return False

    def simulates_spiketimes(self) -> bool:
        return True

    def simulates_voltage_trace(self) -> bool:
        return True

    def get_recording_times(self):
        # [delay, stimulus_start, stimulus_duration, time_to_end]
        self.stimulus = AbstractStimulus()
        delay = 0
        start = self.stimulus.get_stimulus_start_s()
        duration = self.stimulus.get_stimulus_duration_s()
        total_time = len(self.voltage_trace) / self.parameters["step_size"]

        return [delay, start, duration, total_time]