from . import models
from .utils import load_pickle_from_package, save_pickle_to_user, load_pickle_from_user
from .utils import load_json_from_package, save_json_to_user, load_json_from_user

try:
	from collections import Counter, defaultdict
	from typing import Dict, Iterable, Callable, List, Tuple, Any, Iterator
	from tqdm import tqdm
	import toolz
	import importlib.resources as pkg_resources
	import json
	import pickle
except ImportError:
	pass

DEFAULT_VI_TOKENS_FILENAME = 'vi_tokens.pkl'
DEFAULT_OUT_FILENAME = 'vocabs.json'

DEFAULT_VI_TOKENS = defaultdict(set)
DEFAULT_VI_TOKENS = load_pickle_from_package(models, DEFAULT_VI_TOKENS_FILENAME)

# LOWERCASE = True
DEFAULT_WORD_NGRAM_MAX = 4
DEFAULT_BYTE_NGRAM_MAX = 4
DEFAULT_WORD_VOCAB_SIZE = 8000
DEFAULT_BPE_VOCAB_SIZE = 2000
DEFAULT_SOW = '__sow'
DEFAULT_EOW = '__eow'
DEFAULT_UNK = '__unk'
DEFAULT_PAD = '__pad'
DEFAULT_BRE = '__bre'

class VBPE:
	""" Encodes Vietnamese white-space separated text using byte-pair encoding. """

	def __init__(self, vi_tokens=DEFAULT_VI_TOKENS, strict=False,
			  word_ngram_max=DEFAULT_WORD_NGRAM_MAX, bpe_ngram_max=DEFAULT_BYTE_NGRAM_MAX, 
			  word_vocab_size=DEFAULT_WORD_VOCAB_SIZE, bpe_vocab_size=DEFAULT_BPE_VOCAB_SIZE,
			  SOW=DEFAULT_SOW, EOW=DEFAULT_EOW, UNK=DEFAULT_UNK, PAD=DEFAULT_PAD, BRE=DEFAULT_BRE):
		self.word_vocab = {}  # type: Dict[str, int]
		self.bpe_vocab = {}  # type: Dict[str, int]
		self.inverse_word_vocab = {}  # type: Dict[int, str]
		self.inverse_bpe_vocab = {}  # type: Dict[int, str]
		self.vi_tokens = vi_tokens
		# self.lowercase = lowercase
		self.strict = strict
		self.word_ngram_max = word_ngram_max
		self.bpe_ngram_max = bpe_ngram_max
		self.word_vocab_size = word_vocab_size
		self.bpe_vocab_size = bpe_vocab_size
		self.SOW = SOW
		self.EOW = EOW
		self.UNK = UNK
		self.PAD = PAD
		self.BRE = BRE
		self._progress_bar = iter

	@staticmethod
	def process_text(text) -> str:
		# type: (str) -> str
		""" Filter characters """
		remain_unicodes_dict  = {
			# 'digit': [48, 57],
			# 'upper': [65, 90],
			'lower': [97, 122],
			'diacritic': [192, 7929],
			# 'symbol': [33, 47],
			# 'space': [32, 32],
		}

		processed_text = ''
		for c in text.lower():
			is_w = False
			for x in remain_unicodes_dict.values():
				if ord(c) >= x[0] and ord(c) <= x[1]:
					is_w = True
					break

			if is_w:
				processed_text += c
			elif processed_text[-1:] != ' ':
				processed_text += ' '

		return processed_text

	def word_tokenize(self, text) -> List[str]:
		# type: (str) -> List[str]
		""" Split white-space separated text into list of words """
		return text.split(' ')
	
	def get_word_shingles(self, words) -> List[str]:
		# type: (List[str]) -> List[str]
		shingles = []
		i = 0
		while i < len(words):
			shingle = None
			i_ = i
			shingle = words[i]
			for j in range(i + 2, i + self.word_ngram_max + 1):
				if '_'.join(words[i:j]) in self.word_vocab.keys():
					shingle = '_'.join(words[i:j])
					i_ = j - 1
			if shingle:
				shingles.append(shingle)

			i = i_ + 1

		return shingles
	
	def get_shingles(self, word) -> str:
		# type: (str) -> str
		vi_tokens_size = len(list(toolz.concat(self.vi_tokens.values())))
		shingles = []
		i = 0
		while i < len(word):
			shingle = None
			i_ = i
			shingle = word[i]
			for j in range(i + 2, i + vi_tokens_size + 1):
				if word[i:j] in self.vi_tokens[str(len(word[i:j]))]:
					shingle = word[i:j]
					i_ = j - 1
			if shingle:
				shingles.append(shingle)

			i = i_ + 1

		return '_'.join(shingles)
	
	def bpe_tokenize(self, words) -> List[str]:
		# type: (List[str]) -> List[str]
		tokenized_words = []
		for word in words:
			tokenized_words.append(self.get_shingles(word))

		return tokenized_words
	
	def token_pair_count(self, sentences, word_level=True):
		# type: (List[str], bool) -> Iterable[Counter]
		""" Count token pairs """
		if word_level:
			ls_tokens = map(self.word_tokenize, sentences)
			delimiter = '_'
			ngram_max = self.word_ngram_max
		else:
			ls_tokens = [word.split('_') for word in self.bpe_tokenize(sentences)]
			delimiter = ''
			ngram_max = self.bpe_ngram_max

		token_counts = Counter()
		for tokens in ls_tokens:
			for ngram_size in range(1, min([ngram_max, len(tokens)]) + 1):
				for ngram in toolz.sliding_window(ngram_size, tokens):
					token_counts[delimiter.join(ngram)] += 1
		
		yield token_counts

	def learn_vocab(self, sentences, word_level=True) -> Dict[str, int]:
		# type: (List[str], bool) -> Dict[str, int]
		""" Build vocab from vocab_size most common tokens in provided sentences """
		vocab = Counter()

		if word_level:
			vocab_size = self.word_vocab_size
			start_idx = 0
			
			for token in {self.UNK, self.PAD, self.BRE}:
				vocab[token] = int(2**63)
		else:
			vocab_size = self.bpe_vocab_size
			start_idx = self.word_vocab_size
			
			for token in {self.SOW, self.EOW}:
				vocab[token] = int(2**63)
				
		for idx, token_pair_count in enumerate(self.token_pair_count(sentences, word_level=word_level)):
			for token_pair, count in token_pair_count.items():
				vocab[token_pair] += count

			if (idx + 1) % 10000 == 0:
				self.trim_vocab(10 * vocab_size, vocab)

		sorted_vocab = sorted(vocab.items(), key=lambda p: -p[1])[:vocab_size]
		return {bp: idx + start_idx for idx, (bp, count) in enumerate(sorted_vocab)}

	def fit(self, text):
		# type: (VBPE, str) -> None
		""" Learn vocab from text """
		_text = [self.process_text(t.lower().strip()) for t in text.split('\n')]
		
		self.word_vocab = self.learn_vocab(_text)

		remaining_words = [word for word in toolz.concat(map(self.word_tokenize, _text))
						   if word not in self.word_vocab.keys()]
		
		self.bpe_vocab = self.learn_vocab(remaining_words, word_level=False)

		self.inverse_word_vocab = {idx: token for token, idx in self.word_vocab.items()}
		self.inverse_bpe_vocab = {idx: token for token, idx in self.bpe_vocab.items()}

	@staticmethod
	def trim_vocab(n, vocab):
		# type: (int, Dict[str, int]) -> None
		"""  Deletes all pairs below 10 * vocab size to prevent memory problems """
		pair_counts = sorted(vocab.items(), key=lambda p: -p[1])
		pairs_to_trim = [pair for pair, count in pair_counts[n:]]
		for pair in pairs_to_trim:
			del vocab[pair]

	def tokenize(self, text) -> List[str]:
		# type: (VBPE, str) -> List[str]
		""" Split text into word and subword tokens """
		tokens_ls = []
		for t in text.split('\n'):
			word_tokens = self.get_word_shingles(self.word_tokenize(self.process_text(t.lower().strip())))

			tokens = []
			for word_token in word_tokens:
				if word_token in self.word_vocab:
					tokens.append(word_token)
				else:
					tokens.append(self.SOW)
					tokens.extend(self.get_shingles(word_token).split(' '))
					tokens.append(self.EOW)
		
			tokens_ls.append(tokens)
		
		return tokens_ls
	
	def encode(self, sentences, fixed_length=None):
		# type: (VBPE, Iterable[str], int) -> Iterable[List[int]]
		""" Turns space separated text into vocab idxs """
		encoded = []
		for tokens in list(self.tokenize(sentences)):
			in_subword = False
			for token in tokens:
				if in_subword:
					if token in self.bpe_vocab:
						if token == self.EOW:
							in_subword = False
						encoded.append(self.bpe_vocab[token])
					else:
						encoded.append(self.word_vocab[self.UNK])
				else:
					if token == self.SOW:
						in_subword = True
						encoded.append(self.bpe_vocab[token])
					else:
						if token in self.word_vocab:
							encoded.append(self.word_vocab[token])
						else:
							encoded.append(self.word_vocab[self.UNK])
			encoded.append(self.word_vocab[self.BRE])

		if fixed_length is not None:
			encoded = encoded[:fixed_length]
			while len(encoded) < fixed_length:
				encoded.append(self.word_vocab[self.PAD])

		return encoded
	
	def decode(self, encoded):
		# type: (VBPE, Iterable[List[int]]) -> Iterator[str]
		""" Turns token indexes back into space-joined text """
		words = []

		rebuilding_word = False
		current_word = ''
		for idx in encoded:
			if self.inverse_word_vocab.get(idx) == self.BRE:
				if rebuilding_word and self.strict:
					raise ValueError('Encountered BRE token before EOW.')
				if rebuilding_word:
					words.append(current_word + '_')
					current_word = '_'
					words.append('\n')
				else:
					words.append('\n')

			elif self.inverse_bpe_vocab.get(idx) == self.SOW:
				if rebuilding_word and self.strict:
					raise ValueError('Encountered second SOW token before EOW.')
				rebuilding_word = True

			elif self.inverse_bpe_vocab.get(idx) == self.EOW:
				if not rebuilding_word and self.strict:
					raise ValueError('Encountered EOW without matching SOW.')
				rebuilding_word = False
				words.append(current_word)
				current_word = ''

			elif rebuilding_word and (idx in self.inverse_bpe_vocab):
				current_word += self.inverse_bpe_vocab[idx]

			elif rebuilding_word and (idx in self.inverse_word_vocab):
				current_word += self.inverse_word_vocab[idx]

			elif idx in self.inverse_word_vocab:
				words.append(self.inverse_word_vocab[idx])

			elif idx in self.inverse_bpe_vocab:
				if self.strict:
					raise ValueError("Found BPE index {} when not rebuilding word!".format(idx))
				else:
					words.append(self.inverse_bpe_vocab[idx])

			else:
				raise ValueError("Got index {} that was not in word or BPE vocabs!".format(idx))

		return ' '.join(w for w in words if w != '')

	def vocabs_to_dict(self, dont_warn=False):
		# type: (VBPE, bool) -> Dict[str, Dict[str, int]]
		""" Turns vocab into dict that is json-serializeable """
		# if self.custom_tokenizer and not dont_warn:
		# 	print("WARNING! You've specified a non-default tokenizer.  You'll need to reassign it when you load the "
		# 		  "model!")
		return {
			'bpe_vocab': self.bpe_vocab,
			'word_vocab': self.word_vocab,
			'kwargs': {
				'vi_tokens': self.vi_tokens,
				'strict': self.strict,
				'word_ngram_max': self.word_ngram_max,
				'bpe_ngram_max': self.bpe_ngram_max,
				'word_vocab_size': self.word_vocab_size,
				'bpe_vocab_size': self.bpe_vocab_size,
				'EOW': self.EOW,
				'SOW': self.SOW,
				'UNK': self.UNK,
				'PAD': self.PAD,
				'BRE': self.BRE,
			}
		}

	def save(self, filename=DEFAULT_OUT_FILENAME, dont_warn=False, encoding=None):
		# type: (VBPE, str, bool, str) -> None
		""" Serializes and saves vbpe to provided path """
		# with open(outpath, 'w', encoding=encoding) as outfile:
		# 	json.dump(self.vocabs_to_dict(dont_warn), outfile)
		save_json_to_user(self.vocabs_to_dict(dont_warn), filename)

	@classmethod
	def from_dict(cls, vocabs):
		# type: (Any, Dict[str, Dict[str, int]]) -> VBPE
		""" Load vbpe from dict produced with vocabs_to_dict """
		vbpe = VBPE(**vocabs['kwargs'])
		vbpe.word_vocab = vocabs['word_vocab']
		vbpe.bpe_vocab = vocabs['bpe_vocab']

		vbpe.inverse_bpe_vocab = {v: k for k, v in vbpe.bpe_vocab.items()}
		vbpe.inverse_word_vocab = {v: k for k, v in vbpe.word_vocab.items()}

		return vbpe

	@classmethod
	def load(cls, filename=DEFAULT_OUT_FILENAME):
		# type: (Any, str) -> VBPE
		""" Loads an vbpe from path saved with save """
		# with open(in_path) as infile:
		# 	obj = json.load(infile)
		obj = load_json_from_user(filename)
		return cls.from_dict(obj)