import numpy as np
from AbstractTreeLearner import AbstractTreeLearner


class BagLearner(AbstractTreeLearner):
    def __init__(self, learner, bags=9, boost=False, verbose=False, kwargs={}):
        self.learner = learner
        self.verbose = verbose
        self.bags = bags
        self.learners = [learner(**kwargs) for _ in range(bags)]

    def get_bag(self, data_x, data_y):
        num_items = int(data_x.shape[0] * 0.5) # 50% of samples
        bag_x, bag_y = [], []
        for _ in range(num_items):
            i = np.random.randint(0, data_x.shape[0])
            bag_x.append(data_x[i,:])
            bag_y.append(data_y[i])
        return np.array(bag_x), np.array(bag_y)

    def addEvidence(self, data_x, data_y):
        """
        @summary: Add training data to learner
        @param dataX: X values of data to add
        @param dataY: the Y training values
        """
        for learner in self.learners:
            x, y = self.get_bag(data_x, data_y)
            learner.addEvidence(x, y)

    def query(self, points):
        """
        @summary: Estimate a set of test points given the model we built.
        @param points: numpy array with each row corresponding to a query.
        @returns the estimated values according to the saved model.
        """
        return np.mean([l.query(points) for l in self.learners], axis=0)