import sys
import cPickle
import os
import time
import tensorflow as tf
import numpy as np
from model import Model
from utils.parser_utils import minibatches, load_and_preprocess_data

def xavier_weight_init():
    """Returns function that creates random tensor.

    The specified function will take in a shape (tuple or 1-d array) and
    returns a random tensor of the specified shape drawn from the
    Xavier initialization distribution.

    Hint: You might find tf.random_uniform useful.
    """
    def _xavier_initializer(shape, **kwargs):
        """Defines an initializer for the Xavier distribution.
        Specifically, the output should be sampled uniformly from [-epsilon, epsilon] where
            epsilon = sqrt(6) / <sum of the sizes of shape's dimensions>
        e.g., if shape = (2, 3), epsilon = sqrt(6 / (2 + 3))

        This function will be used as a variable initializer.

        Args:
            shape: Tuple or 1-d array that species the dimensions of the requested tensor.
        Returns:
            out: tf.Tensor of specified shape sampled from the Xavier distribution.
        """
        epsilon = np.sqrt(6 / np.sum(shape))
        out_t = tf.Variable(tf.random_uniform(shape=shape, minval=-epsilon, maxval=epsilon))
        return out_t
    return _xavier_initializer

class Config(object):
    """Holds model hyperparams and data information.

    The config class is used to store various hyperparameters and dataset
    information parameters. Model objects are passed a Config() object at
    instantiation. They can then call self.config.<hyperparameter_name> to
    get the hyperparameter settings.
    """
    n_features = 36
    n_classes = 3
    dropout = 0.5  # (p_drop in the handout)
    embed_size = 50
    hidden_size = 200
    batch_size = 1024
    n_epochs = 10
    lr = 0.0005

class ParserModel(Model):
    """
    Implements a feedforward neural network with an embedding layer and single hidden layer.
    This network will predict which transition should be applied to a given partial parse
    configuration.
    """

    def add_placeholders(self):
        """Generates placeholder variables to represent the input tensors"""
        self.input_placeholder = tf.placeholder(tf.int32, [None, self.config.n_features])
        self.labels_placeholder = tf.placeholder(tf.float32, [None, self.config.n_classes])
        self.droupout_placeholder = tf.placeholder(tf.float32)
        
    def create_feed_dict(self, inputs_batch, labels_batch=None, dropout=0):
        """Creates the feed_dict for the dependency parser.

        A feed_dict takes the form of:

        feed_dict = {
                <placeholder>: <tensor of values to be passed for placeholder>,
                ....
        }


        Hint: The keys for the feed_dict should be a subset of the placeholder tensors created in add_placeholders.
        Hint: When an argument is None, don't add it to the feed_dict.

        Args:
            inputs_batch: A batch of input data.
            labels_batch: A batch of label data.
            dropout: The dropout rate.
        Returns:
            feed_dict: The feed dictionary mapping from placeholders to values.
        """
        feed_dict = {self.input_placeholder: inputs_batch, \
                    self.droupout_placeholder: dropout}
        if labels_batch.any():
            feed_dict[self.labels_placeholder] = labels_batch
        return feed_dict
    
    def add_embedding(self):
        """Adds an embedding layer that maps from input tokens (integers) to vectors and then
        concatenates those vectors:
            - Creates a tf.Variable and initializes it with self.pretrained_embeddings.
            - Uses the input_placeholder to index into the embeddings tensor, resulting in a
              tensor of shape (None, n_features, embedding_size).
            - Concatenates the embeddings by reshaping the embeddings tensor to shape
              (None, n_features * embedding_size).

        Hint: You might find tf.nn.embedding_lookup useful.
        Hint: You can use tf.reshape to concatenate the vectors. 

        Returns:
            embeddings: tf.Tensor of shape (None, n_features*embed_size)
        """
        embedding = tf.Variable(self.pretrained_embeddings)
        embeddings = tf.nn.embedding_lookup(embedding, self.input_placeholder)
        embeddings = tf.reshape(embeddings, [-1, self.config.n_features * self.config.embed_size])
        return embeddings    
    
    def add_prediction_op(self):
        """Adds the 1-hidden-layer NN:
            h = Relu(xW + b1)
            h_drop = Dropout(h, dropout_rate)
            pred = h_dropU + b2

        Note that we are not applying a softmax to pred. The softmax will instead be done in
        the add_loss_op function, which improves efficiency because we can use
        tf.nn.softmax_cross_entropy_with_logits

        Use the initializer from q2_initialization.py to initialize W and U (you can initialize b1
        and b2 with zeros)

        Hint: Note that tf.nn.dropout takes the keep probability (1 - p_drop) as an argument.
              Therefore the keep probability should be set to the value of
              (1 - self.dropout_placeholder)

        Returns:
            pred: tf.Tensor of shape (batch_size, n_classes)
        """
        
        x = self.add_embedding()
        xavier = xavier_weight_init()
        with tf.variable_scope("transformation"):
            b1 = tf.Variable(tf.random_uniform([self.config.hidden_size]))
            b2 = tf.Variable(tf.random_uniform([self.config.n_classes]))
            
            W = xavier([self.config.n_features * self.config.embed_size, self.config.hidden_size])
            U = xavier([self.config.hidden_size, self.config.n_classes])
            
            z1 = tf.matmul(x, W) + b1
            h = tf.nn.relu(z1)
            h_dropout = tf.nn.dropout(h, keep_prob=(1-self.droupout_placeholder))
            pred = tf.matmul(h_dropout, U) + b2
        return pred
    
    def add_loss_op(self, pred):
        """Adds Ops for the loss function to the computational graph.
        In this case we are using cross entropy loss.
        The loss should be averaged over all examples in the current minibatch.

        Hint: You can use tf.nn.softmax_cross_entropy_with_logits to simplify your
                    implementation. You might find tf.reduce_mean useful.
        Args:
            pred: A tensor of shape (batch_size, n_classes) containing the output of the neural
                  network before the softmax layer.
        Returns:
            loss: A 0-d tensor (scalar)
        """
        loss = tf.nn.softmax_cross_entropy_with_logits(logits=pred, labels=self.labels_placeholder)
        loss = tf.reduce_mean(loss)
        return loss
    
    def add_training_op(self, loss):
        adam_optim = tf.train.AdamOptimizer(self.config.lr)
        train_op = adam_optim.minimize(loss)
        return train_op
    
    def train_on_batch(self, sess, inputs_batch, labels_batch):
        feed = self.create_feed_dict(inputs_batch, labels_batch=labels_batch,
                                     dropout=self.config.dropout)
        _, loss = sess.run([self.train_op, self.loss], feed_dict=feed)
        return loss
    
    def run_epoch(self, sess, parser, train_examples, dev_set):
        n_minibatches = 1 + len(train_examples) / self.config.batch_size
        prog = tf.keras.utils.Progbar(target=n_minibatches)
        for i, (train_x, train_y) in enumerate(minibatches(train_examples, self.config.batch_size)):
            loss = self.train_on_batch(sess, train_x, train_y)
            prog.update(i + 1, [("train loss", loss)])

        print "Evaluating on dev set",
        dev_UAS, _ = parser.parse(dev_set)
        print "- dev UAS: {:.2f}".format(dev_UAS * 100.0)
        return dev_UAS
    
    def fit(self, sess, saver, parser, train_examples, dev_set):
        best_dev_UAS = 0
        for epoch in range(self.config.n_epochs):
            print "Epoch {:} out of {:}".format(epoch + 1, self.config.n_epochs)
            dev_UAS = self.run_epoch(sess, parser, train_examples, dev_set)
            if dev_UAS > best_dev_UAS:
                best_dev_UAS = dev_UAS
                if saver:
                    print "New best dev UAS! Saving model in ./data/weights/parser.weights"
                    saver.save(sess, './data/weights/parser.weights')
            print
            
    def __init__(self, config, pretrained_embeddings):
        self.pretrained_embeddings = pretrained_embeddings
        self.config = config
        self.build()   

def main(debug=True):
    print 80 * "="
    print "INITIALIZING"
    print 80 * "="
    config = Config()
    parser, embeddings, train_examples, dev_set, test_set = load_and_preprocess_data(debug)
    if not os.path.exists('./data/weights/'):
        os.makedirs('./data/weights/')

    with tf.Graph().as_default() as graph:
        print "Building model...",
        start = time.time()
        model = ParserModel(config, embeddings)
        parser.model = model
        init_op = tf.global_variables_initializer()
        saver = None if debug else tf.train.Saver()
        print "took {:.2f} seconds\n".format(time.time() - start)
    graph.finalize()

    with tf.Session(graph=graph) as session:
        parser.session = session
        session.run(init_op)

        print 80 * "="
        print "TRAINING"
        print 80 * "="
        model.fit(session, saver, parser, train_examples, dev_set)
        
        if not debug:
            print 80 * "="
            print "TESTING"
            print 80 * "="
            print "Restoring the best model weights found on the dev set"
            saver.restore(session, './data/weights/parser.weights')
            print "Final evaluation on test set",
            UAS, dependencies = parser.parse(test_set)
            print "- test UAS: {:.2f}".format(UAS * 100.0)
            print "Writing predictions"
            with open('q2_test.predicted.pkl', 'w') as f:
                cPickle.dump(dependencies, f, -1)
            print "Done!"
class PartialParse(object):
    def __init__(self, sentence):
        """Initializes this partial parse.

        Your code should initialize the following fields:
            self.stack: The current stack represented as a list with the top of the stack as the
                        last element of the list.
            self.buffer: The current buffer represented as a list with the first item on the
                         buffer as the first item of the list
            self.dependencies: The list of dependencies produced so far. Represented as a list of
                    tuples where each tuple is of the form (head, dependent).
                    Order for this list doesn't matter.

        The root token should be represented with the string "ROOT"

        Args:
            sentence: The sentence to be parsed as a list of words.
                      Your code should not modify the sentence.
        """
        # The sentence being parsed is kept for bookkeeping purposes. Do not use it in your code.
        self.sentence = sentence

        self.stack = ['ROOT']
        self.buffer = sentence[:]
        self.dependencies = []
        
    def parse_step(self, transition):
        """Performs a single parse step by applying the given transition to this partial parse

        Args:
            transition: A string that equals "S", "LA", or "RA" representing the shift, left-arc,
                        and right-arc transitions. You can assume the provided transition is a legal
                        transition.
        """
        if transition == 'S':
            self.stack.append(self.buffer[0])
            self.buffer.pop(0)
        elif transition == 'LA':
            self.dependencies.append((self.stack[-1], self.stack[-2]))
            self.stack.pop(-2)
        else:
            self.dependencies.append((self.stack[-2], self.stack[-1]))
            self.stack.pop(-1)
            
    def parse(self, transitions):
        """Applies the provided transitions to this PartialParse

        Args:
            transitions: The list of transitions in the order they should be applied
        Returns:
            dependencies: The list of dependencies produced when parsing the sentence. Represented
                          as a list of tuples where each tuple is of the form (head, dependent)
        """
        for transition in transitions:
            self.parse_step(transition)
        return self.dependencies

def test_step(name, transition, stack, buf, deps,
              ex_stack, ex_buf, ex_deps):
    """Tests that a single parse step returns the expected output"""
    pp = PartialParse([])
    pp.stack, pp.buffer, pp.dependencies = stack, buf, deps

    pp.parse_step(transition)
    stack, buf, deps = (tuple(pp.stack), tuple(pp.buffer), tuple(sorted(pp.dependencies)))
    assert stack == ex_stack, \
        "{:} test resulted in stack {:}, expected {:}".format(name, stack, ex_stack)
    assert buf == ex_buf, \
        "{:} test resulted in buffer {:}, expected {:}".format(name, buf, ex_buf)
    assert deps == ex_deps, \
        "{:} test resulted in dependency list {:}, expected {:}".format(name, deps, ex_deps)
    print "{:} test passed!".format(name)

import copy
def minibatch_parse(sentences, model, batch_size):
    """Parses a list of sentences in minibatches using a model.

    Args:
        sentences: A list of sentences to be parsed (each sentence is a list of words)
        model: The model that makes parsing decisions. It is assumed to have a function
               model.predict(partial_parses) that takes in a list of PartialParses as input and
               returns a list of transitions predicted for each parse. That is, after calling
                   transitions = model.predict(partial_parses)
               transitions[i] will be the next transition to apply to partial_parses[i].
        batch_size: The number of PartialParses to include in each minibatch
    Returns:
        dependencies: A list where each element is the dependencies list for a parsed sentence.
                      Ordering should be the same as in sentences (i.e., dependencies[i] should
                      contain the parse for sentences[i]).
    """

    partial_parses = [PartialParse(s) for s in sentences]

    unfinished_parse = copy.copy(partial_parses)

    while len(unfinished_parse) > 0:
        minibatch = unfinished_parse[0:batch_size]
        # perform transition and single step parser on the minibatch until it is empty
        while len(minibatch) > 0:
            transitions = model.predict(minibatch)
            for index, action in enumerate(transitions):
                minibatch[index].parse_step(action)
            minibatch = [parse for parse in minibatch if len(parse.stack) > 1 or len(parse.buffer) > 0]

        # move to the next batch
        unfinished_parse = unfinished_parse[batch_size:]

    dependencies = []
    for n in range(len(sentences)):
        dependencies.append(partial_parses[n].dependencies)

    return dependencies

class DummyModel(object):
    """Dummy model for testing the minibatch_parse function
    First shifts everything onto the stack and then does exclusively right arcs if the first word of
    the sentence is "right", "left" if otherwise.
    """
    def predict(self, partial_parses):
        return [("RA" if pp.stack[1] is "right" else "LA") if len(pp.buffer) == 0 else "S" for pp in partial_parses]

def test_dependencies(name, deps, ex_deps):
    """Tests the provided dependencies match the expected dependencies"""
    deps = tuple(sorted(deps))
    assert deps == ex_deps, \
        "{:} test resulted in dependency list {:}, expected {:}".format(name, deps, ex_deps)

from nltk.corpus import reuters
from collections import Counter
import random

counts = Counter(reuters.words())
total_count = len(reuters.words())

# Compute the frequencies
for word in counts:
    counts[word] /= float(total_count)

# Generate 100 words of language
text = []

for _ in range(100):
    r = random.random()
    accumulator = .0

    for word, freq in counts.iteritems():
        accumulator += freq

        if accumulator >= r:
            text.append(word)
            break

print ' '.join(text)

from nltk.corpus import reuters
from nltk import bigrams, trigrams
from collections import Counter, defaultdict
import random

model = defaultdict(lambda: defaultdict(lambda: 0))
for sentence in reuters.sents():
    for w1, w2, w3 in trigrams(sentence, pad_right=True, pad_left=True):
        model[(w1, w2)][w3] += 1

for w1_w2 in model:
    total_count = float(sum(model[w1_w2].values()))
    for w3 in model[w1_w2]:
        model[w1_w2][w3] /= total_count



text = [None, None]
sentence_finished = False

while not sentence_finished:
    r = random.random()
    accumulator = .0

    for word in model[tuple(text[-2:])].keys():
        accumulator += model[tuple(text[-2:])][word]

        if accumulator >= r:
            text.append(word)
            break

    if text[-2:] == [None, None]:
        sentence_finished = True


# add the probability computation
text = [None, None]
prob = 1.0  # <- Init probability

sentence_finished = False

while not sentence_finished:
    r = random.random()
    accumulator = .0

    for word in model[tuple(text[-2:])].keys():
        accumulator += model[tuple(text[-2:])][word]

        if accumulator >= r:
            prob *= model[tuple(text[-2:])][word]  # <- Update the probability with the conditional probability of the new word
            text.append(word)
            break

    if text[-2:] == [None, None]:
        sentence_finished = True

print "Probability of text=", prob  # <- Print the probability of the text
print ' '.join([t for t in text if t])

# coding: utf-8

import numpy as np
import scipy as sp
import argparse
import math
import sys
import operator


class FeatureIndex():
    def __init__(self):
        self.word_to_index = {}
        self.index_to_word = {}

    def save(self, fout):
        with open(fout, 'w') as fw:
            for word, idx in self.vocab.iteritems():
                output = '\t'.join([word, str(idx)])
                fw.writelines(output+'\n')
    def read_word_idx(self, fin):
        with open(fin, 'r') as fr:
            lines = fr.readlines()
            for line in lines:
                line = line.strip().split("\t")
                self.index_to_word[line[1]] = line[0]
    def read_word(self, fn):
        self.word_to_index = self.vocab
        uniq_words = set()
        with open('./data/data_note2.txt', 'r') as fr:
            lines = fr.readlines()
            for line in lines:
                words = line.strip().split()
                for word in words:
                    uniq_words.add(word)
        for idx, word in enumerate(uniq_words):
            self.word_to_index[word] = idx

def read_corpus(corpusfile, window = 1):
    vocabulary = FeatureIndex()
    traindata = []
    with open(corpusfile) as corpusinput:
        for line in corpusinput:
            words = line.strip().split()
            wordid = []
            for word in words:
                index = vocabulary.word_to_index[word]
                if index >= 0:
                    wordid.append(index)
                lastword = -1
            for i in range(len(wordid)):
                # 以当前词作为预测目标，前一个词作为输入构造一条样本
                if i > 0:
                    traindata.append((wordid[i - 1], wordid[i]))
                    # 以当前词作为预测目标，下一个词作为输入构造一条样本
                if i < len(wordid) - 1:
                    traindata.append((wordid[i + 1], wordid[i]))
    return vocabulary, traindata

def init_vectors(V, N): 
    """
        Argvs:
            V: Int, dims of the vocabulary
            N: Out, dims of the word vector
    """
    input_vectors = np.random.rand(V, N)
    output_vectors = np.random.rand(N, V)
    return input_vectors, output_vectors

def save_vectors(input_vectors, vectorfile):
    """ save vectors info file, each row is a series float seperate by space
        Argvs:
            input_vectors: A numpy array, each row is a distribution represatation of a word 
            vectorfile: A string, file name
    """
    
    with open(vectorfile, 'w') as vectorfileoutput:
        for row in range(input_vectors.shape[0]):
            vectorfileoutput.write("%s\n" % (" ".join([str(s) for s in input_vectors[row]])))

def compute_loss(V, traindata, input_vectors, output_vectors):
    """
        Argvs:
            V: Int, dims of the vocabulary
            traindata: List of pair(cur_idx, next_idx), (last_idx, cur_idx)
            input_vectors: Numpy Array with shape (V, N)
            output_vectors: Numpy Array with shape (N, V)
    """
    loss = 0.0
    for pair in traindata:
        context_word_onehot = np.zeros((V, 1)) # (V, 1)
        context_word_onehot[pair[0]] = 1
        output_word_onehot = np.zeros((V, 1)) # (V, 1)
        output_word_onehot[pair[1]] = 1

        hidden_output = np.dot(np.transpose(input_vectors), context_word_onehot) # (V, N)
        output_log_prob = np.dot(np.transpose(output_vectors), hidden_output) #(V, 1)
        exp_log_prob = np.exp(output_log_prob)
        sum_exp_prob = np.sum(exp_log_prob)
        output_word_prob = exp_log_prob / sum_exp_prob #(V, 1)
        prob = output_word_prob[pair[1]]
        loss += math.log(1.0/prob)
    return loss



def run_word2vec(corpusfile, vectorfile, vocabulary_file, N=5, learning_rate=0.05, epoch=100):
    """ Get word distribution presatation
        Argvs:
            corpusfile:  words have been splited by space
            vectorfile: file to save vectors
            vocabulary_file: file to save words
            N: Int, vector's dimention, default 5
            learning_rate: Float, defalut 0.05
            epoch: Float, iteration times default 100
    """

    vocabulary, traindata = read_corpus(corpusfile)
    input_vectors, output_vectors = init_vectors(len(vocabulary.vocab), N)
    for t in range(epoch):
        print('epoch = %s, loss = %f' % (t, compute_loss(len(vocabulary.vocab), traindata, input_vectors, output_vectors)))
        for pair in traindata:
            # predict
            context_word_onehot = np.zeros((len(vocabulary.vocab), 1))
            context_word_onehot[pair[0]] = 1
            output_word_onehot = np.zeros((len(vocabulary.vocab), 1))
            output_word_onehot[pair[1]] = 1
            hidden_output = np.dot(np.transpose(input_vectors), context_word_onehot)
            output_log_prob = np.dot(np.transpose(output_vectors), hidden_output)
            exp_log_prob = np.exp(output_log_prob)
            sum_exp_prob = np.sum(exp_log_prob)
            output_word_prob = exp_log_prob / sum_exp_prob

            d_output_log_prob = output_word_prob - output_word_onehot
            d_output_vectors = np.dot(hidden_output, np.transpose(d_output_log_prob))
            output_vectors = output_vectors - learning_rate * d_output_vectors

            d_hidden_to_output_log_prob = output_vectors
            d_hidden_output = np.dot(d_hidden_to_output_log_prob, d_output_log_prob)
            d_input_vectors = np.dot(context_word_onehot, np.transpose(d_hidden_output))
            input_vectors = input_vectors - learning_rate * d_input_vectors
    # save
    save_vectors(input_vectors, vectorfile)
    vocabulary.save(vocabulary_file)

#Huffman Encoding

#Tree-Node Type
class Node:
    def __init__(self,freq):
        self.left = None
        self.right = None
        self.father = None
        self.freq = freq

    def isLeft(self):
        return self.father.left == self

#create nodes创建叶子节点
def createNodes(freqs):
    return [Node(freq) for freq in freqs]

#create Huffman-Tree创建Huffman树
def createHuffmanTree(nodes):
    queue = nodes[:]
    while len(queue) > 1:
        queue.sort(key=lambda item:item.freq)
        node_left = queue.pop(0)
        node_right = queue.pop(0)
        node_father = Node(node_left.freq + node_right.freq)
        node_father.left = node_left
        node_father.right = node_right
        node_left.father = node_father
        node_right.father = node_father
        queue.append(node_father)
    queue[0].father = None
    return queue[0]


#Huffman编码
def huffmanEncoding(nodes,root):
    codes = [''] * len(nodes)
    for i in range(len(nodes)):
        node_tmp = nodes[i]
        while node_tmp != root:
            if node_tmp.isLeft():
                codes[i] = '0' + codes[i]
            else:
                codes[i] = '1' + codes[i]
            node_tmp = node_tmp.father
    return codes

if __name__ == '__main__':

    chars_freqs = [('我', 15), ('喜欢', 8), ('看', 6), ('巴西', 5), ('足球', 3), ('世界杯', 1)]
    nodes = createNodes([item[1] for item in chars_freqs])
    root = createHuffmanTree(nodes)
    codes = huffmanEncoding(nodes,root)
    for item in zip(chars_freqs,codes):
        print 'Character:%s freq:%-2d   encoding: %s' % (item[0][0],item[0][1],item[1])