from pydesim import Atomic, Content, INF, NEG_INF, Port, Errors
from . import WeightInitializers, Ports
from .Neurotransmitters import Linear as LinearNT, Proportional
from ....Ports import NEUROMODULATORS
from .ActionPotentials import Quadratic2
from .LearningRules import RSTDP
import numpy as np


class Sv2(Atomic):

    class States:
        RESTING = "RESTING"
        PROPAGATING = "PROPAGATING"

    DEFAULT = {
        "et_decay_slope": 1,
        "learning_rule": RSTDP,
        "learning_rule_params": {
            "Ap": 0.6,
            "An": 0.3,
            "Tp": 0.001,
            "Tn": 0.003,
            "lr": 0.01,
            "lr_stdp": 0.001,
            "dt_to_et_rate": 1,
            "upper_limit": 1.0,
            "lower_limit": 0.2
        },
        "weight_init": WeightInitializers.uniform,
        "AP_type": Quadratic2,
        "AP_params": {
            "velocity": 40.0,
            "interval": 0.33,
            "min_amplitude": 0.5,
            "ad": 1.5
        },
        "neurotransmitter": Proportional,
        "neurotransmitter_params": {
            "max_limit": 1,
            "recover_rate": 10,
            "dec_per_amplitude": 0.01
        }
    }

    def __init__(
            self,
            name,
            parent,
            pos,
            source_nu,
            d,
            AP_type,
            weight_init,
            learning_rule,
            et_decay_slope=1,
            AP_params={},
            weight_init_params={"low":0.1, "high":0.9},
            learning_rule_params={},
            neurotransmitter=LinearNT,
            neurotransmitter_params={},
            int_transition_callbacks=[],
            ext_transition_callbacks=[],
            state=States.RESTING,
            *args,
            **kwargs
        ):
        """
        d {[float]} -- [length from pre to post neuron]
        et_decay_rate -- in MS scale. ex) decrease 0.3 when elapsed time is 1 ms
        """
        super().__init__(
            name, 
            parent, 
            pos, 
            state=state,
            in_ports={
                Ports.AP: Port(self, Ports.AP),
                Ports.BP: Port(self, Ports.BP),
                Ports.ARRIVE: Port(self, Ports.ARRIVE),
                NEUROMODULATORS: Port(self, NEUROMODULATORS)
            },
            out_ports={
                Ports.STIMULI: Port(self, Ports.STIMULI)
            },
            int_transition_callbacks=int_transition_callbacks,
            ext_transition_callbacks=ext_transition_callbacks
        )

        self.source_nu = source_nu
        if d == 0:
            self.d = 1e-50
        else:
            self.d = d
        self.AP_type = AP_type
        self.AP_params = AP_params
        self.w = weight_init(self, **weight_init_params)
        self.learning_rule = learning_rule(self, **learning_rule_params)

        self.nt = neurotransmitter(self, **neurotransmitter_params)

        self.last_input = NEG_INF
        self.last_bp = NEG_INF

        self.ap_count = 0
        self.eligibility_trace = 0
        self.et_decay_slope = et_decay_slope * 1000

        self.initial_state.update({
            "w": self.w,
            "last_input": self.last_input,
            "last_bp": self.last_bp,
            "_visualizing": False,
            "eligibility_trace": self.eligibility_trace
        })

    def initialize(self):
        super().initialize()
        self.w = self.initial_state["w"]
        self.last_input = self.initial_state["last_input"]
        self.last_bp = self.initial_state["last_bp"]
        self._visualizing = self.initial_state["_visualizing"]
        self.eligibility_trace = self.initial_state["eligibility_trace"]
        self.nt.initialize()

    @Atomic.ext_transition_wrapper
    def ext_transition(self, content, time):
        if self.eligibility_trace < 0:
            self.eligibility_trace += self.et_decay_slope * self.elapsed_time
            if self.eligibility_trace > 0:
                self.eligibility_trace = 0
        elif self.eligibility_trace > 0:
            self.eligibility_trace -= self.et_decay_slope * self.elapsed_time
            if self.eligibility_trace < 0:
                self.eligibility_trace = 0
        self.nt.recover(self.elapsed_time)

        if self.state is self.States.RESTING:
            if content.port == self.in_ports[Ports.AP]:
                self.hold_in(self.States.RESTING, INF)
                value = self.nt.release(content.value)
                self.create_AP(value, time)

            elif content.port == self.in_ports[Ports.BP]:
                self.hold_in(self.States.RESTING, INF)
                self.last_bp = time
                self.learning_rule(time, content)

            elif content.port == self.in_ports[Ports.ARRIVE]:
                self.hold_in(self.States.RESTING, INF)
                self.last_input = time
                self.learning_rule(time, content)
                
            elif content.port == self.in_ports[NEUROMODULATORS]:
                self.hold_in(self.States.RESTING, INF)
                self.learning_rule(time, content)
        else:
            raise Errors.StateError(
                "No transition in current state defined: " + self.state)

    def create_AP(self, A, time):
        if not self._visualizing:
            self.AP_type(
                self.name + "-AP-" + str(self.ap_count), self.parent, self,
                self.d, A * self.w, time,
                **self.AP_params)
        self.ap_count += 1

    @staticmethod
    def connect(name, source, target, params):
        """
        Args:
        name (str): [Name of the synapse]\n
        source (NU): [NU that is connecting to this NU]\n
        target (NU): [NU that is being connected to]\n
        params (dict): [Parameters used in creating new synapse]\n

        Returns:
        [Synapse]: [The created synapse]
        """
        length = np.linalg.norm(
            np.array(target.cell_pos) - np.array(source.cell_pos))
        synapse = Sv2(
            name,
            target,
            target.cell_pos,
            source,
            length,
            **params)
            
        target.in_ports[synapse.name] = Port(target, synapse.name)
        target.couple(target.in_ports[synapse.name],
                      synapse.in_ports[Ports.AP])
        target.couple(
            target.membrane.out_ports[Ports.AP],
            synapse.in_ports[Ports.BP])
        return synapse
