__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_setblocks2():

    """
    wholewises
    """

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

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

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

    print block

    imb.setCoefBlocks(ima.getCoefBlocks())

    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)


def test_split():
    imb = jpegObj.Jpeg("res/test3.jpg")
    c = imb.getCoefMatrix(channel='Y')
    print type(c[0, 0])
    d = c.ravel()
    c1 = np.array([np.split(arr, arr.shape[1] / 8, axis=1) for arr in np.split(c, c.shape[0] / 8)])
    c2 = np.array([np.hsplit(arr, arr.shape[1] / 8) for arr in np.vsplit(c, c.shape[0] / 8)])

    c3 = imb.getCoefBlocks()

    print c1.shape, c2.shape
    if np.array_equal(c1, c2) and np.array_equal(c1, c3):
        print("blocks match")

def _convert_bits(hid_data):
    hid_data_bits = [map(int, '{0:08b}'.format(byte)) for byte in hid_data]
    return np.array(hid_data_bits).ravel()

def test_rawfile():

    raw = [0, 0, 0, 0] + np.fromfile("res/test4.jpg", np.uint8).tolist()
    raw_size = len(raw)
    print raw_size
    for i in xrange(4):
        raw[i] = raw_size % 256
        raw_size /= 256
    raw = np.array(raw)
    print raw.shape,raw

    # print raw.size

    # print bytes2bits(raw)

def test_bitbyte():
    timer.mark()
    raw = np.fromfile("res/test4.jpg", np.uint8)
    timer.report()
    print raw

    bitsraw =  bytes2bits(raw)
    # bitsraw = bitsraw[:24]
    timer.report()
    print bitsraw

    bytesraw = bits2bytes(bitsraw)
    timer.report()
    print bytesraw

def test_iter():
    imb = jpegObj.Jpeg("res/test4.jpg")
    blocks = imb.getCoefBlocks(channel='Y')
    cnt = 0
    for x in np.nditer(blocks, op_flags=['readwrite']):
        if x == 0:
            continue
        x[...] = 2
        print x
    for x in np.nditer(blocks):
        if x == 0:
            continue
        cnt += 1
        print x
    print cnt


if __name__ == '__main__':
    # timer.mark()

    # test_setblock()

    # timer.report()

    #

    # timer.mark()
    # test_setblocks()

    # timer.report()

    # test_split()

    # test_iter()

    # test_rawfile()

    # test_setblocks2()

    # test_bitbyte()

    ima = jpegObj.Jpeg("res/test3.jpg")
    imb = jpegObj.Jpeg("res/steged.jpg")
    print ima.Jgetcompdim(0)
    diffblocks(ima, imb)

    c1 = ima.getCoefBlocks()
    c2 =  imb.getCoefBlocks()

    print c1[0],c2[0]

    pass