import numpy as np
from collections import Counter


class DummyDecisionTree:
    def __init__(self, max_features=None):
        self._class_weights = None
        self._Ys = None

    def fit(self, X, Y):
        total = 0
        self._class_weights = np.zeros(len(set(Y)))
        self._Ys = []
        for i, (v, weight) in enumerate(Counter(Y).items()):
            total += weight
            self._Ys.append(v)
            self._class_weights[i] = weight
        self._class_weights /= total

    def predict(self, X):
        assert self._Ys is not None
        total = X.shape[0]
        return np.random.choice(self._Ys, size=(total, 1),
                                p=self._class_weights)

    def score(self, X, Y):
        Yhat = self.predict(X)
        right = (Y.reshape(-1, 1) == Yhat).sum()
        return right / total

if __name__ == '__main__':
    from sklearn.datasets import load_iris

    dataset = load_iris()
    X, Y = dataset['data'], dataset['target']
    total = X.shape[0]
    train_size = int(total * 0.7)
    indexes = np.arange(total)
    np.random.shuffle(indexes)
    train_index, test_index = indexes[:train_size], indexes[train_size:]
    X_train, Y_train = X[train_index, :], Y[train_index]
    X_test, Y_test = X[test_index, :], Y[test_index]
    tree = DummyDecisionTree()
    tree.fit(X_train, Y_train)
    print(tree.score(X_test, Y_test))