__author__ = 'chunk'

import numpy as np

from ..mjpeg import Jpeg
from ..common import *

__all__ = ['StegBase', 'sample_key']

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


class StegBase(object):
    """
    This is the base class for some of the JPEG-algorithms that behave
    similarly such as JSteg, OutGuess and F3.
    """

    def __init__(self, key=sample_key):
        """
        Constructor of the JPEGSteg class.
        """
        self.t0 = None
        self.cov_jpeg = None
        self.cov_data = None
        self.hid_data = None
        self.key = key

    def _get_cov_data(self, img_path):
        """
        Returns DCT coefficients of the cover image.
        """
        self.cov_jpeg = Jpeg(img_path)
        self.key = self.cov_jpeg.getkey()
        self.cov_data = self.cov_jpeg.getCoefBlocks('Y')
        # self.capacity = self.cov_jpeg.getCapacity('Y')
        self.capacity = np.sum(self.cov_data != 0) - np.size(self.cov_data) / 64
        return self.cov_data


    def _get_hid_data(self, src_hidden, frommem=False):
        """
        Returnsthe secret data as byte sequence.
        """
        if not frommem:
            raw = [0, 0, 0, 0] + np.fromfile(src_hidden, np.uint8).tolist()
        else:
            raw = [0, 0, 0, 0] + np.fromstring(src_hidden, dtype=np.uint8).tolist()

        raw_size = len(raw)
        for i in xrange(4):
            raw[i] = raw_size % 256
            raw_size /= 256
        self.hid_data = np.array(raw, dtype=np.uint8)

        if np.size(self.hid_data) * 8 > self.capacity:
            raise Exception("Cover image is too small to embed data.Cannot fit %d bits in %d NZ-DCT coefficients" % (
                np.size(self.hid_data) * 8, self.capacity))
        return self.hid_data


    def _post_embed_actions(self, src_cover, src_hidden, tgt_stego):
        """
        This function isn't named very accurately. It actually calls the
        _raw_embed function in inherited classes.
        """
        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.setCoefBlocks(cov_data)
            self.cov_jpeg.Jwrite(tgt_stego)

            self._display_rate('embed', self.capacity, bits_cnt)

            # # 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 _post_extract_actions(self, src_steg, tgt_hidden):
        """
        This function isn't named very accurately. It actually calls the
        _raw_extract function in inherited classes.
        """
        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_rate('extract', self.capacity, bits_cnt)

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


    def _looks_like_jpeg(self, path):
        try:
            with open(path, 'r') as f:
                return f.read(2) == '\xff\xd8'
        except IOError:
            return False

    def _display_stats(self, verb, cov_size, emb_size, duration):
        print(
            "%dB %s in %.2fs (%.2f kBps), embedding ratio: %.4f" %
            (emb_size, verb, duration, (emb_size / duration) / 1000.,
             float(emb_size) / cov_size))

    def _display_rate(self, verb, cov_bits, emb_bits):
        print "%s: %d bits\trate: %.4f bpc" % (verb, emb_bits, float(emb_bits) / cov_bits)

    # set & get
    def set_key(self, key):
        assert key != None
        self.key = key

    def get_key(self):
        return self.key

    # dummy functions to please pylint
    def _raw_embed(self, cov_data, hid_data):
        pass

    def _raw_extract(self, steg_data, num_bits):
        pass

    def _dummy_embed_hook(self, cov_data, hid_data):
        pass

    def _dummy_extract_hook(self, steg_data, num_bits):
        pass