from nltk.classifier import maxent
from nltk.classifier import LabeledText
from nltk.token import Token
from nltk.classifier.feature import FunctionFeatureDetector


training_set = [
    (0, 0, 0, 0),
    (0, 1, 0, 0),
    (0, 0, 1, 0),
    (0, 0, 0, 1),
    (0, 1, 1, 0),
    (1, 0, 0, 0),
    (1, 1, 0, 0),
    (1, 0, 1, 0),
    (1, 1, 1, 0),
    ]
results = [0, 0, 0, 0, 0, 1, 1, 1, 1]

num_features = len(training_set[0])

# Make a FeatureDetectorList
class my_feature:
    def __init__(self, dimension, klass):
        self.dimension = dimension
        self.klass = klass
    def __call__(self, token):
        vector, klass = token.text(), token.label()
        return vector[self.dimension] and klass == self.klass

feature_detectors = []
for i in range(num_features):
    f = my_feature(i, 1)
    f = FunctionFeatureDetector(f)
    feature_detectors.append(f)
base_fd_list = maxent.SimpleFDList(feature_detectors)
fd_list = maxent.GIS_FDList(base_fd_list)

train_toks = []
for i in range(len(training_set)):
    x, y = training_set[i], results[i]
    tok = Token(LabeledText(x, y))
    train_toks.append(tok)

trainer = maxent.IISMaxentClassifierTrainer(fd_list)
classifier = trainer.train(train_toks)

# Weights for each of my features, the 1-feature, and the correction
# feature.
weights = classifier.weights()

labels = [0, 1]
classifier = maxent.ConditionalExponentialClassifier(fd_list, labels, weights)
for i in range(len(training_set)):
    token = Token(training_set[i])
    labeled_token = classifier.classify(token)
    label = labeled_token.type().label()
    print label




######################################## Naive Bayes

print "Naive Bayes"

from nltk.classifier import naivebayes

# Make a FeatureDetectorList
class nb_feature:
    def __init__(self, dimension):
        self.dimension = dimension
    def __call__(self, token):
        vector = token.text()
        return vector[self.dimension]

feature_detectors = []
for i in range(num_features):
    f = nb_feature(i)
    f = FunctionFeatureDetector(f)
    feature_detectors.append(f)
fd_list = maxent.SimpleFDList(feature_detectors)

train_toks = []
for i in range(len(training_set)):
    x, y = training_set[i], results[i]
    tok = Token(LabeledText(x, y))
    train_toks.append(tok)

trainer = naivebayes.NBClassifierTrainer(fd_list)
classifier = trainer.train(train_toks)


from nltk import probability
label_fdist = probability.FreqDist()
for y in results:
    label_fdist.inc(y)
label_pdist = probability.MLEProbDist(label_fdist)
fval_pdist = classifier.fval_probdist()
labels = [0, 1]

classifier = naivebayes.NBClassifier(fd_list, labels, label_pdist, fval_pdist)

for i in range(len(training_set)):
    token = Token(training_set[i])
    labeled_token = classifier.classify(token)
    label = labeled_token.type().label()
    print label
    print classifier.distribution_dictionary(token)