#!/usr/bin/env python2
# -*- coding: utf-8 -*-

from __future__ import print_function

from itertools import chain

import data, datetime, argparse, sys

import numpy as np
import tensorflow as tf

from tensorflow.python.ops.rnn_cell import *
from utils import *
from kneser_ney import KneserNeyLM
from nltk.util import ngrams as ngram_generator
from collections import Counter

parser = argparse.ArgumentParser(description='Neural language model')
parser.add_argument("-n", "--run-name", type=str, default=None)
parser.add_argument("-g", "--generate", type=str, default=None)
parser.add_argument("-a", "--learning-rate", type=float, default=0.01)
parser.add_argument("-i", "--b
ootstrap-in", type=bool, default=False)
parser.add_argument("-o", "--bootstrap-out", type=bool, default=False)

parser.add_argument("-m", "--bootstrap-mem", type=bool, default=False)
args = parser.parse_args()

if args.run_name == 'date':
  args.run_name = datetime.datetime.now().strftime("%A %H%M")

with status("Reading NKJP..."):
  arrays, lengths, symbol_map, symbol_count = data.read_nkjp_simple()

dataset_np = np.concatenate([np.concatenate(([data.GO], arr, [data.EOS])) for arr in arrays])
train_test_divider = int(len(lengths) * 0.80)
sequence_count = len(lengths)
test_count = sequence_count - train_test_divider

with status("Loading n-gram model..."):
  n = 5
  ngram_db_path = "ngram_db_{}.npz".format(n)
  try:
    npzfile = np.load(ngram_db_path)
    ngram_dataset = npzfile["ngram_dataset"]
    ngram_probability_table = npzfile["ngram_probability_table"]
  except IOError as e:
    log("Building n-gram model from scratch...")

    ngrams = tuple(Counter() for i in range(n+1))
    ngrams[n].update(ngram for array in arrays[:train_test_divider] for ngram in ngram_generator(chain((data.GO,)*n, array, (data.EOS,)), n))
    for i in range(n - 1, 0, -1):
      for ngram, count in ngrams[i+1].items():
        ngrams[i][ngram[1:]] += count

    log("Precomputing unique prefixes/suffixes")
    # unique_prefixes[i][ngram] where len(ngram) == i contains number of different symbols that proceed given ngram
    unique_prefixes = tuple(Counter() for i in range(n))
    # unique_suffixes[i][ngram] where len(ngram) == i contains number of different symbols that follow given ngram
    unique_suffixes = tuple(Counter() for i in range(n))
    for i in range(n, 0, -1):
      unique_prefixes[i-1].update(ngram[1:] for ngram in ngrams[i].keys())
      unique_suffixes[i-1].update(ngram[:-1] for ngram in ngrams[i].keys())

    log("Indexing ngrams")
    all_ngrams = tuple(set() for i in range(n+1))
    for array in arrays:
      for ngram in ngram_generator(chain((data.GO,)*n, array, (data.EOS,)), n):
        all_ngrams[n].add(ngram)
    for i in range(n - 1, 0, -1):
      for ngram in all_ngrams[i+1]:
        all_ngrams[i].add(ngram[1:])
    # maps ngram tuple to ngram number
    ngram_idx = tuple(dict() for i in range(n))
    for i in range(n, 0, -1):
      for num, ngram in enumerate(all_ngrams[i]):
        ngram_idx[i-1][ngram] = num

    discount = (1.0, 0.5, 0.75, 0.75, 0.75, 0.75)

    def prob(full_ngram):
      current_p = 0.0
      p_multiplier = 1.0
      estimation_base = ngrams
      for i in range(n, 0, -1):
        ngram = full_ngram[-i:]
        prefix = ngram[:-1]
        if estimation_base[i-1][prefix]:
          #print("i", i)
          #print("full ngram", full_ngram)
          #print("ngram", ngram)
          #print("prefix", prefix)
          #print("estamition_base", estimation_base[i][ngram])
          p = max(0, estimation_base[i][ngram] - discount[i]) / estimation_base[i-1][prefix]
          current_p += p * p_multiplier
          p_multiplier *= discount[i] / estimation_base[i-1][prefix] * unique_suffixes[i-1][prefix]
          estimation_base = unique_prefixes
      current_p += p_multiplier / symbol_count # probability of an unseen symbol
      #print(u"Prob of {}: {}".format(''.join(symbol_map[c] for c in ngram), prob_cache[ngram]))
      return current_p

    precomputing_started_clock = log("Precomputing successor probabilities")
    # probabilities for the next symbol based on the last (n-1)-gram
    ngram_probability_table = np.zeros((len(ngram_idx[-2]), symbol_count), dtype=np.float32)
    progress = 0
    for ngram, idx in ngram_idx[-2].items():
      if progress % 500 == 0:
        if progress:
          time_left = (clock() - precomputing_started_clock) / progress * (len(ngram_idx[-2]) - progress)
        else:
          time_left = float("inf")
        log("Precomputing successor probabilities: {:.1f}% ({:.0f} seconds left)".format(100.0 * progress / len(ngram_idx[-2]), time_left))
      probs = np.array([prob(ngram + (i,)) for i in range(symbol_count)])
      probs = probs / max(np.sum(probs), 0.0001)
      #pred_symbol = np.argmax(ngram_probability_table[ngram_idx[n1][prefix]])
      #print(u"Prefix '{}', prediction {:3d} '{}'".format(''.join(symbol_map[c] for c in prefix), pred_symbol, symbol_map[pred_symbol]))
      ngram_probability_table[idx, :] = probs
      progress += 1

    log("Building ngram sequence")
    ngram_dataset = np.zeros(dataset_np.shape + (n,), dtype=np.int32)
    offset = 0
    for array_index, arr in enumerate(arrays):
      for i in range(n):
        ngram_length = i + 1
        for pos, ngram in enumerate(ngram_generator(chain((data.GO,) * ngram_length, arr), ngram_length)):
          ngram_dataset[offset + pos, i] = ngram_idx[i][ngram]
      offset += len(arr) + 2

    np.savez(ngram_db_path, ngram_dataset=ngram_dataset, ngram_probability_table=ngram_probability_table)

def translate1(array):
  return ''.join(symbol_map[i] for i in array)

def translate2(array):
  return '\n'.join(translate1(row) for row in array)

batch_size = 100
check_indices = True
unrolled_iterations = 300
network_size = 500
layer_count = 2
ngram_count = len(ngram_probability_table)

# [sum(len(a) + 2 for a in arrays)] : int32
dataset = tf.Variable(dataset_np,
                      trainable=False, dtype=tf.int32, name="dataset")

ngram_dataset = tf.Variable(ngram_dataset, trainable=False, dtype=tf.int32, name="ngram_dataset")
ngram_probability_table = tf.Variable(ngram_probability_table, trainable=False, dtype=tf.float32, name="ngram_probability_table")

# [num_sequences] : int32
sequence_lengths = tf.Variable([l+2 for l in lengths], trainable=False, name="sequence_lengths")

# [num_sequences] : int32
sequence_offsets = tf.cumsum(sequence_lengths, exclusive=True, name="sequence_offsets")

ngram_embedding_size = 10
ngram_embeddings = tf.Variable(tf.random_normal([ngram_count, ngram_embedding_size], mean=0.0, stddev=0.02),
                               trainable=True, name="ngram_embeddings")


cudnn_cell = tf.contrib.cudnn_rnn.CudnnLSTM(layer_count, network_size, network_size, dropout=0)
cudnn_train_cell = tf.contrib.cudnn_rnn.CudnnLSTM(layer_count, network_size, network_size, dropout=0.5)
cudnn_param_size = cudnn_cell.params_size()
print("Params:", cudnn_param_size)
cudnn_params = tf.Variable(tf.random_normal([cudnn_param_size], mean=0.0, stddev=0.02),
                               trainable=False, name="cudnn_params")


def make_cell(dropout):
  cell = LSTMCell(network_size, use_peepholes=True, cell_clip=10, state_is_tuple=True)
  if dropout:
    cell = DropoutWrapper(cell, output_keep_prob=.5)
  cell = MultiRNNCell([cell] * layer_count, state_is_tuple=True)
  cell = InputProjectionWrapper(cell, num_proj=network_size)
  cell = OutputProjectionWrapper(cell, output_size=symbol_count)
  return cell

def build_input_vector(input_sequence, ngram_input_sequence, ngram_predictions):
  input_list = []
  input_list.append(tf.one_hot(input_sequence, symbol_count))
  if args.bootstrap_in:
    mean, variance = tf.nn.moments(ngram_predictions, [2], keep_dims=True)
    input_list.append(tf.nn.batch_normalization(ngram_predictions, mean, variance, 0, 1, 0.00001))
  if args.bootstrap_mem:
    input_list.append(tf.gather(ngram_embeddings, tf.squeeze(tf.slice(ngram_input_sequence, [0,0,3], [-1,-1,1]))))
  return input_list[0] if len(input_list) == 1 else tf.concat(2, input_list)


def get_total_loss(input_sequence, ngram_predictions, outputs, expected_sequen