import numpy as np 
from collections import Counter
import pandas as pd 
import re
import nltk
import string
import torch
import torch.nn as nn
from nltk import sent_tokenize, word_tokenize
from nltk.stem import WordNetLemmatizer
from nltk.corpus import stopwords

stop_words = set(stopwords.words('english')).union(set(string.punctuation)) 

def preprocessing_rnn(text):
  words=word_tokenize(text)
  filtered_sentence = [] 
  # remove stopwords
  for word in words: 
    if word not in stop_words: 
        filtered_sentence.append(word) 
  
  # lemmatize
  lemma_word = []
  wordnet_lemmatizer = WordNetLemmatizer()
  for w in filtered_sentence:
    word1 = wordnet_lemmatizer.lemmatize(w, pos = "n")
    word2 = wordnet_lemmatizer.lemmatize(word1, pos = "v")
    word3 = wordnet_lemmatizer.lemmatize(word2, pos = ("a"))
    lemma_word.append(word3)
  return lemma_word


def count_vocab_index(train_df, test_df):
    df = pd.concat([train_df, test_df]).sample(frac=1).reset_index(drop=True)
    counts = Counter()
    for _, row in df.iterrows():
        counts.update(preprocessing_rnn(row['discharge_diagnosis']))
    
    # removing the words that have frequency less than 2
    for word in list(counts):
        if counts[word] < 2:
            del counts[word]
    
    vocab2index = {"":0, "UNKNOWN":1}
    words = ["", "UNKNOWN"]
    for word in counts:
        vocab2index[word] = len(words)
        words.append(word)

    return counts, vocab2index


def encode_sentence(text, vocab2index, N = 50):
  tokenized = preprocessing_rnn(text)
  encoded = np.zeros(N, dtype=int)
  enc1 = np.array([vocab2index.get(word, vocab2index["UNKNOWN"]) for word in tokenized])
  length = min(N, len(enc1))
  encoded[:length] = enc1[:length]
  return encoded
        


def get_emb_matrix(w2vmodel, word_counts):
  """ Creates embedding matrix from word vectors"""
  vocab_size = len(word_counts) + 2
  emb_size = w2vmodel.vector_size

  W = np.zeros((vocab_size, emb_size), dtype="float32")
  W[0] = np.zeros(emb_size, dtype='float32') # adding a vector for padding
  W[1] = np.random.uniform(-0.25, 0.25, emb_size) # adding a vector for unknown words 
  
  i = 2
  for word in word_counts:
    if word in w2vmodel.wv:
      W[i] = w2vmodel.wv[word]
    else:
      W[i] = np.random.uniform(-0.25,0.25, emb_size)
    i += 1   
  return W

def create_emb_layer(weights_matrix, non_trainable=False):
  num_embeddings, embedding_dim = weights_matrix.shape
  emb_layer = nn.Embedding(num_embeddings, embedding_dim, padding_idx = 0)
  emb_layer.load_state_dict({'weight': torch.from_numpy(weights_matrix)})
  if non_trainable:
    emb_layer.weight.requires_grad = False

  return emb_layer, num_embeddings, embedding_dim