import scipy as sp
import scipy.stats as st
from brian2 import ms, mV, NeuronGroup, Network, StateMonitor, SpikeMonitor, Synapses
from connectome.model import SYN
from util.timer import Timer
from util.representation import ReprMixin
from scipy.interpolate import interp1d
import brian2
brian2.prefs["codegen.target"] = "numpy"

LUT = {
    "pool_size": [  80,   100,    120,    150,   200,    250,   300],
    "log10_ge": [ -2.1, -2.25,  -2.28, -2.365,  -2.6, -2.625, -2.75],
    "log10_gi": [-0.45,  -0.7, -0.763, -0.894, -1.25,  -1.25,  -1.5]
}

class Efficacy:
    def predict(self, which, pool_size):
        pass

    def ge(self, pool_size):
        return self.predict("ge", pool_size)

    def gi(self, pool_size):
        return self.predict("gi", pool_size)


class LinearEfficacy(Efficacy):
    def __init__(self, ge_slope, ge_intercept, gi_slope, gi_intercept):
        self.models = {"ge": (ge_slope, ge_intercept),
                       "gi": (gi_slope, gi_intercept)}

    def predict(self, which, pool_size):
        slope, intercept = self.models[which]
        log10_pred = slope * pool_size + intercept
        return 10**log10_pred


class EfficacyReg(LinearEfficacy):
    def __init__(self):
        ge = st.linregress(LUT["pool_size"],
                           LUT["log10_ge"])[:2]
        gi = st.linregress(LUT["pool_size"],
                           LUT["log10_gi"])[:2]
        super().__init__(*ge, *gi)

default_reg = EfficacyReg()


class EfficacyInterp(Efficacy):
    def __init__(self):
        self.models = {"ge": interp1d(LUT["pool_size"],
                                     LUT["log10_ge"],
                                     fill_value="extrapolate"),
                       "gi":  interp1d(LUT["pool_size"],
                                     LUT["log10_gi"],
                                     fill_value="extrapolate")}

    def predict(self, which, pool_size):
        return 10**self.models[which](pool_size)

default_interp = EfficacyInterp()


class DynamicSYN(ReprMixin):
    def __init__(self, *,
                 static_syn: SYN,
                 v_resting=-70*mV,
                 v_threshold=-55*mV,
                 pool_delay=st.uniform(.5, 1.5),
                 pool_jitter=st.uniform(0, .3),
                 inh_jitter=st.uniform(.3, .6),
                 exc_refractory=2*ms,
                 g_ratio_jitter=0.5,
                 inh_refractory=1*ms,
                 efficacy=None):
        self.v_resting = v_resting
        self.v_threshold = v_threshold
        self.pool_delay = pool_delay
        self.pool_jitter = pool_jitter
        self.inh_jitter = inh_jitter
        self.exc_refractory = exc_refractory
        self.inh_refractory = inh_refractory
        self.static_syn = static_syn
        self.g_ratio_jitter = g_ratio_jitter
        self.efficacy = efficacy if efficacy is not None else EfficacyInterp()
        self.neuron_group: NeuronGroup = None


    @property
    def ge(self):
        return self.efficacy.ge(self.static_syn.pool_size)

    @property
    def gi(self):
        return self.efficacy.gi(self.static_syn.pool_size)

    def g_cv(self):
        return self.g_ratio_jitter / 2 / sp.sqrt(3)

    @property
    def nr_exc(self):
        return self.static_syn.nr_exc

    @property
    def nr_inh(self):
        return self.static_syn.nr_inh

    @property
    def nr_neurons(self):
        return self.static_syn.nr_neurons

    def ei_property(self, e, i):
        return [e] * self.nr_exc + [i] * self.nr_inh

    def run(self, simulation_time=100*ms):
        with Timer(name="make"):
            self._make_dynamic_net()

        with Timer(name="run"):
            self.dynamic_net.run(simulation_time)

    def _make_dynamic_net(self):
        inh_pools = self._make_static_network()
        self._make_neuron_group()
        synapses = self._make_synapses(inh_pools)
        self.state_monitor_ = StateMonitor(self.neuron_group, 'v', record=True)
        self.spike_monitor_ = SpikeMonitor(self.neuron_group)
        self.dynamic_net = Network(self.neuron_group, self.state_monitor_, self.spike_monitor_, synapses)
        self.apply_efficacies(self.ge, self.gi)
        self._init_neurons()

    def _make_static_network(self):
        self.net_ = self.static_syn()
        self.adjacency_ = self.net_.adjacency_matrix
        pool_connector = self.static_syn.pool_connector_
        self.pools_ = pool_connector.pools_exc
        self.pools_ = [sp.arange(self.nr_exc, self.nr_neurons)] + self.pools_
        self.pool_sets_ = [set(pool) for pool in self.pools_]
        self.n_pools_ = len(self.pool_sets_)
        return pool_connector.pools_inh

    def _make_neuron_group(self):
        taup = 20 * ms  # is used
        v_reversal_i = -80 * mV
        v_reversal_e = 0 * mV

        eqs = '''
        dv/dt = (v_resting-v)/taup: volt
        v_reversal : volt
        v_old = v : volt
        g : 1
        refractory_period : second
        '''
        namespace = {"taup": taup, "v_resting": self.v_resting,
                     "v_threshold": self.v_threshold, "v_reset": self.v_resting}
        neuron_group = NeuronGroup(self.nr_neurons, eqs, threshold='v > v_threshold', reset='v = v_reset',
                                   namespace=namespace,
                                   refractory="refractory_period", name="AllNeurons")
        neuron_group.v_reversal = self.ei_property(v_reversal_e, v_reversal_i)
        neuron_group.refractory_period = self.ei_property(self.exc_refractory,
                                                          self.inh_refractory)
        self.neuron_group = neuron_group

    def apply_efficacies(self, ge, gi):
        g = self.ei_property(ge, gi)
        g += (sp.rand(self.nr_neurons) - .5) * g * self.g_ratio_jitter
        self.g_ = g
        self.neuron_group.g = g

    def _init_neurons(self):
        initial_v = sp.ones(self.nr_neurons) * (
        self.v_resting + (self.v_threshold - self.v_resting) * 0)
        initial_v[self.pools_[
            1]] = self.v_threshold + 1 * mV  # make sure spike emitted at start
        self.neuron_group.v = initial_v

    def _make_synapses(self, inh_pools):
        synapses = Synapses(self.neuron_group, self.neuron_group,
                            on_pre='v_post += g_pre * (v_reversal_pre - v_old_post)',
                            name="AllSynapses")
        post, pre = sp.nonzero(self.adjacency_)
        synapses.connect(i=pre, j=post)

        # I -> E, I delays
        synapses.delay[self.pools_[0], sp.arange(self.nr_neurons)] = self.inh_jitter.rvs(
            len(synapses.delay[self.pools_[0], sp.arange(self.nr_neurons)])) * ms
        # E -> E, I delays
        for pre_pool_exc, post_pool_exc, post_pool_inh in zip(self.pools_[1:-1], self.pools_[2:], inh_pools[1:]):
            post_neurons = sp.concatenate((post_pool_exc, post_pool_inh))
            synapses.delay[pre_pool_exc, post_neurons] = self.pool_delay.rvs() * ms
            synapses.delay[pre_pool_exc, post_neurons] += self.pool_jitter.rvs(
                len(synapses.delay[pre_pool_exc, post_neurons])) * ms
        assert synapses.N == (self.adjacency_ != 0).sum(), "Nr synapse problem"
        self.delays_ = synapses.delay
        return synapses