__author__ = 'hadoop'

from ..common import *

from ..mjpeg import *
from ..msteg import *
from ..msteg.steganography import LSB, F3, F4, F5
from ..mfeat import IntraBlockDiff
from ..mmodel.svm import SVM

from numpy import array
import json
import pickle
import tempfile

import numpy as np
from scipy import stats
from hashlib import md5

np.random.seed(sum(map(ord, "whoami")))
package_dir = os.path.dirname(os.path.abspath(__file__))
classifier = SVM.ModelSVM(toolset='sklearn')

def rddparse_data_CV(raw_row):
    """
    input: (u'key0',u'cf_feat:hog:[0.056273,...]--%--cf_pic:data:\ufffd\ufffd\...--%--cf_tag:hog:True')
    return: ([0.056273,...],1)
    """
    data = raw_row[1].split('--%--')
    feat = json.loads(data[0].split(':')[-1])
    tag = 1 if data[-1].split(':')[-1] == 'True' else 0
    return (feat, tag)


def rddparse_data_ILS(raw_row):
    """
    input: (u'key0',u'cf_feat:hog:[0.056273,...]--%--cf_pic:data:\ufffd\ufffd\...--%--cf_tag:hog:True')
    return: ([0.056273,...],1)

    In fact we can also use mapValues.
    """
    key = raw_row[0]
    # if key == '04650c488a2b163ca8a1f52da6022f03.jpg':
    # with open('/tmp/hhhh','wb') as f:
    # f.write(raw_row[1].decode('unicode-escape')).encode('latin-1')
    items = raw_row[1].decode('unicode-escape').encode('latin-1').split('--%--')
    data = items[0].split('cf_pic:data:')[-1]
    return (key, data)


def rddparse_all_ILS(raw_row):
    """
    Deprecated
    """
    key = raw_row[0]
    items = raw_row[1].decode('unicode-escape').encode('latin-1').split('--%--')

    # @TODO
    # N.B "ValueError: No JSON object could be decoded" Because the spark-hbase IO is based on strings.
    # And the order of items is not as expected. See ../res/row-sample.txt or check in hbase shell for that.

    data = [items[0].split('cf_pic:data:')[-1]] + [json.loads(item.split(':')[-1]) for item in
                                                   items[1:]]

    return (key, data)


def rddparse_dataset_ILS(raw_row):
    if raw_row[0] == '04650c488a2b163ca8a1f52da6022f03.jpg':
        print raw_row
    items = raw_row[1].decode('unicode-escape').encode('latin-1').split('--%--')
    # tag = int(items[-2].split('cf_tag:' + tagtype)[-1])
    # feat = [item for sublist in json.loads(items[-1].split('cf_feat:' + feattype)[-1]) for subsublist in sublist for item in subsublist]
    tag = int(items[-1].split(':')[-1])
    feat = [item for sublist in json.loads(items[0].split(':')[-1]) for subsublist in sublist for
            item in subsublist]

    return (tag, feat)


def rddinfo_ILS(img, info_rate=None, tag_chosen=None, tag_class=None):
    """
    Tempfile is our friend. (?)
    """
    info_rate = info_rate if info_rate != None else 0.0
    tag_chosen = tag_chosen if tag_chosen != None else stats.bernoulli.rvs(0.8)
    tag_class = tag_class if tag_class != None else 0
    try:
        tmpf = tempfile.NamedTemporaryFile(suffix='.jpg', mode='w+b', delete=True)
        tmpf.write(img)
        tmpf.seek(0)
        im = Jpeg(tmpf.name, key=sample_key)
        info = [
            im.image_width,
            im.image_height,
            im.image_width * im.image_height,
            im.getCapacity(),
            im.getQuality(),
            info_rate,
            tag_chosen,
            tag_class
        ]
        return info
    except Exception as e:
        print e
        raise
    finally:
        tmpf.close()


def rddembed_ILS(row, rate=None):
    """
    input:
        e.g. row =('row1',[1,3400,'hello'])
    return:
        newrow = ('row2',[34,5400,'embeded'])
    """
    items = row[1]
    capacity, chosen = int(items[4]), int(items[7])
    if chosen == 0:
        return None
    try:
        tmpf_src = tempfile.NamedTemporaryFile(suffix='.jpg', mode='w+b')
        tmpf_src.write(items[0])
        tmpf_src.seek(0)
        tmpf_dst = tempfile.NamedTemporaryFile(suffix='.jpg', mode='w+b')

        steger = F5.F5(sample_key, 1)

        if rate == None:
            embed_rate = steger.embed_raw_data(tmpf_src.name,
                                               os.path.join(package_dir, '../res/toembed'),
                                               tmpf_dst.name)
        else:
            assert (rate >= 0 and rate < 1)
            # print capacity
            hidden = np.random.bytes(int(int(capacity) * rate) / 8)
            embed_rate = steger.embed_raw_data(tmpf_src.name, hidden, tmpf_dst.name, frommem=True)

        tmpf_dst.seek(0)
        raw = tmpf_dst.read()
        index = md5(raw).hexdigest()

        return (index + '.jpg', [raw] + rddinfo_ILS(raw, embed_rate, 0, 1))

    except Exception as e:
        print e
        raise
    finally:
        tmpf_src.close()
        tmpf_dst.close()


def rddembed_ILS_EXT(row, rate=None):
    """
    input:
        e.g. row =('row1',[1,3400,'hello'])
    return:
        newrow = ('row2',[34,5400,'embeded']) or NULL
        [row,newrow]
    """
    items = row[1]
    capacity, chosen = int(items[4]), int(items[7])
    if chosen == 0:
        return [row]
    try:
        tmpf_src = tempfile.NamedTemporaryFile(suffix='.jpg', mode='w+b')
        tmpf_src.write(items[0])
        tmpf_src.seek(0)
        tmpf_dst = tempfile.NamedTemporaryFile(suffix='.jpg', mode='w+b')

        steger = F5.F5(sample_key, 2)

        if rate == None:
            embed_rate = steger.embed_raw_data(tmpf_src.name,
                                               os.path.join(package_dir, '../res/toembed'),
                                               tmpf_dst.name)
        else:
            assert (rate >= 0 and rate < 1)
            # print capacity
            hidden = np.random.bytes(int(int(capacity) * rate) / 8)
            embed_rate = steger.embed_raw_data(tmpf_src.name, hidden, tmpf_dst.name, frommem=True)

        tmpf_dst.seek(0)
        raw = tmpf_dst.read()
        index = md5(raw).hexdigest()

        return [row, (index + '.jpg', [raw] + rddinfo_ILS(raw, embed_rate, 0, 1))]

    except Exception as e:
        print e
        raise
    finally:
        tmpf_src.close()
        tmpf_dst.close()


def _get_feat(image, feattype='ibd', **kwargs):
    if feattype == 'ibd':
        feater = IntraBlockDiff.FeatIntraBlockDiff()
    else:
        raise Exception("Unknown feature type!")

    desc = feater.feat(image)

    return desc


def rddfeat_ILS(items, feattype='ibd', **kwargs):
    try:
        tmpf_src = tempfile.NamedTemporaryFile(suffix='.jpg', mode='w+b')
        tmpf_src.write(items[0])
        tmpf_src.seek(0)

        desc = json.dumps(_get_feat(tmpf_src.name, feattype=feattype).tolist())
        # print 'desccccccccccccccccccc',desc
        return items + [desc]

    except Exception as e:
        print e
        raise
    finally:
        tmpf_src.close()


def rddanalysis_ILS(items, feattype='ibd', **kwargs):
    head = np.fromstring(items[0][:2], dtype=np.uint8)
    if not np.array_equal(head, [255, 216]):
        return items + [0]
    try:
        tmpf_src = tempfile.NamedTemporaryFile(suffix='.jpg', mode='w+b')
        tmpf_src.write(items[0])
        tmpf_src.seek(0)

        desc = _get_feat(tmpf_src.name, feattype=feattype)
        tag = classifier.predict(desc.ravel())[0]
        # print 'desccccccccccccccccccc',desc
        return items + [tag]

    except Exception as e:
        print e
        raise
    finally:
        tmpf_src.close()

        # return items + classifier.predict(items[-1])


def format_out(row, cols, withdata=False):
    """
    input:
        e.g. row =('row1',[1,3400,'hello'])
            cols = [['cf_info', 'id'], ['cf_info', 'size'], ['cf_tag', 'desc']]
    return:
        [('row1',['row1', 'cf_info', 'id', '1']),('row1',['row1', 'cf_info', 'size', '3400']),('row1',['row1', 'cf_tag', 'desc', 'hello'])]
    """
    puts = []
    key = row[0]
    # if key == '04650c488a2b163ca8a1f52da6022f03.jpg':
    # print row
    if not withdata:
        for data, col in zip(row[1][1:], cols[1:]):
            puts.append((key, [key] + col + [str(data)]))
    else:
        for data, col in zip(row[1], cols):
            puts.append((key, [key] + col + [str(data)]))
    return puts