"""This lib calculates the probability of a given student needs the substitute
test"""

# Third part imports
from pandas import get_dummies
from sklearn.naive_bayes import MultinomialNB
from sklearn.ensemble import AdaBoostClassifier
from sklearn.tree import DecisionTreeClassifier

# Local imports
from mlstudents.dataset import DataSet

class Classify(object):
    """ This class is used to handle classification algorithms. """

    def __init__(self):
        self.data = DataSet()
        self.Xtrainning, self.Ytrainning = self.data.training_data()
        self.Xtesting, self.Ytesting = self.data.testing_data()
        self.Xfinal, self.Yfinal = self.data.final_data()

    def __calculate_difference(self, list1, list2):
        """ Calculates the amount of equals elements in a list. """

        count = 0
        results = list2 - list1
        for result in results:
            if result == 0:
                count += 1
        return count

    def calculate_nb(self, dataset="test"):
        """ Calculates the hits of NaiveBayes. """

        if dataset == "test":
            X = self.Xtesting
            Y = self.Ytesting
        elif dataset == "final":
            X = self.Xfinal
            Y = self.Yfinal

        modelo = MultinomialNB()
        modelo.fit(self.Xtrainning, self.Ytrainning)
        predict_NB = modelo.predict(X)
        res = self.__calculate_difference(predict_NB, Y)
        total = len(Y)

        return res/total

    def calculate_ab(self, dataset="test"):
        """ Calculates the hits of AdaBoost. """

        if dataset == "test":
            X = self.Xtesting
            Y = self.Ytesting
        elif dataset == "final":
            X = self.Xfinal
            Y = self.Yfinal

        modelo = AdaBoostClassifier()
        modelo.fit(self.Xtrainning, self.Ytrainning)
        predict_AB = modelo.predict(X)
        res = self.__calculate_difference(predict_AB, Y)
        total = len(Y)

        return res/total

    def calculate_dt(self, dataset="test"):
        """ Calculates the hits of Decision Tree. """

        if dataset == "test":
            X = self.Xtesting
            Y = self.Ytesting
        elif dataset == "final":
            X = self.Xfinal
            Y = self.Yfinal

        modelo = DecisionTreeClassifier()
        modelo.fit(self.Xtrainning, self.Ytrainning)
        predict_DT = modelo.predict(X)
        res = self.__calculate_difference(predict_DT, Y)
        total = len(Y)

        return res/total

    def calculate_dumb(self):
        """ Calculates the hits of a Dump Classifier. """

        # Creates two arrays
        ones_list = [1] * len(self.Xtesting)
        zeros_list = [0] * len(self.Xtesting)

        res_one = self.__calculate_difference(ones_list, self.Ytesting)
        res_zero = self.__calculate_difference(zeros_list, self.Ytesting)
        total = len(self.Ytesting)

        return max(res_one, res_zero)/total

    def prediction (self, alg=None, csvfile=None):
        """ This method predicts using a new dataset.
        Campus      | Disciplina   | Nota_N1
        centro      | redes        | 0 to 10
        paulista1   | arquitetura  |
        paulista2   |              |
        olimpia     |              |

        """

        df = self.data.discretize(csvfile)
        Xdummies = get_dummies(df)

        if alg == "NaiveBayes":
            modelo = MultinomialNB()
            modelo.fit(self.Xtrainning, self.Ytrainning)
            prediction = modelo.predict(Xdummies.values)

        elif alg == "AdaBoost":
            modelo = AdaBoostClassifier()
            modelo.fit(self.Xtrainning, self.Ytrainning)
            prediction = modelo.predict(Xdummies.values)

        elif alg == "DecisionTree":
            modelo = DecisionTreeClassifier()
            modelo.fit(self.Xtrainning, self.Ytrainning)
            prediction = modelo.predict(Xdummies.values)

        else:
            prediction = None

        return prediction
