from PIL import Image
import math
import numpy
import struct
import sys
import PIL

def calc_DCT(N):
    DCT = numpy.zeros([N, N])
    for i in range(N):
        for j in range(N):
            if i == 0:
                DCT[i, j] = 1 / math.sqrt(N)
            else:
                DCT[i, j] = math.sqrt(2 / N) * math.cos((2 * j + 1) * i * math.pi / (2 * N))
    return numpy.matrix(DCT)

def calc_Q(N, q):
    Q = numpy.zeros([N, N])
    for i in range(N):
        for j in range(N):
            Q[i, j] = 1 + ((1 + i + j) * q)
    return numpy.matrix(Q)

def expand(A, m ,n):
    col = numpy.ones((A.shape[0], n - A.shape[1])) * 255
    A = numpy.append(A, col, 1)
    row = numpy.ones((m - A.shape[0], A.shape[1])) * 255
    A = numpy.append(A, row, 0)
    return A

def quantization(A, Q):
    B = numpy.zeros([A.shape[0], A.shape[1]])
    for i in range(0, A.shape[0]):
        for j in range(0, A.shape[1]):
            B[i, j] = A[i, j] / Q[i, j]
    return numpy.matrix(B)

def InverseQuantization(A, Q):
    B = numpy.zeros([A.shape[0], A.shape[1]])
    for i in range(0, A.shape[0]):
        for j in range(0, A.shape[1]):
            B[i, j] = A[i, j] * Q[i, j]
    return numpy.matrix(B)

def ZigZag(A):
    def move(i, j):
        if j < (A.shape[0] - 1):
            return max(0, i-1), j+1
        else:
            return i+1, j
    S = list()
    x, y = 0, 0
    for n in range(A.shape[0] * A.shape[0]):
        S.append(A[y, x])
        if (x + y) & 1:
            x, y = move(x, y)
        else:
            y, x = move(y, x)
    return S

def ZigZagInverse(S, size):
    def move(i, j):
        if j < (A.shape[0] - 1):
            return max(0, i-1), j+1
        else:
            return i+1, j
    A = numpy.zeros([size, size])
    x, y = 0, 0
    for k in range(size * size):
        A[y, x] = S.pop(0)
        if (x + y) & 1:
            x, y = move(x, y)
        else:
            y, x = move(y, x)
    return numpy.matrix(A)

def DCT_Image(A):
    S = list()
    for i in range(0, A.shape[0], 8):
        for j in range(0, A.shape[1], 8):
            IMG = in_data[i:i + 8, j:j + 8]
            IMG = IMG - 128
            RES = DCT * IMG * DCT_T
            RES = quantization(RES, Q)
            RES = numpy.int8(RES)
            S = S + ZigZag(RES)
    return S

def InverseDCT(S, m, n, Q):
    img = numpy.zeros([0, n])
    for i in range(0, m, 8):
        row = numpy.zeros([8, 0])
        for j in range(0, n, 8):
            A = ZigZagInverse(S, 8)
            A = InverseQuantization(A, Q)
            A = DCT_T * A * DCT
            A = A + 128
            A = numpy.uint8(A)
            row = numpy.append(row, A, 1)
        img = numpy.append(img, row, 0)
    return numpy.matrix(img)

def RLE_Coder(S):
    RES = list()
    counter = 0
    INPUT = list(S)
    n = 0
    while INPUT:
        byte = INPUT.pop(0)
        if byte == 0:
            n = 1
            while True:
                if len(INPUT) == 0:
                    RES.append(0)
                    RES.append(n)
                    break
                byte = INPUT.pop(0)
                if byte == 0:
                    n = n + 1
                    if n == 128:
                        RES.append(0)
                        RES.append(n - 1)
                        n = 1
                else:
                    RES.append(0)
                    RES.append(n)
                    RES.append(byte)
                    break
        else:
            RES.append(byte)

    return RES

def RLE_Decoder(S):
    RES = list()
    INPUT = list(S)
    while INPUT:
        byte = INPUT.pop(0)
        if byte == 0:
            n = INPUT.pop(0)
            for i in range(n):
                RES.append(0)
        else:
            RES.append(byte)
    return RES

def Save(S, imgWidth, imgHeight, filename):
    file = open(filename, "wb")
    file.write(struct.pack("I", imgWidth))
    file.write(struct.pack("I", imgHeight))
    for byte in S:
        file.write(struct.pack("b", byte))
    file.close()
    return

def Open(filename):
    file = open(filename, "rb")
    imgWidth = struct.unpack("I", file.read(4))[0]
    imgHeight = struct.unpack("I", file.read(4))[0]
    S = numpy.fromfile(file, dtype = numpy.int8).tolist()
    file.close()
    return S, imgWidth, imgHeight

DCT = calc_DCT(8)
DCT_T = numpy.transpose(DCT)
Q = calc_Q(8, 5)
Q[0, 0] = 8

print("Compress: comp inputFilename outputFilename")
print("Decompress: decomp inputFilename outputFilename")
print("Exit: q")

while True:
    text = input("> ").split()
    if (len(text) < 1):
        print("wrong command")
    else:
        if len(text) >= 1:
            command = text.pop(0)
            if command == "q":
                sys.exit(0)
        if len(text) != 2:
            print("wrong command")
        else:
            inputFilename = text.pop(0)
            outputFilename = text.pop(0)
            if command == "comp":
                img = Image.open(inputFilename).convert('L')
                imgWidth = img.size[0]
                imgHeight = img.size[1]
                in_data = numpy.matrix(expand(numpy.asarray(img, dtype = numpy.uint8),
                     imgHeight if (imgHeight % 8 == 0) else imgHeight + 8 - (imgHeight % 8),
                     imgWidth if (imgWidth % 8 == 0) else imgWidth + 8 - (imgWidth % 8)))
                serial_out = DCT_Image(in_data)
                print("DCT max = ", max(serial_out))
                print("DCT min = ",min(serial_out))
                RLE_Out = RLE_Coder(serial_out)
                Save(RLE_Out, imgWidth, imgHeight, outputFilename,)
            else:
                if command == "decomp":
                    S, imgWidth, imgHeight = Open(inputFilename)
                    out_data = RLE_Decoder(S)
                    out_data = InverseDCT(out_data, imgHeight, imgWidth, Q)
                    PIL.Image.fromarray(numpy.uint8(out_data)).save(outputFilename)
                else:
                    print("wrong command")