class SeedCLF:
    def __init__(self, id2word, aspects_ids, seed_weights=None, verbose=0, general_ind=4, general_thres=0,
                 hard_pred=False):
        self.id2word = id2word
        self.aspects_ids = aspects_ids
        self.num_aspects = len(aspects_ids)
        self.seed_weights = seed_weights
        self.general_ind = general_ind
        self.general_thres = general_thres
        self.seed_dict = self.create_seed_dict()

    def create_seed_dict(self):
        # Dict: seed_word: (aspect_id, seed_weight)
        seed_dict = {}
        for i, aspect_seeds in enumerate(self.aspects_ids):
            for j, word_id in enumerate(aspect_seeds):
                if self.seed_weights is None:
                    seed_dict[word_id] = (i, 1)
                else:
                    seed_dict[word_id] = (i, self.seed_weights[i][j])
        return seed_dict

    def predict(self, seg):
        seg = list(seg)
        aspect_scores = [0] * self.num_aspects

        for word_id in seg:
            if word_id in self.seed_dict:
                aspect_id, seed_weight = self.seed_dict[word_id]
                aspect_scores[aspect_id] += seed_weight

        if sum(aspect_scores) <= self.general_thres:
            aspect_scores[self.general_ind] = 1