import torch
import os
from torchtext.legacy import data
from torchtext.legacy import datasets
from torchtext.legacy.data import Field, TabularDataset, BucketIterator

import torch.nn as nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
import torch.nn.functional as F
import torch.optim as optim

# Evaluation

from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from sklearn.metrics import f1_score, recall_score, precision_score
from collections import Counter

SEED = 1234

def generate_bigrams(x):
    n_grams = set(zip(*[x[i:] for i in range(2)]))
    for n_gram in n_grams:
        x.append(' '.join(n_gram))
    return x

torch.manual_seed(SEED)
torch.backends.cudnn.deterministic = True
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print (device)
# device = torch.device('cpu')

TEXT = Field(tokenize='spacy', tokenizer_language = 'de_core_news_sm', lower=True, batch_first = True, include_lengths=True,  preprocessing = generate_bigrams)

LABEL = Field(sequential=False, use_vocab=False, dtype=torch.long)
fields = [('GermanText', TEXT), ('GermanGlossToken', TEXT), ('label', LABEL)]

train, valid, test = TabularDataset.splits(path = "../../data/gloss-classification/", train='train.tsv', validation='dev.tsv', test='test.tsv',
                                           format='TSV', fields=fields, skip_header=True)

train_iterator = BucketIterator(train, batch_size=32, sort_key=lambda x: len(x.GermanText),
                            device=device, sort=True, sort_within_batch=True)
valid_iterator = BucketIterator(valid, batch_size=32, sort_key=lambda x: len(x.GermanText),
                            device=device, sort=True, sort_within_batch=True)
test_iterator = BucketIterator(test, batch_size=32, sort_key=lambda x: len(x.GermanText),
                            device=device, sort=True, sort_within_batch=True)

TEXT.build_vocab(train, min_freq=3)


class LSTM(nn.Module):        
    #define all the layers used in model
    def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, 
                 bidirectional, dropout, pad_idx):
        
        #Constructor
        super().__init__()          
        
        #embedding layer
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=pad_idx)
        self.g_embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=pad_idx)
        
        #lstm layer
        self.lstm = nn.LSTM(embedding_dim, 
                           hidden_dim, 
                           num_layers=n_layers, 
                           bidirectional=bidirectional, 
                           dropout=dropout,
                           batch_first=True)
        
        #dense layer
        self.fc = nn.Linear(hidden_dim * 2, embedding_dim)

        self.rel = nn.ReLU()

        self.last_linear = nn.Linear(embedding_dim*2, output_dim)
        
        
    def forward(self, text, gloss):
        
        #text = [batch size,sent_length]
        embedded = self.embedding(text[0])
        # print (gloss)
        g_embedded = self.embedding(gloss[0])
        # g_embedded = g_embedded.permute(1, 0, 2)

        #embedded = [batch size, sent_len, emb dim]
      
        #packed sequence
        packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, text[1].cpu(), batch_first=True)
        
        packed_output, (hidden, cell) = self.lstm(packed_embedded)
        #hidden = [batch size, num layers * num directions,hid dim]
        #cell = [batch size, num layers * num directions,hid dim]
        
        #concat the final forward and backward hidden state
        hidden = torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim = 1)
                
        #hidden = [batch size, hid dim * num directions]
        dense_outputs=self.fc(hidden)

        # print (dense_outputs.shape)
        # print (g_embedded.shape)
        # g_embeeded = g_embedded.squeeze(1)
        g_embedded = torch.flatten(g_embedded, start_dim = 1)
        # print (g_embedded.shape)


        catted = torch.cat((dense_outputs, g_embedded),1)

        relled = self.rel(catted)


        #Final activation function
        outputs=self.last_linear(relled)
        
        #pooled = [batch size, embedding_dim]
                
        return outputs


# vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, 
                 # bidirectional, dropout, pad_idx
#define hyperparameters
INPUT_DIM = len(TEXT.vocab)
HIDDEN_DIM = 100
EMBEDDING_DIM = 100
n_layers = 2
bidirectional = True
dropout = 0.3
OUTPUT_DIM = 3
PAD_IDX = TEXT.vocab.stoi[TEXT.pad_token]

model = LSTM(INPUT_DIM, EMBEDDING_DIM, HIDDEN_DIM, OUTPUT_DIM, n_layers, bidirectional, dropout, PAD_IDX)


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f'The model has {count_parameters(model):,} trainable parameters')

optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()
model = model.to(device)
criterion = criterion.to(device)


def print_evaluation(gold_labels, predicted_labels):

    '''Prints accuracy, precision, recall, f1 score'''

    accuracy = accuracy_score(gold_labels, predicted_labels) * 100
    f1 = f1_score(gold_labels, predicted_labels, average = "macro") * 100
    recall = recall_score(gold_labels, predicted_labels, average = "macro") * 100
    precision = precision_score(gold_labels, predicted_labels, average = "macro") * 100
  

    a = [accuracy, precision,  recall, f1]
    for i in range (4):
        a[i] = round(a[i],2)

    return a

def binary_accuracy(preds, y):
    """
    Returns accuracy per batch, i.e. if you get 8/10 right, this returns 0.8, NOT 8
    """

    #round predictions to the closest integer
    rounded_preds = preds.argmax(1)
    # print (rounded_preds)
    np_preds = rounded_preds.cpu().detach().numpy()
    np_y = y.cpu().detach().numpy()

    # print (print_evaluation(np_y, np_preds))

    correct = (rounded_preds == y).float() #convert into float for division 
    acc = correct.sum() / len(correct)
    return acc

def train(model, iterator, optimizer, criterion):
    
    epoch_loss = 0
    epoch_acc = 0
    
    model.train()
    
    for batch in iterator:
        
        optimizer.zero_grad()
        
        predictions = model(batch.GermanText, batch.GermanGlossToken).squeeze(1)
        
        loss = criterion(predictions, batch.label)
        
        acc = binary_accuracy(predictions, batch.label)
        
        loss.backward()
        
        optimizer.step()
        
        epoch_loss += loss.item()
        epoch_acc += acc.item()
        
    return epoch_loss / len(iterator), epoch_acc / len(iterator)

def evaluate(model, iterator, criterion):
    
    epoch_loss = 0
    epoch_acc = 0
    
    model.eval()
    
    with torch.no_grad():
    
        for batch in iterator:

            predictions = model(batch.GermanText, batch.GermanGlossToken).squeeze(1)
            
            loss = criterion(predictions, batch.label)
            
            acc = binary_accuracy(predictions, batch.label)

            epoch_loss += loss.item()
            epoch_acc += acc.item()
        
    return epoch_loss / len(iterator), epoch_acc / len(iterator)

import time

def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs


print ("begin training")
N_EPOCHS = 20

best_valid_loss = float('inf')

for epoch in range(N_EPOCHS):

    start_time = time.time()
    
    train_loss, train_acc = train(model, train_iterator, optimizer, criterion)
    valid_loss, valid_acc = evaluate(model, valid_iterator, criterion)
    
    end_time = time.time()

    epoch_mins, epoch_secs = epoch_time(start_time, end_time)
    
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), 'models/fast-text.pt')
    
    print(f'Epoch: {epoch+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')
    print(f'\tTrain Loss: {train_loss:.3f} | Train Acc: {train_acc*100:.2f}%')
    print(f'\t Val. Loss: {valid_loss:.3f} |  Val. Acc: {valid_acc*100:.2f}%')


    all_preds = []
    all_y = []
    model.eval()
    
    with torch.no_grad():
    
        for batch in test_iterator:

            predictions = model(batch.GermanText, batch.GermanGlossToken).squeeze(1)
            rounded_preds = predictions.argmax(1)
            np_preds = list(rounded_preds.cpu().detach().numpy())
            np_y = list(batch.label.cpu().detach().numpy())

            all_preds += np_preds
            all_y += np_y

    print ("FINAL INFO")
    print (len(all_preds), len(all_y))
    print (Counter(all_preds))
    print (Counter(all_y))
    print (print_evaluation(all_y, all_preds))

