# this class trains a classifier and returns it
from __future__ import division
from __future__ import unicode_literals
import nltk
import codecs
import time
import os

###################################################
# Trainer Class using the Decision Tree Algorithm #
# Contains also a function to access accuracy and #
# errors.                                         #
###################################################

class GenericTrainer:
    
    # calculate and return the accuracy on the development set
    # and the baseline accuracy
    def return_accuracy(self, dev_features):
        dev_features_without_word, words = self.make_wordfree_feature_dict(dev_features)
        corr_with_MB = 0
        for word in dev_features_without_word:
            if word[1] == "O":
                corr_with_MB += 1
        accuracy_base = corr_with_MB / len(dev_features_without_word)
        return nltk.classify.accuracy(self.classifier, dev_features_without_word), accuracy_base
        
    # remove feature "word" from featureset to exclude from training (was added by default)
    # at the same time, assemble the list of words (used below to output mislabeled words)
    def make_wordfree_feature_dict(self, train_features):
        words = []
        train_features_ww = []
        for tup in train_features:
            features = {}
            for f in tup[0]:
                if f == "word":
                    words.append(tup[0]["word"])
                else:
                    features[f] = tup[0][f]
            train_features_ww.append((features,tup[1]))
        return train_features_ww, words

class DTTrainer(GenericTrainer):
    
    def __init__(self, train_features):
        # train the classifier right at its initialization
        print("  Decision Tree Classifier initialized")
        train_features_without_word, words =  self.make_wordfree_feature_dict(train_features)
        self.classifier = self.train(train_features_without_word)

    

    # use the nltk implementation of the Decision Tree algorithm to train 
    def train(self, train_features):
        classifier = nltk.DecisionTreeClassifier.train(train_features)
        return classifier
    
    

    # function to
    # pretty print the decision tree
    #def pretty_format():
    
    # print the following info to the screen
    # and write to log file:
    # accuracy
    # the incorrectly classified words along with the gold label and the predicted label
    def analyze_classifier(self, bestclassifier, dev_features):
        try:
            os.makedirs('logs')
        except OSError:
            pass
        logname = "logs/log.decisiontree" + time.strftime('%Y-%m-%d-%H-%M-%S') + ".txt"
        logfile = codecs.open(logname,'w','utf8')
        dev_features_without_word, words = self.make_wordfree_feature_dict(dev_features)
        corr_with_MB = 0
        for word in dev_features_without_word:
            if word[1] == "O":
                corr_with_MB += 1
        acc_base= corr_with_MB / len(dev_features_without_word)
        errors = []
        i = 0
        # for each token in the development set
        for i, tok in enumerate(dev_features_without_word):
            # word
            word = words[i]
            # gold label
            tag = tok[1]
            # predicted label
            guess = bestclassifier.classify(tok[0])
            
            # if the gold label is not the same as the predicted label
            if guess != tag:
                errors.append( (tag, guess, word) )

        print("\n=== Summary best classifier ===")
        print('{:>6,}'.format(len(dev_features_without_word)) + " words were in the development set")
        print('{:>6,}'.format(len(dev_features_without_word)-len(errors)) +
              " of these were labeled correctly by your classifier, and")
        print('{:>6,}'.format(len(errors)) + " were mislabeled")
        print('{:<8.5f}'.format(self.return_accuracy(dev_features)[0]) + "is the resulting accuracy")
        print('{:<8.5f}'.format(acc_base) +
              "is the baseline accuracy (if all words were labeled as non-entity)")
        print("An overview of the errors is in " + logname + '\n')
        
        logfile.write('{:>6,}'.format(len(dev_features_without_word)) +
                      " words were in the development set\r\n")
        logfile.write('{:>6,}'.format(len(errors)) + " were mislabeled\r\n\r\n")
        logfile.write('{:>6,}'.format(len(dev_features_without_word) - corr_with_MB) +
                      " would have been mislabeled if all words were labeled as non-entity" +
                      '\r\n\r\n')
        logfile.write("The following features have been used:\r\n\r\n")
        features = ''
        for f in dev_features_without_word[0][0]:
            features += f + '\r\n'
        logfile.write(features + '\r\n\r\n')

        # print decision tree
        # note: non-str feature values (Ture/Fales/None)
        #       produce a Type Error in Py3 (but not in Py2)
        try:
            logfile.write('Decision tree:\r\n\r\n' + str(bestclassifier) + '\r\n\r\n')
        except TypeError:
            logfile.write('Decision tree not printed due to non-str feature values\r\n')

        # print/write all the errors
        logfile.write('Errors:\r\n\r\n')
        for (tag, guess, word) in sorted(errors):
            logfile.write('correct={:<8} guess={:<8s} word={:<30}\r\n'.format(tag, guess, word))
        logfile.close()
        
        

###################################################
# Trainer Class using the Naive Bayes Algorithm   #
# Contains also a function to access accuracy and #
# errors.                                         #
###################################################
        
class NBTrainer(GenericTrainer):

    def __init__(self, train_features):
        # train the classifier right at its initialization
        print("  Naive Bayes Classifier initialized")
        train_features_without_word,words = self.make_wordfree_feature_dict(train_features)
        self.classifier = self.train(train_features_without_word)

    # use the nltk implementation of the Naive Bayes algorithm to train
    def train(self, train_features):
        classifier = nltk.NaiveBayesClassifier.train(train_features)
        return classifier
        
    # print the following info to the screen and write to log file:
    # accuracy
    # the incorrectly classified words along with the gold label and the predicted label
    def analyze_classifier(self, bestclassifier, dev_features):
        try:
            os.makedirs('logs')
        except OSError:
            pass
        logname="logs/log.naivebayes" + time.strftime('%Y-%m-%d-%H-%M-%S') + ".txt"
        logfile = codecs.open(logname,'w','utf8')
        
        # collect the features that were most informative for the algorithm to make the decision
        # it gives the two labels it had to make a decision between the probability distribution of 
        # the feature 1 for class1 and feature value for class2
        n = 10
        summary_mostinf_features = ""
        cpdist = bestclassifier._feature_probdist
        for (fname, fval) in bestclassifier.most_informative_features(n):
            def labelprob(l):
                return cpdist[l, fname].prob(fval)

            labels = sorted([l for l in bestclassifier._labels
                             if fval in cpdist[l, fname].samples()],
                            key=labelprob)
            if len(labels) == 1:
                continue
            l0 = labels[0]
            l1 = labels[-1]
            if cpdist[l0, fname].prob(fval) == 0:
                ratio = 'INF'
            else:
                ratio = '%8.1f' % (cpdist[l1, fname].prob(fval) /
                                   cpdist[l0, fname].prob(fval))
            summary_mostinf_features += ('%24s = %-14r %11s : %-11s = %s : 1.0\r\n\r\n' %
                   (fname, fval, ("%s" % l1)[:11], ("%s" % l0)[:11], ratio))
      
        dev_features_without_word, words = self.make_wordfree_feature_dict(dev_features)
        corr_with_MB = 0
        for word in dev_features_without_word:
            if word[1] == "O":
                corr_with_MB += 1
        acc_base = corr_with_MB / len(dev_features_without_word)
        errors = []
        i = 0
        # for each token in the development set
        for i, tok in enumerate(dev_features_without_word):
            # word
            word = words[i]
            # gold label
            tag =  tok[1]
            # predicted label
            guess = bestclassifier.classify(tok[0])
            
            # if the gold label is not the same as the predicted label
            if guess != tag:
                errors.append( (tag, guess, word) )
        # print/write the information
        print("\n=== Summary best classifier ===")
        print('{:>6,}'.format(len(dev_features_without_word)) +
              " words were in the development set")
        print('{:>6,}'.format(len(dev_features_without_word)-len(errors)) +
              " of these were labeled correctly by your classifier, and")
        print('{:>6,}'.format(len(errors)) + " were mislabeled")
        print('{:<8.5f}'.format(self.return_accuracy(dev_features)[0]) + "is the resulting accuracy")
        print('{:<8.5f}'.format(acc_base) +
              "is the baseline accuracy (if all words were labeled as non-entity)")
        print("An overview of the errors is in " + logname + '\n')

        logfile.write('{:>6,}'.format(len(dev_features_without_word)) +
                      " words were in the development set\r\n")
        logfile.write('{:>6,}'.format(len(errors)) + " were mislabeled\r\n\r\n")
        logfile.write('{:>6,}'.format(len(dev_features_without_word)-corr_with_MB) +
                      " would have been mislabeled if all words were labeled as non-entity" +
                      '\r\n\r\n')

        logfile.write("The following features have been used:\r\n\r\n")
        features = ''
        for f in dev_features_without_word[0][0]:
            features += f + '\r\n'
        logfile.write(features + '\r\n\r\n')

        logfile.write("Most informative features:\r\n\r\n")
        # bestclassifier.show_most_informative_features(5)
        logfile.write(summary_mostinf_features + '\r\n\r\n')

        # print/write all the errors
        logfile.write('Errors:\r\n\r\n')
        for (tag, guess, word) in sorted(errors):
            logfile.write('correct={:<8} guess={:<8s} word={:<30}\r\n'.format(tag, guess, word))
        logfile.close()
        

