import numpy as np
import matplotlib.pyplot as plt
from collections import Counter
from string import punctuation

class Tokenizer():
  def __init__(self):

    self.vocab = None
    self.pad_idx = 0
    self.unk_idx = 1

  def preprocessing(self, texts: list):

    '''
    Método que pre procesa un documento, transformando las palabras a minúsculas, eliminando puntuaciones y caracterés "\n". Devuelve los textos pre procesados.
    texts: documentos a ser procesados (list)
    '''

    texts = [text.lower() for text in texts] # lowercase
    texts = [''.join([c for c in text if c not in punctuation]) for text in texts] # delete punctuation
    texts = [text.split('\n') for text in texts] # eliminate \n
    texts = [' '.join(text) for text in texts]

    return texts

  def train(self, texts: list):

    '''
    Método que entrena el tokenizador, construyendo un vocabulario de tokens y su codificación respectiva.
    texts: documentos que el tokenizador usa para construir el vocabulario (list)
    '''

    # preprocessing
    texts = self.preprocessing(texts)

    # joint text
    megadoc = ' '.join(texts)
    words = megadoc.split()

    self.counts = Counter(words) # Construye un diccionario de palabras. Las claves son las palabras y los valores son la frecuencia
    self.vocab = sorted(self.counts, key = self.counts.get, reverse = True) # Ordenamos la palabras por frecuencia
    self.vocab_to_int = {word: ii for ii, word in enumerate(self.vocab, 2)} # Construimos diccionario para mapear palabra a número entero. Empezamos los índices en 2
    self.vocab_to_int[self.unk_idx] = '<unk>' # token para palabras no reconocidas
    self.vocab_to_int[self.pad_idx] = '<pad>' # token para padding

    self.int_to_vocab = {value: key for key, value in self.vocab_to_int.items()}

  def encode(self, texts: list):

    '''
    Método que usa el vocabulario construido para codificar textos. Devuelve los textos codificados.
    texts: Documentos a ser codificados (list)
    '''

    if self.vocab_to_int is None:
      raise ValueError('Debes entrenar el tokenizador primero')

    encoded_text = [[self.vocab_to_int[word] if word in self.vocab_to_int.keys() else 1 for word in text.split()] for text in texts]

    return encoded_text

  def decode(self, texts: list):

    '''
    Método que usa el vocabulario construido para decodificar textos. Devuelve los textos decodificados.
    texts: Documentos a ser decodificados (list)
    '''

    if self.vocab is None:
      raise ValueError('Debes entrenar el tokenizador primero')
    
    decoded_text = [[self.int_to_vocab[word] if word in self.int_to_vocab.keys() else 'unk' for word in text] for text in texts]

    return decoded_text

  def filter_text(self, encoded_text: list, encoded_labels: list, min_tokens = 1, max_tokens = 1e6):

    '''
    Método que filtra una colección de documentos en función de la cantidad de tokens. Devuelve la coleccion de documentos filtrados.
    encoded_text: Textos codificados a ser filtrados (list)
    encoded_labels: Etiquetas a filtrar en función del texto (list)
    min_tokens: Cantidad mínima de tokens permitida (int)
    max_tokens: Cantidad máxima de tokens permitida (int)
    '''

    print('Documentos antes de eliminación:', len(encoded_text))

    #Extraemos los índices de todos los reviews que tienen longitud cumpliendo los filtros
    filter_idx = [ii for ii, text in enumerate(encoded_text) if min_tokens <= len(text) <= max_tokens]

    #Nos quedamos solo con los reviews con longitud que cumplen los filtros
    encoded_text = [encoded_text[ii] for ii in filter_idx]

    #Lo mismo con los labels
    encoded_labels = np.array([encoded_labels[ii] for ii in filter_idx])

    print('Documentos después de eliminación:', len(encoded_text))

    return encoded_text, encoded_labels

  def padding(self, encoded_text: list, vector_size: int):
    
    '''
    Método que hace padding a una secuencia, fijando el largo de las secuencias en un número determinado vector_size:
      Las secuencias con largo mayor a vector_size, son acortadas por la derecha hasta ese valor.
      Las secuencias con largo menor a vector_size, son llenadas con 0s por la izquierda hasta completar ese valor.
    Retorna la secuencia modificada.

    encoded_text: lista con los textos a modificar (list)
    vector_size: largo de los documentos a fijar (int)
    '''

    features = np.zeros((len(encoded_text), vector_size), dtype = int)

    for i, row in enumerate(encoded_text):
      features[i, -len(row):] = np.array(row)[:vector_size]

    return features

  def tokenize(self, texts: list, 
               #vector_size: int
               ):

    '''
    Método que tokeniza documentos a partir del vocabulario construido. Devuelve los documentos codificados en un largo de tamaño vector_size
    texts: Textos a ser tokenizados (list)
    vector_size: Largo de los textos de salida (int)
    '''

    if self.vocab is None:
      raise ValueError('Debes entrenar el tokenizador primero')

    if self.vector_size is None:
      raise ValueError('Debes especificar vector_size en objeto Tokenizer (tokenizer.vector_size = x)')

    # preprocessing
    texts = self.preprocessing(texts)

    # encode
    encoded_text = self.encode(texts)

    # padding
    features = self.padding(encoded_text, self.vector_size)
  
    return features

  def graph_distribution(self, encoded_text):

    '''
    Método que grafica la distribución del largo de los documentos de entrenamiento.
    '''
    
    if self.vocab is None:
      raise ValueError('Debes entrenar el tokenizador primero')

    text_lens = Counter([len(text) for text in encoded_text]) #Contamos cuantas palabras hay en cada review

    plt.figure(figsize = (12, 6))
    plt.bar(text_lens.keys(), text_lens.values())
    plt.title('Distribución del largo de documentos en el Dataset')
    plt.xlabel('Cantidad de tokens')
    plt.ylabel('Frecuencia')
    plt.show()

  def __len__(self):
    return len(self.vocab_to_int) if self.vocab is not None else 0