# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

# pylint:skip-file
import logging
import sys, random, time, math
import mxnet as mx
import numpy as np
from nce import *
from operator import itemgetter
from optparse import OptionParser
from collections import Counter

import logging
head = head = '%(asctime)-15s %(message)s'
logging.basicConfig(level=logging.INFO, format=head)


# ----------------------------------------------------------------------------------------
EMBEDDING_SIZE = 100
BATCH_SIZE = 256
NUM_LABEL = 5
NUM_EPOCH = 20
MIN_COUNT = 5  # only works when doing nagative sampling, keep it same as nce-loss
GRAMS = 3      # here we use triple-letter representation
MAX_SUBWORDS = 10
PADDING_CHAR = '</s>'


# ----------------------------------------------------------------------------------------
def get_net(vocab_size, num_input, num_label):
    data = mx.sym.Variable('data')
    mask = mx.sym.Variable('mask')  # use mask to handle variable-length input.
    label = mx.sym.Variable('label')
    label_mask = mx.sym.Variable('label_mask')
    label_weight = mx.sym.Variable('label_weight')
    embed_weight = mx.sym.Variable('embed_weight')

    # Get embedding for one-hot input.
    # get sub-word units input.
    unit_embed = mx.sym.Embedding(data=data, input_dim=vocab_size,
                                  weight=embed_weight,
                                  output_dim=EMBEDDING_SIZE)

    # mask embedding_output to get summation of sub-word units'embedding.
    unit_embed = mx.sym.broadcast_mul(lhs=unit_embed, rhs=mask, name='data_units_embed')

    # sum over all these words then you get word-embedding.
    data_embed = mx.sym.sum(unit_embed, axis=2)

    # Slice input equally along specified axis.
    datavec = mx.sym.SliceChannel(data=data_embed,
                                  num_outputs=num_input,
                                  squeeze_axis=1, name='data_slice')
    pred = datavec[0]
    for i in range(1, num_input):
        pred = pred + datavec[i]

    return nce_loss_subwords(data=pred,
                             label=label,
                             label_mask=label_mask,
                             label_weight=label_weight,
                             embed_weight=embed_weight,
                             vocab_size=vocab_size,
                             num_hidden=EMBEDDING_SIZE,
                             num_label=num_label)


def get_subword_units(token, gram=GRAMS):
    """Return subword-units presentation, given a word/token.
    """
    if token == '</s>':  # special token for padding purpose.
        return [token]
    t = '#' + token + '#'
    return [t[i:i + gram] for i in range(0, len(t) - gram + 1)]


def get_subword_representation(wid, vocab_inv, units_vocab, max_len):
    token = vocab_inv[wid]
    units = [units_vocab[unit] for unit in get_subword_units(token)]
    weights = [1] * len(units) + [0] * (max_len - len(units))
    units = units + [units_vocab[PADDING_CHAR]] * (max_len - len(units))
    return units, weights


def prepare_subword_units(tks):
    # statistics on units
    units_vocab = {PADDING_CHAR:1}
    max_len = 0
    unit_set = set()
    logging.info('grams: %d', GRAMS)
    logging.info('counting max len...')
    for tk in tks:
        res = get_subword_units(tk)
        unit_set.update(i for i in res)
        if max_len < len(res):
            max_len = len(res)
    logging.info('preparing units vocab...')
    for unit in unit_set:
        if len(unit) == 0:
            continue
        if unit not in units_vocab:
            units_vocab[unit] = len(units_vocab)
        uid = units_vocab[unit]
    return units_vocab, max_len


def load_data_as_subword_units(name):
    tks = []
    fread = open(name, 'r')
    logging.info('reading corpus from file...')
    for line in fread:
        line = line.strip().decode('utf-8')
        tks.extend(line.split(' '))

    logging.info('Total tokens: %d', len(tks))

    tks = [tk for tk in tks if len(tk) <= MAX_SUBWORDS]
    c = Counter(tks)

    logging.info('Total vocab: %d', len(c))

    vocab = {}
    vocab_inv = {}
    freq = [0]
    data = []
    for tk in tks:
        if len(tk) == 0:
            continue
        if tk not in vocab:
            vocab[tk] = len(vocab)
            freq.append(0)
        wid = vocab[tk]
        vocab_inv[wid] = tk
        data.append(wid)
        freq[wid] += 1

    negative = []
    for i, v in enumerate(freq):
        if i == 0 or v < MIN_COUNT:
            continue
        v = int(math.pow(v * 1.0, 0.75))  # sample negative w.r.t. its frequency
        negative += [i for _ in range(v)]

    logging.info('counting subword units...')
    units_vocab, max_len = prepare_subword_units(tks)
    logging.info('vocabulary size: %d', len(vocab))
    logging.info('subword unit size: %d', len(units_vocab))

    logging.info('generating input data...')
    units = []
    weights = []
    for wid in data:
        word_units, weight = get_subword_representation(wid, vocab_inv, units_vocab, max_len)
        units.append(word_units)
        weights.append(weight)

    negative_units = []
    negative_weights = []
    for wid in negative:
        word_units, weight = get_subword_representation(wid, vocab_inv, units_vocab, max_len)
        negative_units.append(word_units)
        negative_weights.append(weight)

    return data, units, weights, negative_units, negative_weights, vocab, units_vocab, freq, max_len


class SimpleBatch(object):
    def __init__(self, data_names, data, label_names, label):
        self.data = data
        self.label = label
        self.data_names = data_names
        self.label_names = label_names

    @property
    def provide_data(self):
        return [(n, x.shape) for n, x in zip(self.data_names, self.data)]

    @property
    def provide_label(self):
        return [(n, x.shape) for n, x in zip(self.label_names, self.label)]


class DataIter(mx.io.DataIter):
    def __init__(self, fname, batch_size, num_label):
        super(DataIter, self).__init__()
        self.batch_size = batch_size
        self.data, self.units, self.weights, self.negative_units, self.negative_weights, \
        self.vocab, self.units_vocab, self.freq, self.max_len = load_data_as_subword_units(fname)
        self.vocab_size = len(self.units_vocab)
        self.num_label = num_label
        self.provide_data = [('data', (batch_size, num_label - 1, self.max_len)),
                             ('mask', (batch_size, num_label - 1, self.max_len, 1))]
        self.provide_label = [('label', (self.batch_size, num_label, self.max_len)),
                              ('label_weight', (self.batch_size, num_label)),
                              ('label_mask', (self.batch_size, num_label, self.max_len, 1))]

    def sample_ne(self):
        # a negative sample.
        return self.negative_units[random.randint(0, len(self.negative_units) - 1)]

    def sample_ne_indices(self):
        return [random.randint(0, len(self.negative_units) - 1) for _ in range(self.num_label - 1)]

    def __iter__(self):
        logging.info('DataIter start.')
        batch_data = []
        batch_data_mask = []
        batch_label = []
        batch_label_mask = []
        batch_label_weight = []
        start = random.randint(0, self.num_label - 1)
        for i in range(start, len(self.units) - self.num_label - start, self.num_label):
            context_units = self.units[i: i + self.num_label / 2] + \
                            self.units[i + 1 + self.num_label / 2: i + self.num_label]
            context_mask = self.weights[i: i + self.num_label / 2] + \
                           self.weights[i + 1 + self.num_label / 2: i + self.num_label]
            target_units = self.units[i + self.num_label / 2]
            target_word = self.data[i + self.num_label / 2]
            if self.freq[target_word] < MIN_COUNT:
                continue
            indices = self.sample_ne_indices()
            target = [target_units] + [self.negative_units[i] for i in indices]
            target_weight = [1.0] + [0.0 for _ in range(self.num_label - 1)]
            target_mask = [self.weights[i + self.num_label / 2]] + [self.negative_weights[i] for i in indices]

            batch_data.append(context_units)
            batch_data_mask.append(context_mask)
            batch_label.append(target)
            batch_label_mask.append(target_mask)
            batch_label_weight.append(target_weight)

            if len(batch_data) == self.batch_size:
                # reshape for broadcast_mul
                batch_data_mask = np.reshape(batch_data_mask, (batch_size, num_label - 1, self.max_len, 1))
                batch_label_mask = np.reshape(batch_label_mask, (batch_size, num_label, self.max_len, 1))
                data_all = [mx.nd.array(batch_data), mx.nd.array(batch_data_mask)]
                label_all = [mx.nd.array(batch_label), mx.nd.array(batch_label_weight), mx.nd.array(batch_label_mask)]
                data_names = ['data', 'mask']
                label_names = ['label', 'label_weight', 'label_mask']
                # clean up
                batch_data = []
                batch_data_mask = []
                batch_label = []
                batch_label_weight = []
                batch_label_mask = []
                yield SimpleBatch(data_names, data_all, label_names, label_all)

    def reset(self):
        pass


if __name__ == '__main__':
    head = '%(asctime)-15s %(message)s'
    logging.basicConfig(level=logging.DEBUG, format=head)

    parser = OptionParser()
    parser.add_option("-g", "--gpu", action="store_true", dest="gpu", default=False,
                      help="use gpu")

    batch_size = BATCH_SIZE
    num_label = NUM_LABEL

    data_train = DataIter("./data/text8", batch_size, num_label)

    network = get_net(data_train.vocab_size, num_label - 1, num_label)

    options, args = parser.parse_args()
    # devs = mx.cpu()
    devs = [mx.cpu(i) for i in range(4)]
    if options.gpu == True:
        devs = mx.gpu()
    model = mx.model.FeedForward(ctx=devs,
                                 symbol=network,
                                 num_epoch=NUM_EPOCH,
                                 learning_rate=0.3,
                                 momentum=0.9,
                                 wd=0.0000,
                                 initializer=mx.init.Xavier(factor_type="in", magnitude=2.34))

    metric = NceAuc()
    model.fit(X=data_train,
              eval_metric=metric,
              batch_end_callback=mx.callback.Speedometer(batch_size, 50), )
