#!/usr/bin/env python
# coding: utf-8

# In[3]:


from sklearn.model_selection import train_test_split

from sklearn.datasets import fetch_20newsgroups
# https://scikit-learn.org/0.19/datasets/twenty_newsgroups.html

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import accuracy_score


# In[71]:


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


# In[84]:


train_x = pd.read_csv('train/in.tsv', header=None, sep='\t', quoting=csv.QUOTE_NONE, error_bad_lines=False)
train_y = pd.read_csv('train/expected.tsv', header=None, sep='\t', quoting=csv.QUOTE_NONE, error_bad_lines=False)
dev_x = pd.read_csv('dev-0/in.tsv', header=None, sep='\t', quoting=csv.QUOTE_NONE, error_bad_lines=False)
dev_y = pd.read_csv('dev-0/expected.tsv', header=None, sep='\t', quoting=csv.QUOTE_NONE, error_bad_lines=False)
test_x = pd.read_csv('test-A/in.tsv', header=None, sep='\t',quoting=csv.QUOTE_NONE, error_bad_lines=False)


# In[85]:


print(len(train_x))


# In[86]:


print(len(train_y))


# In[87]:


train_y = train_y[0]


# In[100]:


dev_y = dev_y[0]


# In[88]:


print(type(train_y))


# In[89]:


train_y = train_y.to_numpy()


# In[102]:


dev_y = dev_y.to_numpy()


# In[90]:


train_x.head


# In[91]:


dev_x.head()


# In[92]:


train_x = train_x[0]


# In[93]:


vec_model = Word2Vec(train_x, vector_size=100, window=5, min_count=1, workers=4)


# In[ ]:


def w2v(model, data):
    return np.array([np.mean([model.wv[word] if word in model.wv.key_to_index else np.zeros(100, dtype=float) for word in doc], axis=0) for doc in data])
    


# In[ ]:


w2v()


# In[96]:


dev_x = dev_x[0]
test_x = test_x[0]


# In[95]:


vec_x_train = np.array([np.mean([vec_model.wv[word] if word in vec_model.wv.key_to_index else np.zeros(100, dtype=float) for word in doc], axis=0) for doc in train_x])


# In[97]:


vec_x_dev = np.array([np.mean([vec_model.wv[word] if word in vec_model.wv.key_to_index else np.zeros(100, dtype=float) for word in doc], axis=0) for doc in dev_x])
vec_x_test = np.array([np.mean([vec_model.wv[word] if word in vec_model.wv.key_to_index else np.zeros(100, dtype=float) for word in doc], axis=0) for doc in test_x])


# In[36]:


X_dev0_w2v = vectorize(vec_model,dev_x)
X_test_w2v = vectorize(vec_model,test_x)


# In[7]:


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


# In[37]:


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


# In[8]:


def get_loss_acc(model, X_dataset, Y_dataset):
    loss_score = 0
    acc_score = 0
    items_total = 0
    model.eval()
    for i in range(0, Y_dataset.shape[0], BATCH_SIZE):
        X = np.array(X_dataset[i:i+BATCH_SIZE]).astype(np.float32)
        X = torch.tensor(X)
        Y = Y_dataset[i:i+BATCH_SIZE]
        Y = torch.tensor(Y.astype(np.float32)).reshape(-1,1)
        Y_predictions = model(X)
        acc_score += torch.sum((Y_predictions > 0.5) == Y).item()
        items_total += Y.shape[0]

        loss = criterion(Y_predictions, Y)

        loss_score += loss.item() * Y.shape[0]
    return (loss_score / items_total), (acc_score / items_total)


# In[9]:


def predict(model, data):
    model.eval()
    predictions = []
    for x in data:
        X = torch.tensor(np.array(x).astype(np.float32))
        Y_predictions = model(X)
        if Y_predictions[0] > 0.5:
            predictions.append("1")
        else:
            predictions.append("0")
    return predictions


# In[18]:


FEAUTERES = 100


# In[62]:


BATCH_SIZE = 5


# In[58]:


nn_model = NeuralNetworkModel()


# In[103]:


for epoch in range(7):
    loss_score = 0
    acc_score = 0
    items_total = 0
    nn_model.train()
    for i in range(0, train_y.shape[0], BATCH_SIZE):
        X = vec_x_train[i:i+BATCH_SIZE]
        X = torch.tensor(X)
        Y = train_y[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]

    display(epoch)
    display(get_loss_acc(nn_model,vec_x_train, train_y))
    display(get_loss_acc(nn_model, vec_x_dev, dev_y))


# In[104]:


dev_pred = predict(nn_model, vec_x_dev)
test_pred = predict(nn_model, vec_x_test)


# In[105]:


dev_pred


# In[119]:


dev_pred = [int(i) for i in dev_pred]
test_pred = [int(i) for i in test_pred]


# In[120]:


dev_pred = np.array(dev_pred)
test_pred = np.array(test_pred)


# In[117]:


dev_pred


# In[121]:


np.savetxt("dev-0/out.tsv",dev_pred, delimiter="\t", fmt='%d')


# In[122]:


np.savetxt("test-A/out.tsv",test_pred, delimiter="\t", fmt='%d')


# In[ ]: