"""
This script is _not_ to be viewed as a high quality reference solution, but 
rather a simple hack to illustrate the efficiency of different
approaches. Timings are done with the time module, but seem to be consistent. 
"""

from PIL import Image
import numpy as np
import time
from scipy import weave

#parameters used by all versions:
iter = 10
kappa = 0.1
eps = 0.001
test = True

############################
#Python loop version
############################
#from image to Numpy:
data = np.array(Image.open("disasterbefore.jpg"),dtype=float)
n, m = data.shape[:2]
data_new = np.zeros(data.shape)
data_new[:] = data 

t0 = time.clock()
for q in xrange(iter):
    for i in xrange(1,n-1):
        for j in xrange(1,m-1):
            data_new[i][j] = data[i][j] +kappa*(data[i-1][j]
                                        +data[i][j-1]-4*data[i][j]
                                        +data[i][j+1]+data[i+1][j]) ;
    data[:] = data_new    
print "Python loop version: ", time.clock()-t0
#from array back to image
Image.fromarray(data.astype('uint8')).save("disaster_after.jpg")
#save Py data as reference
data_ref = np.zeros(data.shape)
data_ref[:] = data

#########################
#Vectorized numpy version
#########################
data = np.array(Image.open("disasterbefore.jpg"),dtype=float)

def laplacian(v):
    return (v[0:-2,1:-1] + v[1:-1,0:-2] - 4*v[1:-1,1:-1] 
            + v[1:-1,2:] + v[2:,1:-1] )

t0 = time.clock()
for q in xrange(iter):
    new_data = laplacian(data)
    data[1:-1,1:-1] += kappa*new_data

print "Numpy version: ", time.clock()-t0
Image.fromarray(data.astype('uint8')).save("disaster_after2.jpg")

if test:
    print np.sum(abs(data-data_ref) >= eps)


#########################
#Weave 1: additional C loop for data copy
#########################
data = np.array(Image.open("disasterbefore.jpg"),dtype=float)
data_new = np.zeros_like(data)
data_new[:] = data

code = r"""
for (int q=0; q<iter; q++)
{
   for (int i=1; i<n-1; i++)
      {
          for (int j=1; j<m-1; j++)
            {   
                data_new[i*m+j] = data[i*m+j]+kappa*(data[(i-1)*m+j]
                       +data[i*m + j-1] -4*data[i*m+j] +data[i*m +j+1]
                       +data[(i+1)*m+j]);
            }
        }
  
   for (int i=1; i<n-1; i++)
      for (int j=1; j<m-1; j++)
         data[i*m+j] = data_new[i*m+j];
}
"""
t0 = time.clock()
io = ["data", "data_new","iter","kappa","m","n"]
weave.inline(code, io)
print "Weave version 1 (extra loop): ", time.clock()-t0
Image.fromarray(data.astype('uint8')).save("disaster_after_w1.jpg")
if test:
    print np.sum(data-data_ref >= eps)


#########################
#Weave 2: pointer swap
#########################
data = np.array(Image.open("disasterbefore.jpg"),dtype=float)
data_new = np.zeros_like(data)
data_new[:] = data
tmp = np.zeros_like(data)
code = r"""
for (int q=0; q<iter; q++)
{
   for (int i=1; i<n-1; i++)
      {
          for (int j=1; j<m-1; j++)
            {
                data_new[i*m+j] = data[i*m+j]+kappa*(data[(i-1)*m+j]
                       +data[i*m + j-1] -4*data[i*m+j] +data[i*m +j+1]
                       +data[(i+1)*m+j]);
            }
        }

tmp = data; 
data = data_new;
data_new = tmp;
}
"""
io = ["data", "data_new","tmp","iter","kappa","m","n"]
t0 = time.clock()
weave.inline(code, io)
print "Weave version 2 (pointer swap): ", time.clock()-t0
Image.fromarray(data.astype('uint8')).save("disaster_after_w2.jpg")
if test:
    print np.sum(data-data_ref >= eps)

#########################
#Weave 3: outer loop in Python
#########################
data = np.array(Image.open("disasterbefore.jpg"),dtype=float)
data_new = np.zeros_like(data)
data_new[:] = data

code = r"""
   for (int i=1; i<n-1; i++)
      {
          for (int j=1; j<m-1; j++)
            {
                data_new[i*m+j] = data[i*m+j]+kappa*(data[(i-1)*m+j]
                       +data[i*m + j-1] -4*data[i*m+j] +data[i*m +j+1]
                       +data[(i+1)*m+j]);
            }
        }
"""
io = ["data", "data_new","kappa","m","n"]
t0 = time.clock()
for q in range(iter):
    weave.inline(code, io)
    data[:] = data_new
print "Weave version 3 (outer Py loop): ", time.clock()-t0
Image.fromarray(data.astype('uint8')).save("disaster_after_w3.jpg")
if test:
    print np.sum(data-data_ref >= eps)