#!/usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt
import pdb

def lms(u, d, M, step, leak=0, initCoeffs=None, N=None, returnCoeffs=False):

    if N is None:
        N = len(u)-M+1

    initCoeffs = np.zeros(M)

    # Initialization
    y = np.zeros(N) # Filter output
    e = np.zeros(N) # Error signal
    w = initCoeffs  # Initialise equaliser
    leakstep = (1 - step*leak)
    if returnCoeffs:
        W = np.zeros((N, M)) # Matrix to hold coeffs for each equaliser step

    # Equalise
    for n in xrange(N):

        x = np.flipud(u[n:n+M]) #

        y[n] = np.dot(x, w)

        e[n] = d[n+M-1] - y[n]

        w = leakstep * w + step * x * e[n]

        y[n] = np.dot(x, w)

        if returnCoeffs:
            W[n] = w

        if returnCoeffs:
            w = W

    return y, e, w



np.random.seed(1337)
ulen    = 2000
coeff   = np.concatenate(([1], np.zeros(10), [-0.9], np.zeros(7), [0.1]))
u       = np.random.randn(ulen)
d       = np.convolve(u, coeff)
M       = 20 # No. of taps
step    = 0.003 # Step size
y, e, w = lms(u, d, M, step)
print np.allclose(w, coeff)



plt.figure()
plt.subplot(1,1,1)
plt.plot(np.abs(e[0:400]))

plt.show()

#pdb.set_trace()