#!/usr/bin/python

import math
from PIL import Image

Z = math.sqrt(0.5)

# one pass of the forward transform
def dwt(vec):
    N = len(vec)
    low = [Z * (vec[i*2] + vec[i*2+1]) for i in range(N/2)]
    high = [Z * (vec[i*2+1] - vec[i*2]) for i in range(N/2)]
    return low + high

# one pass of the inverse transform
def idwt(vec):
    N = len(vec)
    low = vec[0:N/2]
    high = vec[N/2:N]
    return [Z * (low[i/2] + ((i & 1)*2-1) * high[i/2]) for i in range(N)]

# recursive 2d forward transform, lev interations
# image size should be more than 2^lev

def dwt2d(arr, lev):
    for iter in range(lev):
        rows = len(arr) >> iter
        cols = len(arr[0]) >> iter
        for row in range(rows):
            vec = [arr[row][col] for col in range(cols)]
            vec = dwt(vec)
            for col in range(cols):
                arr[row][col] = int(vec[col])
        for col in range(cols):
            vec = [arr[row][col] for row in range(rows)]
            vec = dwt(vec)
            for row in range(rows):
                arr[row][col] = vec[row]

# recursive 2d inverse transform, lev interations
# image size should be more than 2^lev

def idwt2d(arr, lev):
    for iter in range(lev-1, -1, -1):
        rows = len(arr) >> iter
        cols = len(arr[0]) >> iter
        for col in range(cols):
            vec = [arr[row][col] for row in range(rows)]
            vec = idwt(vec)
            for row in range(rows):
                arr[row][col] = vec[row]
        for row in range(rows):
            vec = [arr[row][col] for col in range(cols)]
            vec = idwt(vec)
            for col in range(cols):
                arr[row][col] = int(vec[col])

# read image

im = Image.open("kuva.png")
rows, cols = im.size
levels = int(math.log(rows,2))-2

imi = im.load()    # pixels in PixelArray
imf = [[float(imi[row, col]) for col in range(cols)] for row in range(rows)]

dwt2d(imf, levels) # forward transform

T = 0.0            # set threshold

for row in range(rows):
    for col in range(cols):
        c = imf[row][col] # coefficient

        # add your magic here
        
        imf[row][col] = c

idwt2d(imf, levels) # inverse transform

# update pixels
for row in range(rows):
    for col in range(cols):
        imi[row, col] = int(imf[row][col]) 

# display image
im.show()