__author__ = 'chunk'

import numpy as np
import mjsteg
import jpegObj
from common import *

timer = Timer()

sample = [[7, 12, 14, -12, 1, 0, -1, 0],
          [6, 5, -10, 0, 6, 0, 0, 0],
          [0, 6, -5, 4, 0, -1, 0, 0],
          [0, -3, 0, 1, -1, 0, 0, 0],
          [-3, 5, 0, 0, 0, 0, 0, 0],
          [2, -1, 0, 0, 0, 0, 0, 0],
          [0, 0, 0, 0, 0, 0, 0, 0],
          [0, 0, 0, 0, 0, 0, 0, 0]]


def diffblock(c1, c2):
    diff = False
    if np.array_equal(c1, c2):
        print("blocks match")
    else:
        print("blocks not match")
        diff = True

    return diff


def diffblocks(a, b):
    diff = False
    for comp in range(a.image_components):
        xmax, ymax = a.Jgetcompdim(comp)
        for y in range(ymax):
            for x in range(xmax):
                if a.Jgetblock(x, y, comp) != b.Jgetblock(x, y, comp):
                    print("blocks({},{}) in component {} not match".format(y, x, comp))
                    diff = True
    return diff


def test_setblocks():
    """
    wholewise
    """
    imb = jpegObj.Jpeg("res/test4.jpg")
    block = imb.getCoefBlock(channel='Y', loc=(-1, 2))
    print block

    imb.setCoefMatrix(np.array([[0] * 800 for i in range(600)]), channel='Y')

    block = imb.getCoefBlock(channel='Y', loc=(-1, 2))
    print block

    imb.Jwrite("res/test4.jpg")

    ima = jpegObj.Jpeg("res/test3.jpg")
    imb = jpegObj.Jpeg("res/test4.jpg")
    diffblocks(ima, imb)


def test_setblock():
    """
    blockwise
    """
    imb = jpegObj.Jpeg("res/test4.jpg")
    block = imb.getCoefBlock(channel='Y', loc=(0, 1))
    print block

    # print b.setCoefBlock(np.array([[27] + [0] * 7] + [[0] * 8 for i in range(7)]), loc=(0, 1))
    imb.setCoefBlock(np.array(sample), loc=(0, 1))
    blocks1 = imb.getCoefBlock(channel='Y', loc=(0, 1))
    print blocks1
    blocks2 = imb.Jgetblock(1, 0, 0)
    block_to_show = np.frombuffer(blocks2, dtype=np.int16, count=-1, offset=0).reshape(8, 8)
    print block_to_show
    diffblock(blocks1, block_to_show)


if __name__ == '__main__':
    timer.mark()
    test_setblock()
    timer.report()

    timer.mark()
    test_setblocks()
    timer.report()