#!/usr/bin/env python
from __future__ import division
from numpy import *
import scipy.linalg as la
from matplotlib import pyplot as plt
from matplotlib import animation
#Daubechie's fourth-order coefficients for discrete wavelet transform

c = array([1+sqrt(3), 3+sqrt(3),3-sqrt(3),1-sqrt(3)])/(4*sqrt(2))
c2 = array([c[3],-c[2],c[1],-c[0]])
def e(i,n):
	e = zeros(n); e[i]=1; return e
def step_back(x):
	N = shape(x)[-1]; assert(N%2==0)
	M = zeros((N,N))
	v = zeros(N); v[:4] = c
	v2 = zeros(N); v2[:4] = c2
	M[::2,::2] = la.circulant(v)[::2,::2] #don't need transpose because linalg.circulant is column based
	M[1::2,1::2] = la.circulant(v2)[::2,::2]
	S = zeros((N,N)) #separation matrix
	S[::2,:N/2] = la.circulant(e(0,N/2))
	S[1::2,N/2:] = la.circulant(e(0,N/2))
	return dot(M,dot(S.T,x))

def main():
	x = zeros((12,2**12))
	x[0,4] = 1
	x[0,29] = 1
	#x[0,:] = 1-linspace(0,1,2**12) #linear ramp
	for i in range(11):
		end = 4*2**i
		x[i+1] = x[i]
		x[i+1,:end] = step_back(x[i,:end])

	fig = plt.figure()
	ax = plt.axes(xlim=(0, 2**12), ylim=(-1e-3, 1e-2))
	plt.title('Preimage of $e_5 + e_{30}$ under discrete wavelet transform')
	line, = ax.plot([], [], lw=2)
	def init():
		line.set_data([], []); return line,
	def animate(i):
		if i>11: i=11
		ax.set_ylim(0,amax(1.1*x[i])) 
		line.set_data(arange(2**12), x[i]); return line,
	anim = animation.FuncAnimation(fig, animate, init_func=init,frames=25, interval=20)
	#anim.save('wavelet.gif', writer='imagemagick', fps=5)
	plt.show()

if __name__ == '__main__':
	main()