import numpy as np
import torch
from gensim import downloader
import gensim

with open(r"C:\Users\Kacper Dudzic\Desktop\paranormal-or-skeptic\train\in.tsv", "r", encoding="utf-8") as f:
    X_train = [line.strip() for line in f.readlines()]
with open(r"C:\Users\Kacper Dudzic\Desktop\paranormal-or-skeptic\train\expected.tsv", "r", encoding="utf-8") as f:
    Y_train = np.array([int(line.strip()) for line in f.readlines()])

glove_vectors = downloader.load("glove-wiki-gigaword-100")
X_train = [np.mean([glove_vectors[tk] for tk in gensim.utils.tokenize(text, lowercase=True) if tk in glove_vectors] or
                   [np.zeros(100)], axis=0) for text in X_train]


class NeuralNetworkModel(torch.nn.Module):

    def __init__(self):
        super(NeuralNetworkModel, self).__init__()
        self.fc1 = torch.nn.Linear(100, 500)
        self.fc2 = torch.nn.Linear(500, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        x = torch.sigmoid(x)
        return x


nn_model = NeuralNetworkModel()
criterion = torch.nn.BCELoss()
optimizer = torch.optim.Adam(nn_model.parameters())

BATCH_SIZE = 5
for epoch in range(5):
    loss_score = 0
    acc_score = 0
    items_total = 0
    nn_model.train()
    for i in range(0, Y_train.shape[0], BATCH_SIZE):
        X = X_train[i:i + BATCH_SIZE]
        X = torch.tensor(np.array(X).astype(np.float32))
        Y = Y_train[i:i + BATCH_SIZE]
        Y = torch.tensor(Y.astype(np.float32)).reshape(-1, 1)
        Y_predictions = nn_model(X)
        acc_score += torch.sum((Y_predictions > 0.5) == Y).item()
        items_total += Y.shape[0]
        optimizer.zero_grad()
        loss = criterion(Y_predictions, Y)
        loss.backward()
        optimizer.step()

        loss_score += loss.item() * Y.shape[0]

    print(epoch)
    print((loss_score / items_total), (acc_score / items_total))

PATH = "new_model_full.pt"
torch.save(nn_model, PATH)