import lzma

from naivebayes import NaiveBayesTextClassifier
import nltk
from nltk.corpus import stopwords

nltk.download("stopwords")

# Read train files
with lzma.open("train/in.tsv.xz", "rt", encoding="utf-8") as train_file:
    x_train = [x.strip().lower() for x in train_file.readlines()]

with open("train/expected.tsv", "r", encoding="utf-8") as train_file:
    y_train = [int(x.strip()) for x in train_file.readlines()]

nbc = NaiveBayesTextClassifier(
    categories=[0, 1],
    stop_words=stopwords.words("english"),
    min_df=1
)
step = 15000
for i in range(0, len(x_train), step):
    nbc.train(x_train[i:min(i+step, len(x_train))], y_train[i:min(i+step, len(x_train))])

# Read dev files
with lzma.open("dev-0/in.tsv.xz", "rt", encoding="utf-8") as dev_file:
    x_dev = [x.strip().lower() for x in dev_file.readlines()]

# Read test file
with lzma.open("test-A/in.tsv.xz", "rt", encoding="utf-8") as test_file:
    x_test = [x.strip().lower() for x in test_file.readlines()]

# Predict dev
pred_dev = [str(x) + "\n" for x in nbc.classify(x_dev)]
with open("dev-0/out.tsv", "w", encoding="utf-8") as dev_out_file:
    dev_out_file.writelines(pred_dev)

# Predict dev
pred_test = [str(x) + "\n" for x in nbc.classify(x_test)]
with open("test-A/out.tsv", "w", encoding="utf-8") as test_out_file:
    test_out_file.writelines(pred_test)