# -*- coding: utf-8 -*-
__author__ = 'chunk'

from ..common import *
from .dependencies import *
from . import *
# from ..mdata import MSR, CV, ILSVRC, ILSVRC_S

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

import sys
from pyspark import RDD
from pyspark import SparkConf, SparkContext
from pyspark.mllib.classification import LogisticRegressionWithSGD, SVMWithSGD
from pyspark.mllib.regression import LabeledPoint
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__))


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, 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 [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 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


# scconf = SparkConf()
# scconf.setSparkHome("HPC-server") \
# .setMaster("spark://HPC-server:7077") \
#     .setAppName("example")
# sc = SparkContext(conf=scconf)
#
#
# def read_hbase(table_name, func=None, collect=False):
#     """
#     ref - http://happybase.readthedocs.org/en/latest/user.html#retrieving-data
#
#     Filter format:
#         columns=['cf1:col1', 'cf1:col2']
#         or
#         columns=['cf1']
#
#     """
#
#     hconf = {
#         "hbase.zookeeper.quorum": "HPC-server, HPC, HPC2",
#         # "hbase.zookeeper.quorum": self.host,
#         "hbase.mapreduce.inputtable": table_name,
#     }
#
#     hbase_rdd = sc.newAPIHadoopRDD(inputFormatClass=hparams["inputFormatClass"],
#                                            keyClass=hparams["readKeyClass"],
#                                            valueClass=hparams["readValueClass"],
#                                            keyConverter=hparams["readKeyConverter"],
#                                            valueConverter=hparams["readValueConverter"],
#                                            conf=hconf)
#
#     parser = func if func != None else rddparse_data_CV
#     hbase_rdd = hbase_rdd.map(lambda x: parser(x))
#
#     if collect:
#         return hbase_rdd.collect()
#     else:
#         return hbase_rdd
#
#
# def write_hbase(table_name, data, fromrdd=False, columns=None, withdata=False):
#     """
#     Data Format: (Deprecated)
#         e.g. [["row8", "f1", "", "caocao cao"], ["row9", "f1", "c1", "asdfg hhhh"]]
#
#     Data(from dictionary):
#         e.g. data ={'row1':[1,3400,'hello'], 'row2':[34,5000,'here in mine']},
#             cols = ['cf_info:id', 'cf_info:size', 'cf_tag:desc']
#     Data(from Rdd):
#         e.g. data =[('row1',[1,3400,'hello']), ('row2',[34,5000,'here in mine'])],
#             cols = ['cf_info:id', 'cf_info:size', 'cf_tag:desc']
#     """
#     hconf = {
#         "hbase.zookeeper.quorum": "HPC-server, HPC, HPC2",  # "hbase.zookeeper.quorum": self.host,
#         "hbase.mapreduce.inputtable": table_name,
#         "hbase.mapred.outputtable": table_name,
#         "mapreduce.outputformat.class": hparams["outputFormatClass"],
#         "mapreduce.job.output.key.class": hparams["writeKeyClass"],
#         "mapreduce.job.output.value.class": hparams["writeValueClass"],
#     }
#     cols = [col.split(':') for col in columns]
#     if not fromrdd:
#         rdd_data = sc.parallelize(data)
#     else:
#         rdd_data = data
#
#     rdd_data.flatMap(lambda x: format_out(x, cols, withdata=withdata)).saveAsNewAPIHadoopDataset(
#         conf=hconf,
#         keyConverter=hparams["writeKeyConverter"],
#         valueConverter=hparams["writeValueConverter"])


class Sparker(object):
    def __init__(self, host='HPC-server', appname='NewPySparkApp', **kwargs):
        load_env()
        self.host = host
        self.appname = appname
        self.master = kwargs.get('master', 'spark://%s:7077' % self.host)
        self.conf = SparkConf()
        self.conf.setSparkHome(self.host) \
            .setMaster(self.master) \
            .setAppName(self.appname)

        # self.conf.set("spark.akka.frameSize","10685760")
        # self.conf.set("spark.driver.extraClassPath", extraClassPath) \
        # .set("spark.executor.extraClassPath", extraClassPath) \
        # .set("SPARK_CLASSPATH", extraClassPath) \
        # .set("spark.driver.memory", "1G") \
        # .set("spark.yarn.jar", sparkJar)

        self.sc = SparkContext(conf=self.conf)

        self.model = None

    def read_hbase(self, table_name, func=None, collect=False, parallelism=40):
        """
        ref - http://happybase.readthedocs.org/en/latest/user.html#retrieving-data

        Filter format:
            columns=['cf1:col1', 'cf1:col2']
            or
            columns=['cf1']

        """

        hconf = {
            "hbase.zookeeper.quorum": "HPC-server, HPC, HPC2",
            # "hbase.zookeeper.quorum": self.host,
            "hbase.mapreduce.inputtable": table_name,
        }

        hbase_rdd = self.sc.newAPIHadoopRDD(inputFormatClass=hparams["inputFormatClass"],
                                            keyClass=hparams["readKeyClass"],
                                            valueClass=hparams["readValueClass"],
                                            keyConverter=hparams["readKeyConverter"],
                                            valueConverter=hparams["readValueConverter"],
                                            conf=hconf)

        parser = func if func != None else rddparse_data_CV
        hbase_rdd = hbase_rdd.map(lambda x: parser(x))

        if collect:
            return hbase_rdd.collect()
        else:
            """
            RDD-hbase bug fixed.(with 'repartition()')
            <http://stackoverflow.com/questions/29011574/how-is-spark-partitioning-from-hdfs>

            When Spark reads a file from HDFS, it creates a single partition for a single input split. Input split is set by the Hadoop InputFormat used to read this file. For instance, if you use textFile() it would be TextInputFormat in Hadoop, which would return you a single partition for a single block of HDFS (but the split between partitions would be done on line split, not the exact block split), unless you have a compressed text file. In case of compressed file you would get a single partition for a single file (as compressed text files are not splittable).
            When you call rdd.repartition(x) it would perform a shuffle of the data from N partititons you have in rdd to x partitions you want to have, partitioning would be done on round robin basis.
            If you have a 30GB uncompressed text file stored on HDFS, then with the default HDFS block size setting (128MB) it would be stored in 235 blocks, which means that the RDD you read from this file would have 235 partitions. When you call repartition(1000) your RDD would be marked as to be repartitioned, but in fact it would be shuffled to 1000 partitions only when you will execute an action on top of this RDD (lazy execution concept)

            """
            return hbase_rdd.repartition(parallelism)

    def write_hbase(self, table_name, data, fromrdd=False, columns=None, withdata=False):
        """
        Data Format: (Deprecated)
            e.g. [["row8", "f1", "", "caocao cao"], ["row9", "f1", "c1", "asdfg hhhh"]]

        Data(from dictionary):
            e.g. data ={'row1':[1,3400,'hello'], 'row2':[34,5000,'here in mine']},
                cols = ['cf_info:id', 'cf_info:size', 'cf_tag:desc']
        Data(from Rdd):
            e.g. data =[('row1',[1,3400,'hello']), ('row2',[34,5000,'here in mine'])],
                cols = ['cf_info:id', 'cf_info:size', 'cf_tag:desc']
        """
        hconf = {
            "hbase.zookeeper.quorum": "HPC-server, HPC, HPC2",
        # "hbase.zookeeper.quorum": self.host,
            "hbase.mapreduce.inputtable": table_name,
            "hbase.mapred.outputtable": table_name,
            "mapreduce.outputformat.class": hparams["outputFormatClass"],
            "mapreduce.job.output.key.class": hparams["writeKeyClass"],
            "mapreduce.job.output.value.class": hparams["writeValueClass"],
        }
        cols = [col.split(':') for col in columns]
        if not fromrdd:
            rdd_data = self.sc.parallelize(data)
        else:
            rdd_data = data

        rdd_data.flatMap(
            lambda x: format_out(x, cols, withdata=withdata)).saveAsNewAPIHadoopDataset(
            conf=hconf,
            keyConverter=hparams["writeKeyConverter"],
            valueConverter=hparams["writeValueConverter"])

    def train_svm(self, X, Y=None):

        if Y == None:
            # From rdd_labeled
            assert isinstance(X, RDD)
            svm = SVMWithSGD.train(X)
        else:
            # data = []
            # for feat, tag in zip(X, Y):
            # data.append(LabeledPoint(tag, feat))
            # svm = SVMWithSGD.train(self.sc.parallelize(data))
            hdd_data = self.sc.parallelize(zip(X, Y), 20).map(lambda x: LabeledPoint(x[1], x[0]))
            svm = SVMWithSGD.train(hdd_data)
        self.model = svm
        # with open('res/svm_spark.model', 'wb') as modelfile:
        # model = pickle.dump(svm, modelfile)

        return self.model

    def predict_svm(self, x, collect=False, model=None):
        """
        From pyspark.mlib.classification.py:

            >> svm.predict([1.0])
            1
            >> svm.predict(sc.parallelize([[1.0]])).collect()
            [1]
            >> svm.clearThreshold()
            >> svm.predict(array([1.0]))
            1.25...
        """
        if model is None:
            if self.model != None:
                model = self.model
            else:
                # with open('res/svm_spark.model', 'rb') as modelfile:
                # model = pickle.load(modelfile)
                raise Exception("No model available!")

        res = model.predict(x)
        if collect:
            return res.collect()
        else:
            return res

    def test_svm(self, X, Y=None, model=None):
        if model is None:
            if self.model != None:
                model = self.model
            else:
                # with open('res/svm_spark.model', 'rb') as modelfile:
                # model = pickle.load(modelfile)
                raise Exception("No model available!")

        if Y == None:
            assert isinstance(X, RDD)
            pass
        else:
            result_Y = np.array(self.predict_svm(X, collect=True))
            return np.mean(Y == result_Y)