import pandas as pd
import numpy as np
import torch
from nltk.tokenize import word_tokenize
import gensim.downloader
from csv import QUOTE_NONE


print('initialization')
word2vec = gensim.downloader.load('word2vec-google-news-300')


def get_word2vec(document):
    return np.mean([word2vec[token] for token in document if token in word2vec] or [np.zeros(300)], axis=0)


class MyNeuralNetwork(torch.nn.Module):

    def __init__(self, input_size, hidden_size, num_classes):
        super(MyNeuralNetwork, self).__init__()
        self.fc1 = torch.nn.Linear(input_size, hidden_size)
        self.fc2 = torch.nn.Linear(hidden_size, num_classes)

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


# wczytanie danych
print('loading data')
train_x = pd.read_table('train/in.tsv.xz', error_bad_lines=False, header=None, quoting=QUOTE_NONE,
                        names=['content', 'id'])
train_y = pd.read_table('train/expected.tsv', error_bad_lines=False, header=None, quoting=QUOTE_NONE,
                        names=['label'])['label']
dev_x = pd.read_table('dev-0/in.tsv.xz', error_bad_lines=False, header=None, quoting=QUOTE_NONE,
                      names=['content', 'id'])
test_x = pd.read_table('test-A/in.tsv.xz', error_bad_lines=False, header=None, quoting=QUOTE_NONE,
                       names=['content', 'id'])

# preprocessing danych
print('word tokenize')
train_x = [word_tokenize(row) for row in train_x['content'].str.lower()]
dev_x = [word_tokenize(row) for row in dev_x['content'].str.lower()]
test_x = [word_tokenize(row) for row in test_x['content'].str.lower()]

print('word2vec')
train_x = [get_word2vec(document) for document in train_x]
dev_x = [get_word2vec(document) for document in dev_x]
test_x = [get_word2vec(document) for document in test_x]

# trenowanie
print('model training')
network = MyNeuralNetwork(300, 600, 1)
criterion = torch.nn.BCELoss()
optimizer = torch.optim.SGD(network.parameters(), lr=0.02)
epochs = 15
batch_size = 15

for epoch in range(epochs):
    network.train()
    for i in range(0, train_y.shape[0], batch_size):
        x = train_x[i:i + batch_size]
        x = torch.tensor(x)
        y = train_y[i:i + batch_size]
        y = torch.tensor(y.astype(np.float32).to_numpy()).reshape(-1, 1)
        outputs = network(x.float())
        loss = criterion(outputs, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# ewaluacja
print('evaluation')
dev_y_prediction = []
test_y_prediction = []
with torch.no_grad():
    for i in range(0, len(dev_x), batch_size):
        x = dev_x[i:i + batch_size]
        x = torch.tensor(x)
        outputs = network(x.float())
        prediction = outputs > 0.5
        dev_y_prediction += prediction.tolist()
    for i in range(0, len(test_x), batch_size):
        x = test_x[i:i + batch_size]
        x = torch.tensor(x)
        outputs = network(x.float())
        prediction = outputs > 0.5
        test_y_prediction += prediction.tolist()

# zapisanie danych
print('saving data')
np.asarray(dev_y_prediction, dtype=np.int32).tofile('./dev-0/out.tsv', sep='\n')
np.asarray(test_y_prediction, dtype=np.int32).tofile('./test-A/out.tsv', sep='\n')

print('done')