__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 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")
    jpegObj.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")
    jpegObj.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

    jpegObj.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 jpegObj.diffblocks(ima, imc)

    # c1 = ima.getCoefBlocks()
    # c2 =  imb.getCoefBlocks()
    #
    # # print c1[0],c2[0]
    # s1 = imb.getsignal(channel='Y')
    # s2 = ima.getsignal(channel='Y')
    # imb.setsignal(s2,channel='Y')
    # imb.Jwrite('res/new.jpg')


    # print base.acMask(8,16)
    # mmask = base.acMaskBlock
    # print mmask
    # sample = np.array(sample)[mmask]
    # print np.hstack([[],sample])

    pass