__author__ = 'chunk'

import numpy as np

import mjsteg
import jpegObj

from jpegObj import base

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]]

sample_key = [46812L, 20559L, 31360L, 16681L, 27536L, 39553L, 5427L, 63029L, 56572L, 36476L, 25695L, 61908L, 63014L,
              5908L, 59816L, 56765L]

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
    cnt = 0
    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
                    cnt += 1
    return diff, cnt

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", key=sample_key)

    # imb = jpegObj.Jpeg("res/new.jpg",key=sample_key)
    imc = jpegObj.Jpeg("res/steged.jpg", key=sample_key)
    print ima.Jgetcompdim(0)
    print ima.getkey(), imc.getkey()
    print diffblocks(ima, imc)

    # c1 = ima.getCoefBlocks()
    # c2 =  imb.getCoefBlocks()
    #
    # # print c1[0],c2[0]
    # s1 = imb.getsignal(channel='Y')
    # s2 = ima.getsignal(channel='Y')