import math

from sklearn import tree
from sklearn.tree import _tree, DecisionTreeRegressor

from learner import Learner
from model import Model


class DecisionTreeLearner(Learner):

    def __init__(self, models, feature_names):
        super().__init__(models, feature_names)
        self.models = [DecisionTreeModel(m, feature_names, index) for index, m in enumerate(self.models)]

    @staticmethod
    def name():
        return "DecisionTreeLearner"

    @staticmethod
    def short_name():
        return "dtl"


class DecisionTreeModel(Model):

    def __init__(self, model_args, feature_names, ID):
        model = DecisionTreeRegressor(max_depth=model_args[0],
                                      min_samples_leaf=model_args[1])
        super().__init__(model, feature_names)
        self.ID = "dtl_" + str(ID)

    def __str__(self):
        return str(self.model)

    def name(self):
        return "dt (max_depth=" + str(self.model.max_depth) + \
               ", min_leaf=" + str(self.model.min_samples_leaf) + ")"

    def feature_importances(self, _):
        return [(self.ID, self.model.feature_importances_)]

    def custom_output(self, model_index, outfolder):
        # Visualize trees
        tree.export_graphviz(self.model, feature_names=self.feature_names,
                             out_file=outfolder + "tree" + str(model_index) + ".dot")

        # Output Go code
        self.tree_to_go_code(model_index, outfolder)

    def tree_to_go_code(self, model_index, outfolder):
        decision_tree = self.model.tree_
        with open(outfolder + "tree" + str(model_index) + "code.go", "w") as code_file:
            def subtree_to_go_code(node, depth):
                indent = "\t" * depth
                if decision_tree.feature[node] != _tree.TREE_UNDEFINED:
                    feature = self.feature_names[decision_tree.feature[node]]
                    threshold = decision_tree.threshold[node]
                    code_file.write("{}if s.{} <= {} {{\n".format(indent, feature, math.floor(threshold)))
                    subtree_to_go_code(decision_tree.children_left[node], depth + 1)
                    code_file.write("{}}}\n".format(indent))
                    subtree_to_go_code(decision_tree.children_right[node], depth)
                else:
                    code_file.write("{}return {}\n".format(indent, decision_tree.value[node][0][0]))

            code_file.write("// Package ab (Code generated by a Python script)\n")
            code_file.write("package ab\n\n")
            code_file.write("func getEstimatedValueDT(s *Sample) float64 {\n")
            subtree_to_go_code(0, 1)
            code_file.write("}\n")

    def score(self, X, y):
        m = [4, 10, 20, 40, 60, 80]

        def group_func(X, ind_X):
            for split in m:
                if X["num_stones"].loc[ind_X] <= split:
                    return split
            return 121

        subsets = X.groupby(lambda i: group_func(X, i))

        s = {-1: (self.model.score(X, y), len(X))}
        for key in subsets.groups.keys():
            inp = X.loc[subsets.groups[key]]
            out = y.loc[subsets.groups[key]]
            s[key] = (self.model.score(inp, out), len(inp))
        return s
