__author__ = 'chunk'

import time

import numpy as np

import numpy.random as rnd

from msteg.StegBase import *

import mjsteg
import jpegObj
from common import *

class F4(StegBase):
    """ This module has two methods: <i>embed_raw_data</i> to embed data
    with the F3 algorithm and <i>extract_raw_data</i> to extract data

    which was embedded previously. """

    def __init__(self, key=sample_key):
        """

        Constructor of the F3 class.

        """

        StegBase.__init__(self, key)

    def _get_cov_data(self, img_path):
        """

        Returns DCT coefficients of the cover image.

        """

        self.cov_jpeg = jpegObj.Jpeg(img_path, key=self.key)

        cov_data = self.cov_jpeg.getsignal(channel='Y')

        self.cov_data = np.array(cov_data, dtype=np.int16)

        return self.cov_data

    def embed_raw_data(self, src_cover, src_hidden, tgt_stego):

        self.t0 = time.time()

        try:
            cov_data = self._get_cov_data(src_cover)

            hid_data = self._get_hid_data(src_hidden)

            # print hid_data.dtype,type(hid_data),hid_data.tolist()

            cov_data, bits_cnt = self._raw_embed(cov_data, hid_data)

            if bits_cnt != np.size(hid_data) * 8:
                raise Exception("Expected embedded size is %db but actually %db." % (
                    np.size(hid_data) * 8, bits_cnt))

            self.cov_jpeg.setsignal(cov_data, channel='Y')

            self.cov_jpeg.Jwrite(tgt_stego)

            # size_cov = os.path.getsize(tgt_stego)
            size_cov = np.size(cov_data) / 8
            size_embedded = np.size(hid_data)

            self._display_stats("embedded", size_cov, size_embedded,
                                time.time() - self.t0)

        except TypeError as e:
            raise e
        except Exception as expt:
            print "Exception when embedding!"
            raise

    def extract_raw_data(self, src_steg, tgt_hidden):

        self.t0 = time.time()

        try:

            steg_data = self._get_cov_data(src_steg)
            # emb_size = os.path.getsize(src_steg)
            emb_size = np.size(steg_data) / 8


            # recovering file size
            header_size = 4 * 8
            size_data, bits_cnt = self._raw_extract(steg_data, header_size)
            if bits_cnt < header_size:

                raise Exception("Expected embedded size is %db but actually %db." % (
                    header_size, bits_cnt))

            size_data = bits2bytes(size_data)

            size_hd = 0
            for i in xrange(4):
                size_hd += size_data[i] * 256 ** i

            raw_size = size_hd * 8

            if raw_size > np.size(steg_data):
                raise Exception("Supposed secret data too large for stego image.")

            hid_data, bits_cnt = self._raw_extract(steg_data, raw_size)

            if bits_cnt != raw_size:
                raise Exception("Expected embedded size is %db but actually %db." % (
                    raw_size, bits_cnt))

            hid_data = bits2bytes(hid_data)
            # print hid_data.dtype,type(hid_data),hid_data.tolist()
            hid_data[4:].tofile(tgt_hidden)

            self._display_stats("extracted", emb_size,
                                np.size(hid_data),
                                time.time() - self.t0)
        except Exception as expt:
            print "Exception when extracting!"
            raise

    def _raw_embed(self, cov_data, hid_data):
        """

        cov_data - 1-D numpy.int16 array (permunated)

        hid_data - 1-D numpy.uint8 array

        """
        hid_data = bytes2bits(hid_data)
        i = 0
        for x in np.nditer(cov_data, op_flags=['readwrite']):

            if x == 0: continue

            m = (hid_data[i] & 1)
            if x > 0 and x & 1 != m:
                x[...] -= 1
            elif x < 0 and x & 1 == m:
                x[...] += 1
            if x == 0: continue
            i += 1
            if i == hid_data.size: break

        return cov_data, i

    def _raw_extract(self, steg_data, num_bits):
        """
        Just a small helper function to extract hidden data.
        steg_data - 1-D numpy.int16 array (permunated)

        """

        hid_data = np.zeros(num_bits, np.uint8)

        j = 0
        for x in steg_data:

            if x == 0: continue
            if j >= num_bits: break

            if x > 0:
                hid_data[j] = x & 1
            else:
                hid_data[j] = (x & 1) ^ 1

            j = j + 1

        return hid_data, j

    def __str__(self):
        return "F4"