import numpy as np
import math,random
import matplotlib.pyplot as plt
import pygame
from pygame.locals import *
 #filter coefficients
c0 = (1+math.sqrt(3))/(4*math.sqrt(2))
c1 = (3+math.sqrt(3))/(4*math.sqrt(2))
c2 = (3-math.sqrt(3))/(4*math.sqrt(2))
c3 = (1-math.sqrt(3))/(4*math.sqrt(2))

#initialize filters
filter_a = [c0,c1,c2,c3]
filter_b = [c3,-c2,c1,-c0]


def wavelet2signal_step(signal):
    signal_dim = len(signal)
    
    #construct filterbank
    filter_bank = np.zeros((signal_dim,signal_dim))
    filter_bank[::2,:4] = filter_a
    filter_bank[1::2,:4] = filter_b
    #sheer the filtermatrix
    for i in xrange(signal_dim):
        filter_bank[i,:] = np.roll(filter_bank[i,:],2*(i/2))
    #make split_matrix
    split_matrix = np.zeros((signal_dim,signal_dim))
    split_matrix[:,:1] = 1
    for i in xrange(signal_dim):
        split_matrix[i,:] = np.roll(split_matrix[i,:],2*i)
    split_matrix[signal_dim/2:,:] = np.roll(split_matrix[signal_dim/2:,:],1)
    
    signal =np.matrix(filter_bank).transpose()* np.matrix(split_matrix).transpose()* signal
    return signal
    
def signal2wavelet_step(signal):
    signal_dim = len(signal)
    
    #construct filterbank
    filter_bank = np.zeros((signal_dim,signal_dim))
    filter_bank[::2,:4] = filter_a
    filter_bank[1::2,:4] = filter_b
    #sheer the filtermatrix
    for i in xrange(signal_dim):
        filter_bank[i,:] = np.roll(filter_bank[i,:],2*(i/2))
    #make split_matrix
    split_matrix = np.zeros((signal_dim,signal_dim))
    split_matrix[:,:1] = 1
    for i in xrange(signal_dim):
        split_matrix[i,:] = np.roll(split_matrix[i,:],2*i)
    split_matrix[signal_dim/2:,:] = np.roll(split_matrix[signal_dim/2:,:],1)
    
    
    signal = np.matrix(split_matrix)*np.matrix(filter_bank)*signal
    return signal


signal_dim = 2**12
signal = np.zeros(signal_dim)


signal[4] = 1.0
signal[29] = 1.0

    
#wavelet 2 signal
signal = np.matrix(signal).transpose() 

for i in xrange(2,int(math.log(signal_dim,2)+1)):
    signal[:2**i] = wavelet2signal_step(signal[:2**i])

signal = np.array(signal).flatten()
plt.plot(xrange(len(signal)),signal)


#signal 2 wavelet

signal = np.matrix(signal).transpose() 
for i in reversed(xrange(2,int(math.log(signal_dim,2)+1))):
    signal[:2**i] = signal2wavelet_step(signal[:2**i])

#only for display purposes hack
signal[4] = .3
signal[29] = .3
### hack end
signal = np.array(signal).flatten()
plt.plot(xrange(len(signal)),signal)
plt.show()