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

# In[1]:


import csv


# In[2]:


get_ipython().system('pip install gensim')


# In[17]:


import nltk
nltk.download('punkt')


# In[9]:


get_ipython().system('pip install nltk')


# In[3]:


get_ipython().system('pip install torch')


# In[4]:


import gensim.downloader
import numpy as np
import pandas as pd
import torch


# In[5]:


import torch.nn as nn
from nltk import word_tokenize


# In[13]:


header_names = ["content", "id", "label"]


# In[23]:


class FF(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(FF, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu1(out)
        out = self.relu2(out)
        out = self.fc3(out)
        return torch.sigmoid(out)

train_set_labels = pd.read_table(
    "train/expected.tsv",
    error_bad_lines=False,
    quoting=csv.QUOTE_NONE,
    header=None,
    names=header_names[2:],
)

train_set_features = pd.read_table(
    "train/in.tsv.xz",
    error_bad_lines=False,
    quoting=csv.QUOTE_NONE,
    header=None,
    names=header_names[:2],
)


test_set = pd.read_table(
    "test-A/in.tsv.xz",
    error_bad_lines=False,
    header=None,
    quoting=csv.QUOTE_NONE,
    names=header_names[:2],
)

dev_set = pd.read_table(
    "dev-0/in.tsv.xz",
    error_bad_lines=False,
    header=None,
    quoting=csv.QUOTE_NONE,
    names=header_names[:2],
)
X_train = train_set_features["content"].str.lower()
y_train = train_set_labels["label"]

X_dev = dev_set["content"].str.lower()
X_test = test_set["content"].str.lower()
X_train = [word_tokenize(content) for content in X_train]
X_dev = [word_tokenize(content) for content in X_dev]
X_test = [word_tokenize(content) for content in X_test]
word2vec = gensim.downloader.load("word2vec-google-news-300")


# In[24]:


X_train = [
    np.mean(
        [word2vec[word] for word in content if word in word2vec] or [np.zeros(300)],
        axis=0,
    )
    for content in X_train
]
X_dev = [
    np.mean(
        [word2vec[word] for word in content if word in word2vec] or [np.zeros(300)],
        axis=0,
    )
    for content in X_dev
]

X_test = [
    np.mean(
        [word2vec[word] for word in content if word in word2vec] or [np.zeros(300)],
        axis=0,
    )
    for content in X_test
]
hidden_layer = 650
epochs = 15
batch_size = 10


# In[27]:


output_dim = 1


input_dim =300
model = FF(input_dim, hidden_layer, output_dim)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
criterion = torch.nn.BCELoss()


# In[28]:


for epoch in range(epochs):
    model.train()
    for i in range(0, y_train.shape[0], batch_size):
        X = X_train[i : i + batch_size]
        X = torch.tensor(X)
        y = y_train[i : i + batch_size]
        y = torch.tensor(y.astype(np.float32).to_numpy()).reshape(-1, 1)

        outputs = model(X.float())
        loss = criterion(outputs, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
test_prediction = []
dev_prediction = []
model.eval()
with torch.no_grad():
    for i in range(0, len(X_test), batch_size):
        X = X_test[i : i + batch_size]
        X = torch.tensor(X)

        outputs = model(X.float())

        prediction = outputs > 0.5
        test_prediction += prediction.tolist()

    for i in range(0, len(X_dev), batch_size):
        X = X_dev[i : i + batch_size]
        X = torch.tensor(X)

        outputs = model(X.float())

        prediction = outputs > 0.5
        dev_prediction += prediction.tolist()

test_prediction = np.asarray(test_prediction, dtype=np.int32)
dev_prediction = np.asarray(dev_prediction, dtype=np.int32)
test_prediction.tofile("./test-A/out.tsv", sep="\n")
dev_prediction.tofile("./dev-0/out.tsv", sep="\n")


# In[ ]:





# In[ ]:





# In[ ]: