__author__ = 'chunk'

import numpy as np
from .. import mjpeg
from ..mjpeg import base
from ..msteg.steganography import LSB, F3, F4, F5

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]


package_dir = os.path.dirname(os.path.abspath(__file__))


def test_setblocks():
    """
    wholewise
    """
    imb = mjpeg.Jpeg(os.path.join(package_dir, "../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(os.path.join(package_dir, "../res/test4.jpg"))

    ima = mjpeg.Jpeg(os.path.join(package_dir, "../res/test3.jpg"))
    imb = mjpeg.Jpeg(os.path.join(package_dir, "../res/test4.jpg"))
    mjpeg.diffblocks(ima, imb)


def test_setblocks2():
    """
    wholewises
    """
    ima = mjpeg.Jpeg(os.path.join(package_dir, "../res/test3.jpg"))
    imb = mjpeg.Jpeg(os.path.join(package_dir, "../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(os.path.join(package_dir, "../res/test4.jpg"))

    ima = mjpeg.Jpeg(os.path.join(package_dir, "../res/test3.jpg"))
    imb = mjpeg.Jpeg(os.path.join(package_dir, "../res/test4.jpg"))
    mjpeg.diffblocks(ima, imb)


def test_setblock():
    """
    blockwise
    """
    imb = mjpeg.Jpeg(os.path.join(package_dir, "../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
    mjpeg.diffblock(blocks1, block_to_show)


def test_split():
    imb = mjpeg.Jpeg(os.path.join(package_dir, "../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(os.path.join(package_dir, "../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(os.path.join(package_dir, "../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 = mjpeg.Jpeg(os.path.join(package_dir, "../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


def test_jpeg():
    ima = mjpeg.Jpeg(os.path.join(package_dir, "../res/cropped_pil.jpg"), key=sample_key)
    print ima.getQuality()
    print ima.getCapacity('All')
    """
    55df71faed43b37a8ac242120a72d 74903
    988d826212a6eed023ab99284971d 113279
    221379062154e32d0efa5cce32f1b 97137
    cecdb4b3fef5433faa3d24259e3e4 140687
    e0255871e984f61ee8dc1800b9184 81294
    tmp000 84057
    """

    ima = mjpeg.Jpeg(os.path.join(package_dir, "../res/166500/cecdb4b3fef5433faa3d24259e3e4.jpg"), key=sample_key)
    print ima.getQuality()
    print ima.getCapacity()

    # imb = mjpeg.Jpeg("res/new.jpg",key=sample_key)
    imc = mjpeg.Jpeg(os.path.join(package_dir, "../res/steged.jpg"), key=sample_key)
    print ima.Jgetcompdim(0)
    print ima.getkey(), imc.getkey()
    print mjpeg.diffblocks(ima, imc)



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 = mjpeg.Jpeg("res/166500/cecdb4b3fef5433faa3d24259e3e4.jpg", key=sample_key)
    print ima.getQuality()
    print ima.getCapacity()
    sys.exit(0)

    # imb = mjpeg.Jpeg("res/new.jpg",key=sample_key)
    imc = mjpeg.Jpeg("res/steged.jpg", key=sample_key)
    print ima.Jgetcompdim(0)
    print ima.getkey(), imc.getkey()
    print mjpeg.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