class RectifiedLinearWithThreshold(nengo.RectifiedLinear):
    threshold = nengo.params.NumberParam('threshold')

    def __init__(self, amplitude=1, threshold=0):
        super().__init__(amplitude=amplitude)
        self.threshold = threshold

    def step_math(self, dt, J, output):
        output[...] = self.amplitude * np.where(J>self.threshold, J, 0)
        


"""
Explanation: First, we need to add a new neuron model to nengo.  It has a RectifiedLinear neuron already, but it doesn't have an option to set a threshold.  So here we just subclass that model and create that parameter.
Note that we could also do something similar to set a maximum output for the ReLU as well.
End of explanation
"""


class LeakyIntegrator(nengo.LIF):
    threshold = nengo.params.NumberParam('threshold')
    
    def __init__(self, amplitude=1, threshold=0, tau_rc=0.02, min_voltage=0):
        super().__init__(amplitude=amplitude, tau_rc=tau_rc, tau_ref=0, min_voltage=min_voltage)
        self.threshold = threshold
        
    def step_math(self, dt, J, output, voltage, refractory_time):
        # to map to the paper: output is y, voltage is u, and J is the inputs times the dictionary matrix
        
        # y = T(u)
        output[:] = np.where(voltage>self.threshold, voltage, 0)
        
        # du/dt = -u + y + phi*r
        dv = -voltage + output + J
        
        # perform the voltage update
        voltage += dv*(dt/self.tau_rc)
        
        # limit the voltage to positive only
        voltage[voltage < self.min_voltage] = self.min_voltage

"""
Explanation: Now we make a new neuron model: the (non-spiking) leaky integrator.  We base this on the existing LIF model (since it already has a voltage as an internal state variable).
End of explanation
"""


class NegOja(nengo.Oja):
    pass

from nengo.builder.learning_rules import SimOja, get_pre_ens, get_post_ens, build_or_passthrough

@nengo.builder.Builder.register(NegOja)
def build_negoja(model, oja, rule):
    conn = rule.connection
    pre_activities = model.sig[get_pre_ens(conn).neurons]['out']
    post_activities = model.sig[get_post_ens(conn).neurons]['out']
    pre_filtered = build_or_passthrough(model, oja.pre_synapse, pre_activities)
    post_filtered = build_or_passthrough(
        model, oja.post_synapse, post_activities)

    model.add_op(SimOja(pre_filtered,
                        post_filtered,
                        model.sig[conn]['weights'],
                        model.sig[rule]['delta'],
                        learning_rate=-oja.learning_rate,
                        beta=oja.beta))

    # expose these for probes
    model.sig[rule]['pre_filtered'] = pre_filtered
    model.sig[rule]['post_filtered'] = post_filtered

"""
Explanation: For the V1->Residual connection, we need a sign-flipped Oja rule.  This can be done by setting a negative learning rate for the normal Oja rule.  However, nengo currently disallows a negative value for a learning rate.  So, we implement a negative variant of the Oja rule.
End of explanation
"""


seed = 10
np.random.seed(seed=seed)

N = 10    # size of the input
J = 50    # number of neurons

# generate an example image
span = np.linspace(0, 1, N)
image = (np.sin(2*np.pi*span)+1)/2

model = nengo.Network(seed=seed)
with model:
    # the input, which will be N-dimensional 
    #  (note that we can replace this with input that changes over time to give a more real-world example)
    stim = nengo.Node(image)
    
    # for the residual, we use N neurons which are ReLU neurons with a threshold of -infinity.  That is, 
    # there is no nonlinearity here at all.
    residual = nengo.Ensemble(n_neurons=N, dimensions=1,
                          neuron_type=RectifiedLinearWithThreshold(threshold=-np.inf),
                          gain=nengo.dists.Choice([1]),
                          bias=nengo.dists.Choice([0]))
    
    v1 = nengo.Ensemble(n_neurons=J, dimensions=1,
                          neuron_type=LeakyIntegrator(threshold=0.1, tau_rc=0.001),
                          gain=nengo.dists.Choice([1]),
                          bias=nengo.dists.Choice([0]))

    # the initial dictionary
    w = np.random.uniform(-0.1, 0.1, (J, N))
    
    # now we make the learning connections.  Nengo has Oja built in, which is Hebbian with a 
    # stabilizing factor.  We can set beta=0 to get rid of that factor, to make it just Hebbian.
    learning_rate=1e-1
        
    tau = 0        # low-pass filter to use when using the w matrix
    tau_learn = 0  # low-pass filter to use when computing the Hebbian update
        
    c1 = nengo.Connection(residual.neurons, v1.neurons, transform=w, synapse=tau,
                     learning_rule_type=nengo.Oja(beta=0, pre_synapse=tau_learn, post_synapse=tau_learn, learning_rate=learning_rate))
        
    c2 = nengo.Connection(v1.neurons, residual.neurons, transform=-w.T, synapse=tau,
                     learning_rule_type=NegOja(beta=0, pre_synapse=tau_learn, post_synapse=tau_learn, learning_rate=learning_rate))
    
    # subtract the reconstruction from the stim to get the residual layer
    nengo.Connection(stim, residual.neurons, synapse=0)
    
    # record data from each layer
    p_v1 = nengo.Probe(v1.neurons)
    p_res = nengo.Probe(residual.neurons)
    p_c1 = nengo.Probe(c1, 'weights')
    p_c2 = nengo.Probe(c2, 'weights')
    
    
sim = nengo.Simulator(model)
sim.run(0.4)
                        

plt.plot(sim.trange(), sim.data[p_v1])
plt.title('activity of V1 neurons')
plt.show()

"""
Explanation: Now we build our model.  We have three basic "layers":  the image, the residual, and v1
End of explanation
"""


w = sim.data[p_c1]    # the weight matrix at each point in time
v1 = sim.data[p_v1]   # the activity of V1 over time



recon = np.zeros((w.shape[0], w.shape[2]))
for i in range(w.shape[0]):
    recon[i] = np.dot(v1[i], w[i])
# TODO: figure out a nicer way to do the above math in numpy    
    
plt.imshow(recon, aspect='auto')
plt.ylabel('time step')
plt.title('reconstruction')
plt.colorbar()
plt.figure()
plt.plot(image, label='original')
plt.plot(recon[-1], label='final reconstruction')
plt.legend()
plt.show()

plt.imshow(sim.data[p_res][:,:], aspect='auto')
plt.ylabel('time step')
plt.title('residual')
plt.colorbar()


plt.imshow(sim.data[p_v1], aspect='auto')
plt.colorbar()
plt.title('V1')
plt.show()

"""
Explanation: Now what does the reconstruction look like?
End of explanation
"""


