import numpy as np
import torch
import gensim
from sklearn import preprocessing
import pandas as pd
from gensim.test.utils import common_texts
from gensim.models import Word2Vec

with open("train/in.tsv") as f:
    X_train = f.readlines()
    with open("train/expected.tsv") as ff:
        Y_train = ff.readlines() 
        with open("test-A/in.tsv") as d:
            X = d.readlines() 

model = Word2Vec(X_train, min_count=1,size= 50,workers=3, window =3, sg = 1)
X_train=model[X_train]
X=model[X]
FEAUTERES=10000
class NeuralNetworkModel(torch.nn.Module):

    def __init__(self):
        super(NeuralNetworkModel, self).__init__()
        self.fc1 = torch.nn.Linear(FEAUTERES,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()

BATCH_SIZE = 5

criterion = torch.nn.BCELoss()

optimizer = torch.optim.SGD(nn_model.parameters(), lr = 0.1)

for epoch in range(10):
    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(X.astype(np.float32).todense())
        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] 
        with open('test-A/out1.tsv', 'w') as file:
            for e in Y_predictions:
                file.write("%f\n" % e)