import numpy as np
import math


class ID3Classifier:

    def __init__(self):
        self.tree = None
        self.target = None
        self.features = None


    # ===========================
    # ENTROPY
    # ===========================
    def entropy(self, data):

        values, counts = np.unique(data, return_counts=True)

        ent = 0

        for count in counts:

            prob = count / sum(counts)

            if prob > 0:
                ent -= prob * math.log2(prob)

        return ent


    # ===========================
    # INFORMATION GAIN
    # ===========================
    def information_gain(self, df, feature):

        total_entropy = self.entropy(df[self.target])

        values, counts = np.unique(df[feature], return_counts=True)

        weighted_entropy = 0

        for i in range(len(values)):

            subset = df[df[feature] == values[i]]

            subset_entropy = self.entropy(subset[self.target])

            weight = counts[i] / sum(counts)

            weighted_entropy += weight * subset_entropy

        return total_entropy - weighted_entropy


    # ===========================
    # ID3 TREE BUILDER
    # ===========================
    def _id3(self, data, features):

        # Pure node
        if len(np.unique(data[self.target])) == 1:

            return np.unique(data[self.target])[0]

        # No features left
        if len(features) == 0:

            return data[self.target].mode()[0]


        igs = [self.information_gain(data, f) for f in features]

        best = features[np.argmax(igs)]

        tree = {best: {}}

        remaining = [f for f in features if f != best]


        for value in np.unique(data[best]):

            subset = data[data[best] == value]

            if len(subset) == 0:

                tree[best][value] = data[self.target].mode()[0]

            else:

                tree[best][value] = self._id3(subset, remaining)

        return tree


    # ===========================
    # FIT MODEL
    # ===========================
    def fit(self, df, features, target):

        self.target = target
        self.features = features

        self.tree = self._id3(df, features)


    # ===========================
    # PRINT TREE
    # ===========================
    def print_tree(self, tree=None, indent=""):

        if tree is None:
            tree = self.tree

        if not isinstance(tree, dict):

            print(indent + "→", tree)

            return

        for key in tree:

            print(indent + key)

            for value in tree[key]:

                print(indent + " ├─", value)

                self.print_tree(tree[key][value], indent + " │   ")


    # ===========================
    # SINGLE PREDICT
    # ===========================
    def _predict(self, tree, sample):

        if not isinstance(tree, dict):

            return tree

        root = list(tree.keys())[0]

        value = sample[root]

        if value not in tree[root]:

            return None

        return self._predict(tree[root][value], sample)


    # ===========================
    # MULTIPLE PREDICT
    # ===========================
    def predict(self, df):

        preds = []

        for _, row in df.iterrows():

            preds.append(self._predict(self.tree, row))

        return np.array(preds)