from data_process import getData
import nltk 
import numpy as np
import pandas as pd
import random
import os
from utils import seed_all
import pickle
import torch
import torch.nn as nn
from torchtext.data.utils import get_tokenizer
import gensim.downloader as api
from gensim.models import KeyedVectors
from keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences
from keras.utils.np_utils import to_categorical
from sklearn.model_selection import train_test_split
from model import *
from tqdm import tqdm
import torch.optim as optim 
import sys
from sklearn import metrics
from sklearn.metrics import fbeta_score

MAX_SEQ_LENGTH = 100  # greater than 100 in length will be truncated
loss_function = nn.CrossEntropyLoss()
HIDDEN_DIM = 64

#seed_all()

def evalFunction(model, X, Y, device):
    model.to(device)
    cpu_target = None
    cpu_output = None
    eval_loss =  0
    eval_acc = 0
    cor = 0
    tot = 0
    for sentence, tags in tqdm(zip(X,Y),total=len(X)):
        with torch.no_grad():
            mask_index = np.where(sentence == 0)[0]
            mask_start_index = mask_index[0] if mask_index.any() else -1

            if cpu_target is None:
                cpu_target = np.array(tags[:mask_start_index])
            else:
                cpu_target = np.append(cpu_target, np.array(tags[:mask_start_index]), axis=0)


            sentence = sentence.to(device,dtype=torch.long)
            tags = tags.to(device, dtype=torch.float32)

            tag_scores = model(sentence)
            eval_loss += loss_function(tag_scores, tags)

            if cpu_output is None:
                cpu_output = tag_scores[:mask_start_index].cpu().numpy()
            else:
                cpu_output = np.append(cpu_output, tag_scores[:mask_start_index].cpu().numpy(), axis=0)
            
    cor = sum(np.argmax(cpu_output, axis=1) == np.argmax(cpu_target, axis=1))
    tot = len(cpu_output)
    acc = cor/tot
    print(f"Evaluation -> loss:{eval_loss}, accuracy:{cor/tot}")
    return eval_loss, acc



def trainer(model, X, Y, X_val, Y_val, epochs, device, optimizer):
    model.to(device)
    model.train()
    val_loss = []
    best_accuracy = 0
    for epoch in tqdm(range(epochs)):
        epoch_loss = 0
        for sentence, tags in tqdm(zip(X,Y),total=len(X)):

            sentence = sentence.to(device,dtype=torch.long)
            tags = tags.to(device, dtype=torch.float32)

            model.zero_grad()
            tag_scores = model(sentence)
            # print(tags.shape, tag_scores.shape)
            # sys.exit()
            loss = loss_function(tag_scores, tags)
            epoch_loss += loss.item()
            loss.backward()
            optimizer.step()
        eval_loss, accuracy = evalFunction(model, X_val, Y_val, device)
        val_loss.append(eval_loss)
        if accuracy > best_accuracy:
            print("Saving model")
            torch.save(model.state_dict(), "model1.bin")
            best_accuracy = accuracy
        print(f"Epoch {epoch+1} Loss-> ", epoch_loss)

    # with torch.no_grad():
    #     inputs = X[0][0].to(device, torch.long)
    #     tags = Y[0][0].to(device, dtype=torch.float32)
    #     tag_scores = model(inputs)
    #     print(loss_function(tag_scores, tags))

if __name__ == "__main__":
    words, pos_tags = getData()
    num_vocab_words = len(set(word.lower() for sentence in words for word in sentence))
    num_vocab_tags = len(set(tag.lower() for sentence in pos_tags for tag in sentence))

    print("Total number of tagged sentences: {}".format(len(words)))
    print("Vocabulary size: {}".format(num_vocab_words))
    print("Total number of tags: {}".format(num_vocab_tags))
    print('sample X: ', words[0], '\n')
    print('sample Y: ', pos_tags[0], '\n')


    # Tokenize input sentence
    word_tokenizer = Tokenizer()                     
    word_tokenizer.fit_on_texts(words)                    
    words_encoded = word_tokenizer.texts_to_sequences(words)

    # Tokenize pos tags
    tag_tokenizer = Tokenizer()
    tag_tokenizer.fit_on_texts(pos_tags)
    pos_tags_encoded = tag_tokenizer.texts_to_sequences(pos_tags)
    with open('word_tokenizer.pickle', 'wb') as handle:
        pickle.dump(word_tokenizer, handle)

    with open('pos_tokenizer.pickle', 'wb') as handle:
        pickle.dump(tag_tokenizer, handle)

    print("** Raw data point **", "\n", "-"*100, "\n")
    print('X: ', words[0], '\n')
    print('Y: ', pos_tags[0], '\n')
    print()
    print("** Encoded data point **", "\n", "-"*100, "\n")
    print('X: ', words_encoded[0], '\n')
    print('Y: ', pos_tags_encoded[0], '\n')

    lens = [len(seq) for seq in words_encoded]
    print("Length of longest sentence: {}".format(max(lens)))

    # truncate amd pad sentence to equal length
    words_padded = pad_sequences(words_encoded, maxlen=MAX_SEQ_LENGTH, padding="post", truncating="post")
    pos_tags_padded = pad_sequences(pos_tags_encoded, maxlen=MAX_SEQ_LENGTH, padding="post", truncating="post")

    print(words_padded[0], "\n"*3)
    print(pos_tags_padded[0])

    X, Y = words_padded, pos_tags_padded


    # load word2vec from the gensim library
    word2vec = api.load("glove-wiki-gigaword-100")
    # result = word2vec.most_similar(positive = ["king", "woman"], negative = ["man"])
    # print(result)
    # print()

    EMBEDDING_SIZE = len(word2vec['man']) # here 100
    VOCABULARY_SIZE = len(word_tokenizer.word_index) + 1
    embedding_weights = np.zeros((VOCABULARY_SIZE, EMBEDDING_SIZE))
    word2id = word_tokenizer.word_index

    # copy vectors from word2vec model to the words present in corpus
    for word, index in word2id.items():
        try:
            embedding_weights[index, :] = word2vec[word]
        except Exception as e:
            pass
    
    print("Embeddings shape: {}".format(embedding_weights.shape))
    #open("embeddings.npy")
    np.save("embddings.npy", embedding_weights)
    #embedding_weights = np.load(f)
    Y = to_categorical(Y)
    print(Y.shape)
    NUM_CLASSES = Y.shape[2]

    # Data conversion for torch
    embedding_weights = torch.tensor(embedding_weights,  dtype=torch.float32)
    X = torch.tensor(X, dtype=torch.long)
    Y = torch.tensor(Y, dtype=torch.float32)

    # Split data into train, test, val
    TEST_SIZE = 0.15
    VALID_SIZE = 0.15
    X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=TEST_SIZE, random_state=4)
    X_train, X_validation, Y_train, Y_validation = train_test_split(X_train, Y_train, test_size=VALID_SIZE, random_state=4)

    print("TRAINING DATA")
    print('Shape of input sequences: {}'.format(X_train.shape))
    print('Shape of output sequences: {}'.format(Y_train.shape))
    print("-"*50)
    print("VALIDATION DATA")
    print('Shape of input sequences: {}'.format(X_validation.shape))
    print('Shape of output sequences: {}'.format(Y_validation.shape))
    print("-"*50)
    print("TESTING DATA")
    print('Shape of input sequences: {}'.format(X_test.shape))
    print('Shape of output sequences: {}'.format(Y_test.shape))


    device = torch.device('cpu')
    embedding_weights = embedding_weights.to(device)
    model = PosTaggerBi(EMBEDDING_SIZE, HIDDEN_DIM, VOCABULARY_SIZE, NUM_CLASSES, embedding_weights)
    optimizer = optim.Adam(model.parameters(),lr=5e-5)
    trainer(model, X_train, Y_train, X_validation, Y_validation, 20, device, optimizer=optimizer)