import pandas as pd
import numpy as np
import nltk
nltk.download('punkt')
import torch
import torch.nn as nn
import torch.optim as optim
from nltk.corpus import stopwords
#from torchtext import data
import re
import random
import spacy
#library imports
import torch
import torch.nn as nn
import pandas as pd
import numpy as np
import re
import spacy
import jovian
from collections import Counter
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F
import string
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.metrics import make_scorer, accuracy_score, precision_score, recall_score, f1_score

SEED=42
torch.manual_seed(SEED)





from nltk.tokenize import word_tokenize
text = "God is Great! I won a lottery."


def spacy_token(x):
  x= re.sub(r'[^a-zA-Z\s]','',x)
  x= re.sub(r'[\n]',' ',x)
  return word_tokenize(x)

print(spacy_token(text))
stop_words = set(stopwords.words('english'))


reviews = pd.read_csv("data/Data.csv")
print(reviews.shape)
reviews.head()


testData = pd.read_csv("data/Test_set_paper_for_testing.csv")
print(testData.shape)
testData.head()

def tokenize (text):
    text = re.sub(r"[^\x00-\x7F]+", " ", text)
    regex = re.compile('[' + re.escape(string.punctuation) + '0-9\\r\\t\\n]') # remove punctuation and numbers
    nopunct = regex.sub(" ", text.lower())
    return word_tokenize(nopunct)

counts = Counter()
for index, row in reviews.iterrows():
    counts.update(tokenize(row['claims']))

#creating vocabulary
vocab2index = {"":0, "UNK":1}
words = ["", "UNK"]
for word in counts:
    vocab2index[word] = len(words)
    words.append(word)


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

reviews['encoded'] = reviews['claims'].apply(lambda x: np.array(encode_sentence(x,vocab2index )))
print(reviews.head())

testData['encoded'] = testData['claims'].apply(lambda x: np.array(encode_sentence(x,vocab2index )))
print(testData.head())


print(Counter(reviews['classes']))
X_train = list(reviews['encoded'])
y_train = list(reviews['classes'])

X_valid = list(testData['encoded'])
y_valid = list(testData['classes'])


for i in range(len(y_train)):
    if(y_train[i] == -1):
        y_train[i] = 2
        
for i in range(len(y_valid)):
    if(y_valid[i] == -1):
        y_valid[i] = 2

#X_train, X_valid, y_train, y_valid = train_test_split(X, y, test_size=0.2)


class ReviewsDataset(Dataset):
    def __init__(self, X, Y):
        self.X = X
        self.y = Y

    def __len__(self):
        return len(self.y)

    def __getitem__(self, idx):
        return torch.from_numpy(self.X[idx][0].astype(np.int32)), self.y[idx], self.X[idx][1]

train_ds = ReviewsDataset(X_train, y_train)
valid_ds = ReviewsDataset(X_valid, y_valid)

def train_model(model, epochs=10, lr=0.001):
    parameters = filter(lambda p: p.requires_grad, model.parameters())
    optimizer = torch.optim.Adam(parameters, lr=lr)
    for i in range(epochs):
        model.train()
        sum_loss = 0.0
        total = 0
        for x, y, l in train_dl:
            x = x.long()
            y = y.long()
            y_pred = model(x, l)
            optimizer.zero_grad()
            loss = F.cross_entropy(y_pred, y)
            loss.backward()
            optimizer.step()
            sum_loss += loss.item()*y.shape[0]
            total += y.shape[0]
        val_loss, val_acc, val_rmse, f1 = validation_metrics(model, val_dl)
        if i % 5 == 1:
            print("train loss %.3f, val loss %.3f, val accuracy %.3f, val rmse %.3f and F1 score %.3f" % (sum_loss/total, val_loss, val_acc, val_rmse, f1))
            
def validation_metrics (model, valid_dl):
    model.eval()
    correct = 0
    total = 0
    sum_loss = 0.0
    sum_rmse = 0.0
    for x, y, l in valid_dl:
        x = x.long()
        y = y.long()
        y_hat = model(x, l)
        loss = F.cross_entropy(y_hat, y)
        pred = torch.max(y_hat, 1)[1]
        f1 = f1_score(y, pred, average='macro')
        correct += (pred == y).float().sum()
        total += y.shape[0]
        sum_loss += loss.item()*y.shape[0]
        sum_rmse += np.sqrt(mean_squared_error(pred, y.unsqueeze(-1)))*y.shape[0]
    return sum_loss/total, correct/total, sum_rmse/total, f1

batch_size = 5
vocab_size = len(words)
train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
val_dl = DataLoader(valid_ds, batch_size=batch_size)


class LSTM_fixed_len(torch.nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.dropout = nn.Dropout(0.3)
        self.embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)
        self.linear = nn.Linear(hidden_dim, 5)

    def forward(self, x, s):
        x = self.embeddings(x)
        x = self.dropout(x)
        x_pack = pack_padded_sequence(x, s, batch_first=True, enforce_sorted=False)
        out_pack, (ht, ct) = self.lstm(x_pack)
        out = self.linear(ht[-1])
        return out

model_fixed =  LSTM_fixed_len(vocab_size, 50, 50)
train_model(model_fixed, epochs=30, lr=0.01)
