import nltk
from estnltk import Text
import urllib.request
from collections import Counter

A2 = Text(urllib.request.urlopen("http://www.tlu.ee/~kais/Digihum_tehnoloogiad/Nadal6/A2_2018_I.txt").read().decode("utf8").lower()).postags
B1 = Text(urllib.request.urlopen("http://www.tlu.ee/~kais/Digihum_tehnoloogiad/Nadal6/B1_2018_I.txt").read().decode("utf8").lower()).postags
B1 = B1[0: len(A2)]
pikkus = 3
A2kolmikud = ["-".join(A2[arv:arv+pikkus]) for arv in range(len(A2)-(pikkus-1))]
B1kolmikud = ["-".join(B1[arv:arv+pikkus]) for arv in range(len(B1)-(pikkus-1))]

def tunnused(kolmik):
	return{'sõnaliigikolmik': kolmik}

andmed = [[tunnused(kolmik), 'A2'] for kolmik in A2kolmikud]
andmed+= [[tunnused(kolmik), 'B1'] for kolmik in B1kolmikud]

mudel = nltk.NaiveBayesClassifier.train(andmed)
print("Model Accuracy")
print(nltk.classify.accuracy(mudel, andmed))
mudel.show_most_informative_features(20)

A2_2 = Text(urllib.request.urlopen("http://www.tlu.ee/~kais/Digihum_tehnoloogiad/Nadal6/A2_2018_II.txt").read().decode("utf8").lower()).postags
A2_kolmikud2 = ["-".join(A2_2[arv:arv+pikkus]) for arv in range(len(A2_2)-(pikkus-1))]
print("Classification of New Text")
print(Counter([mudel.classify(tunnused(kolmik)) for kolmik in A2_kolmikud2]))