from scipy.fftpack import dct
from scipy.fftpack import idct
from scipy.misc import imread

import matplotlib.pyplot as plt
import numpy as np

__author__ = "Wei-Chung Huang"
__copyright__ = "Copyright 2018, The SJSU MSCS Master project"
__license__ = "MIT"
__version__ = "1.0.0"


class JPGEncode:

    qt_matrix = np.asarray([[16, 11, 10, 16, 24, 40, 51, 61],
                            [12, 12, 14, 19, 26, 58, 60, 55],
                            [14, 13, 16, 24, 40, 57, 69, 56],
                            [14, 17, 22, 29, 51, 87, 80, 62],
                            [18, 22, 37, 56, 68, 109, 103, 77],
                            [24, 35, 55, 64, 81, 104, 113, 92],
                            [49, 64, 78, 87, 103, 121, 120, 101],
                            [72, 92, 95, 98, 112, 100, 103, 99]])

    def __init__(self):
        pass

    @staticmethod
    def dct2(x):
        """
        Perform a 2D Discrete Cosine Transform.
        :param x: The input array
        :return: An array of DCT coefficients
        """
        return dct(dct(x, axis=0, norm='ortho'), axis=1, norm='ortho')

    @staticmethod
    def idct2(x):
        """
        Perform a 2D Inverse Discrete Cosine Transform.
        :param x: The input DCT coefficients
        :return: The transformed input array
        """
        return idct(idct(x, axis=0, norm='ortho'), axis=1, norm='ortho')

    def encode(self, img, bk_size=8, thresh=0.01, use_qtm=False, debug=False):
        """
        Encode an image with block-wise DCT transform with a scalar uniform quantization
        :param img: An 2D input image
        :param bk_size: Block size for DCT transform
        :param thresh: A threshold that specify the percentage of coefficient should be reduced
        :param use_qtm: A flag to use quantization matrix as a compression
        :param debug: A flag to dump the compressed image
        :return: A codeword after a zig-zag scan
        """
        assert use_qtm is False or use_qtm is True and bk_size == 8, \
            "Block-size should be 8 when using QTM"

        # centerize the image by shifting 128 from [0, 255]
        img_shift = img - 128

        # Do 8x8 DCT on image (in-place)
        imsize = img_shift.shape
        dct_coeff = np.zeros(imsize)
        dct_thresh = np.zeros(imsize)
        if use_qtm:
            for i in np.r_[:imsize[0]:bk_size]:
                for j in np.r_[:imsize[1]:bk_size]:
                    dct_coeff[i:(i + bk_size), j:(j + bk_size)] = self.dct2(img_shift[i:(i + bk_size), j:(j + bk_size)])
                    dct_thresh[i:(i + bk_size), j:(j + bk_size)] \
                        = np.round(dct_coeff[i:(i + bk_size), j:(j + bk_size)] / self.qt_matrix)
        else:
            for i in np.r_[:imsize[0]:bk_size]:
                for j in np.r_[:imsize[1]:bk_size]:
                    dct_coeff[i:(i + bk_size), j:(j + bk_size)] = self.dct2(img_shift[i:(i + bk_size), j:(j + bk_size)])
            dct_thresh = np.round(dct_coeff * (abs(dct_coeff) > (thresh * np.max(dct_coeff))))

        # Encode each DCT coefficient block
        codewords = []
        for i in np.r_[:imsize[0]:bk_size]:
            for j in np.r_[:imsize[1]:bk_size]:
                codewords.append(self.zigzagScan(dct_thresh[i:(i + bk_size), j:(j + bk_size)]))

        if debug:
            plt.figure()
            plt.imshow(img, cmap='gray')

            plt.figure()
            plt.imshow(dct_coeff, cmap='gray', vmax=np.max(dct_coeff) * 0.01, vmin=0)
            plt.title("8x8 DCTs of the image")

            plt.figure()
            plt.imshow(dct_thresh, cmap='gray', vmax=np.max(dct_coeff) * 0.01, vmin=0)
            plt.title("Thresholded 8x8 DCTs of the image")

            im_dct = np.zeros(imsize)
            if use_qtm:
                for i in np.r_[:imsize[0]:bk_size]:
                    for j in np.r_[:imsize[1]:bk_size]:
                        reversed_coeff = dct_thresh[i:(i + bk_size), j:(j + bk_size)] * self.qt_matrix
                        im_dct[i:(i + bk_size), j:(j + bk_size)] = np.round(self.idct2(reversed_coeff)) + 128
            else:
                for i in np.r_[:imsize[0]:bk_size]:
                    for j in np.r_[:imsize[1]:bk_size]:
                        im_dct[i:(i + bk_size), j:(j + bk_size)] \
                            = np.round(self.idct2(dct_thresh[i:(i + bk_size), j:(j + bk_size)])) + 128

            plt.figure()
            plt.imshow(np.hstack((img, im_dct)), cmap='gray')
            plt.title("Comparison between original and DCT compressed images")

            plt.show()
            percent_nonzeros = np.sum(dct_thresh != 0.0) / (imsize[0] * imsize[1] * 1.0)
            print("Keeping only %f%% of the DCT coefficients" % (percent_nonzeros * 100.0))

        return codewords

    def encode_with_qtm_dithering(self, img, bk_size=8):
        """
        Encode an image with block-wise DCT transform with a scalar uniform quantization and Quantization Matrix.
        Also, the quantized cells are dithered.
        :param img: An 2D input image
        :param bk_size: Block size for DCT transform
        :return: A codeword after a zig-zag scan
        """
        assert bk_size == 8, "Block-size should be 8 when using QTM"

        # centerize the image by shifting 128 from [0, 255]
        img_shift = img - 128

        # Do 8x8 DCT on image (in-place)
        imsize = img_shift.shape
        dct_coeff = np.zeros(imsize)
        dct_thresh = np.zeros(imsize)
        for i in np.r_[:imsize[0]:bk_size]:
            for j in np.r_[:imsize[1]:bk_size]:
                dct_coeff[i:(i + bk_size), j:(j + bk_size)] = self.dct2(img_shift[i:(i + bk_size), j:(j + bk_size)])
                dct_thresh[i:(i + bk_size), j:(j + bk_size)] = self.dithering_8x8(dct_coeff[i:(i + bk_size), j:(j + bk_size)])

        # Encode each DCT coefficient block
        codewords = []
        for i in np.r_[:imsize[0]:bk_size]:
            for j in np.r_[:imsize[1]:bk_size]:
                codewords.append(self.zigzagScan(dct_thresh[i:(i + bk_size), j:(j + bk_size)]))

        return codewords

    def dithering_8x8(self, src):
        """
        Dither an image by using Floyd Steinberg.
        (In general, dithering can be viewed as an error diffusion process.)
        :param src: 2D ndarray with size 8x8.
        :return: A dithered block.
        """
        dst = np.zeros([8, 8])
        for y in range(src.shape[0] - 1):
            for x in range(src.shape[1] - 1):
                old_pixel = src[y][x]
                new_pixel = old_pixel / self.qt_matrix[y][x]
                quant_error = old_pixel - new_pixel
                dst[y][x] = round(new_pixel)
                dst[y][x + 1] += quant_error * 7 / 16
                dst[y + 1][x - 1] += quant_error * 3 / 16
                dst[y + 1][x] += quant_error * 5 / 16
                dst[y + 1][x + 1] += quant_error * 1 / 16
        return dst

    @staticmethod
    def zigzagScan(matrix):
        if matrix is None or len(matrix) is 0 or len(matrix[0]) is 0:
            return []

        m = len(matrix)
        n = len(matrix[0])
        i = j = 0
        res = [0] * (m * n)
        counter = 0
        res[counter] = matrix[i][j]
        counter += 1
        if j + 1 < n:
            j += 1
        else:
            i += 1

        while counter < m * n:
            while i < m and j >= 0:
                res[counter] = matrix[i][j]
                counter += 1
                i += 1
                j -= 1
            i -= 1
            j += 1
            if i + 1 < m:
                i += 1
            else:
                j += 1

            while i >= 0 and j < n:
                res[counter] = matrix[i][j]
                counter += 1
                i -= 1
                j += 1
            i += 1
            j -= 1
            if j + 1 < n:
                j += 1
            else:
                i += 1

        return res


if __name__ == '__main__':
    """
    Example of encoding an image by the JPEG compression.
    """
    codec = JPGEncode()
    img = imread('../lena.png').astype(float)

    print(type(img))
    codeword = codec.encode(img, thresh=0.05, use_qtm=False, debug=True)
    print(len(codeword))
    print(codeword)
