__author__ = 'chunk'

from . import *

from ..mfeat import IntraBlockDiff

from ..mspark import rdd, SC

from pyspark.mllib.regression import LabeledPoint

from ..common import *

import os, sys

from hashlib import md5
import csv

import json

import happybase

from ..mjpeg import *
from ..msteg import *
from ..msteg.steganography import LSB, F3, F4, F5

import numpy as np

from scipy import stats

import tempfile

np.random.seed(sum(map(ord, "whoami")))

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

class DataILSVRC_S(DataDumperBase):

    """
    This module is specially for ILSVRC data processing under spark & hbase.

    We posit that the DB(e.g. HBase) has only the images data with md5 name as id.
    The task is to gennerate info(size,capacity,quality,etc.) and class & chosen tags, and then to perform embedding and finally to calcculate ibd features.

    Each step includes reading from & writing to Hbase (though PC).
    And each step must have a 'spark' mode option, which means that the operation is performed by spark with reading & wrting through RDDs.

    copyright(c) 2015 chunkplus@gmail.com

    """

    def __init__(self, base='ILSVRC2013_DET_val', category='Train_1', tablename=None):

        DataDumperBase.__init__(self, base, category)

        self.base = base

        self.category = category

        self.dict_data = {}

        self.rdd_data = None

        if tablename is None:
            self.table_name = self.base.strip('/').split('/')[-1]
            if self.category != None:
                self.table_name += ('-' + self.category)
        else:
            self.table_name = tablename

        self.sparker = SC.Sparker(host='HPC-server', appname='ImageILSVRC-S',
                                  master='spark://HPC-server:7077')

        self.steger = F5.F5(sample_key, 1)

    def get_table(self):

        print "getting table [%s]..." % self.table_name

        if self.table != None:
            return self.table

        if self.connection is None:
            c = happybase.Connection('HPC-server')
            self.connection = c

        tables = self.connection.tables()
        if self.table_name not in tables:
            families = {'cf_pic': dict(),
                        'cf_info': dict(max_versions=10),
                        'cf_tag': dict(),
                        'cf_feat': dict(),
                        }
            self.connection.create_table(name=self.table_name, families=families)

        table = self.connection.table(name=self.table_name)

        self.table = table

        return table

    def delete_table(self, table_name=None, disable=True):

        print "deleting table..."

        if table_name == None:
            table_name = self.table_name

        if self.connection is None:
            c = happybase.Connection('HPC-server')
            self.connection = c

        tables = self.connection.tables()
        if table_name not in tables:
            return False
        else:
            try:
                self.connection.delete_table(table_name, disable)
            except:
                print 'Exception when deleting table.'
                raise
        return True

    def _get_info(self, 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')
            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
        finally:
            tmpf.close()

    def _get_feat(self, image, feattype='ibd', **kwargs):

        # size = kwargs.get('size', (48, 48))
        #
        # if feattype == 'hog':
        # feater = HOG.FeatHOG(size=size)
        if feattype == 'ibd':

            feater = IntraBlockDiff.FeatIntraBlockDiff()
        else:
            raise Exception("Unknown feature type!")

        desc = feater.feat(image)

        return desc

    def _rddparse_data(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(raw_row):

        key = raw_row[0]

        items = raw_row[1].decode('unicode-escape').encode('latin-1').split('--%--')

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

        return (key, data)

    def _rdd_embed(self, row):
        """
        input:
            e.g. row =('row1',[1,3400,'hello'])
        return:
            newrow = ('row2',[34,5400,'embeded'])
        """
        items = row[1]
        capacity, rate, chosen = items[4], items[6], 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')

            if rate == None:

                embed_rate = self.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 = self.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] + self._get_info(raw, embed_rate, 0, 1))

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

    def _extract_data(self, mode='hbase', writeback=False, withdata=True):

        """
        Get info barely out of image data.
        """

        print "extracting data..."

        if mode == 'hbase':
            if self.table == None:
                self.table = self.get_table()

            cols = ['cf_pic:data']

            for key, data in self.table.scan(columns=cols):
                data = data['cf_pic:data']

                self.dict_data[key] = [data] + self._get_info(data)

            if not writeback:
                return self.dict_data
            else:
                try:
                    with self.table.batch(batch_size=5000) as b:
                        for imgname, imginfo in self.dict_data.items():
                            b.put(imgname,
                                  {
                                      # 'cf_pic:data': imginfo[0],

                                      'cf_info:width': str(imginfo[1]),
                                      'cf_info:height': str(imginfo[2]),
                                      'cf_info:size': str(imginfo[3]),
                                      'cf_info:capacity': str(imginfo[4]),
                                      'cf_info:quality': str(imginfo[5]),
                                      'cf_info:rate': str(imginfo[6]),
                                      'cf_tag:chosen': str(imginfo[7]),
                                      'cf_tag:class': str(imginfo[8]),

                                  })

                except ValueError:
                    raise


        elif mode == 'spark':

            if self.sparker == None:
                self.sparker = SC.Sparker(host='HPC-server', appname='ImageILSVRC-S',

                                          master='spark://HPC-server:7077')

            cols = [
                'cf_pic:data',
                'cf_info:width',
                'cf_info:height',
                'cf_info:size',
                'cf_info:capacity',
                'cf_info:quality',
                'cf_info:rate',
                'cf_tag:chosen',
                'cf_tag:class'
            ]

            # # Debug

            # tmp_data = self.sparker.read_hbase(self.table_name, func=rdd.rddparse_data_ILS,

            # collect=False)
            # # tmp_data = tmp_data.mapValues(lambda data: [data] + rdd.rddinfo_ILS(data))
            # print tmp_data.collect()[0][1]
            # return

            self.rdd_data = self.sparker.read_hbase(self.table_name, func=rdd.rddparse_data_ILS,

                                                    collect=False).mapValues(

                lambda data: [data] + rdd.rddinfo_ILS(data))

            if not writeback:
                return self.rdd_data
            else:

                self.sparker.write_hbase(self.table_name, self.rdd_data, fromrdd=True, columns=cols,

                                         withdata=withdata)

        elif mode == 'analysis':
            if self.sparker == None:
                self.sparker = SC.Sparker(host='HPC-server', appname='ImageILSVRC-S',
                                          master='spark://HPC-server:7077')

            cols = [
                'cf_pic:data',
            ]

            # # Debug
            # tmp_data = self.sparker.read_hbase(self.table_name, func=rdd.rddparse_data_ILS,
            # collect=False)
            # # tmp_data = tmp_data.mapValues(lambda data: [data] + rdd.rddinfo_ILS(data))
            # print tmp_data.collect()[0][1]
            # return


            self.rdd_data = self.sparker.read_hbase(self.table_name, func=rdd.rddparse_data_ILS,
                                                    collect=False).mapValues(
                lambda data: [data])

            if not writeback:
                return self.rdd_data
            else:
                self.sparker.write_hbase(self.table_name, self.rdd_data, fromrdd=True, columns=cols,
                                         withdata=withdata)

        else:

            raise Exception("Unknown mode!")

    def _embed_data(self, mode='hbase', rate=None, readforward=False, writeback=False, withdata=True):

        print "embedding data..."

        if mode == 'hbase':
            if self.table == None:
                self.table = self.get_table()

            if readforward:
                self.dict_data = {}

                cols = [
                    'cf_pic:data',
                    'cf_info:width',
                    'cf_info:height',
                    'cf_info:size',
                    'cf_info:capacity',
                    'cf_info:quality',
                    'cf_info:rate',
                    'cf_tag:chosen',
                    'cf_tag:class'
                ]

                for key, data in self.table.scan(columns=cols):
                    data = [data[k] for k in cols]

                    self.dict_data[key] = data

            dict_data_ext = {}

            for imgname, imgdata in self.dict_data.items():

                capacity, chosen = int(imgdata[4]), int(imgdata[7])

                if chosen == 0:
                    continue

                try:
                    tmpf_src = tempfile.NamedTemporaryFile(suffix='.jpg', mode='w+b')
                    tmpf_src.write(imgdata[0])

                    tmpf_src.seek(0)

                    tmpf_dst = tempfile.NamedTemporaryFile(suffix='.jpg', mode='w+b')

                    if rate == None:

                        embed_rate = self.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(capacity * rate) / 8)

                        embed_rate = self.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()
                    dict_data_ext[index + '.jpg'] = [raw] + self._get_info(raw, embed_rate, 0, 1)

                except Exception as e:
                    print e

                    raise

                finally:
                    tmpf_src.close()
                    tmpf_dst.close()

            self.dict_data.update(dict_data_ext)

            if not writeback:
                return self.dict_data
            else:
                try:
                    with self.table.batch(batch_size=5000) as b:
                        for imgname, imginfo in dict_data_ext.items():
                            b.put(imgname,
                                  {
                                      'cf_pic:data': imginfo[0],

                                      'cf_info:width': str(imginfo[1]),
                                      'cf_info:height': str(imginfo[2]),
                                      'cf_info:size': str(imginfo[3]),
                                      'cf_info:capacity': str(imginfo[4]),
                                      'cf_info:quality': str(imginfo[5]),
                                      'cf_info:rate': str(imginfo[6]),
                                      'cf_tag:chosen': str(imginfo[7]),
                                      'cf_tag:class': str(imginfo[8]),
                                  })

                except ValueError:
                    raise

        elif mode == 'spark':

            if self.sparker == None:
                self.sparker = SC.Sparker(host='HPC-server', appname='ImageILSVRC-S',

                                          master='spark://HPC-server:7077')

            cols = [
                'cf_pic:data',
                'cf_info:width',
                'cf_info:height',
                'cf_info:size',
                'cf_info:capacity',
                'cf_info:quality',
                'cf_info:rate',
                'cf_tag:chosen',
                'cf_tag:class'
            ]

            if readforward:

                self.rdd_data = self.sparker.read_hbase(self.table_name, func=rdd.rddparse_all_ILS, collect=False)

            # rdd_data_ext = self.rdd_data.map(lambda x: rdd.rddembed_ILS(x, rate=rate)).filter(lambda x: x != None)
            # self.rdd_data = self.rdd_data.union(rdd_data_ext)

            self.rdd_data = self.rdd_data.flatMap(lambda x: rdd.rddembed_ILS_EXT(x, rate=rate))

            if not writeback:

                return self.rdd_data

            else:

                self.sparker.write_hbase(self.table_name, self.rdd_data, fromrdd=True, columns=cols,

                                         withdata=withdata)

        else:

            raise Exception("Unknown mode!")

    def _extract_feat(self, mode='hbase', feattype='ibd', readforward=False, writeback=False, withdata=False):

        print "extracting feat..."

        if mode == 'hbase':
            if self.table == None:
                self.table = self.get_table()

            if readforward:
                self.dict_data = {}

                cols = [
                    'cf_pic:data',
                    'cf_info:width',
                    'cf_info:height',
                    'cf_info:size',
                    'cf_info:capacity',
                    'cf_info:quality',
                    'cf_info:rate',
                    'cf_tag:chosen',
                    'cf_tag:class'
                ]

                for key, data in self.table.scan(columns=cols):
                    data = [data[k] for k in cols]

                    self.dict_data[key] = data

            for imgname, imgdata in self.dict_data.items():
                try:
                    tmpf_src = tempfile.NamedTemporaryFile(suffix='.jpg', mode='w+b')
                    tmpf_src.write(imgdata[0])

                    tmpf_src.seek(0)

                    desc = json.dumps(self._get_feat(tmpf_src.name, feattype=feattype).tolist())

                    self.dict_data[imgname].append(desc)

                except Exception as e:
                    print e

                    raise

                finally:
                    tmpf_src.close()

            if not writeback:
                return self.dict_data
            else:
                try:
                    with self.table.batch(batch_size=5000) as b:
                        for imgname, imginfo in self.dict_data.items():
                            b.put(imgname,
                                  {
                                      'cf_pic:data': imginfo[0],

                                      'cf_info:width': str(imginfo[1]),
                                      'cf_info:height': str(imginfo[2]),
                                      'cf_info:size': str(imginfo[3]),
                                      'cf_info:capacity': str(imginfo[4]),
                                      'cf_info:quality': str(imginfo[5]),
                                      'cf_info:rate': str(imginfo[6]),
                                      'cf_tag:chosen': str(imginfo[7]),
                                      'cf_tag:class': str(imginfo[8]),
                                      'cf_feat:' + feattype: imginfo[9],
                                  })

                except ValueError:
                    raise

        elif mode == 'spark':

            if self.sparker == None:
                self.sparker = SC.Sparker(host='HPC-server', appname='ImageILSVRC-S',

                                          master='spark://HPC-server:7077')

            cols = [
                'cf_pic:data',
                'cf_info:width',
                'cf_info:height',
                'cf_info:size',
                'cf_info:capacity',
                'cf_info:quality',
                'cf_info:rate',
                'cf_tag:chosen',
                'cf_tag:class',
                'cf_feat:' + feattype,
            ]

            if readforward:

                self.rdd_data = self.sparker.read_hbase(self.table_name, func=rdd.rddparse_all_ILS, collect=False)

            self.rdd_data = self.rdd_data.mapValues(lambda items: rdd.rddfeat_ILS(items, feattype))

            # print self.rdd_data.collect()[0]
            # return

            if not writeback:

                return self.rdd_data

            else:

                print "writing back..."

                self.sparker.write_hbase(self.table_name, self.rdd_data, fromrdd=True, columns=cols,

                                         withdata=withdata)

        else:
            raise Exception("Unknown mode!")

    def _analysis(self, mode='analysis', feattype='ibd', readforward=False, writeback=True, withdata=False):

        if mode == 'analysis':
            if self.sparker == None:
                self.sparker = SC.Sparker(host='HPC-server', appname='ImageILSVRC-S',
                                          master='spark://HPC-server:7077')

            cols = [
                'cf_pic:data',

                'cf_tag:class',
            ]

            if readforward:
                self.rdd_data = self.sparker.read_hbase(self.table_name, func=rdd.rddparse_all_ILS, collect=False)

            self.rdd_data = self.rdd_data.mapValues(lambda items: rdd.rddanalysis_ILS(items))

            # print self.rdd_data.collect()[0]
            # return

            if not writeback:
                return self.rdd_data
            else:
                print "writing back..."
                self.sparker.write_hbase(self.table_name, self.rdd_data, fromrdd=True, columns=cols,
                                         withdata=withdata)

        else:
            raise Exception("Unknown mode!")

    def format(self):

        print "formatting..."

        self._extract_data(mode='hbase', writeback=False)
        self._embed_data(mode='hbase', rate=0.1, readforward=False, writeback=False)
        self._extract_feat(mode='hbase', feattype='ibd', readforward=False, writeback=True)

    def load_data(self, mode='hbase', feattype='ibd', tagtype='class', collect=False):

        print "loading data..."

        INDEX = []
        X = []
        Y = []

        if mode == "local":
            dict_dataset = {}

            if feattype == 'coef':  # raw
                with open(self.list_file, 'rb') as tsvfile:
                    tsvfile = csv.reader(tsvfile, delimiter='\t')
                    for line in tsvfile:
                        hash = line[0]
                        tag = line[-1]
                        image = os.path.join(self.img_dir, hash[:3], hash[3:] + '.jpg')
                        if image:
                            im = Jpeg(image, key=sample_key)
                            dict_dataset[hash] = (tag, im.getCoefMatrix(channel='Y'))

                for tag, feat in dict_dataset.values():

                    feat.ravel()[[i * 304 + j for i in range(0, 304, 8) for j in range(0, 304, 8)]] = 0

                    X.append(feat.tolist())
                    Y.append(int(tag))

            else:
                with open(self.list_file, 'rb') as tsvfile:
                    tsvfile = csv.reader(tsvfile, delimiter='\t')
                    for line in tsvfile:
                        hash = line[0]
                        tag = line[-1]
                        path_feat = os.path.join(self.feat_dir, hash[:3], hash[3:] + '.' + feattype)
                        if path_feat:
                            with open(path_feat, 'rb') as featfile:
                                dict_dataset[hash] = (tag, json.loads(featfile.read()))

                for tag, feat in dict_dataset.values():
                    # X.append([item for sublist in feat for subsublist in sublist for item in subsublist])
                    X.append(np.array(feat).ravel().tolist())
                    Y.append(int(tag))

        elif mode == "hbase":

            if self.table == None:
                self.table = self.get_table()

            col_feat, col_tag = 'cf_feat:' + feattype, 'cf_tag:' + tagtype
            for key, data in self.table.scan(columns=[col_feat, col_tag]):

                # X.append(
                # [item for sublist in json.loads(data[col_feat]) for subsublist in sublist for item in subsublist])
                X.append(np.array(json.loads(data[col_feat])).ravel().tolist())

                Y.append(int(data[col_tag]))

        elif mode == "spark" or mode == "cluster":

            if self.sparker == None:
                self.sparker = SC.Sparker(host='HPC-server', appname='ImageILSVRC-S',

                                          master='spark://HPC-server:7077')

            rdd_dataset = self.sparker.read_hbase(self.table_name, func=rdd.rddparse_dataset_ILS, collect=False)

            if not collect:
                rdd_dataset = rdd_dataset.map(lambda x: LabeledPoint(x[0], x[1]))
                return rdd_dataset

            for tag, feat in rdd_dataset.collect():

                X.append(feat)
                Y.append(tag)

        else:
            raise Exception("Unknown mode!")

        return X, Y