__author__ = 'chunk'

"""
ref - http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.115.3651&rep=rep1&type=pdf

<p>This module implements the rather sophisticated F5 algorithm which was
invented by Andreas Westfeld.</p>

Unlike its vastly inferior predecessors, namely F3 and F4, it features matrix
encoding which makes it possible to embed a chunk of k bits within 2^k - 1
bits of the cover data and only change one bit (at most). A bit change is
done by subtracting the absolute value of the corresponding DCT coefficient.
When the embedding process begins, the parameter k is computed based on
the capacity of the cover image and the prospective embedding ratio.
With small amount of hidden data k becomes large which leads to a greater
embedding efficiency (embedded information per bit change).<br />

A permutation (initialized by a user-supplied seed) of the DCT coefficients
helps to scatter each chunk across the entire image.
F5 can be seen as meta-algorithm as it uses a coding scheme to change
as little data as possible and then applies a simpler algorithm (such as F3)
to actually embed data. That is why this module allows the user to specify
which embedding function (one of JSteg, F3, F4) should be used.
"""

import time
import math
import numpy as np
import numpy.random as rnd
from msteg import *
from F4 import F4
import mjsteg
import pyjpegobj
from common import *


class F5(StegBase):
    """ This module has two methods: <i>embed_raw_data</i> to embed data
    with the F5 algorithm and <i>extract_raw_data</i> to extract data
    which was embedded previously. """

    def __init__(self, key=sample_key, k=None):
        """
        Constructor of the F5 class.
        """
        StegBase.__init__(self, key)
        self._embed_fun = None
        self.default_embedding = True

        # needed because k is embedded separately
        self.k_coeff = k


    def _get_cov_data(self, img_path):
        """
        Returns DCT coefficients of the cover image.
        """
        self.cov_jpeg = pyjpegobj.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, embed_fun='Default'):
        """This method embeds arbitrary data into a cover image.
        The cover image must be a JPEG.

        @param embed_fun:
        Specifies which embedding function should be used. Must be one of
        'Default', 'F3', 'Jsteg'. If 'Default' is selected, the algorithm uses
        the same behavior as Westfeld's implementation, i.e. decrementing
        absolute values for n > 1 (F3) and using F4 in the special case n = 1.
        Selecting F3 or JSteg results in using that scheme for all n.
        """
        self.t0 = time.time()

        if embed_fun == 'F3':
            self._embed_fun = self._f3_embed
            self.default_embedding = False
        elif embed_fun == 'JSteg' or embed_fun == 'LSB':
            self._embed_fun = self._jsteg_embed
            self.default_embedding = False
        else:
            self._embed_fun = self._f3_embed
            self.default_embedding = True

        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, embed_fun='Default'):
        self.t0 = time.time()

        if embed_fun == 'F3':
            self._embed_fun = self._f3_embed
            self.default_embedding = False
        elif embed_fun == 'JSteg' or embed_fun == 'LSB':
            self._embed_fun = self._jsteg_embed
            self.default_embedding = False
        else:
            self._embed_fun = self._f3_embed
            self.default_embedding = True

        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[:header_size])
            print 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 _embed_k(self, cov_data, hid_data):
        np.random.seed(self.seed)
        self.dct_p = np.random.permutation(cov_data.size)
        self.k_coeff = self._find_max_k(cov_data, hid_data)
        self.ui.display_status('setting k = %d' % self.k_coeff)
        k_split = self.lookup_tab.split_byte(self.k_coeff, 1)[-4:]
        # embed k in F3-like style
        for m in k_split:
            success = False
            while not success:
                self.cov_ind += 1
                while cov_data[self.dct_p[self.cov_ind]] == 0 or \
                                        self.dct_p[self.cov_ind] % 64 == 0:
                    self.cov_ind += 1
                if m != cov_data[self.dct_p[self.cov_ind]] & 1:
                    cov_data[self.dct_p[self.cov_ind]] -= \
                        math.copysign(1, cov_data[self.dct_p[self.cov_ind]])
                success = cov_data[self.dct_p[self.cov_ind]] != 0

    def _extract_k(self, steg_data):
        # initializing the MT is done only once in order to retain the state
        self.dct_p = np.random.seed(self.seed)
        self.dct_p = np.random.permutation(self.steg_data.size)
        k_split = np.zeros(4, np.uint8)
        for i in xrange(k_split.size):
            self.steg_ind += 1
            while self.steg_data[self.dct_p[self.steg_ind]] == 0 or \
                                    self.dct_p[self.steg_ind] % 64 == 0:
                self.steg_ind += 1
            k_split[i] = self.steg_data[self.dct_p[self.steg_ind]] & 1
        self.k_coeff = self.lookup_tab.merge_words(tuple([0, 0, 0, 0] +
                                                         list(k_split)), 1)

    def _find_max_k(self, cov_data, hid_data):
        cnt = 4  # information about k take up 4 bits
        # find number of DCT coefficients
        update_cnt = 10000
        for i, c in enumerate(cov_data):
            if update_cnt == 0:
                self._set_progress(
                    int(30 * (float(i) / float(cov_data.size))))
                update_cnt = 10000
            update_cnt -= 1
            # pessimistic, but accurate estimation of the capacity of the image
            ci = int(c)
            if (not (ci is 0)) and (not ((i % 64) is 0)) \
                    and (not (ci is 1)) and (not (ci is -1)):
                cnt += 1
        hid_size = hid_data.size
        cov_size = cnt
        if cov_size < hid_size:
            raise Exception("Cannot fit %d bits in %d DCT coefficients. \
                    Cover image is too small." % (hid_size, cov_size))
        self.ui.display_status('DCT embedding ratio = %f' \
                               % (float(hid_size) / float(cov_size)))
        k = 1
        while True:
            k += 1
            n = (1 << k) - 1
            num_chunks = cov_size / n
            num_emb_bits = num_chunks * k
            if num_emb_bits < hid_size:
                return min(k - 1, 15)

    # low level embedding functions
    def _f3_embed(self, cov_data, ind):
        cov_data[ind] -= math.copysign(1, cov_data[ind])

    def _jsteg_embed(self, cov_data, ind):
        m = 1 ^ (cov_data[ind] & 1)
        cov_data[ind] = (cov_data[ind] & 0xfffffffe) | m

    def _raw_embed(self, cov_data, hid_data):
        k = self.k_coeff
        n = (1 << k) - 1

        if n == 1 and self.default_embedding:
            # in case k = n = 1, Westfeld's implementation uses F4 for embedding.
            f4 = F4(key=self.key)
            return f4._raw_embed(cov_data, hid_data)

        hid_data = bytes2bits(hid_data)
        if len(hid_data) % k != 0:
            hid_data = list(hid_data) + [0 for x in range(k - len(hid_data) % k)]

        ind_nonzero = np.nonzero(cov_data)[0]

        if np.size(ind_nonzero) * k < len(hid_data) * n:
            raise Exception("Supposed secret data too large for stego image.")

        ind_cov = 0
        for ind_hid in range(0, len(hid_data), k):
            msg_chunk = hid_data[ind_hid:ind_hid + k]
            cov_chunk = ind_nonzero[ind_cov:ind_cov + n]
            ind_cov += n

            success = False
            while not success:
                h = 0
                for i in xrange(n):
                    h ^= ((cov_data[cov_chunk[i]] & 1) * (i + 1))
                scalar_x = 0
                for i in xrange(k):
                    scalar_x = (scalar_x << 1) + msg_chunk[
                        i]  # N.B. hid_data[0]:high (that is x2), hid_data[1]:low (that is x1)
                s = scalar_x ^ h
                if s != 0:
                    self._embed_fun(cov_data, cov_chunk[s - 1])
                else:
                    break

                if cov_data[cov_chunk[s - 1]] == 0:  # shrinkage
                    cov_chunk[s - 1:-1] = cov_chunk[s:]
                    cov_chunk[-1] = ind_nonzero[ind_cov]
                    ind_cov += 1
                else:
                    success = True

        return cov_data, ind_hid + k

    def _raw_extract(self, steg_data, num_bits):
        k = self.k_coeff
        n = (1 << k) - 1

        if n == 1 and self.default_embedding:
            f4 = F4(key=self.key)
            return f4._raw_extract(steg_data, num_bits)

        num_bits_ceil = num_bits
        if num_bits % k != 0:
            num_bits_ceil = k * (num_bits / k + 1)

        hid_data = np.zeros(num_bits_ceil, np.uint8)
        curr_chunk = np.zeros(k, np.uint8)
        steg_data = steg_data[np.nonzero(steg_data)]
        ind_hid = 0
        for ind_cov in range(0, len(steg_data), n):
            steg_chunk = steg_data[ind_cov:ind_cov + n]

            h = 0  # hash value
            for i in xrange(n):
                h ^= ((steg_chunk[i] & 1) * (i + 1))

            for i in xrange(k):
                curr_chunk[k - i - 1] = h & 1  # N.B. hid_data[0]:high (that is x2), hid_data[1]:low (that is x1)
                h >>= 1

            hid_data[ind_hid:ind_hid + k] = curr_chunk[0:k]
            ind_hid += k

            if ind_hid >= num_bits_ceil: break

        return hid_data, num_bits_ceil

    def __str__(self):
        return 'F5'