import numpy as np
import matplotlib.pyplot as plt
from PIL import Image as im
import math


original_image = im.open("DSC_0222.png")
width, height = original_image.size

'''Checking the dimension of image. If it is a perfect multiple of 4, subtract 1 from height and width to
adjust the first row and first column for prediction'''
if width%4 == 0:
    new_width = int(width / 4) - 1
else:
    new_width = int(width / 4)
new_width *= 4
if height%4 == 0:
    new_height = int(height / 4) - 1
else:
    new_height = int(height / 4)
new_height *= 4

'''Reading image for RGB values and converting to Luma(Y)'''
R = np.zeros([height, width])
G = np.zeros([height, width])
B = np.zeros([height, width])
Y = np.zeros([height, width])
predicted_Y = np.zeros([height, width])
predicted_image = np.zeros((height, width, 3), dtype=np.uint8)

best_pred = [0] * 10

'''Prediction block arrays'''
P = np.zeros([4, 4])
P0 = np.zeros([4, 4])
P1 = np.zeros([4, 4])
P2 = np.zeros([4, 4])
P3 = np.zeros([4, 4])
P4 = np.zeros([4, 4])
P5 = np.zeros([4, 4])
P6 = np.zeros([4, 4])
P7 = np.zeros([4, 4])
P8 = np.zeros([4, 4])
P9 = np.zeros([4, 4])

pixel_values = original_image.load()

for x in range(height):
    for y in range(width):
        R[x, y], G[x, y], B[x, y] = pixel_values[y, x]
        Y[x, y] = int((R[x, y] + G[x, y] + B[x, y]) / 3)


'''Assigning Prediction matrix the first row and first column'''
for x in range(height):
    predicted_image[x, 0] = Y[x, 0]
    predicted_Y[x, 0] = Y[x, 0]
for y in range(width):
    predicted_image[0, y] = Y[0, y]
    predicted_Y[0, y] = Y[0, y]

'''Calculating remainder dimension if not a perfect multiple of 4'''
height_remainder = height - new_height - 1
width_remainder = width - new_width - 1

'''Assigning Prediction matrix the last row/s and last column/s if any'''

for x in range(height_remainder):
    predicted_image[height - x - 1, 0] = Y[height_remainder - x, 0]
    predicted_Y[height - x - 1, 0] = Y[height_remainder - x, 0]
for y in range(width_remainder):
    predicted_image[0, width - y - 1] = Y[0, width_remainder - y]
    predicted_Y[0, width - y - 1] = Y[0, width_remainder - y]

'''
----------------------------------------------------------------------------------------------------------
Taken from WebP (averaging)
4x4 Luma Prediction Modes
'''
def AVG3(a, b, c):
    return int(a + 2 * b + c + 2)>>2


def AVG2(a, b):
    return int(a + b + 1)>>1


'''Vertical Mode'''
def VE4(a, b, c, d):
    X = a
    A = b[0]
    B = b[1]
    C = b[2]
    D = b[3]
    E = c[0]
    for i in range(4):
        P0[i, 0] = AVG3(X, A, B)
        P0[i, 1] = AVG3(A, B, C)
        P0[i, 2] = AVG3(B, C, D)
        P0[i, 3] = AVG3(C, D, E)


'''Horizontal Mode'''
def HE4(a, b, c, d):
    X = a
    I = d[0]
    J = d[1]
    K = d[2]
    L = d[3]
    for i in range(4):
        P1[0, i] = AVG3(X, I, J)
        P1[1, i] = AVG3(I, J, K)
        P1[2, i] = AVG3(J, K, L)
        P1[3, i] = AVG3(K, L, L)

'''DC Mode'''
def DC4(A, B, C, D):
    dc = 0
    for i in range(4):
        dc += B[i] + D[3 - i]
    P3.fill(int(dc)>>3)
    # print P

'''Right Down Mode'''
def RD4(a, b, c, d):
    X = a
    I = d[0]
    J = d[1]
    K = d[2]
    L = d[3]
    A = b[0]
    B = b[1]
    C = b[2]
    D = b[3]
    P4[0, 3] = AVG3(D, C, B)
    P4[0, 2] = P4[1, 3] = AVG3(A, B, C)
    P4[0, 1] = P4[1, 2] = P4[2, 3] = AVG3(X, A, B)
    P4[0, 0] = P4[1, 1] = P4[2, 2] = P4[3, 3] = AVG3(A, X, I)
    P4[1, 0] = P4[2, 1] = P4[3, 2] = AVG3(X, I, J)
    P4[2, 0] = P4[3, 1] = AVG3(I, J, K)
    P4[3, 0] = AVG3(J, K, L)

'''Left Down Mode'''
def LD4(a, b, c, d):
    X = a
    A = b[0]
    B = b[1]
    C = b[2]
    D = b[3]
    E = c[0]
    F = c[1]
    G = c[2]
    H = c[3]
    P5[0, 0] = AVG3(A, B, C)
    P5[1, 0] = P5[0, 1] = AVG3(B, C, D)
    P5[2, 0] = P5[1, 1] = P5[0, 2] = AVG3(C, D, E)
    P5[3, 0] = P5[2, 1] = P5[1, 2] = P5[0, 3] = AVG3(D, E, F)
    P5[3, 1] = P5[2, 2] = P5[1, 3] = AVG3(E, F, G)
    P5[3, 2] = P5[2, 3] = AVG3(F, G, H)
    P5[3, 3] = AVG3(G, H, H)

'''Vertical Right Mode'''
def VR4(a, b, c, d):
    X = a
    I = d[0]
    J = d[1]
    K = d[2]
    A = b[0]
    B = b[1]
    C = b[2]
    D = b[3]
    P6[0, 0] = P6[2, 1] = AVG2(X, A)
    P6[0, 1] = P6[2, 2] = AVG2(A, B)
    P6[0, 2] = P6[2, 3] = AVG2(B, C)
    P6[0, 3] = AVG2(C, D)

    P6[3, 0] = AVG3(K, J, I)
    P6[2, 0] = AVG3(J, I, X)
    P6[1, 0] = P6[3, 1] = AVG3(I, X, A)
    P6[1, 1] = P6[3, 2] = AVG3(X, A, B)
    P6[1, 2] = P6[3, 3] = AVG3(A, B, C)
    P6[1, 3] = AVG3(B, C, D)

'''Vertical Left Mode'''
def VL4(a, b, c, d):
    A = b[0]
    B = b[1]
    C = b[2]
    D = b[3]
    E = c[0]
    F = c[1]
    G = c[2]
    H = c[3]
    P7[0, 0] = AVG2(A, B)
    P7[0, 1] = P7[2, 0] = AVG2(B, C)
    P7[0, 2] = P7[2, 1] = AVG2(C, D)
    P7[0, 3] = P7[2, 2] = AVG2(D, E)
    P7[1, 0] = AVG3(A, B, C)
    P7[1, 1] = P7[3, 0] = AVG3(B, C, D)
    P7[1, 2] = P7[3, 1] = AVG3(C, D, E)
    P7[1, 3] = P7[3, 2] = AVG3(D, E, F)
    P7[2, 3] = AVG2(E, F)
    P7[3, 3] = AVG3(E, F, G)

'''Horizontal Up Mode'''
def HU4(a, b, c, d):
    I = d[0]
    J = d[1]
    K = d[2]
    L = d[3]
    P8[0, 0] = AVG2(I, J)
    P8[0, 2] = P8[1, 0] = AVG2(J, K)
    P8[1, 2] = P8[2, 0] = AVG2(K, L)
    P8[0, 1] = AVG3(I, J, K)
    # P8[2, 2] = P8[0, 3] = AVG2(K, L)
    P8[0, 3] = P8[1, 1] = AVG3(J, K, L)
    P8[1, 3] = P8[2, 1] = AVG3(K, L, L)
    P8[2, 2] = P8[2, 3] = P8[3, 0] = P8[3, 1] = P8[3, 2] = P8[3, 3] = L

'''Horizontal Down Mode'''
def HD4(a, b, c, d):
    X = a
    I = d[0]
    J = d[1]
    K = d[2]
    L = d[3]
    A = b[0]
    B = b[1]
    C = b[2]
    P9[0, 0] = P9[1, 2] = AVG2(I, X)
    P9[1, 0] = P9[2, 2] = AVG2(J, I)
    P9[2, 0] = P9[3, 2] = AVG2(K, J)
    P9[3, 0] = AVG2(L, K)
    P9[0, 3] = AVG3(A, B, C)
    P9[0, 2] = AVG3(X, A, B)
    P9[0, 1] = P9[1, 3] = AVG3(I, X, A)
    P9[1, 1] = P9[2, 3] = AVG3(J, I, X)
    P9[2, 1] = P9[3, 3] = AVG3(K, J, I)
    P9[3, 1] = AVG3(L, K, J)

'''True Motion Prediction'''
def TM4(a, b, c, d):
    for i in range(4):
        for j in range(4):
            P2[i, j] = d[i] + b[j] - a
            if P2[i, j] < 0:
                P2[i, j] = 0
            if P2[i, j] > 255:
                P2[i, j] = 255

'''
-----------------------------------------------------------------------------------------------------------
Best Fit only using SSE as I am just checking the quality of predicted image and not the size
'''
def best_fit(a, b, cc, d):
    cnt = 0
    for i in range(4):
        for j in range(4):
            diff = Y[a + i, b + j] - cc[i, j]
            cnt += diff**2
    best_pred[d] = cnt


'''
Main entry point is here
D - Stores elements to the left of the 4x4 micro block (I, J, K, L in webp source code)
B - Stores elements to the top of the 4x4 micro block (A, B, C, D in webp source code)
C - Stores elements to the top right of the 4x4 micro block (E, F, G, H in webp source code)
A - The element to the left top corner of the 4x4 micro block (Q in webp source code)
'''
D = [0] * 4
B = [0] * 4
C = [0] * 4
for x in range(1, new_height - 3, 4):
    for y in range(1, new_width - 3, 4):
        A = predicted_Y[x - 1, y - 1]

        if not y == new_width - 4:
            for i in range(4):
                D[i] = predicted_Y[x + i, y - 1]
                B[i] = predicted_Y[x - 1, y + i]
                C[i] = predicted_Y[x - 1, y + i + 4]

            '''Function calls for the different prediction modes'''
            VE4(A, B, C, D)
            best_fit(x, y, P0, 0)
            HE4(A, B, C, D)
            best_fit(x, y, P1, 1)
            TM4(A, B, C, D)
            best_fit(x, y, P2, 2)
            DC4(A, B, C, D)
            best_fit(x, y, P3, 3)
            RD4(A, B, C, D)
            best_fit(x, y, P4, 4)
            LD4(A, B, C, D)
            best_fit(x, y, P5, 5)
            VR4(A, B, C, D)
            best_fit(x, y, P6, 6)
            VL4(A, B, C, D)
            best_fit(x, y, P7, 7)
            HU4(A, B, C, D)
            best_fit(x, y, P8, 8)
            HD4(A, B, C, D)
            best_fit(x, y, P9, 9)
            select = best_pred.index(min(best_pred))

            '''Choosing the predicted block of the best mode'''
            if select == 0:
                P = P0
            elif select == 1:
                P = P1
            elif select == 2:
                P = P2
            elif select == 3:
                P = P3
            elif select == 4:
                P = P4
            elif select == 5:
                P = P5
            elif select == 6:
                P = P6
            elif select == 7:
                P = P7
            elif select == 8:
                P = P8
            elif select == 9:
                P = P9

        else:
            '''For the rightmost column of 4x4 blocks, the modes which use 'C' will not be used'''
            for i in range(4):
                D[i] = predicted_Y[x + i, y - 1]
                B[i] = predicted_Y[x - 1, y + i]

                HE4(A, B, C, D)
                best_fit(x, y, P1, 1)
                TM4(A, B, C, D)
                best_fit(x, y, P2, 2)
                DC4(A, B, C, D)
                best_fit(x, y, P3, 3)
                RD4(A, B, C, D)
                best_fit(x, y, P4, 4)
                VR4(A, B, C, D)
                best_fit(x, y, P6, 6)
                HU4(A, B, C, D)
                best_fit(x, y, P8, 8)
                HD4(A, B, C, D)
                best_fit(x, y, P9, 9)
                select = best_pred.index(min(best_pred))

                if select == 1:
                    P = P1
                elif select == 2:
                    P = P2
                elif select == 3:
                    P = P3
                elif select == 4:
                    P = P4
                elif select == 6:
                    P = P6
                elif select == 7:
                    P = P7
                elif select == 8:
                    P = P8
                elif select == 9:
                    P = P9

        for k in range(4):
            for l in range(4):
                pred = int((P[k, l]))
                predicted_Y[x + k, y + l] = pred
                predicted_image[x + l, k + y] = (pred, pred, pred)

'''Calculation of Sum of squares of errors'''
sse = 0
for i in range(new_height-4):
    for j in range(new_width-4):
        sse += (Y[i,j]-predicted_Y[i, j])**2
sse /= (new_width*new_height)
sse = math.sqrt(sse)
print 'SSE', sse


'''Displaying the original image'''
imz = im.fromarray(Y)
plt.imshow(imz)
plt.show()

'''Displaying the predicted image'''
imz = im.fromarray(predicted_image)
plt.imshow(imz)
plt.show()